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}