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