use std::collections::BTreeSet;
use std::fmt;
use std::num::NonZeroU32;
use std::pin::Pin;
use std::task::{Context, Poll};
use futures_util::Stream;
use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::Error as _};
use serde_json::Value;
use url::Url;
use crate::KitError;
use crate::core::content::base64_bytes;
use crate::core::{GenerateRequestView, selected_tools_for_choice};
use crate::error::ValidationViolation;
use crate::hosted::{HostedTool, HostedToolOutput};
use crate::models::ModelId;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum FileReference {
Provider {
id: String,
media_type: Option<String>,
},
Url {
url: Url,
media_type: Option<String>,
},
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum Artifact {
Inline {
media_type: String,
file_name: Option<String>,
description: Option<String>,
#[serde(with = "base64_bytes")]
data: Vec<u8>,
},
File {
reference: FileReference,
file_name: Option<String>,
description: Option<String>,
},
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum ContentPart {
Text {
text: String,
},
Thinking {
text: Option<String>,
},
ImageUrl {
url: Url,
},
Image {
media_type: String,
#[serde(with = "base64_bytes")]
data: Vec<u8>,
},
Audio {
media_type: String,
#[serde(with = "base64_bytes")]
data: Vec<u8>,
},
Document {
media_type: String,
#[serde(with = "base64_bytes")]
data: Vec<u8>,
file_name: Option<String>,
},
File {
reference: FileReference,
},
Json {
value: Value,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Role {
System,
Developer,
User,
Assistant,
Tool,
}
#[derive(Clone, PartialEq)]
pub struct GenerationContinuation {
pub(crate) model_id: ModelId,
pub(crate) inner: NativeContinuation,
}
#[derive(Clone, PartialEq)]
pub(crate) enum NativeContinuation {
OpenAi { output_items: Vec<Value> },
Anthropic { content_blocks: Vec<Value> },
Gemini { signed_parts: Vec<Value> },
Ollama { thinking: String },
}
impl fmt::Debug for GenerationContinuation {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("GenerationContinuation")
.field("model_id", &self.model_id)
.field("state", &"<redacted>")
.finish()
}
}
#[derive(Clone, PartialEq)]
pub struct Message<C = GenerationContinuation> {
role: Role,
content: Vec<ContentPart>,
function_calls: Vec<FunctionCall>,
function_results: Vec<FunctionResult>,
continuation: Option<C>,
}
impl<C> fmt::Debug for Message<C> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("Message")
.field("role", &self.role)
.field("content", &self.content)
.field("function_calls", &self.function_calls)
.field("function_results", &self.function_results)
.field(
"continuation",
&self.continuation.is_some().then_some("<redacted>"),
)
.finish()
}
}
impl<C> Message<C> {
fn from_role(role: Role, content: Vec<ContentPart>) -> Self {
Self {
role,
content,
function_calls: Vec::new(),
function_results: Vec::new(),
continuation: None,
}
}
pub fn system(content: Vec<ContentPart>) -> Result<Self, KitError> {
Self::instruction(Role::System, content)
}
pub fn developer(content: Vec<ContentPart>) -> Result<Self, KitError> {
Self::instruction(Role::Developer, content)
}
fn instruction(role: Role, content: Vec<ContentPart>) -> Result<Self, KitError> {
let message = Self::from_role(role, content);
message.validate_shape()?;
Ok(message)
}
pub fn user(content: Vec<ContentPart>) -> Result<Self, KitError> {
let message = Self::from_role(Role::User, content);
message.validate_shape()?;
Ok(message)
}
pub fn assistant(
content: Vec<ContentPart>,
function_calls: Vec<FunctionCall>,
) -> Result<Self, KitError> {
let message = Self {
function_calls,
..Self::from_role(Role::Assistant, content)
};
message.validate_shape()?;
Ok(message)
}
pub fn tool(results: Vec<FunctionResult>) -> Result<Self, KitError> {
let message = Self {
function_results: results,
..Self::from_role(Role::Tool, Vec::new())
};
message.validate_shape()?;
Ok(message)
}
pub fn user_text(text: impl Into<String>) -> Result<Self, KitError> {
Self::user(vec![ContentPart::Text { text: text.into() }])
}
pub fn assistant_text(text: impl Into<String>) -> Result<Self, KitError> {
Self::assistant(vec![ContentPart::Text { text: text.into() }], Vec::new())
}
pub fn with_continuation(mut self, continuation: C) -> Result<Self, KitError> {
if self.role != Role::Assistant {
return Err(KitError::invalid_input(
ValidationViolation::Incompatible {
field: "message.continuation",
other: "message.role",
},
"continuation can be attached only to assistant messages",
));
}
self.continuation = Some(continuation);
Ok(self)
}
#[must_use]
pub const fn role(&self) -> Role {
self.role
}
#[must_use]
pub fn content(&self) -> &[ContentPart] {
&self.content
}
#[must_use]
pub fn function_calls(&self) -> &[FunctionCall] {
&self.function_calls
}
#[must_use]
pub fn function_results(&self) -> &[FunctionResult] {
&self.function_results
}
#[must_use]
pub const fn continuation(&self) -> Option<&C> {
self.continuation.as_ref()
}
pub(crate) fn try_map_continuation<D>(
self,
map: &mut impl FnMut(C) -> Result<D, KitError>,
) -> Result<Message<D>, KitError> {
Ok(Message {
role: self.role,
content: self.content,
function_calls: self.function_calls,
function_results: self.function_results,
continuation: self.continuation.map(map).transpose()?,
})
}
pub(crate) fn returned_assistant(
content: Vec<ContentPart>,
function_calls: Vec<FunctionCall>,
continuation: Option<C>,
) -> Self {
Self {
role: Role::Assistant,
content,
function_calls,
function_results: Vec::new(),
continuation,
}
}
pub(crate) fn validate_for_request(&self) -> Result<(), KitError> {
self.validate_shape()?;
if self
.content
.iter()
.any(|part| matches!(part, ContentPart::Thinking { .. }))
&& self.continuation.is_none()
{
return Err(KitError::invalid_input(
ValidationViolation::MissingField {
field: "message.continuation",
},
"thinking input requires assistant continuation state",
));
}
Ok(())
}
fn validate_shape(&self) -> Result<(), KitError> {
for part in &self.content {
part.validate()?;
}
match self.role {
Role::System | Role::Developer => {
if self.content.is_empty()
|| self
.content
.iter()
.any(|part| !matches!(part, ContentPart::Text { .. }))
{
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "message.content",
rule: "instructions require non-empty text-only content",
},
"instruction messages require non-empty text-only content",
));
}
if !self.function_calls.is_empty()
|| !self.function_results.is_empty()
|| self.continuation.is_some()
{
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "message.role",
rule: "instructions cannot carry functions or continuation",
},
"instruction messages cannot carry functions or continuation state",
));
}
}
Role::User => {
if self.content.is_empty() {
return Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "message.content",
},
"user message content must not be empty",
));
}
if self.content.iter().any(|part| {
matches!(
part,
ContentPart::Thinking { .. } | ContentPart::Json { .. }
)
}) {
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "message.content",
rule: "user content cannot include thinking or raw JSON",
},
"user messages cannot contain thinking or raw JSON content",
));
}
if !self.function_calls.is_empty()
|| !self.function_results.is_empty()
|| self.continuation.is_some()
{
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "message.role",
rule: "user messages cannot carry functions or continuation",
},
"user messages cannot carry functions or continuation state",
));
}
}
Role::Assistant => {
if self.content.is_empty()
&& self.function_calls.is_empty()
&& self.continuation.is_none()
{
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "message",
rule: "assistant requires content, calls, or continuation",
},
"assistant message requires content, a function call, or retained continuation",
));
}
if self.content.iter().any(|part| {
!matches!(
part,
ContentPart::Text { .. } | ContentPart::Thinking { .. }
)
}) {
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "message.content",
rule: "assistant input requires text or continued thinking",
},
"assistant input supports only text and continued thinking content",
));
}
if !self.function_results.is_empty() {
return Err(KitError::invalid_input(
ValidationViolation::Incompatible {
field: "message.function_results",
other: "message.role",
},
"assistant messages cannot carry function results",
));
}
}
Role::Tool => {
if !self.content.is_empty()
|| !self.function_calls.is_empty()
|| self.function_results.is_empty()
|| self.continuation.is_some()
{
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "message",
rule: "tool messages require only function results",
},
"tool messages require function results and no direct content, calls, or continuation",
));
}
}
}
Ok(())
}
}
impl GenerationContinuation {
pub(crate) fn open_ai(model_id: ModelId, output_items: Vec<Value>) -> Self {
Self {
model_id,
inner: NativeContinuation::OpenAi { output_items },
}
}
pub(crate) fn anthropic(model_id: ModelId, content_blocks: Vec<Value>) -> Self {
Self {
model_id,
inner: NativeContinuation::Anthropic { content_blocks },
}
}
pub(crate) fn gemini(model_id: ModelId, signed_parts: Vec<Value>) -> Self {
Self {
model_id,
inner: NativeContinuation::Gemini { signed_parts },
}
}
pub(crate) fn ollama(model_id: ModelId, thinking: String) -> Self {
Self {
model_id,
inner: NativeContinuation::Ollama { thinking },
}
}
}
impl<C> Serialize for Message<C> {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
if self.continuation.is_some() {
return Err(S::Error::custom(
"provider continuation cannot be serialized; retain the returned message in memory",
));
}
#[derive(Serialize)]
struct Wire<'a> {
role: Role,
content: &'a [ContentPart],
function_calls: &'a [FunctionCall],
function_results: &'a [FunctionResult],
}
Wire {
role: self.role,
content: &self.content,
function_calls: &self.function_calls,
function_results: &self.function_results,
}
.serialize(serializer)
}
}
impl<'de, C> Deserialize<'de> for Message<C> {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Wire {
role: Role,
#[serde(default)]
content: Vec<ContentPart>,
#[serde(default)]
function_calls: Vec<FunctionCall>,
#[serde(default)]
function_results: Vec<FunctionResult>,
}
let wire = Wire::deserialize(deserializer)?;
let message = Self {
role: wire.role,
content: wire.content,
function_calls: wire.function_calls,
function_results: wire.function_results,
continuation: None,
};
message
.validate_for_request()
.map_err(serde::de::Error::custom)?;
Ok(message)
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct JsonSchema(Value);
impl JsonSchema {
pub fn new(schema: Value) -> Result<Self, KitError> {
let object = schema.as_object().ok_or_else(|| {
KitError::invalid_input(
ValidationViolation::InvalidFormat {
field: "json_schema",
rule: "JSON object",
},
"JSON schema must be encoded as an object",
)
})?;
if !matches!(object.get("type"), Some(Value::String(kind)) if kind == "object") {
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "json_schema.type",
rule: "must declare object",
},
"JSON schema root must declare `type: object`",
));
}
Ok(Self(schema))
}
#[must_use]
pub const fn schema(&self) -> &Value {
&self.0
}
#[must_use]
pub fn into_value(self) -> Value {
self.0
}
}
impl Serialize for JsonSchema {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
self.0.serialize(serializer)
}
}
impl<'de> Deserialize<'de> for JsonSchema {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
Self::new(Value::deserialize(deserializer)?).map_err(serde::de::Error::custom)
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum ResponseFormat {
Text,
JsonObject,
JsonSchema(JsonSchema),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReasoningEffort {
Minimal,
Low,
Medium,
High,
Maximum,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReasoningSummaryMode {
Auto,
Hidden,
Summary,
Verbose,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ReasoningConfig {
pub effort: Option<ReasoningEffort>,
pub max_reasoning_tokens: Option<u32>,
pub summary_mode: ReasoningSummaryMode,
}
impl Default for ReasoningConfig {
fn default() -> Self {
Self {
effort: None,
max_reasoning_tokens: None,
summary_mode: ReasoningSummaryMode::Auto,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AllowedToolMode {
Auto,
Required,
}
impl Serialize for AllowedToolMode {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
self.as_wire().serialize(serializer)
}
}
impl<'de> Deserialize<'de> for AllowedToolMode {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
Self::from_wire(&value)
}
}
impl AllowedToolMode {
pub(crate) fn as_wire(self) -> &'static str {
match self {
Self::Auto => "auto",
Self::Required => "required",
}
}
fn from_wire<E>(value: &str) -> Result<Self, E>
where
E: serde::de::Error,
{
match value {
"required" => Ok(Self::Required),
"auto" => Ok(Self::Auto),
other => Err(E::custom(format!("unknown allowed tool mode `{other}`"))),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ToolChoice {
Auto,
None,
Required,
Any,
Tool {
name: String,
},
Allowed {
mode: AllowedToolMode,
names: Vec<String>,
},
}
impl Serialize for ToolChoice {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
self.as_wire().serialize(serializer)
}
}
impl<'de> Deserialize<'de> for ToolChoice {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
match Value::deserialize(deserializer)? {
Value::String(value) => match value.as_str() {
"auto" => Ok(Self::Auto),
"none" => Ok(Self::None),
"required" => Ok(Self::Required),
"any" => Ok(Self::Any),
other => Err(serde::de::Error::custom(format!(
"unknown tool choice variant `{other}`"
))),
},
Value::Object(object) => {
serde_json::from_value::<ToolChoiceObjectWire>(Value::Object(object))
.map_err(serde::de::Error::custom)?
.into_tool_choice()
}
_ => Err(serde::de::Error::custom(
"tool choice must be encoded as a string or object",
)),
}
}
}
impl ToolChoice {
fn as_wire(&self) -> ToolChoiceWireRef<'_> {
match self {
Self::Auto => ToolChoiceWireRef::String("auto"),
Self::None => ToolChoiceWireRef::String("none"),
Self::Required => ToolChoiceWireRef::String("required"),
Self::Any => ToolChoiceWireRef::String("any"),
Self::Tool { name } => ToolChoiceWireRef::Object(ToolChoiceObjectWireRef {
kind: "tool",
mode: None,
name: Some(name),
names: None,
}),
Self::Allowed { mode, names } => ToolChoiceWireRef::Object(ToolChoiceObjectWireRef {
kind: "allowed",
mode: Some(mode.as_wire()),
name: None,
names: Some(names.as_slice()),
}),
}
}
}
#[derive(Serialize)]
#[serde(untagged)]
enum ToolChoiceWireRef<'a> {
String(&'static str),
Object(ToolChoiceObjectWireRef<'a>),
}
#[derive(Serialize)]
struct ToolChoiceObjectWireRef<'a> {
#[serde(rename = "type")]
kind: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
mode: Option<&'static str>,
#[serde(skip_serializing_if = "Option::is_none")]
name: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
names: Option<&'a [String]>,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ToolChoiceObjectWire {
#[serde(rename = "type", default)]
kind: Option<String>,
#[serde(default)]
mode: Option<String>,
#[serde(default)]
name: Option<String>,
#[serde(default)]
names: Option<Value>,
}
impl ToolChoiceObjectWire {
fn into_tool_choice<E>(self) -> Result<ToolChoice, E>
where
E: serde::de::Error,
{
let kind = self
.kind
.ok_or_else(|| E::custom("tool choice objects require a string `type` field"))?;
match kind.as_str() {
"tool" => {
let name = self
.name
.filter(|name| !name.trim().is_empty())
.ok_or_else(|| E::custom("tool choice objects require a non-empty `name`"))?;
Ok(ToolChoice::Tool { name })
}
"allowed" => Ok(ToolChoice::Allowed {
mode: self
.mode
.as_deref()
.map(AllowedToolMode::from_wire)
.transpose()?
.unwrap_or(AllowedToolMode::Auto),
names: parse_allowed_tool_names(self.names)?,
}),
other => Err(E::custom(format!(
"unknown tool choice object type `{other}`"
))),
}
}
}
fn parse_allowed_tool_names<E>(value: Option<Value>) -> Result<Vec<String>, E>
where
E: serde::de::Error,
{
match value {
None => Ok(Vec::new()),
Some(Value::Array(items)) => items
.into_iter()
.map(|item| match item {
Value::String(name) if !name.trim().is_empty() => Ok(name),
Value::String(_) => {
Err(E::custom("allowed tool entries must be non-empty strings"))
}
_ => Err(E::custom("allowed tool entries must be strings")),
})
.collect(),
Some(_) => Err(E::custom("allowed tool names must be encoded as an array")),
}
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct ProviderFunctionTool {
name: String,
description: Option<String>,
input_schema: JsonSchema,
}
impl ProviderFunctionTool {
pub fn new(
name: impl Into<String>,
description: Option<String>,
input_schema: JsonSchema,
) -> Result<Self, KitError> {
let name = name.into();
FunctionCall::validate_name("tools.function.name", &name)?;
if description
.as_ref()
.is_some_and(|value| value.trim().is_empty())
{
return Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "tools.function.description",
},
"function tool description must be omitted or non-empty",
));
}
Ok(Self {
name,
description,
input_schema,
})
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn description(&self) -> Option<&str> {
self.description.as_deref()
}
#[must_use]
pub const fn input_schema(&self) -> &JsonSchema {
&self.input_schema
}
#[must_use]
pub fn into_parts(self) -> (String, Option<String>, JsonSchema) {
(self.name, self.description, self.input_schema)
}
}
impl<'de> Deserialize<'de> for ProviderFunctionTool {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Wire {
name: String,
description: Option<String>,
input_schema: JsonSchema,
}
let wire = Wire::deserialize(deserializer)?;
Self::new(wire.name, wire.description, wire.input_schema).map_err(serde::de::Error::custom)
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", content = "declaration", rename_all = "snake_case")]
pub enum ProviderToolSpec {
Function(ProviderFunctionTool),
Hosted(HostedTool),
}
impl ProviderToolSpec {
pub(crate) const fn hosted_tool(&self) -> Option<&HostedTool> {
match self {
Self::Function(_) => None,
Self::Hosted(tool) => Some(tool),
}
}
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct FunctionCall {
id: Option<String>,
name: String,
arguments: Value,
}
impl<'de> Deserialize<'de> for FunctionCall {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Wire {
id: Option<String>,
name: String,
arguments: Value,
}
let wire = Wire::deserialize(deserializer)?;
Self::new(wire.id, wire.name, wire.arguments).map_err(serde::de::Error::custom)
}
}
impl FunctionCall {
pub fn new(
id: Option<String>,
name: impl Into<String>,
arguments: Value,
) -> Result<Self, KitError> {
let name = name.into();
Self::validate_name("function_call.name", &name)?;
if let Some(id) = id.as_deref() {
Self::validate_id("function_call.id", id)?;
}
if !arguments.is_object() {
return Err(KitError::invalid_input(
ValidationViolation::InvalidFormat {
field: "function_call.arguments",
rule: "JSON object",
},
"function arguments must be a JSON object",
));
}
Ok(Self {
id,
name,
arguments,
})
}
#[must_use]
pub fn id(&self) -> Option<&str> {
self.id.as_deref()
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn arguments(&self) -> &Value {
&self.arguments
}
pub(crate) fn validate_name(label: &'static str, value: &str) -> Result<(), KitError> {
let mut bytes = value.bytes();
let first = bytes.next();
if value.len() > 64
|| !first.is_some_and(|byte| byte.is_ascii_alphabetic() || byte == b'_')
|| !bytes.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
{
Err(KitError::invalid_input(
ValidationViolation::InvalidFormat {
field: label,
rule: "1..=64 ASCII letters, digits, underscores, or hyphens; first character letter or underscore",
},
format!(
"{label} must use 1..=64 portable ASCII letters, digits, underscores, or hyphens and start with a letter or underscore"
),
))
} else {
Ok(())
}
}
fn validate_id(label: &'static str, value: &str) -> Result<(), KitError> {
if value.is_empty()
|| value.len() > 256
|| value.trim() != value
|| value.chars().any(char::is_control)
{
Err(KitError::invalid_input(
ValidationViolation::InvalidFormat {
field: label,
rule: "1..=256 unpadded UTF-8 bytes without control characters",
},
format!(
"{label} must contain 1..=256 unpadded UTF-8 bytes without control characters"
),
))
} else {
Ok(())
}
}
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct FunctionResult {
call_id: Option<String>,
name: String,
content: Vec<ContentPart>,
is_error: bool,
}
impl<'de> Deserialize<'de> for FunctionResult {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Wire {
call_id: Option<String>,
name: String,
content: Vec<ContentPart>,
is_error: bool,
}
let wire = Wire::deserialize(deserializer)?;
Self::new(wire.call_id, wire.name, wire.content, wire.is_error)
.map_err(serde::de::Error::custom)
}
}
impl FunctionResult {
pub fn new(
call_id: Option<String>,
name: impl Into<String>,
content: Vec<ContentPart>,
is_error: bool,
) -> Result<Self, KitError> {
let name = name.into();
FunctionCall::validate_name("function_result.name", &name)?;
if let Some(call_id) = call_id.as_deref() {
FunctionCall::validate_id("function_result.call_id", call_id)?;
}
if content.is_empty() {
return Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "function_result.content",
},
"function result content must not be empty",
));
}
for part in &content {
part.validate()?;
}
if content
.iter()
.any(|part| matches!(part, ContentPart::Thinking { .. }))
{
return Err(KitError::invalid_input(
ValidationViolation::Unsupported {
field: "function_result.content",
rule: "thinking content is not permitted",
},
"function results cannot contain provider thinking state",
));
}
Ok(Self {
call_id,
name,
content,
is_error,
})
}
#[must_use]
pub fn call_id(&self) -> Option<&str> {
self.call_id.as_deref()
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn content(&self) -> &[ContentPart] {
&self.content
}
#[must_use]
pub const fn is_error(&self) -> bool {
self.is_error
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Default)]
#[serde(deny_unknown_fields)]
pub struct GenerationConfig {
temperature: Option<f32>,
top_p: Option<f32>,
top_k: Option<NonZeroU32>,
max_output_tokens: Option<NonZeroU32>,
stop_sequences: Vec<String>,
presence_penalty: Option<f32>,
frequency_penalty: Option<f32>,
seed: Option<u64>,
tool_choice: Option<ToolChoice>,
}
impl GenerationConfig {
pub fn with_temperature(mut self, value: f32) -> Result<Self, KitError> {
if !value.is_finite() || value < 0.0 {
return Err(KitError::invalid_input(
ValidationViolation::OutOfRange {
field: "generation.temperature",
rule: "finite and non-negative",
},
"temperature must be finite and non-negative",
));
}
self.temperature = Some(value);
Ok(self)
}
pub fn with_top_p(mut self, value: f32) -> Result<Self, KitError> {
if !(value.is_finite() && 0.0 < value && value <= 1.0) {
return Err(KitError::invalid_input(
ValidationViolation::OutOfRange {
field: "generation.top_p",
rule: "finite and in (0, 1]",
},
"top_p must be in (0, 1]",
));
}
self.top_p = Some(value);
Ok(self)
}
#[must_use]
pub fn with_top_k(mut self, value: NonZeroU32) -> Self {
self.top_k = Some(value);
self
}
#[must_use]
pub fn with_max_output_tokens(mut self, value: NonZeroU32) -> Self {
self.max_output_tokens = Some(value);
self
}
pub fn with_stop_sequences(mut self, values: Vec<String>) -> Result<Self, KitError> {
let mut seen = BTreeSet::new();
for value in &values {
if value.trim().is_empty() {
return Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "generation.stop_sequences",
},
"stop sequences must not be empty",
));
}
if !seen.insert(value) {
return Err(KitError::invalid_input(
ValidationViolation::Duplicate {
field: "generation.stop_sequences",
},
"stop sequences must be unique",
));
}
}
self.stop_sequences = values;
Ok(self)
}
pub fn with_presence_penalty(mut self, value: f32) -> Result<Self, KitError> {
self.presence_penalty = Some(Self::penalty("generation.presence_penalty", value)?);
Ok(self)
}
pub fn with_frequency_penalty(mut self, value: f32) -> Result<Self, KitError> {
self.frequency_penalty = Some(Self::penalty("generation.frequency_penalty", value)?);
Ok(self)
}
fn penalty(name: &'static str, value: f32) -> Result<f32, KitError> {
if value.is_finite() && (-2.0..=2.0).contains(&value) {
Ok(value)
} else {
Err(KitError::invalid_input(
ValidationViolation::OutOfRange {
field: name,
rule: "finite and in [-2, 2]",
},
format!("{name} must be finite and in [-2, 2]"),
))
}
}
#[must_use]
pub fn with_seed(mut self, value: u64) -> Self {
self.seed = Some(value);
self
}
#[must_use]
pub fn with_tool_choice(mut self, value: ToolChoice) -> Self {
self.tool_choice = Some(value);
self
}
#[must_use]
pub const fn temperature(&self) -> Option<f32> {
self.temperature
}
#[must_use]
pub const fn top_p(&self) -> Option<f32> {
self.top_p
}
#[must_use]
pub const fn top_k(&self) -> Option<NonZeroU32> {
self.top_k
}
#[must_use]
pub const fn max_output_tokens(&self) -> Option<NonZeroU32> {
self.max_output_tokens
}
#[must_use]
pub fn stop_sequences(&self) -> &[String] {
&self.stop_sequences
}
#[must_use]
pub const fn presence_penalty(&self) -> Option<f32> {
self.presence_penalty
}
#[must_use]
pub const fn frequency_penalty(&self) -> Option<f32> {
self.frequency_penalty
}
#[must_use]
pub const fn seed(&self) -> Option<u64> {
self.seed
}
#[must_use]
pub const fn tool_choice(&self) -> Option<&ToolChoice> {
self.tool_choice.as_ref()
}
}
impl<'de> Deserialize<'de> for GenerationConfig {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize, Default)]
#[serde(deny_unknown_fields)]
struct Wire {
temperature: Option<f32>,
top_p: Option<f32>,
top_k: Option<NonZeroU32>,
max_output_tokens: Option<NonZeroU32>,
#[serde(default)]
stop_sequences: Vec<String>,
presence_penalty: Option<f32>,
frequency_penalty: Option<f32>,
seed: Option<u64>,
tool_choice: Option<ToolChoice>,
}
let wire = Wire::deserialize(deserializer)?;
let mut config = Self::default();
if let Some(value) = wire.temperature {
config = config
.with_temperature(value)
.map_err(serde::de::Error::custom)?;
}
if let Some(value) = wire.top_p {
config = config.with_top_p(value).map_err(serde::de::Error::custom)?;
}
config.top_k = wire.top_k;
config.max_output_tokens = wire.max_output_tokens;
config = config
.with_stop_sequences(wire.stop_sequences)
.map_err(serde::de::Error::custom)?;
if let Some(value) = wire.presence_penalty {
config = config
.with_presence_penalty(value)
.map_err(serde::de::Error::custom)?;
}
if let Some(value) = wire.frequency_penalty {
config = config
.with_frequency_penalty(value)
.map_err(serde::de::Error::custom)?;
}
config.seed = wire.seed;
config.tool_choice = wire.tool_choice;
Ok(config)
}
}
#[derive(Clone, Debug, PartialEq, Serialize)]
#[serde(bound(serialize = ""), deny_unknown_fields)]
pub struct GenerateRequest<C = GenerationContinuation> {
messages: Vec<Message<C>>,
tools: Vec<ProviderToolSpec>,
response_format: Option<ResponseFormat>,
reasoning: Option<ReasoningConfig>,
generation: GenerationConfig,
}
impl<C> GenerateRequest<C> {
pub fn new(messages: Vec<Message<C>>) -> Result<Self, KitError> {
let request = Self {
messages,
tools: Vec::new(),
response_format: None,
reasoning: None,
generation: GenerationConfig::default(),
};
request.validate()?;
Ok(request)
}
pub fn with_tools(mut self, tools: Vec<ProviderToolSpec>) -> Result<Self, KitError> {
self.tools = tools;
self.validate_tools()?;
Ok(self)
}
#[must_use]
pub fn with_response_format(mut self, value: ResponseFormat) -> Self {
self.response_format = Some(value);
self
}
#[must_use]
pub fn with_reasoning(mut self, value: ReasoningConfig) -> Self {
self.reasoning = Some(value);
self
}
pub fn with_generation(mut self, value: GenerationConfig) -> Result<Self, KitError> {
self.generation = value;
self.validate_tools()?;
Ok(self)
}
#[must_use]
pub fn messages(&self) -> &[Message<C>] {
&self.messages
}
#[must_use]
pub fn tools(&self) -> &[ProviderToolSpec] {
&self.tools
}
#[must_use]
pub const fn response_format(&self) -> Option<&ResponseFormat> {
self.response_format.as_ref()
}
#[must_use]
pub const fn reasoning(&self) -> Option<&ReasoningConfig> {
self.reasoning.as_ref()
}
#[must_use]
pub const fn generation(&self) -> &GenerationConfig {
&self.generation
}
pub fn try_map_continuation<D>(
self,
mut map: impl FnMut(C) -> Result<D, KitError>,
) -> Result<GenerateRequest<D>, KitError> {
Ok(GenerateRequest {
messages: self
.messages
.into_iter()
.map(|message| message.try_map_continuation(&mut map))
.collect::<Result<_, _>>()?,
tools: self.tools,
response_format: self.response_format,
reasoning: self.reasoning,
generation: self.generation,
})
}
fn validate(&self) -> Result<(), KitError> {
if self.messages.is_empty() {
return Err(KitError::invalid_input(
ValidationViolation::Empty {
field: "generation.messages",
},
"generation request messages must not be empty",
));
}
let mut saw_ordinary = false;
for message in &self.messages {
match message.role() {
crate::generation::Role::System | crate::generation::Role::Developer
if saw_ordinary =>
{
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "generation.messages",
rule: "instructions must form a leading prefix",
},
"system and developer instructions must form one leading prefix",
));
}
crate::generation::Role::System | crate::generation::Role::Developer => {}
_ => saw_ordinary = true,
}
message.validate_for_request()?;
}
if !saw_ordinary {
return Err(KitError::invalid_input(
ValidationViolation::InvalidValue {
field: "generation.messages",
rule: "requires a non-instruction message",
},
"generation request requires at least one non-instruction message",
));
}
self.validate_tools()
}
fn validate_tools(&self) -> Result<(), KitError> {
let mut names = BTreeSet::new();
for tool in &self.tools {
if let ProviderToolSpec::Function(function) = tool
&& !names.insert(function.name())
{
return Err(KitError::invalid_input(
ValidationViolation::Duplicate {
field: "tools.function.name",
},
format!("duplicate function tool name `{}`", function.name()),
));
}
}
selected_tools_for_choice(&self.tools, self.generation.tool_choice())?;
Ok(())
}
#[must_use]
pub(crate) fn as_view(&self) -> GenerateRequestView<'_, C> {
GenerateRequestView {
messages: &self.messages,
tools: &self.tools,
response_format: self.response_format.as_ref(),
reasoning: self.reasoning.as_ref(),
generation: &self.generation,
}
}
}
impl<'de, C> Deserialize<'de> for GenerateRequest<C> {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(bound(deserialize = ""), deny_unknown_fields)]
struct Wire<C> {
messages: Vec<Message<C>>,
#[serde(default)]
tools: Vec<ProviderToolSpec>,
response_format: Option<ResponseFormat>,
reasoning: Option<ReasoningConfig>,
#[serde(default)]
generation: GenerationConfig,
}
let wire = Wire::deserialize(deserializer)?;
let mut request = Self::new(wire.messages).map_err(serde::de::Error::custom)?;
request.tools = wire.tools;
request.response_format = wire.response_format;
request.reasoning = wire.reasoning;
request.generation = wire.generation;
request.validate().map_err(serde::de::Error::custom)?;
Ok(request)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {
Stop,
Length,
ToolCall,
ContentFilter,
Error,
Paused,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum Citation {
Url {
url: String,
title: Option<String>,
snippet: Option<String>,
start_index: Option<usize>,
end_index: Option<usize>,
},
File {
reference: FileReference,
title: Option<String>,
snippet: Option<String>,
start_index: Option<usize>,
end_index: Option<usize>,
},
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(bound(serialize = "", deserialize = ""))]
pub struct GenerateResponse<C = GenerationContinuation> {
pub model: ModelId,
pub message: Message<C>,
pub finish_reason: Option<FinishReason>,
pub citations: Vec<Citation>,
pub artifacts: Vec<Artifact>,
pub hosted_tool_outputs: Vec<HostedToolOutput>,
pub usage: Usage,
}
impl<C> GenerateResponse<C> {
pub fn try_map_continuation<D>(
self,
mut map: impl FnMut(C) -> Result<D, KitError>,
) -> Result<GenerateResponse<D>, KitError> {
Ok(GenerateResponse {
model: self.model,
message: self.message.try_map_continuation(&mut map)?,
finish_reason: self.finish_reason,
citations: self.citations,
artifacts: self.artifacts,
hosted_tool_outputs: self.hosted_tool_outputs,
usage: self.usage,
})
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(
bound(serialize = "", deserialize = ""),
tag = "event",
rename_all = "snake_case"
)]
pub enum GenerateEvent<C = GenerationContinuation> {
Started {
model: ModelId,
},
ContentDelta {
delta: String,
},
ThinkingDelta {
delta: String,
},
HostedTool(Box<HostedToolOutput>),
Usage(Usage),
Completed(Box<GenerateResponse<C>>),
}
impl<C> GenerateEvent<C> {
pub fn try_map_continuation<D>(
self,
map: impl FnMut(C) -> Result<D, KitError>,
) -> Result<GenerateEvent<D>, KitError> {
Ok(match self {
Self::Started { model } => GenerateEvent::Started { model },
Self::ContentDelta { delta } => GenerateEvent::ContentDelta { delta },
Self::ThinkingDelta { delta } => GenerateEvent::ThinkingDelta { delta },
Self::HostedTool(output) => GenerateEvent::HostedTool(output),
Self::Usage(usage) => GenerateEvent::Usage(usage),
Self::Completed(response) => {
GenerateEvent::Completed(Box::new(response.try_map_continuation(map)?))
}
})
}
}
pub struct GenerateEventStream<C = GenerationContinuation> {
inner: Pin<Box<dyn Stream<Item = Result<GenerateEvent<C>, KitError>> + Send + 'static>>,
}
impl<C: 'static> GenerateEventStream<C> {
pub fn new<S>(stream: S) -> Self
where
S: Stream<Item = Result<GenerateEvent<C>, KitError>> + Send + 'static,
{
Self {
inner: Box::pin(stream),
}
}
}
impl<C> Stream for GenerateEventStream<C> {
type Item = Result<GenerateEvent<C>, KitError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.inner.as_mut().poll_next(cx)
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub reasoning_tokens: Option<u64>,
pub cached_input_tokens: Option<u64>,
pub total_tokens: Option<u64>,
}
impl Usage {
pub(crate) fn reported_tokens(input_tokens: Option<u64>, output_tokens: Option<u64>) -> Self {
Self {
input_tokens,
output_tokens,
..Self::default()
}
}
pub(crate) fn with_reported_total(mut self, total_tokens: Option<u64>) -> Self {
self.total_tokens = total_tokens;
self
}
pub(crate) fn with_cached_input(mut self, cached_input_tokens: Option<u64>) -> Self {
self.cached_input_tokens = cached_input_tokens;
self
}
pub(crate) fn with_reasoning(mut self, reasoning_tokens: Option<u64>) -> Self {
self.reasoning_tokens = reasoning_tokens;
self
}
}