quorum_rs/llms/
mod.rs

1//! Defines the low-level traits for interacting with language models.
2//!
3//! This module provides a generic `AiModel` trait that abstracts the specific
4//! details of communicating with different LLM providers. The goal is to have a
5//! single, consistent interface that the high-level agent logic can use.
6//!
7//! For a quick-start model implementation, see [`SimpleOpenAIModel`].
8
9pub mod error;
10pub mod openai_compatible;
11pub mod rate_limiter;
12pub mod simple_model;
13pub mod simulated;
14pub mod span;
15pub mod strategies;
16pub mod stub;
17pub use error::LlmError;
18pub use openai_compatible::OpenAICompatibleModel;
19pub use rate_limiter::RateLimiter;
20pub use simple_model::SimpleOpenAIModel;
21pub use span::LlmRequestSpan;
22pub use stub::StubModel;
23
24use crate::agents::AgentConfig;
25use async_openai::types::{
26    ChatCompletionRequestMessage, ChatCompletionTool, ChatCompletionToolChoiceOption,
27    CreateChatCompletionResponse,
28};
29use async_trait::async_trait;
30use dyn_clone::DynClone;
31use std::fmt::Debug;
32
33/// A "partial" request struct that contains the universal parts of a request
34/// (messages and tools) that are assembled by the agent's ReAct loop.
35#[derive(Debug, Clone)]
36pub struct RequestConfig {
37    pub messages: Vec<ChatCompletionRequestMessage>,
38    pub tools: Option<Vec<ChatCompletionTool>>,
39    pub tool_choice: Option<ChatCompletionToolChoiceOption>,
40    pub presence_penalty: Option<f32>,
41}
42
43/// Result of a chat completion including the response, raw request, and timing.
44pub struct ChatCompletionResult {
45    pub response: CreateChatCompletionResponse,
46    pub raw_request: String,
47    pub timing: TimingMetadata,
48    pub provider_backend: Option<String>,
49    /// `Some` when the SDK shrink-guard rewrote `max_tokens`.
50    pub shrink_info: Option<ShrinkInfo>,
51}
52
53/// Timing metadata for an LLM call.
54pub struct TimingMetadata {
55    pub ttft_ms: Option<u64>,
56    pub generation_ms: Option<u64>,
57}
58
59/// Shrink-guard state for one LLM call. Consumed by
60/// `LlmRequestSpan::complete` to populate
61/// [`LlmRequestComplete`] and emit [`ContextEmergencyShrink`].
62#[derive(Debug, Clone)]
63pub struct ShrinkInfo {
64    /// `true` when `available < floor` — the floor-clamp case,
65    /// distinct from healthy `available > floor` adaptive shrinks.
66    pub floor_used: bool,
67    /// Raw headroom (saturating non-negative). Pre-clamp, so
68    /// `available = 0` and `available = 199` are reported distinctly.
69    pub available_space: u32,
70    pub requested_max: u32,
71    pub floor: u32,
72    pub estimated_input: u32,
73    pub context_window: u32,
74}
75
76/// A trait for a client that can interact with a language model.
77/// This is the lowest level of abstraction for model interaction.
78#[async_trait]
79pub trait AiModel: Send + Sync + DynClone + Debug {
80    /// The core function for interacting with a model.
81    /// It takes an agent's configuration and a partial request config,
82    /// combines them into a full request, and returns the model's response.
83    async fn chat_completion(
84        &self,
85        agent: &AgentConfig,
86        request_config: RequestConfig,
87    ) -> Result<ChatCompletionResult, LlmError>;
88}
89
90// Make the trait cloneable for dependency injection.
91dyn_clone::clone_trait_object!(AiModel);
92
93/// Overrides that strategies can apply to the request.
94#[derive(Debug, Clone, Default)]
95pub struct RequestOverrides {
96    pub max_tokens: Option<u32>,
97}
98
99/// A trait for transforming LLM requests/responses for specific providers.
100///
101/// Different providers (OpenAI, vLLM, Harmony) require different request
102/// formats and return responses in different structures. The `ChatStrategy`
103/// trait allows the `OpenAICompatibleModel` to delegate provider-specific
104/// transformations to pluggable strategy implementations.
105#[async_trait]
106pub trait ChatStrategy: Send + Sync {
107    /// Prepares the JSON body for the HTTP request.
108    async fn prepare_request(
109        &self,
110        agent: &AgentConfig,
111        request: &RequestConfig,
112        overrides: &RequestOverrides,
113    ) -> Result<serde_json::Value, LlmError>;
114
115    /// Parses the raw response body into a standardized OpenAI response.
116    async fn parse_response(
117        &self,
118        response_body: &str,
119    ) -> Result<CreateChatCompletionResponse, LlmError>;
120
121    /// Returns the API endpoint suffix (e.g., "/chat/completions").
122    fn endpoint_suffix(&self) -> &str {
123        "/chat/completions"
124    }
125
126    /// Returns whether this strategy supports streaming responses.
127    fn supports_streaming(&self) -> bool {
128        true
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    /// Minimal strategy that only implements the required methods,
137    /// leaving `endpoint_suffix` and `supports_streaming` at their defaults.
138    #[derive(Debug)]
139    struct DefaultStrategy;
140
141    #[async_trait]
142    impl ChatStrategy for DefaultStrategy {
143        async fn prepare_request(
144            &self,
145            _agent: &AgentConfig,
146            _request: &RequestConfig,
147            _overrides: &RequestOverrides,
148        ) -> Result<serde_json::Value, LlmError> {
149            Ok(serde_json::json!({}))
150        }
151
152        async fn parse_response(
153            &self,
154            _response_body: &str,
155        ) -> Result<CreateChatCompletionResponse, LlmError> {
156            Err(LlmError::Other("not implemented".into()))
157        }
158    }
159
160    #[test]
161    fn chat_strategy_default_endpoint_suffix() {
162        let strategy = DefaultStrategy;
163        assert_eq!(strategy.endpoint_suffix(), "/chat/completions");
164    }
165
166    #[test]
167    fn chat_strategy_default_supports_streaming() {
168        let strategy = DefaultStrategy;
169        assert!(strategy.supports_streaming());
170    }
171
172    #[test]
173    fn request_overrides_default() {
174        let overrides = RequestOverrides::default();
175        assert!(overrides.max_tokens.is_none());
176    }
177}
quorum_rs/llms/mod.rs

quorum_rs/llms/
mod.rs
