Skip to main content

ic_rig/
completion.rs

1//! Completion model trait and associated request/response types.
2//!
3//! [`CompletionModel`] is the single interface that provider implementations
4//! must satisfy. The [`Agent`](crate::agent::Agent) and any other high-level
5//! construct in this crate are generic over it.
6//!
7//! # Implementing a provider
8//!
9//! ```rust,ignore
10//! use irig::completion::{
11//!     CompletionModel, CompletionRequest, CompletionResponse, CompletionError,
12//! };
13//!
14//! pub struct MyModel<H> {
15//!     client: H,
16//!     model: String,
17//!     api_key: String,
18//! }
19//!
20//! impl<H: HttpClient> CompletionModel for MyModel<H> {
21//!     type Error = CompletionError;
22//!
23//!     async fn complete(
24//!         &self,
25//!         request: CompletionRequest,
26//!     ) -> Result<CompletionResponse, Self::Error> {
27//!         // 1. Serialise `request` into the provider's API format.
28//!         // 2. Call `self.client.post(...)`.
29//!         // 3. Deserialise the response into `CompletionResponse`.
30//!         todo!()
31//!     }
32//! }
33//! ```
34
35use crate::message::{Message, ToolCall};
36use crate::tool::ToolDefinition;
37use thiserror::Error;
38
39// ── Request ───────────────────────────────────────────────────────────────────
40
41/// Everything a completion model needs to generate a response.
42#[derive(Debug, Clone)]
43pub struct CompletionRequest {
44    /// Conversation history, including the latest user turn.
45    pub messages: Vec<Message>,
46    /// Tools the model is allowed to call.
47    pub tools: Vec<ToolDefinition>,
48    /// Optional sampling temperature (provider-specific range).
49    pub temperature: Option<f64>,
50    /// Optional maximum number of tokens to generate.
51    pub max_tokens: Option<u32>,
52    /// Provider-specific extra parameters (passed through as-is).
53    pub extra: Option<serde_json::Value>,
54}
55
56impl CompletionRequest {
57    pub fn new(messages: Vec<Message>) -> Self {
58        Self {
59            messages,
60            tools: Vec::new(),
61            temperature: None,
62            max_tokens: None,
63            extra: None,
64        }
65    }
66}
67
68// ── Response ──────────────────────────────────────────────────────────────────
69
70/// The model's choice of response for a single completion turn.
71#[derive(Debug, Clone)]
72pub enum ModelChoice {
73    /// The model produced a text reply.
74    Message(String),
75    /// The model wants to call one or more tools.
76    ToolCall(Vec<ToolCall>),
77}
78
79/// Token usage reported by the provider (optional — not all providers expose this).
80#[derive(Debug, Clone, Default)]
81pub struct Usage {
82    pub prompt_tokens: u32,
83    pub completion_tokens: u32,
84}
85
86/// The full response returned by a [`CompletionModel`].
87#[derive(Debug, Clone)]
88pub struct CompletionResponse {
89    /// What the model decided to do this turn.
90    pub choice: ModelChoice,
91    /// Token usage, if the provider reported it.
92    pub usage: Option<Usage>,
93}
94
95// ── Error ─────────────────────────────────────────────────────────────────────
96
97/// Errors that can occur during a completion request.
98#[derive(Debug, Error)]
99pub enum CompletionError {
100    /// The HTTP transport returned an error.
101    #[error("HTTP error: {0}")]
102    Http(String),
103
104    /// The response body could not be deserialised.
105    #[error("JSON error: {0}")]
106    Json(#[from] serde_json::Error),
107
108    /// The provider returned a non-2xx status with a message.
109    #[error("Provider error ({status}): {message}")]
110    Provider { status: u16, message: String },
111
112    /// The response structure was unexpected or missing required fields.
113    #[error("Response error: {0}")]
114    Response(String),
115}
116
117// ── Trait ─────────────────────────────────────────────────────────────────────
118
119/// The core abstraction for any text-generation model.
120///
121/// Implement this for each provider (OpenAI, Anthropic, …) or for your own
122/// local model backend.
123pub trait CompletionModel {
124    type Error: std::error::Error + 'static;
125
126    /// Send a completion request and return the model's response.
127    fn complete(
128        &self,
129        request: CompletionRequest,
130    ) -> impl std::future::Future<Output = Result<CompletionResponse, Self::Error>>;
131}