Skip to main content

infinity_provider_protocol/
completion.rs

1//! Completion request, response streaming, and error types.
2
3use std::pin::Pin;
4
5use futures_util::Stream;
6use serde::{Deserialize, Serialize};
7
8use crate::message::{Message, Reasoning, ToolCall};
9
10/// A request to stream one completion from a model.
11///
12/// This is the entire API surface providers must support; anything
13/// provider-specific (thinking configuration, beta flags, ...) travels in
14/// `additional_params` or is applied by the provider itself.
15#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
16pub struct CompletionRequest {
17    /// The system prompt.
18    pub preamble: Option<String>,
19    /// The conversation so far, ending with the content to complete.
20    pub chat_history: Vec<Message>,
21    /// Definitions of the tools available to the model.
22    pub tools: Vec<ToolDefinition>,
23    /// Maximum number of output tokens. `None` uses the provider's default.
24    pub max_tokens: Option<u64>,
25    /// Extra provider-specific request parameters, merged into the request
26    /// by the provider.
27    pub additional_params: Option<serde_json::Value>,
28}
29
30/// Definition of one tool the model may call.
31#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
32pub struct ToolDefinition {
33    pub name: String,
34    pub description: String,
35    /// JSON schema of the tool's arguments.
36    pub parameters: serde_json::Value,
37}
38
39/// Token usage reported by a provider for one completion.
40#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
41pub struct Usage {
42    /// The number of input ("prompt") tokens used in a given request.
43    #[serde(default)]
44    pub input_tokens: u64,
45    /// The number of output ("completion") tokens used in a given request.
46    #[serde(default)]
47    pub output_tokens: u64,
48    /// Stored separately as some providers only report one number.
49    #[serde(default)]
50    pub total_tokens: u64,
51    /// The number of cached input tokens (from prompt caching). 0 if not
52    /// reported by the provider.
53    #[serde(default)]
54    pub cached_input_tokens: u64,
55}
56
57/// Provider-declared classification of a completion error, telling callers
58/// how to react. Classification is the provider's job — it knows its own
59/// failure modes — so the agent runtime never has to parse error message
60/// strings.
61#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
62pub enum ErrorClass {
63    /// Transient failure (dropped stream, backend hiccup, request timeout);
64    /// retrying the same request may succeed.
65    Transient,
66    /// Rate limited / throttled; retrying the same request may succeed
67    /// after a longer backoff.
68    Throttled,
69    /// The request's input does not fit the model's context window;
70    /// retrying the same request can never succeed. Callers must shrink the
71    /// input (drop or truncate oversized messages, compact history).
72    ContextOverflow,
73    /// Permanent failure (bad request, unknown model, access denied); do
74    /// not retry.
75    Fatal,
76}
77
78/// Error invoking a completion model.
79#[derive(Debug, thiserror::Error)]
80pub enum CompletionError {
81    /// Error building the completion request.
82    #[error("RequestError: {0}")]
83    RequestError(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
84
85    /// Error parsing the completion response.
86    #[error("ResponseError: {0}")]
87    ResponseError(String),
88
89    /// Error returned by the completion model provider, with the provider's
90    /// retry classification.
91    #[error("ProviderError: {message}")]
92    ProviderError { message: String, class: ErrorClass },
93
94    /// JSON (de)serialization error.
95    #[error("JsonError: {0}")]
96    JsonError(#[from] serde_json::Error),
97}
98
99impl CompletionError {
100    /// A provider error with an explicit retry classification.
101    pub fn provider(class: ErrorClass, message: impl Into<String>) -> Self {
102        Self::ProviderError {
103            message: message.into(),
104            class,
105        }
106    }
107
108    /// How callers should treat this error.
109    ///
110    /// * [`RequestError`](Self::RequestError) and
111    ///   [`JsonError`](Self::JsonError) are our side failing to build or
112    ///   parse — retrying the same request cannot help ([`ErrorClass::Fatal`]).
113    /// * [`ResponseError`](Self::ResponseError) is a transport/framing
114    ///   failure while reading the response — typically transient.
115    /// * [`ProviderError`](Self::ProviderError) carries the provider's own
116    ///   classification.
117    pub fn class(&self) -> ErrorClass {
118        match self {
119            Self::RequestError(_) | Self::JsonError(_) => ErrorClass::Fatal,
120            Self::ResponseError(_) => ErrorClass::Transient,
121            Self::ProviderError { class, .. } => *class,
122        }
123    }
124}
125
126/// One streamed item of a completion response.
127///
128/// Deltas stream incremental UI-facing content; the non-delta variants
129/// (`ToolCall`, `Reasoning`) carry the complete block to be recorded in
130/// history. A well-formed stream ends with `Final` (though providers may
131/// terminate without one, e.g. on network errors).
132#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
133pub enum StreamChunk {
134    /// A chunk of assistant text.
135    Text(String),
136    /// A complete tool call (all arguments accumulated).
137    ToolCall(ToolCall),
138    /// An incremental piece of an in-progress tool call.
139    ToolCallDelta {
140        /// Provider-supplied tool call id.
141        id: String,
142        content: ToolCallDeltaContent,
143    },
144    /// A complete reasoning block.
145    Reasoning(Reasoning),
146    /// A chunk of reasoning text.
147    ReasoningDelta { id: Option<String>, text: String },
148    /// The final chunk, carrying the completion's token usage (if the
149    /// provider reports it).
150    Final(FinalResponse),
151}
152
153/// The content of a [`StreamChunk::ToolCallDelta`].
154#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
155pub enum ToolCallDeltaContent {
156    /// The name of the tool being called (start of the call).
157    Name(String),
158    /// A fragment of the call's JSON arguments.
159    Delta(String),
160}
161
162/// Payload of [`StreamChunk::Final`].
163#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
164pub struct FinalResponse {
165    pub usage: Option<Usage>,
166}
167
168/// The streaming response returned by a model invocation.
169pub type ModelStream = Pin<Box<dyn Stream<Item = Result<StreamChunk, CompletionError>> + Send>>;