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