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    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 tokio::sync::mpsc;
38use tokio_util::sync::CancellationToken;
39
40/// LLM client trait
41#[async_trait]
42pub trait LlmClient: Send + Sync {
43    /// Report the client's explicitly supported active-generation capacity.
44    ///
45    /// The conservative default is single-flight. Providers that can safely
46    /// serve more active generations must override this with a typed contract;
47    /// callers must not infer concurrency from provider names or endpoint
48    /// strings.
49    fn model_generation_concurrency(&self) -> ModelGenerationConcurrency {
50        ModelGenerationConcurrency::single_flight()
51    }
52
53    /// Derive a provider client bound to one logical agent session.
54    ///
55    /// Stateless providers can keep the default and share the existing client.
56    /// Account-backed providers whose transport uses a live session identity
57    /// should return an independent client so parallel child agents do not
58    /// contend for the parent's active operation.
59    fn fork_for_session(&self, _session_id: &str) -> Option<std::sync::Arc<dyn LlmClient>> {
60        None
61    }
62
63    /// Complete a conversation (non-streaming)
64    async fn complete(
65        &self,
66        messages: &[Message],
67        system: Option<&str>,
68        tools: &[ToolDefinition],
69    ) -> Result<LlmResponse>;
70
71    /// Complete a conversation with streaming
72    /// Returns a receiver for streaming events.
73    /// The cancel_token is checked during the HTTP request; if cancelled, the request is aborted.
74    async fn complete_streaming(
75        &self,
76        messages: &[Message],
77        system: Option<&str>,
78        tools: &[ToolDefinition],
79        cancel_token: CancellationToken,
80    ) -> Result<mpsc::Receiver<StreamEvent>>;
81
82    /// Report the strongest provider-native structured-output enforcement this
83    /// client supports. Used by [`structured`] to decide whether to force a
84    /// tool call, request a native `response_format`, or fall back to
85    /// prompt-and-parse. Defaults to no native support.
86    fn native_structured_support(&self) -> structured::NativeStructuredSupport {
87        structured::NativeStructuredSupport::None
88    }
89
90    /// Complete a conversation while honoring a structured-output directive
91    /// (forced `tool_choice` and/or native `response_format`).
92    ///
93    /// The default implementation ignores the directive and behaves exactly
94    /// like [`LlmClient::complete`], so existing clients keep working unchanged;
95    /// providers that support native structured output override this.
96    async fn complete_structured(
97        &self,
98        messages: &[Message],
99        system: Option<&str>,
100        tools: &[ToolDefinition],
101        _directive: &structured::StructuredDirective,
102    ) -> Result<LlmResponse> {
103        self.complete(messages, system, tools).await
104    }
105
106    /// Streaming counterpart of [`LlmClient::complete_structured`]. Defaults to
107    /// [`LlmClient::complete_streaming`], ignoring the directive.
108    async fn complete_streaming_structured(
109        &self,
110        messages: &[Message],
111        system: Option<&str>,
112        tools: &[ToolDefinition],
113        _directive: &structured::StructuredDirective,
114        cancel_token: CancellationToken,
115    ) -> Result<mpsc::Receiver<StreamEvent>> {
116        self.complete_streaming(messages, system, tools, cancel_token)
117            .await
118    }
119}
120
121// Include test modules — these reference internal types via crate paths
122#[cfg(test)]
123#[path = "tests.rs"]
124mod tests_file;