Skip to main content

a3s_code_core/llm/
factory.rs

1//! LLM client factory
2
3use super::anthropic::AnthropicClient;
4use super::http::ReqwestHttpClient;
5use super::openai::OpenAiClient;
6use super::structured::NativeStructuredSupport;
7use super::types::SecretString;
8use super::zhipu::ZhipuClient;
9use super::LlmClient;
10use crate::retry::RetryConfig;
11use std::collections::HashMap;
12use std::sync::Arc;
13use std::time::Duration;
14
15/// LLM client configuration
16#[derive(Clone, Default)]
17pub struct LlmConfig {
18    pub provider: String,
19    pub model: String,
20    pub api_key: SecretString,
21    pub base_url: Option<String>,
22    pub headers: HashMap<String, String>,
23    pub session_id_header: Option<String>,
24    pub session_id: Option<String>,
25    pub retry_config: Option<RetryConfig>,
26    /// Per-model API HTTP timeout in milliseconds. None means no HTTP timeout.
27    pub api_timeout_ms: Option<u64>,
28    /// Sampling temperature (0.0–1.0). None uses the provider default.
29    pub temperature: Option<f32>,
30    /// Maximum tokens to generate. None uses the client default.
31    pub max_tokens: Option<usize>,
32    /// Extended thinking budget in tokens (Anthropic only).
33    pub thinking_budget: Option<usize>,
34    /// Request token-level log probabilities from OpenAI-compatible providers.
35    pub logprobs: Option<bool>,
36    /// Number of alternative logprobs per token when logprobs are requested.
37    pub top_logprobs: Option<usize>,
38    /// When true, temperature is never sent to the API (e.g., o1 models).
39    pub disable_temperature: bool,
40    /// Explicit native structured-output support. None means infer from provider
41    /// and base URL; custom OpenAI-compatible endpoints default to prompt mode.
42    pub native_structured_support: Option<NativeStructuredSupport>,
43}
44
45impl std::fmt::Debug for LlmConfig {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        f.debug_struct("LlmConfig")
48            .field("provider", &self.provider)
49            .field("model", &self.model)
50            .field("api_key", &"[REDACTED]")
51            .field("base_url", &self.base_url)
52            .field("headers", &self.headers.keys().collect::<Vec<_>>())
53            .field("session_id_header", &self.session_id_header)
54            .field(
55                "session_id",
56                &self.session_id.as_ref().map(|_| "[REDACTED]"),
57            )
58            .field("retry_config", &self.retry_config)
59            .field("api_timeout_ms", &self.api_timeout_ms)
60            .field("temperature", &self.temperature)
61            .field("max_tokens", &self.max_tokens)
62            .field("thinking_budget", &self.thinking_budget)
63            .field("logprobs", &self.logprobs)
64            .field("top_logprobs", &self.top_logprobs)
65            .field("disable_temperature", &self.disable_temperature)
66            .field("native_structured_support", &self.native_structured_support)
67            .finish()
68    }
69}
70
71impl LlmConfig {
72    pub fn new(
73        provider: impl Into<String>,
74        model: impl Into<String>,
75        api_key: impl Into<String>,
76    ) -> Self {
77        Self {
78            provider: provider.into(),
79            model: model.into(),
80            api_key: SecretString::new(api_key.into()),
81            base_url: None,
82            headers: HashMap::new(),
83            session_id_header: None,
84            session_id: None,
85            retry_config: None,
86            api_timeout_ms: None,
87            temperature: None,
88            max_tokens: None,
89            thinking_budget: None,
90            logprobs: None,
91            top_logprobs: None,
92            disable_temperature: false,
93            native_structured_support: None,
94        }
95    }
96
97    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
98        self.base_url = Some(base_url.into());
99        self
100    }
101
102    pub fn with_headers(mut self, headers: HashMap<String, String>) -> Self {
103        self.headers = headers;
104        self
105    }
106
107    pub fn with_session_id_header(mut self, header_name: impl Into<String>) -> Self {
108        self.session_id_header = Some(header_name.into());
109        self
110    }
111
112    pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
113        self.session_id = Some(session_id.into());
114        self
115    }
116
117    pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
118        self.retry_config = Some(retry_config);
119        self
120    }
121
122    pub fn with_api_timeout(mut self, timeout_ms: u64) -> Self {
123        self.api_timeout_ms = Some(timeout_ms);
124        self
125    }
126
127    pub fn with_temperature(mut self, temperature: f32) -> Self {
128        self.temperature = Some(temperature);
129        self
130    }
131
132    pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
133        self.max_tokens = Some(max_tokens);
134        self
135    }
136
137    pub fn with_thinking_budget(mut self, budget: usize) -> Self {
138        self.thinking_budget = Some(budget);
139        self
140    }
141
142    pub fn with_logprobs(mut self, enabled: bool) -> Self {
143        self.logprobs = Some(enabled);
144        self
145    }
146
147    pub fn with_top_logprobs(mut self, top_logprobs: usize) -> Self {
148        self.logprobs = Some(true);
149        self.top_logprobs = Some(top_logprobs);
150        self
151    }
152
153    pub fn with_native_structured_support(mut self, support: NativeStructuredSupport) -> Self {
154        self.native_structured_support = Some(support);
155        self
156    }
157
158    pub(crate) fn resolved_headers(&self) -> HashMap<String, String> {
159        let mut headers = self.headers.clone();
160        if let (Some(header_name), Some(session_id)) = (&self.session_id_header, &self.session_id) {
161            headers.insert(header_name.clone(), session_id.clone());
162        }
163        headers
164    }
165}
166
167fn default_openai_native_structured_support(base_url: Option<&str>) -> NativeStructuredSupport {
168    match base_url {
169        None => NativeStructuredSupport::JsonSchema,
170        Some(url) if url.contains("api.openai.com") => NativeStructuredSupport::JsonSchema,
171        Some(_) => NativeStructuredSupport::None,
172    }
173}
174
175/// Create LLM client with full configuration (supports custom base_url)
176pub fn create_client_with_config(config: LlmConfig) -> Arc<dyn LlmClient> {
177    let retry = config.retry_config.clone().unwrap_or_default();
178    let http = config
179        .api_timeout_ms
180        .map(|timeout_ms| {
181            ReqwestHttpClient::with_timeout(Duration::from_millis(timeout_ms))
182                .expect("failed to build LLM HTTP client with API timeout")
183        })
184        .map(|client| Arc::new(client) as Arc<dyn super::http::HttpClient>);
185    let api_key = config.api_key.expose().to_string();
186    let headers = config.resolved_headers();
187
188    match config.provider.as_str() {
189        "anthropic" | "claude" => {
190            let mut client = AnthropicClient::new(api_key, config.model)
191                .with_provider_name(config.provider.clone())
192                .with_retry_config(retry);
193            if let Some(http) = http.clone() {
194                client = client.with_http_client(http);
195            }
196            if let Some(base_url) = config.base_url {
197                client = client.with_base_url(base_url);
198            }
199            if !config.disable_temperature {
200                if let Some(temp) = config.temperature {
201                    client = client.with_temperature(temp);
202                }
203            }
204            if let Some(max) = config.max_tokens {
205                client = client.with_max_tokens(max);
206            }
207            if let Some(budget) = config.thinking_budget {
208                client = client.with_thinking_budget(budget);
209            }
210            Arc::new(client)
211        }
212        "openai" | "gpt" => {
213            let native_structured_support = config.native_structured_support.unwrap_or_else(|| {
214                default_openai_native_structured_support(config.base_url.as_deref())
215            });
216            let mut client = OpenAiClient::new(api_key, config.model)
217                .with_provider_name(config.provider.clone())
218                .with_retry_config(retry)
219                .with_native_structured_support(native_structured_support);
220            if let Some(http) = http.clone() {
221                client = client.with_http_client(http);
222            }
223            if let Some(base_url) = config.base_url {
224                client = client.with_base_url(base_url);
225            }
226            if !headers.is_empty() {
227                client = client.with_headers(headers.clone());
228            }
229            if !config.disable_temperature {
230                if let Some(temp) = config.temperature {
231                    client = client.with_temperature(temp);
232                }
233            }
234            if let Some(max) = config.max_tokens {
235                client = client.with_max_tokens(max);
236            }
237            if let Some(enabled) = config.logprobs {
238                client = client.with_logprobs(enabled);
239            }
240            if let Some(top_logprobs) = config.top_logprobs {
241                client = client.with_top_logprobs(top_logprobs);
242            }
243            Arc::new(client)
244        }
245        "glm" | "zhipu" | "bigmodel" => {
246            let mut client = ZhipuClient::new(api_key, config.model).with_retry_config(retry);
247            if let Some(http) = http.clone() {
248                client = client.with_http_client(http);
249            }
250            if let Some(base_url) = config.base_url {
251                client = client.with_base_url(base_url);
252            }
253            if !config.disable_temperature {
254                if let Some(temp) = config.temperature {
255                    client = client.with_temperature(temp);
256                }
257            }
258            if let Some(max) = config.max_tokens {
259                client = client.with_max_tokens(max);
260            }
261            if let Some(enabled) = config.logprobs {
262                client = client.with_logprobs(enabled);
263            }
264            if let Some(top_logprobs) = config.top_logprobs {
265                client = client.with_top_logprobs(top_logprobs);
266            }
267            Arc::new(client)
268        }
269        // OpenAI-compatible providers (deepseek, groq, together, ollama, etc.)
270        _ => {
271            tracing::info!(
272                "Using OpenAI-compatible client for provider '{}'",
273                config.provider
274            );
275            let mut client = OpenAiClient::new(api_key, config.model)
276                .with_provider_name(config.provider.clone())
277                .with_retry_config(retry);
278            if let Some(http) = http.clone() {
279                client = client.with_http_client(http);
280            }
281            if let Some(base_url) = config.base_url {
282                client = client.with_base_url(base_url);
283            }
284            if !headers.is_empty() {
285                client = client.with_headers(headers.clone());
286            }
287            if !config.disable_temperature {
288                if let Some(temp) = config.temperature {
289                    client = client.with_temperature(temp);
290                }
291            }
292            if let Some(max) = config.max_tokens {
293                client = client.with_max_tokens(max);
294            }
295            if let Some(enabled) = config.logprobs {
296                client = client.with_logprobs(enabled);
297            }
298            if let Some(top_logprobs) = config.top_logprobs {
299                client = client.with_top_logprobs(top_logprobs);
300            }
301            Arc::new(client)
302        }
303    }
304}