a3s_code_core/llm/
factory.rs1use 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#[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 pub api_timeout_ms: Option<u64>,
27 pub temperature: Option<f32>,
29 pub max_tokens: Option<usize>,
31 pub thinking_budget: Option<usize>,
33 pub logprobs: Option<bool>,
35 pub top_logprobs: Option<usize>,
37 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
156pub 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 _ => {
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}