1use std::future::Future;
4use std::pin::Pin;
5use std::time::Duration;
6
7use anyhow::Result;
8use reqwest::Client;
9use serde::{Deserialize, Serialize};
10use tracing::{debug, info};
11
12use super::{AiClient, AiClientCapabilities, AiClientMetadata, RequestOptions};
13use crate::claude::{error::ClaudeError, model_config::get_model_registry};
14
15const PROBE_TIMEOUT: Duration = Duration::from_secs(2);
19
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24pub enum ProbeSource {
25 LmStudio,
27 Ollama,
29}
30
31impl ProbeSource {
32 #[must_use]
34 pub fn as_str(self) -> &'static str {
35 match self {
36 Self::LmStudio => "lmstudio",
37 Self::Ollama => "ollama",
38 }
39 }
40}
41
42#[derive(Serialize, Debug)]
44struct Message {
45 role: String,
46 content: String,
47}
48
49#[derive(Serialize, Debug)]
58struct ResponseFormatField {
59 #[serde(rename = "type")]
60 kind: &'static str,
61 json_schema: JsonSchemaSpec,
62}
63
64#[derive(Serialize, Debug)]
71struct JsonSchemaSpec {
72 name: &'static str,
73 strict: bool,
74 schema: serde_json::Value,
75}
76
77#[derive(Serialize, Debug)]
79struct OpenAiRequest {
80 model: String,
81 messages: Vec<Message>,
82 #[serde(skip_serializing_if = "Option::is_none")]
83 max_tokens: Option<i32>,
84 #[serde(skip_serializing_if = "Option::is_none")]
85 max_completion_tokens: Option<i32>,
86 #[serde(skip_serializing_if = "Option::is_none")]
87 temperature: Option<f32>,
88 stream: bool,
89 #[serde(skip_serializing_if = "Option::is_none")]
90 response_format: Option<ResponseFormatField>,
91}
92
93#[derive(Deserialize, Debug)]
95struct Choice {
96 message: ResponseMessage,
97 #[allow(dead_code)] finish_reason: Option<String>,
99}
100
101#[derive(Deserialize, Debug)]
103struct ResponseMessage {
104 #[allow(dead_code)] role: String,
106 content: String,
107}
108
109#[derive(Deserialize, Debug)]
111struct OpenAiResponse {
112 choices: Vec<Choice>,
113 model: Option<String>,
114 usage: Option<Usage>,
115}
116
117#[derive(Deserialize, Debug)]
119#[allow(dead_code)] struct Usage {
121 prompt_tokens: Option<i32>,
122 completion_tokens: Option<i32>,
123 total_tokens: Option<i32>,
124}
125
126pub struct OpenAiAiClient {
128 client: Client,
130 api_key: Option<String>,
132 model: String,
134 base_url: String,
136 max_tokens: Option<i32>,
138 temperature: Option<f32>,
140 active_beta: Option<(String, String)>,
142 loaded_context_length: Option<usize>,
147}
148
149#[derive(Deserialize, Debug)]
151struct LmStudioModelsResponse {
152 data: Vec<LmStudioModel>,
153}
154
155#[derive(Deserialize, Debug)]
157struct LmStudioModel {
158 id: String,
159 state: Option<String>,
160 loaded_context_length: Option<usize>,
161}
162
163impl OpenAiAiClient {
164 pub fn new(
166 model: String,
167 api_key: Option<String>,
168 base_url: String,
169 max_tokens: Option<i32>,
170 temperature: Option<f32>,
171 active_beta: Option<(String, String)>,
172 ) -> Result<Self> {
173 let client = super::build_http_client()?;
174
175 Ok(Self {
176 client,
177 api_key,
178 model,
179 base_url,
180 max_tokens,
181 temperature,
182 active_beta,
183 loaded_context_length: None,
184 })
185 }
186
187 pub fn new_ollama(
189 model: String,
190 base_url: Option<String>,
191 active_beta: Option<(String, String)>,
192 ) -> Result<Self> {
193 Self::new(
194 model,
195 None, base_url.unwrap_or_else(|| "http://localhost:11434".to_string()),
197 Some(4096), Some(0.1), active_beta,
200 )
201 }
202
203 pub fn new_openai(
205 model: String,
206 api_key: String,
207 active_beta: Option<(String, String)>,
208 ) -> Result<Self> {
209 Self::new(
210 model,
211 Some(api_key),
212 "https://api.openai.com".to_string(),
213 None, Some(0.1), active_beta,
216 )
217 }
218
219 fn get_max_tokens(&self) -> i32 {
221 if let Some(configured_max) = self.max_tokens {
222 return configured_max;
223 }
224 super::registry_max_output_tokens(&self.model, &self.active_beta)
225 }
226
227 fn get_api_url(&self) -> String {
234 let base = self.base_url.trim_end_matches('/');
235 let url = format!("{base}/v1/chat/completions");
236 debug!(base_url = %self.base_url, full_url = %url, "Constructed OpenAI-compatible API URL");
237 url
238 }
239
240 fn is_ollama(&self) -> bool {
242 self.base_url.contains("localhost")
243 || self.base_url.contains("127.0.0.1")
244 || self.api_key.is_none()
245 }
246
247 fn service_tag(&self) -> &'static str {
251 if self.is_ollama() {
252 "ollama"
253 } else {
254 "openai"
255 }
256 }
257
258 fn is_gpt5_series(&self) -> bool {
260 self.model.starts_with("gpt-5") || self.model.starts_with("o1")
261 }
262
263 #[must_use]
265 pub fn loaded_context_length(&self) -> Option<usize> {
266 self.loaded_context_length
267 }
268
269 pub fn set_loaded_context_length(&mut self, value: usize) {
273 self.loaded_context_length = Some(value);
274 }
275
276 fn build_request(
280 &self,
281 system_prompt: &str,
282 user_prompt: &str,
283 response_format: Option<ResponseFormatField>,
284 ) -> OpenAiRequest {
285 let mut messages = Vec::new();
286
287 if !system_prompt.is_empty() {
288 messages.push(Message {
289 role: "system".to_string(),
290 content: system_prompt.to_string(),
291 });
292 }
293
294 messages.push(Message {
295 role: "user".to_string(),
296 content: user_prompt.to_string(),
297 });
298
299 let max_tokens = self.get_max_tokens();
300 if self.is_gpt5_series() {
301 OpenAiRequest {
302 model: self.model.clone(),
303 messages,
304 max_tokens: None,
305 max_completion_tokens: Some(max_tokens),
306 temperature: None,
308 stream: false,
309 response_format,
310 }
311 } else {
312 OpenAiRequest {
313 model: self.model.clone(),
314 messages,
315 max_tokens: Some(max_tokens),
316 max_completion_tokens: None,
317 temperature: self.temperature,
318 stream: false,
319 response_format,
320 }
321 }
322 }
323
324 async fn send_inner(&self, request: OpenAiRequest) -> Result<String> {
332 debug!(
333 max_tokens = ?request.max_tokens,
334 max_completion_tokens = ?request.max_completion_tokens,
335 configured_temperature = ?self.temperature,
336 effective_temperature = ?request.temperature,
337 message_count = request.messages.len(),
338 is_gpt5_series = self.is_gpt5_series(),
339 response_format_set = request.response_format.is_some(),
340 "Built OpenAI-compatible request payload"
341 );
342
343 let api_url = self.get_api_url();
344 info!(url = %api_url, model = %self.model, "Sending request to OpenAI-compatible API");
345
346 let mut req_builder = self
347 .client
348 .post(&api_url)
349 .header("Content-Type", "application/json")
350 .json(&request);
351
352 if let Some(ref api_key) = self.api_key {
353 req_builder = req_builder.header("Authorization", format!("Bearer {api_key}"));
354 }
355
356 let started = std::time::Instant::now();
357 let send_result = req_builder.send().await;
358 super::record_ai_http(self.service_tag(), "POST", &api_url, started, &send_result);
359 let response = send_result.map_err(|e| ClaudeError::NetworkError(e.to_string()))?;
360
361 let response = super::check_error_response(response).await?;
362
363 let openai_response: OpenAiResponse = response
364 .json()
365 .await
366 .map_err(|e| ClaudeError::InvalidResponseFormat(e.to_string()))?;
367
368 debug!(
369 choice_count = openai_response.choices.len(),
370 model = ?openai_response.model,
371 usage = ?openai_response.usage,
372 "Received OpenAI-compatible API response"
373 );
374
375 let result = openai_response
376 .choices
377 .first()
378 .map(|choice| choice.message.content.clone())
379 .ok_or_else(|| {
380 ClaudeError::InvalidResponseFormat("No choices in response".to_string()).into()
381 });
382
383 super::log_response_success("OpenAI-compatible", &result);
384
385 result
386 }
387
388 pub async fn probe_loaded_context_length(&mut self) -> Option<ProbeSource> {
402 let host = host_root(&self.base_url);
403 let service = self.service_tag();
404
405 if let Some(value) = probe_lm_studio(&self.client, &host, &self.model, service).await {
406 self.loaded_context_length = Some(value);
407 return Some(ProbeSource::LmStudio);
408 }
409
410 if let Some(value) = probe_ollama_native(&self.client, &host, &self.model, service).await {
411 self.loaded_context_length = Some(value);
412 return Some(ProbeSource::Ollama);
413 }
414
415 None
416 }
417}
418
419fn host_root(base_url: &str) -> String {
425 let trimmed = base_url.trim_end_matches('/');
426 trimmed
427 .strip_suffix("/v1")
428 .unwrap_or(trimmed)
429 .trim_end_matches('/')
430 .to_string()
431}
432
433async fn probe_lm_studio(client: &Client, host: &str, model: &str, service: &str) -> Option<usize> {
438 let url = format!("{host}/api/v0/models");
439 debug!(url = %url, model = %model, "Probing LM Studio for loaded context length");
440
441 let started = std::time::Instant::now();
442 let result = client.get(&url).timeout(PROBE_TIMEOUT).send().await;
443 super::record_ai_http(service, "GET", &url, started, &result);
444 let response = result.ok()?;
445 if !response.status().is_success() {
446 debug!(status = %response.status(), "LM Studio probe returned non-success");
447 return None;
448 }
449 let body: LmStudioModelsResponse = response.json().await.ok()?;
450 body.data
451 .into_iter()
452 .find(|entry| entry.id == model && entry.state.as_deref() == Some("loaded"))
453 .and_then(|entry| entry.loaded_context_length)
454}
455
456async fn probe_ollama_native(
461 client: &Client,
462 host: &str,
463 model: &str,
464 service: &str,
465) -> Option<usize> {
466 let url = format!("{host}/api/show");
467 debug!(url = %url, model = %model, "Probing Ollama for loaded context length");
468
469 let started = std::time::Instant::now();
470 let result = client
471 .post(&url)
472 .timeout(PROBE_TIMEOUT)
473 .json(&serde_json::json!({ "name": model }))
474 .send()
475 .await;
476 super::record_ai_http(service, "POST", &url, started, &result);
477 let response = result.ok()?;
478 if !response.status().is_success() {
479 debug!(status = %response.status(), "Ollama probe returned non-success");
480 return None;
481 }
482 let body: serde_json::Value = response.json().await.ok()?;
483 let model_info = body.get("model_info")?.as_object()?;
484 for (key, value) in model_info {
485 if key.ends_with(".context_length") {
486 if let Some(n) = value.as_u64() {
487 return usize::try_from(n).ok();
488 }
489 }
490 }
491 None
492}
493
494impl AiClient for OpenAiAiClient {
495 fn send_request<'a>(
496 &'a self,
497 system_prompt: &'a str,
498 user_prompt: &'a str,
499 ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
500 Box::pin(async move {
501 debug!(
502 system_prompt_len = system_prompt.len(),
503 user_prompt_len = user_prompt.len(),
504 model = %self.model,
505 base_url = %self.base_url,
506 is_ollama = self.is_ollama(),
507 "Preparing OpenAI-compatible API request"
508 );
509
510 let request = self.build_request(system_prompt, user_prompt, None);
511 self.send_inner(request).await
512 })
513 }
514
515 fn capabilities(&self) -> AiClientCapabilities {
516 AiClientCapabilities {
517 supports_response_schema: true,
518 }
519 }
520
521 fn send_request_with_options<'a>(
522 &'a self,
523 system_prompt: &'a str,
524 user_prompt: &'a str,
525 options: RequestOptions,
526 ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
527 Box::pin(async move {
528 debug!(
529 system_prompt_len = system_prompt.len(),
530 user_prompt_len = user_prompt.len(),
531 has_schema = options.response_schema.is_some(),
532 model = %self.model,
533 base_url = %self.base_url,
534 is_ollama = self.is_ollama(),
535 "Preparing OpenAI-compatible API request (with options)"
536 );
537
538 let response_format = options.response_schema.map(|schema| ResponseFormatField {
539 kind: "json_schema",
540 json_schema: JsonSchemaSpec {
541 name: "response",
542 strict: true,
543 schema,
544 },
545 });
546
547 let request = self.build_request(system_prompt, user_prompt, response_format);
548 self.send_inner(request).await
549 })
550 }
551
552 fn get_metadata(&self) -> AiClientMetadata {
553 let registry = get_model_registry();
554
555 let max_context_length = if let Some(probed) = self.loaded_context_length {
559 probed
560 } else if registry.get_input_context(&self.model) > 0 {
561 registry.get_input_context(&self.model)
562 } else {
563 32768 };
565
566 let max_response_length = if registry.get_max_output_tokens(&self.model) > 0 {
567 registry.get_max_output_tokens(&self.model)
568 } else {
569 4096 };
571
572 let provider = if self.is_ollama() {
573 "Ollama".to_string()
574 } else {
575 "OpenAI".to_string()
576 };
577
578 AiClientMetadata {
579 provider,
580 model: self.model.clone(),
581 max_context_length,
582 max_response_length,
583 active_beta: self.active_beta.clone(),
584 }
585 }
586}
587
588#[cfg(test)]
589#[allow(clippy::unwrap_used, clippy::expect_used)]
590mod tests {
591 use super::*;
592
593 #[test]
594 fn new_ollama() {
595 let client = OpenAiAiClient::new_ollama("llama2".to_string(), None, None).unwrap();
596 assert_eq!(client.model, "llama2");
597 assert_eq!(client.base_url, "http://localhost:11434");
598 assert!(client.api_key.is_none());
599 assert!(client.is_ollama());
600 }
601
602 #[test]
603 fn new_ollama_custom_url() {
604 let client = OpenAiAiClient::new_ollama(
605 "codellama".to_string(),
606 Some("http://192.168.1.100:11434".to_string()),
607 None,
608 )
609 .unwrap();
610 assert_eq!(client.base_url, "http://192.168.1.100:11434");
611 assert!(client.is_ollama());
612 }
613
614 #[test]
615 fn new_openai() {
616 let client =
617 OpenAiAiClient::new_openai("gpt-4".to_string(), "sk-test123".to_string(), None)
618 .unwrap();
619 assert_eq!(client.model, "gpt-4");
620 assert_eq!(client.base_url, "https://api.openai.com");
621 assert_eq!(client.api_key, Some("sk-test123".to_string()));
622 assert!(!client.is_ollama());
623 }
624
625 #[test]
626 fn get_api_url() {
627 let client = OpenAiAiClient::new_ollama("llama2".to_string(), None, None).unwrap();
628 let url = client.get_api_url();
629 assert_eq!(url, "http://localhost:11434/v1/chat/completions");
630 }
631
632 #[test]
633 fn get_api_url_trailing_slash() {
634 let client = OpenAiAiClient::new(
635 "test-model".to_string(),
636 None,
637 "http://localhost:11434/".to_string(),
638 None,
639 None,
640 None,
641 )
642 .unwrap();
643 let url = client.get_api_url();
644 assert_eq!(url, "http://localhost:11434/v1/chat/completions");
645 }
646
647 #[test]
648 fn is_ollama_detection() {
649 let ollama_client = OpenAiAiClient::new(
651 "llama2".to_string(),
652 None,
653 "http://localhost:11434".to_string(),
654 None,
655 None,
656 None,
657 )
658 .unwrap();
659 assert!(ollama_client.is_ollama());
660
661 let local_client = OpenAiAiClient::new(
663 "llama2".to_string(),
664 Some("fake-key".to_string()),
665 "http://127.0.0.1:11434".to_string(),
666 None,
667 None,
668 None,
669 )
670 .unwrap();
671 assert!(local_client.is_ollama());
672
673 let no_key_client = OpenAiAiClient::new(
675 "llama2".to_string(),
676 None,
677 "http://remote-server.com".to_string(),
678 None,
679 None,
680 None,
681 )
682 .unwrap();
683 assert!(no_key_client.is_ollama());
684
685 let openai_client = OpenAiAiClient::new(
687 "gpt-4".to_string(),
688 Some("sk-real-key".to_string()),
689 "https://api.openai.com".to_string(),
690 None,
691 None,
692 None,
693 )
694 .unwrap();
695 assert!(!openai_client.is_ollama());
696 }
697
698 #[test]
699 fn service_tag_distinguishes_ollama_from_openai() {
700 let ollama = OpenAiAiClient::new(
703 "llama2".to_string(),
704 None,
705 "http://localhost:11434".to_string(),
706 None,
707 None,
708 None,
709 )
710 .unwrap();
711 assert_eq!(ollama.service_tag(), "ollama");
712
713 let openai = OpenAiAiClient::new(
714 "gpt-4".to_string(),
715 Some("sk-real-key".to_string()),
716 "https://api.openai.com".to_string(),
717 None,
718 None,
719 None,
720 )
721 .unwrap();
722 assert_eq!(openai.service_tag(), "openai");
723 }
724
725 #[test]
728 fn gpt5_series_gpt5_models() {
729 let client = OpenAiAiClient::new(
730 "gpt-5-preview".to_string(),
731 Some("key".to_string()),
732 "https://api.openai.com".to_string(),
733 None,
734 None,
735 None,
736 )
737 .unwrap();
738 assert!(client.is_gpt5_series());
739
740 let client2 = OpenAiAiClient::new(
741 "gpt-5".to_string(),
742 Some("key".to_string()),
743 "https://api.openai.com".to_string(),
744 None,
745 None,
746 None,
747 )
748 .unwrap();
749 assert!(client2.is_gpt5_series());
750 }
751
752 #[test]
753 fn gpt5_series_o1_models() {
754 let client = OpenAiAiClient::new(
755 "o1-mini".to_string(),
756 Some("key".to_string()),
757 "https://api.openai.com".to_string(),
758 None,
759 None,
760 None,
761 )
762 .unwrap();
763 assert!(client.is_gpt5_series());
764
765 let client2 = OpenAiAiClient::new(
766 "o1-preview".to_string(),
767 Some("key".to_string()),
768 "https://api.openai.com".to_string(),
769 None,
770 None,
771 None,
772 )
773 .unwrap();
774 assert!(client2.is_gpt5_series());
775 }
776
777 #[test]
778 fn gpt5_series_regular_models_not_matched() {
779 let client = OpenAiAiClient::new(
780 "gpt-4".to_string(),
781 Some("key".to_string()),
782 "https://api.openai.com".to_string(),
783 None,
784 None,
785 None,
786 )
787 .unwrap();
788 assert!(!client.is_gpt5_series());
789
790 let client2 = OpenAiAiClient::new(
791 "gpt-4o-mini".to_string(),
792 Some("key".to_string()),
793 "https://api.openai.com".to_string(),
794 None,
795 None,
796 None,
797 )
798 .unwrap();
799 assert!(!client2.is_gpt5_series());
800 }
801
802 #[test]
805 fn get_max_tokens_configured_value_wins() {
806 let client = OpenAiAiClient::new(
807 "gpt-4".to_string(),
808 Some("key".to_string()),
809 "https://api.openai.com".to_string(),
810 Some(8192),
811 None,
812 None,
813 )
814 .unwrap();
815 assert_eq!(client.get_max_tokens(), 8192);
816 }
817
818 #[test]
819 fn get_max_tokens_from_registry() {
820 let client =
822 OpenAiAiClient::new_openai("gpt-4o".to_string(), "key".to_string(), None).unwrap();
823 let tokens = client.get_max_tokens();
824 assert!(tokens > 0, "expected positive token limit, got {tokens}");
826 }
827
828 #[test]
831 fn get_metadata_openai() {
832 let client =
833 OpenAiAiClient::new_openai("gpt-4o".to_string(), "key".to_string(), None).unwrap();
834 let metadata = client.get_metadata();
835 assert_eq!(metadata.provider, "OpenAI");
836 assert_eq!(metadata.model, "gpt-4o");
837 assert!(metadata.active_beta.is_none());
838 }
839
840 #[test]
841 fn get_metadata_ollama() {
842 let client = OpenAiAiClient::new_ollama("llama2".to_string(), None, None).unwrap();
843 let metadata = client.get_metadata();
844 assert_eq!(metadata.provider, "Ollama");
845 assert_eq!(metadata.model, "llama2");
846 }
847
848 #[test]
849 fn get_metadata_with_beta() {
850 let beta = Some(("anthropic-beta".to_string(), "output-128k".to_string()));
851 let client =
852 OpenAiAiClient::new_openai("gpt-4o".to_string(), "key".to_string(), beta).unwrap();
853 let metadata = client.get_metadata();
854 assert!(metadata.active_beta.is_some());
855 let (key, value) = metadata.active_beta.unwrap();
856 assert_eq!(key, "anthropic-beta");
857 assert_eq!(value, "output-128k");
858 }
859
860 #[test]
863 fn request_gpt5_uses_max_completion_tokens() {
864 let request = OpenAiRequest {
865 model: "gpt-5".to_string(),
866 messages: vec![Message {
867 role: "user".to_string(),
868 content: "hello".to_string(),
869 }],
870 max_tokens: None,
871 max_completion_tokens: Some(4096),
872 temperature: None,
873 stream: false,
874 response_format: None,
875 };
876
877 let json = serde_json::to_string(&request).unwrap();
878 assert!(json.contains("max_completion_tokens"));
879 assert!(!json.contains("\"max_tokens\""));
881 }
882
883 #[test]
884 fn request_regular_model_uses_max_tokens() {
885 let request = OpenAiRequest {
886 model: "gpt-4".to_string(),
887 messages: vec![Message {
888 role: "user".to_string(),
889 content: "hello".to_string(),
890 }],
891 max_tokens: Some(4096),
892 max_completion_tokens: None,
893 temperature: Some(0.1),
894 stream: false,
895 response_format: None,
896 };
897
898 let json = serde_json::to_string(&request).unwrap();
899 assert!(json.contains("\"max_tokens\""));
900 assert!(!json.contains("max_completion_tokens"));
901 assert!(json.contains("\"temperature\""));
902 }
903
904 #[test]
909 fn capabilities_advertise_response_schema_support_openai() {
910 let client =
911 OpenAiAiClient::new_openai("gpt-4o".to_string(), "key".to_string(), None).unwrap();
912 assert!(client.capabilities().supports_response_schema);
913 }
914
915 #[test]
916 fn capabilities_advertise_response_schema_support_ollama() {
917 let client = OpenAiAiClient::new_ollama("llama2".to_string(), None, None).unwrap();
918 assert!(client.capabilities().supports_response_schema);
919 }
920
921 #[test]
929 fn build_request_omits_response_format_without_schema() {
930 let client =
931 OpenAiAiClient::new_openai("gpt-4o".to_string(), "key".to_string(), None).unwrap();
932 let request = client.build_request("sys", "user", None);
933 let body = serde_json::to_value(&request).unwrap();
934 assert!(
935 body.get("response_format").is_none(),
936 "expected response_format to be omitted, got: {body}"
937 );
938 }
939
940 #[test]
943 fn build_request_embeds_response_format_with_schema_regular_model() {
944 let client =
945 OpenAiAiClient::new_openai("gpt-4o".to_string(), "key".to_string(), None).unwrap();
946 let schema = serde_json::json!({
947 "type": "object",
948 "properties": { "answer": { "type": "string" } },
949 "required": ["answer"],
950 "additionalProperties": false,
951 });
952 let response_format = Some(ResponseFormatField {
953 kind: "json_schema",
954 json_schema: JsonSchemaSpec {
955 name: "response",
956 strict: true,
957 schema: schema.clone(),
958 },
959 });
960 let request = client.build_request("sys", "user", response_format);
961 let body = serde_json::to_value(&request).unwrap();
962 assert_eq!(body["response_format"]["type"], "json_schema");
963 assert_eq!(body["response_format"]["json_schema"]["name"], "response");
964 assert_eq!(body["response_format"]["json_schema"]["strict"], true);
965 assert_eq!(body["response_format"]["json_schema"]["schema"], schema);
966 assert!(body.get("max_tokens").is_some());
968 assert!(body.get("max_completion_tokens").is_none());
969 }
970
971 #[test]
974 fn build_request_embeds_response_format_with_schema_gpt5() {
975 let client =
976 OpenAiAiClient::new_openai("gpt-5".to_string(), "key".to_string(), None).unwrap();
977 let schema = serde_json::json!({ "type": "object", "additionalProperties": false });
978 let response_format = Some(ResponseFormatField {
979 kind: "json_schema",
980 json_schema: JsonSchemaSpec {
981 name: "response",
982 strict: true,
983 schema: schema.clone(),
984 },
985 });
986 let request = client.build_request("sys", "user", response_format);
987 let body = serde_json::to_value(&request).unwrap();
988 assert_eq!(body["response_format"]["type"], "json_schema");
989 assert_eq!(body["response_format"]["json_schema"]["schema"], schema);
990 assert!(body.get("max_completion_tokens").is_some());
992 assert!(body.get("max_tokens").is_none());
993 assert!(body.get("temperature").is_none());
994 }
995
996 #[test]
999 fn host_root_strips_trailing_slash() {
1000 assert_eq!(host_root("http://localhost:1234/"), "http://localhost:1234");
1001 }
1002
1003 #[test]
1004 fn host_root_strips_v1_suffix() {
1005 assert_eq!(
1006 host_root("http://localhost:1234/v1"),
1007 "http://localhost:1234"
1008 );
1009 }
1010
1011 #[test]
1012 fn host_root_strips_v1_with_trailing_slash() {
1013 assert_eq!(
1014 host_root("http://localhost:1234/v1/"),
1015 "http://localhost:1234"
1016 );
1017 }
1018
1019 #[test]
1020 fn host_root_passthrough_when_no_v1() {
1021 assert_eq!(
1022 host_root("http://localhost:11434"),
1023 "http://localhost:11434"
1024 );
1025 }
1026
1027 #[test]
1030 fn probe_source_as_str_stable() {
1031 assert_eq!(ProbeSource::LmStudio.as_str(), "lmstudio");
1032 assert_eq!(ProbeSource::Ollama.as_str(), "ollama");
1033 }
1034
1035 #[test]
1038 fn metadata_uses_probed_value_when_set() {
1039 let mut client = OpenAiAiClient::new_ollama("llama2".to_string(), None, None).unwrap();
1040 client.set_loaded_context_length(8192);
1041 let metadata = client.get_metadata();
1042 assert_eq!(metadata.max_context_length, 8192);
1043 }
1044
1045 #[test]
1046 fn metadata_falls_back_to_registry_when_probe_value_absent() {
1047 let client =
1051 OpenAiAiClient::new_ollama("totally-unknown-model".to_string(), None, None).unwrap();
1052 let metadata = client.get_metadata();
1053 let expected = get_model_registry().get_input_context("totally-unknown-model");
1054 assert_eq!(metadata.max_context_length, expected);
1055 }
1056
1057 #[test]
1058 fn loaded_context_length_starts_unset() {
1059 let client = OpenAiAiClient::new_ollama("llama2".to_string(), None, None).unwrap();
1060 assert!(client.loaded_context_length().is_none());
1061 }
1062
1063 #[test]
1064 fn loaded_context_length_round_trips() {
1065 let mut client = OpenAiAiClient::new_ollama("llama2".to_string(), None, None).unwrap();
1066 client.set_loaded_context_length(4096);
1067 assert_eq!(client.loaded_context_length(), Some(4096));
1068 }
1069
1070 use wiremock::matchers::{body_json, method, path};
1073 use wiremock::{Mock, MockServer, ResponseTemplate};
1074
1075 fn ollama_client_pointing_at(server_uri: &str, model: &str) -> OpenAiAiClient {
1076 OpenAiAiClient::new_ollama(model.to_string(), Some(server_uri.to_string()), None).unwrap()
1077 }
1078
1079 #[tokio::test]
1080 async fn probe_returns_lm_studio_value_when_model_loaded() {
1081 let server = MockServer::start().await;
1082 Mock::given(method("GET"))
1083 .and(path("/api/v0/models"))
1084 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1085 "data": [
1086 {
1087 "id": "llama-3.2-3b-instruct",
1088 "state": "loaded",
1089 "loaded_context_length": 4096_u64,
1090 "max_context_length": 131_072_u64,
1091 }
1092 ]
1093 })))
1094 .mount(&server)
1095 .await;
1096
1097 let mut client = ollama_client_pointing_at(&server.uri(), "llama-3.2-3b-instruct");
1098 let source = client.probe_loaded_context_length().await;
1099 assert_eq!(source, Some(ProbeSource::LmStudio));
1100 assert_eq!(client.loaded_context_length(), Some(4096));
1101 assert_eq!(client.get_metadata().max_context_length, 4096);
1103 }
1104
1105 #[tokio::test]
1106 async fn probe_skips_lm_studio_entry_when_model_not_loaded() {
1107 let server = MockServer::start().await;
1108 Mock::given(method("GET"))
1112 .and(path("/api/v0/models"))
1113 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1114 "data": [
1115 { "id": "model-a", "state": "not-loaded", "loaded_context_length": 4096_u64 }
1116 ]
1117 })))
1118 .mount(&server)
1119 .await;
1120 Mock::given(method("POST"))
1121 .and(path("/api/show"))
1122 .and(body_json(serde_json::json!({ "name": "model-a" })))
1123 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1124 "model_info": { "llama.context_length": 8192_u64 }
1125 })))
1126 .mount(&server)
1127 .await;
1128
1129 let mut client = ollama_client_pointing_at(&server.uri(), "model-a");
1130 let source = client.probe_loaded_context_length().await;
1131 assert_eq!(source, Some(ProbeSource::Ollama));
1132 assert_eq!(client.loaded_context_length(), Some(8192));
1133 }
1134
1135 #[tokio::test]
1136 async fn probe_skips_lm_studio_when_model_id_does_not_match() {
1137 let server = MockServer::start().await;
1138 Mock::given(method("GET"))
1139 .and(path("/api/v0/models"))
1140 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1141 "data": [
1142 { "id": "other-model", "state": "loaded", "loaded_context_length": 4096_u64 }
1143 ]
1144 })))
1145 .mount(&server)
1146 .await;
1147 Mock::given(method("POST"))
1148 .and(path("/api/show"))
1149 .respond_with(ResponseTemplate::new(404))
1150 .mount(&server)
1151 .await;
1152
1153 let mut client = ollama_client_pointing_at(&server.uri(), "wanted-model");
1154 let source = client.probe_loaded_context_length().await;
1155 assert!(source.is_none());
1156 assert!(client.loaded_context_length().is_none());
1157 }
1158
1159 #[tokio::test]
1160 async fn probe_falls_back_to_ollama_native() {
1161 let server = MockServer::start().await;
1162 Mock::given(method("GET"))
1163 .and(path("/api/v0/models"))
1164 .respond_with(ResponseTemplate::new(404))
1165 .mount(&server)
1166 .await;
1167 Mock::given(method("POST"))
1168 .and(path("/api/show"))
1169 .and(body_json(serde_json::json!({ "name": "qwen2.5-coder" })))
1170 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1171 "model_info": {
1172 "general.architecture": "qwen2",
1173 "qwen2.context_length": 32768_u64,
1174 "qwen2.embedding_length": 3584_u64
1175 }
1176 })))
1177 .mount(&server)
1178 .await;
1179
1180 let mut client = ollama_client_pointing_at(&server.uri(), "qwen2.5-coder");
1181 let source = client.probe_loaded_context_length().await;
1182 assert_eq!(source, Some(ProbeSource::Ollama));
1183 assert_eq!(client.loaded_context_length(), Some(32768));
1184 }
1185
1186 #[tokio::test]
1187 async fn probe_returns_none_when_neither_endpoint_responds() {
1188 let server = MockServer::start().await;
1189 Mock::given(method("GET"))
1190 .and(path("/api/v0/models"))
1191 .respond_with(ResponseTemplate::new(500))
1192 .mount(&server)
1193 .await;
1194 Mock::given(method("POST"))
1195 .and(path("/api/show"))
1196 .respond_with(ResponseTemplate::new(500))
1197 .mount(&server)
1198 .await;
1199
1200 let mut client = ollama_client_pointing_at(&server.uri(), "anything");
1201 let source = client.probe_loaded_context_length().await;
1202 assert!(source.is_none());
1203 assert!(client.loaded_context_length().is_none());
1204 let registry_value = get_model_registry().get_input_context("anything");
1206 assert_eq!(client.get_metadata().max_context_length, registry_value);
1207 }
1208
1209 #[tokio::test]
1210 async fn probe_returns_none_when_ollama_payload_lacks_context_length_key() {
1211 let server = MockServer::start().await;
1212 Mock::given(method("GET"))
1213 .and(path("/api/v0/models"))
1214 .respond_with(ResponseTemplate::new(404))
1215 .mount(&server)
1216 .await;
1217 Mock::given(method("POST"))
1218 .and(path("/api/show"))
1219 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1220 "model_info": {
1221 "general.architecture": "phantom",
1222 "phantom.embedding_length": 1024_u64
1223 }
1224 })))
1225 .mount(&server)
1226 .await;
1227
1228 let mut client = ollama_client_pointing_at(&server.uri(), "ghost");
1229 let source = client.probe_loaded_context_length().await;
1230 assert!(source.is_none());
1231 }
1232
1233 #[tokio::test]
1234 async fn probe_handles_v1_suffix_in_base_url() {
1235 let server = MockServer::start().await;
1239 Mock::given(method("GET"))
1240 .and(path("/api/v0/models"))
1241 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1242 "data": [
1243 { "id": "lm", "state": "loaded", "loaded_context_length": 2048_u64 }
1244 ]
1245 })))
1246 .mount(&server)
1247 .await;
1248
1249 let base_with_v1 = format!("{}/v1", server.uri());
1250 let mut client = ollama_client_pointing_at(&base_with_v1, "lm");
1251 let source = client.probe_loaded_context_length().await;
1252 assert_eq!(source, Some(ProbeSource::LmStudio));
1253 assert_eq!(client.loaded_context_length(), Some(2048));
1254 }
1255
1256 #[tokio::test]
1257 async fn probe_ignores_lm_studio_entry_with_no_loaded_context_length() {
1258 let server = MockServer::start().await;
1259 Mock::given(method("GET"))
1262 .and(path("/api/v0/models"))
1263 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1264 "data": [ { "id": "x", "state": "loaded", "loaded_context_length": null } ]
1265 })))
1266 .mount(&server)
1267 .await;
1268 Mock::given(method("POST"))
1269 .and(path("/api/show"))
1270 .respond_with(ResponseTemplate::new(404))
1271 .mount(&server)
1272 .await;
1273
1274 let mut client = ollama_client_pointing_at(&server.uri(), "x");
1275 let source = client.probe_loaded_context_length().await;
1276 assert!(source.is_none());
1277 }
1278
1279 #[tokio::test]
1280 async fn probe_returns_none_when_lm_studio_returns_invalid_json() {
1281 let server = MockServer::start().await;
1284 Mock::given(method("GET"))
1285 .and(path("/api/v0/models"))
1286 .respond_with(ResponseTemplate::new(200).set_body_string("<html>not json</html>"))
1287 .mount(&server)
1288 .await;
1289 Mock::given(method("POST"))
1290 .and(path("/api/show"))
1291 .respond_with(ResponseTemplate::new(404))
1292 .mount(&server)
1293 .await;
1294
1295 let mut client = ollama_client_pointing_at(&server.uri(), "anything");
1296 let source = client.probe_loaded_context_length().await;
1297 assert!(source.is_none());
1298 }
1299
1300 #[tokio::test]
1301 async fn probe_returns_none_when_ollama_returns_invalid_json() {
1302 let server = MockServer::start().await;
1305 Mock::given(method("GET"))
1306 .and(path("/api/v0/models"))
1307 .respond_with(ResponseTemplate::new(404))
1308 .mount(&server)
1309 .await;
1310 Mock::given(method("POST"))
1311 .and(path("/api/show"))
1312 .respond_with(ResponseTemplate::new(200).set_body_string("{not json"))
1313 .mount(&server)
1314 .await;
1315
1316 let mut client = ollama_client_pointing_at(&server.uri(), "anything");
1317 let source = client.probe_loaded_context_length().await;
1318 assert!(source.is_none());
1319 }
1320
1321 #[tokio::test]
1322 async fn probe_returns_none_when_ollama_response_lacks_model_info() {
1323 let server = MockServer::start().await;
1326 Mock::given(method("GET"))
1327 .and(path("/api/v0/models"))
1328 .respond_with(ResponseTemplate::new(404))
1329 .mount(&server)
1330 .await;
1331 Mock::given(method("POST"))
1332 .and(path("/api/show"))
1333 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1334 "details": { "family": "llama" }
1335 })))
1336 .mount(&server)
1337 .await;
1338
1339 let mut client = ollama_client_pointing_at(&server.uri(), "anything");
1340 let source = client.probe_loaded_context_length().await;
1341 assert!(source.is_none());
1342 }
1343
1344 #[tokio::test]
1345 async fn probe_returns_none_when_ollama_model_info_is_not_object() {
1346 let server = MockServer::start().await;
1350 Mock::given(method("GET"))
1351 .and(path("/api/v0/models"))
1352 .respond_with(ResponseTemplate::new(404))
1353 .mount(&server)
1354 .await;
1355 Mock::given(method("POST"))
1356 .and(path("/api/show"))
1357 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1358 "model_info": "not an object"
1359 })))
1360 .mount(&server)
1361 .await;
1362
1363 let mut client = ollama_client_pointing_at(&server.uri(), "anything");
1364 let source = client.probe_loaded_context_length().await;
1365 assert!(source.is_none());
1366 }
1367
1368 #[tokio::test]
1369 async fn probe_returns_none_when_ollama_context_length_is_not_u64() {
1370 let server = MockServer::start().await;
1374 Mock::given(method("GET"))
1375 .and(path("/api/v0/models"))
1376 .respond_with(ResponseTemplate::new(404))
1377 .mount(&server)
1378 .await;
1379 Mock::given(method("POST"))
1380 .and(path("/api/show"))
1381 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1382 "model_info": { "llama.context_length": "8192" }
1383 })))
1384 .mount(&server)
1385 .await;
1386
1387 let mut client = ollama_client_pointing_at(&server.uri(), "anything");
1388 let source = client.probe_loaded_context_length().await;
1389 assert!(source.is_none());
1390 }
1391
1392 #[tokio::test]
1393 async fn probe_returns_none_when_server_unreachable() {
1394 let mut client = ollama_client_pointing_at("http://127.0.0.1:1", "anything");
1399 let source = client.probe_loaded_context_length().await;
1400 assert!(source.is_none());
1401 }
1402
1403 async fn mock_chat_completion_ok(server: &MockServer) {
1408 Mock::given(method("POST"))
1409 .and(path("/v1/chat/completions"))
1410 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1411 "choices": [
1412 {
1413 "message": { "role": "assistant", "content": "ok" },
1414 "finish_reason": "stop"
1415 }
1416 ],
1417 "model": "test-model"
1418 })))
1419 .mount(server)
1420 .await;
1421 }
1422
1423 fn openai_client_pointing_at(server_uri: &str, model: &str) -> OpenAiAiClient {
1424 OpenAiAiClient::new(
1425 model.to_string(),
1426 Some("test-key".to_string()),
1427 server_uri.to_string(),
1428 Some(1024),
1429 Some(0.1),
1430 None,
1431 )
1432 .unwrap()
1433 }
1434
1435 #[tokio::test]
1443 async fn send_request_with_options_serializes_response_format_on_the_wire() {
1444 let _ = tracing_subscriber::fmt()
1445 .with_max_level(tracing::Level::DEBUG)
1446 .with_test_writer()
1447 .try_init();
1448
1449 let server = MockServer::start().await;
1450 mock_chat_completion_ok(&server).await;
1451
1452 let client = openai_client_pointing_at(&server.uri(), "gpt-4o");
1453 let schema = serde_json::json!({
1454 "type": "object",
1455 "properties": { "answer": { "type": "string" } },
1456 "required": ["answer"],
1457 "additionalProperties": false,
1458 });
1459 let options = RequestOptions::default().with_response_schema(schema.clone());
1460
1461 let result = client
1462 .send_request_with_options("system", "user", options)
1463 .await
1464 .unwrap();
1465 assert_eq!(result, "ok");
1466
1467 let received = server.received_requests().await.unwrap();
1468 assert_eq!(
1469 received.len(),
1470 1,
1471 "expected exactly one chat-completions request"
1472 );
1473 let body: serde_json::Value = serde_json::from_slice(&received[0].body).unwrap();
1474 assert_eq!(body["response_format"]["type"], "json_schema");
1475 assert_eq!(body["response_format"]["json_schema"]["name"], "response");
1476 assert_eq!(body["response_format"]["json_schema"]["strict"], true);
1477 assert_eq!(body["response_format"]["json_schema"]["schema"], schema);
1478 }
1479
1480 #[tokio::test]
1484 async fn send_request_omits_response_format_on_the_wire() {
1485 let server = MockServer::start().await;
1486 mock_chat_completion_ok(&server).await;
1487
1488 let client = openai_client_pointing_at(&server.uri(), "gpt-4o");
1489 let _ = client.send_request("system", "user").await.unwrap();
1490
1491 let received = server.received_requests().await.unwrap();
1492 assert_eq!(received.len(), 1);
1493 let body: serde_json::Value = serde_json::from_slice(&received[0].body).unwrap();
1494 assert!(
1495 body.get("response_format").is_none(),
1496 "expected response_format to be absent from wire body, got: {body}"
1497 );
1498 }
1499
1500 #[test]
1504 fn build_request_skips_empty_system_prompt() {
1505 let client =
1506 OpenAiAiClient::new_openai("gpt-4o".to_string(), "key".to_string(), None).unwrap();
1507 let request = client.build_request("", "user prompt", None);
1508 assert_eq!(request.messages.len(), 1);
1509 assert_eq!(request.messages[0].role, "user");
1510 assert_eq!(request.messages[0].content, "user prompt");
1511 }
1512
1513 #[tokio::test]
1519 async fn send_request_propagates_network_error_on_unreachable_server() {
1520 let client = openai_client_pointing_at("http://127.0.0.1:1", "gpt-4o");
1521 let err = client
1522 .send_request("system", "user")
1523 .await
1524 .expect_err("expected network error against closed port");
1525 let chain = format!("{err:#}");
1526 assert!(
1527 chain.to_lowercase().contains("network"),
1528 "expected network-error wording in chain, got: {chain}"
1529 );
1530 }
1531
1532 #[tokio::test]
1537 async fn send_request_propagates_http_error_response() {
1538 let server = MockServer::start().await;
1539 Mock::given(method("POST"))
1540 .and(path("/v1/chat/completions"))
1541 .respond_with(ResponseTemplate::new(500).set_body_string("upstream boom"))
1542 .mount(&server)
1543 .await;
1544
1545 let client = openai_client_pointing_at(&server.uri(), "gpt-4o");
1546 let err = client
1547 .send_request("system", "user")
1548 .await
1549 .expect_err("expected error from 500 response");
1550 let chain = format!("{err:#}");
1551 assert!(
1552 chain.contains("HTTP 500"),
1553 "expected 'HTTP 500' in error chain, got: {chain}"
1554 );
1555 }
1556
1557 #[tokio::test]
1560 async fn send_request_propagates_json_parse_error() {
1561 let server = MockServer::start().await;
1562 Mock::given(method("POST"))
1563 .and(path("/v1/chat/completions"))
1564 .respond_with(ResponseTemplate::new(200).set_body_string("{not valid json"))
1565 .mount(&server)
1566 .await;
1567
1568 let client = openai_client_pointing_at(&server.uri(), "gpt-4o");
1569 let err = client
1570 .send_request("system", "user")
1571 .await
1572 .expect_err("expected error from malformed JSON body");
1573 let chain = format!("{err:#}");
1577 assert!(
1578 chain.contains("Invalid response format"),
1579 "expected 'Invalid response format' in error chain, got: {chain}"
1580 );
1581 }
1582
1583 #[tokio::test]
1588 async fn send_request_errors_when_response_has_no_choices() {
1589 let server = MockServer::start().await;
1590 Mock::given(method("POST"))
1591 .and(path("/v1/chat/completions"))
1592 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1593 "choices": [],
1594 "model": "test-model"
1595 })))
1596 .mount(&server)
1597 .await;
1598
1599 let client = openai_client_pointing_at(&server.uri(), "gpt-4o");
1600 let err = client
1601 .send_request("system", "user")
1602 .await
1603 .expect_err("expected error when choices array is empty");
1604 let chain = format!("{err:#}");
1605 assert!(
1606 chain.contains("No choices in response"),
1607 "expected 'No choices in response' in error chain, got: {chain}"
1608 );
1609 }
1610}