Skip to main content

a3s_code_core/llm/
mod.rs

1//! LLM client abstraction layer
2//!
3//! Provides a unified interface for interacting with LLM providers
4//! (Anthropic Claude, OpenAI, Zhipu AI GLM, and OpenAI-compatible providers).
5
6mod admission;
7pub mod anthropic;
8mod error;
9pub mod factory;
10pub mod http;
11pub mod openai;
12pub mod structured;
13mod token_estimation;
14mod types;
15pub mod zhipu;
16
17// Re-export public types
18pub use admission::{
19    ModelGenerationAdmission, ModelGenerationAdmissionError, ModelGenerationConcurrency,
20    ModelGenerationPermit,
21};
22pub use anthropic::AnthropicClient;
23pub(crate) use error::non_retryable_llm_error_message;
24pub use error::NonRetryableLlmError;
25pub use factory::{create_client_with_config, LlmConfig};
26pub use http::{
27    clear_http_metrics_callback, default_http_client, set_http_metrics_callback, HttpClient,
28    HttpClientError, HttpMetricsCallback, HttpMetricsRecord, HttpResponse, StreamingHttpResponse,
29};
30pub use openai::OpenAiClient;
31pub(crate) use token_estimation::{estimate_message_tokens, estimate_prompt_tokens};
32pub use types::*;
33pub use zhipu::ZhipuClient;
34
35use anyhow::Result;
36use async_trait::async_trait;
37use std::time::Duration;
38use tokio::sync::mpsc;
39use tokio_util::sync::CancellationToken;
40
41/// LLM client trait
42#[async_trait]
43pub trait LlmClient: Send + Sync {
44    /// Report the client's explicitly supported active-generation capacity.
45    ///
46    /// The conservative default is single-flight. Providers that can safely
47    /// serve more active generations must override this with a typed contract;
48    /// callers must not infer concurrency from provider names or endpoint
49    /// strings.
50    fn model_generation_concurrency(&self) -> ModelGenerationConcurrency {
51        ModelGenerationConcurrency::single_flight()
52    }
53
54    /// Derive a provider client bound to one logical agent session.
55    ///
56    /// Stateless providers can keep the default and share the existing client.
57    /// Account-backed providers whose transport uses a live session identity
58    /// should return an independent client so parallel child agents do not
59    /// contend for the parent's active operation.
60    fn fork_for_session(&self, _session_id: &str) -> Option<std::sync::Arc<dyn LlmClient>> {
61        None
62    }
63
64    /// Return a view of this client configured for one active generation
65    /// deadline. The caller still owns and enforces the outer deadline.
66    ///
67    /// Composite and account-backed clients can use this budget to configure
68    /// their underlying transport without inferring timeout intent from error
69    /// text. Stateless clients may keep the default.
70    fn with_active_generation_timeout(
71        &self,
72        _timeout: Duration,
73    ) -> Option<std::sync::Arc<dyn LlmClient>> {
74        None
75    }
76
77    /// Complete a conversation (non-streaming)
78    async fn complete(
79        &self,
80        messages: &[Message],
81        system: Option<&str>,
82        tools: &[ToolDefinition],
83    ) -> Result<LlmResponse>;
84
85    /// Complete a conversation with streaming
86    /// Returns a receiver for streaming events.
87    /// The cancel_token is checked during the HTTP request; if cancelled, the request is aborted.
88    async fn complete_streaming(
89        &self,
90        messages: &[Message],
91        system: Option<&str>,
92        tools: &[ToolDefinition],
93        cancel_token: CancellationToken,
94    ) -> Result<mpsc::Receiver<StreamEvent>>;
95
96    /// Report the strongest provider-native structured-output enforcement this
97    /// client supports. Used by [`structured`] to decide whether to force a
98    /// tool call, request a native `response_format`, or fall back to
99    /// prompt-and-parse. Defaults to no native support.
100    fn native_structured_support(&self) -> structured::NativeStructuredSupport {
101        structured::NativeStructuredSupport::None
102    }
103
104    /// Report whether [`LlmClient::complete_structured`] uses a transport that
105    /// is independent from the streaming implementation.
106    ///
107    /// The conservative default is false because several account-backed
108    /// clients implement `complete` by opening a stream and waiting for its
109    /// terminal event. Composite reliability layers use this capability to
110    /// avoid presenting the same streaming failure mode as a non-streaming
111    /// fallback.
112    fn has_distinct_non_streaming_transport(&self) -> bool {
113        false
114    }
115
116    /// Complete a conversation while honoring a structured-output directive
117    /// (forced `tool_choice` and/or native `response_format`).
118    ///
119    /// The default implementation ignores the directive and behaves exactly
120    /// like [`LlmClient::complete`], so existing clients keep working unchanged;
121    /// providers that support native structured output override this.
122    async fn complete_structured(
123        &self,
124        messages: &[Message],
125        system: Option<&str>,
126        tools: &[ToolDefinition],
127        _directive: &structured::StructuredDirective,
128    ) -> Result<LlmResponse> {
129        self.complete(messages, system, tools).await
130    }
131
132    /// Streaming counterpart of [`LlmClient::complete_structured`]. Defaults to
133    /// [`LlmClient::complete_streaming`], ignoring the directive.
134    async fn complete_streaming_structured(
135        &self,
136        messages: &[Message],
137        system: Option<&str>,
138        tools: &[ToolDefinition],
139        _directive: &structured::StructuredDirective,
140        cancel_token: CancellationToken,
141    ) -> Result<mpsc::Receiver<StreamEvent>> {
142        self.complete_streaming(messages, system, tools, cancel_token)
143            .await
144    }
145}
146
147// Include test modules — these reference internal types via crate paths
148#[cfg(test)]
149#[path = "tests.rs"]
150mod tests_file;