lc/
provider.rs

1use anyhow::Result;
2use futures_util::StreamExt;
3use reqwest::Client;
4use serde::{Deserialize, Serialize};
5use std::time::Duration;
6
7use super::template_processor::TemplateProcessor;
8
9#[derive(Debug, Serialize)]
10pub struct ChatRequest {
11    pub model: String,
12    pub messages: Vec<Message>,
13    pub max_tokens: Option<u32>,
14    pub temperature: Option<f32>,
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub tools: Option<Vec<Tool>>,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub stream: Option<bool>,
19}
20
21// Chat request without model field for providers that specify model in URL
22#[derive(Debug, Serialize)]
23pub struct ChatRequestWithoutModel {
24    pub messages: Vec<Message>,
25    pub max_tokens: Option<u32>,
26    pub temperature: Option<f32>,
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub tools: Option<Vec<Tool>>,
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub stream: Option<bool>,
31}
32
33impl From<&ChatRequest> for ChatRequestWithoutModel {
34    fn from(request: &ChatRequest) -> Self {
35        Self {
36            messages: request.messages.clone(),
37            max_tokens: request.max_tokens,
38            temperature: request.temperature,
39            tools: request.tools.clone(),
40            stream: request.stream,
41        }
42    }
43}
44
45
46#[derive(Debug, Serialize)]
47pub struct EmbeddingRequest {
48    pub model: String,
49    pub input: String,
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub encoding_format: Option<String>,
52}
53
54#[derive(Debug, Serialize)]
55pub struct ImageGenerationRequest {
56    pub prompt: String,
57    pub model: Option<String>,
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub n: Option<u32>,
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub size: Option<String>,
62    #[serde(skip_serializing_if = "Option::is_none")]
63    pub quality: Option<String>,
64    #[serde(skip_serializing_if = "Option::is_none")]
65    pub style: Option<String>,
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub response_format: Option<String>,
68}
69
70#[derive(Debug, Deserialize)]
71pub struct ImageGenerationResponse {
72    pub data: Vec<ImageData>,
73}
74
75#[derive(Debug, Deserialize, Clone)]
76pub struct ImageData {
77    #[serde(skip_serializing_if = "Option::is_none")]
78    pub url: Option<String>,
79    #[serde(skip_serializing_if = "Option::is_none")]
80    pub b64_json: Option<String>,
81    #[serde(skip_serializing_if = "Option::is_none")]
82    pub revised_prompt: Option<String>,
83}
84
85#[derive(Debug, Deserialize)]
86pub struct EmbeddingResponse {
87    pub data: Vec<EmbeddingData>,
88    pub usage: EmbeddingUsage,
89}
90
91#[derive(Debug, Deserialize, Clone)]
92pub struct EmbeddingData {
93    pub embedding: Vec<f64>,
94}
95
96#[derive(Debug, Deserialize, Clone)]
97pub struct EmbeddingUsage {
98    pub total_tokens: u32,
99}
100
101#[derive(Debug, Serialize, Clone)]
102pub struct Tool {
103    #[serde(rename = "type")]
104    pub tool_type: String,
105    pub function: Function,
106}
107
108#[derive(Debug, Serialize, Clone)]
109pub struct Function {
110    pub name: String,
111    pub description: String,
112    pub parameters: serde_json::Value,
113}
114
115// Updated Message struct to support multimodal content
116#[derive(Debug, Serialize, Deserialize, Clone)]
117pub struct Message {
118    pub role: String,
119    #[serde(flatten)]
120    pub content_type: MessageContent,
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub tool_calls: Option<Vec<ToolCall>>,
123    #[serde(skip_serializing_if = "Option::is_none")]
124    pub tool_call_id: Option<String>,
125}
126
127// New enum to support both text and multimodal content
128#[derive(Debug, Serialize, Deserialize, Clone)]
129#[serde(untagged)]
130pub enum MessageContent {
131    Text { content: Option<String> },
132    Multimodal { content: Vec<ContentPart> },
133}
134
135// Content part for multimodal messages
136#[derive(Debug, Serialize, Deserialize, Clone)]
137#[serde(tag = "type")]
138pub enum ContentPart {
139    #[serde(rename = "text")]
140    Text { text: String },
141    #[serde(rename = "image_url")]
142    ImageUrl { image_url: ImageUrl },
143}
144
145#[derive(Debug, Serialize, Deserialize, Clone)]
146pub struct ImageUrl {
147    pub url: String,
148    #[serde(skip_serializing_if = "Option::is_none")]
149    pub detail: Option<String>, // "low", "high", or "auto"
150}
151
152impl Message {
153    pub fn user(content: String) -> Self {
154        Self {
155            role: "user".to_string(),
156            content_type: MessageContent::Text {
157                content: Some(content),
158            },
159            tool_calls: None,
160            tool_call_id: None,
161        }
162    }
163
164    #[allow(dead_code)]
165    pub fn user_with_image(text: String, image_data: String, detail: Option<String>) -> Self {
166        Self {
167            role: "user".to_string(),
168            content_type: MessageContent::Multimodal {
169                content: vec![
170                    ContentPart::Text { text },
171                    ContentPart::ImageUrl {
172                        image_url: ImageUrl {
173                            url: image_data,
174                            detail,
175                        },
176                    },
177                ],
178            },
179            tool_calls: None,
180            tool_call_id: None,
181        }
182    }
183
184    pub fn assistant(content: String) -> Self {
185        Self {
186            role: "assistant".to_string(),
187            content_type: MessageContent::Text {
188                content: Some(content),
189            },
190            tool_calls: None,
191            tool_call_id: None,
192        }
193    }
194
195    pub fn assistant_with_tool_calls(tool_calls: Vec<ToolCall>) -> Self {
196        Self {
197            role: "assistant".to_string(),
198            content_type: MessageContent::Text { content: None },
199            tool_calls: Some(tool_calls),
200            tool_call_id: None,
201        }
202    }
203
204    pub fn tool_result(tool_call_id: String, content: String) -> Self {
205        Self {
206            role: "tool".to_string(),
207            content_type: MessageContent::Text {
208                content: Some(content),
209            },
210            tool_calls: None,
211            tool_call_id: Some(tool_call_id),
212        }
213    }
214
215    // Helper method to get text content for backward compatibility
216    pub fn get_text_content(&self) -> Option<&String> {
217        match &self.content_type {
218            MessageContent::Text { content } => content.as_ref(),
219            MessageContent::Multimodal { content } => {
220                // Return the first text content if available
221                content.iter().find_map(|part| match part {
222                    ContentPart::Text { text } => Some(text),
223                    _ => None,
224                })
225            }
226        }
227    }
228}
229
230#[derive(Debug, Deserialize)]
231pub struct ChatResponse {
232    pub choices: Vec<Choice>,
233}
234
235#[derive(Debug, Deserialize)]
236pub struct Choice {
237    pub message: ResponseMessage,
238}
239
240#[derive(Debug, Deserialize)]
241pub struct ResponseMessage {
242    #[allow(dead_code)]
243    pub role: String,
244    pub content: Option<String>,
245    pub tool_calls: Option<Vec<ToolCall>>,
246}
247
248#[derive(Debug, Serialize, Deserialize, Clone)]
249pub struct ToolCall {
250    pub id: String,
251    #[serde(rename = "type")]
252    pub call_type: String,
253    pub function: FunctionCall,
254}
255
256#[derive(Debug, Serialize, Deserialize, Clone)]
257pub struct FunctionCall {
258    pub name: String,
259    pub arguments: String,
260}
261
262
263#[derive(Debug, Deserialize)]
264pub struct ModelsResponse {
265    #[serde(alias = "models")]
266    pub data: Vec<Model>,
267}
268
269#[derive(Debug, Deserialize)]
270pub struct Provider {
271    pub provider: String,
272    #[allow(dead_code)]
273    pub status: String,
274    #[serde(default)]
275    #[allow(dead_code)]
276    pub supports_tools: bool,
277    #[serde(default)]
278    #[allow(dead_code)]
279    pub supports_structured_output: bool,
280}
281
282#[derive(Debug, Deserialize)]
283pub struct Model {
284    pub id: String,
285    #[serde(default = "default_object_type")]
286    pub object: String,
287    #[serde(default)]
288    pub providers: Vec<Provider>,
289}
290
291fn default_object_type() -> String {
292    "model".to_string()
293}
294
295#[derive(Debug, Deserialize)]
296pub struct TokenResponse {
297    pub token: String,
298    pub expires_at: i64, // Unix timestamp
299}
300
301pub struct OpenAIClient {
302    client: Client,
303    streaming_client: Client, // Separate client optimized for streaming
304    base_url: String,
305    api_key: String,
306    models_path: String,
307    chat_path: String,
308    custom_headers: std::collections::HashMap<String, String>,
309    provider_config: Option<crate::config::ProviderConfig>,
310    template_processor: Option<TemplateProcessor>,
311}
312
313impl OpenAIClient {
314    pub fn new_with_headers(
315        base_url: String,
316        api_key: String,
317        models_path: String,
318        chat_path: String,
319        custom_headers: std::collections::HashMap<String, String>,
320    ) -> Self {
321        use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
322
323        // Create default headers including the required tracking headers
324        let mut default_headers = HeaderMap::new();
325        default_headers.insert(
326            HeaderName::from_static("http-referer"),
327            HeaderValue::from_static("https://lc.viwq.dev/"),
328        );
329        default_headers.insert(
330            HeaderName::from_static("x-title"),
331            HeaderValue::from_static("lc"),
332        );
333
334        // Create optimized HTTP client with connection pooling and keep-alive settings
335        // This client keeps compression enabled for regular requests
336        let client = Client::builder()
337            .pool_max_idle_per_host(10) // Keep up to 10 idle connections per host
338            .pool_idle_timeout(Duration::from_secs(90)) // Keep connections alive for 90 seconds
339            .tcp_keepalive(Duration::from_secs(60)) // TCP keep-alive every 60 seconds
340            .timeout(Duration::from_secs(60)) // Total request timeout
341            .connect_timeout(Duration::from_secs(10)) // Connection establishment timeout
342            .user_agent(concat!(
343                env!("CARGO_PKG_NAME"),
344                "/",
345                env!("CARGO_PKG_VERSION")
346            ))
347            .default_headers(default_headers.clone())
348            .build()
349            .expect("Failed to create optimized HTTP client");
350
351        // Create a separate streaming-optimized client
352        let streaming_client = Client::builder()
353            .timeout(Duration::from_secs(300)) // Longer timeout for streaming
354            .default_headers(default_headers)
355            .build()
356            .expect("Failed to create streaming-optimized HTTP client");
357
358        Self {
359            client,
360            streaming_client,
361            base_url: base_url.trim_end_matches('/').to_string(),
362            api_key,
363            models_path,
364            chat_path,
365            custom_headers,
366            provider_config: None,
367            template_processor: None,
368        }
369    }
370
371    pub fn new_with_provider_config(
372        base_url: String,
373        api_key: String,
374        models_path: String,
375        chat_path: String,
376        custom_headers: std::collections::HashMap<String, String>,
377        provider_config: crate::config::ProviderConfig,
378    ) -> Self {
379        use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
380
381        // Create default headers including the required tracking headers
382        let mut default_headers = HeaderMap::new();
383        default_headers.insert(
384            HeaderName::from_static("http-referer"),
385            HeaderValue::from_static("https://lc.viwq.dev/"),
386        );
387        default_headers.insert(
388            HeaderName::from_static("x-title"),
389            HeaderValue::from_static("lc"),
390        );
391
392        // Create optimized HTTP client with connection pooling and keep-alive settings
393        // This client keeps compression enabled for regular requests
394        let client = Client::builder()
395            .pool_max_idle_per_host(10) // Keep up to 10 idle connections per host
396            .pool_idle_timeout(Duration::from_secs(90)) // Keep connections alive for 90 seconds
397            .tcp_keepalive(Duration::from_secs(60)) // TCP keep-alive every 60 seconds
398            .timeout(Duration::from_secs(60)) // Total request timeout
399            .connect_timeout(Duration::from_secs(10)) // Connection establishment timeout
400            .user_agent(concat!(
401                env!("CARGO_PKG_NAME"),
402                "/",
403                env!("CARGO_PKG_VERSION")
404            ))
405            .default_headers(default_headers.clone())
406            .build()
407            .expect("Failed to create optimized HTTP client");
408
409        // Create a separate streaming-optimized client
410        let streaming_client = Client::builder()
411            .timeout(Duration::from_secs(300)) // Longer timeout for streaming
412            .default_headers(default_headers)
413            .build()
414            .expect("Failed to create streaming-optimized HTTP client");
415
416        // Create template processor if any endpoint templates are configured
417        let template_processor = if provider_config.chat_templates.is_some()
418            || provider_config.images_templates.is_some()
419            || provider_config.embeddings_templates.is_some()
420            || provider_config.models_templates.is_some() {
421            match TemplateProcessor::new() {
422                Ok(processor) => Some(processor),
423                Err(e) => {
424                    eprintln!("Warning: Failed to create template processor: {}", e);
425                    None
426                }
427            }
428        } else {
429            None
430        };
431
432        Self {
433            client,
434            streaming_client,
435            base_url: base_url.trim_end_matches('/').to_string(),
436            api_key,
437            models_path,
438            chat_path,
439            custom_headers,
440            provider_config: Some(provider_config),
441            template_processor,
442        }
443    }
444
445    /// Get the chat URL, handling both traditional paths and full URLs with model replacement
446    fn get_chat_url(&self, model: &str) -> String {
447        if let Some(ref config) = self.provider_config {
448            // Use the provider config's URL generation method which handles template variables
449            config.get_chat_url(model)
450        } else {
451            // Fallback to original logic for backward compatibility
452            if self.chat_path.starts_with("https://") {
453                // Full URL with model replacement
454                self.chat_path
455                    .replace("{model_name}", model)
456                    .replace("{model}", model)
457            } else {
458                // Traditional path-based approach
459                format!("{}{}", self.base_url, self.chat_path)
460            }
461        }
462    }
463
464
465
466    pub async fn chat(&self, request: &ChatRequest) -> Result<String> {
467        let url = self.get_chat_url(&request.model);
468
469        let mut req = self
470            .client
471            .post(&url)
472            .header("Content-Type", "application/json");
473
474        // Disable compression for streaming requests
475        if request.stream == Some(true) {
476            req = req.header("Accept-Encoding", "identity");
477        }
478
479        // Add Authorization header only if no custom headers are present
480        // This allows providers like Gemini to use custom authentication headers
481        if self.custom_headers.is_empty() {
482            req = req.header("Authorization", format!("Bearer {}", self.api_key));
483        }
484
485        // Add custom headers
486        for (name, value) in &self.custom_headers {
487            req = req.header(name, value);
488        }
489
490        // Check if we have a template for this provider/model/endpoint
491        let request_body = if let Some(ref config) = &self.provider_config {
492            if let Some(ref processor) = &self.template_processor {
493                // Get template for chat endpoint
494                let template = config.get_endpoint_template("chat", &request.model);
495
496                if let Some(template_str) = template {
497                    // Clone the processor to avoid mutable borrow issues
498                    let mut processor_clone = processor.clone();
499                    // Use template to transform request
500                    match processor_clone.process_request(request, &template_str, &config.vars) {
501                        Ok(json_value) => Some(json_value),
502                        Err(e) => {
503                            eprintln!("Warning: Failed to process request template: {}. Falling back to default.", e);
504                            None
505                        }
506                    }
507                } else {
508                    None
509                }
510            } else {
511                None
512            }
513        } else {
514            None
515        };
516
517        // Send request with template-processed body or fall back to default logic
518        let response = if let Some(json_body) = request_body {
519            req.json(&json_body).send().await?
520        } else {
521            // Fall back to existing logic
522            // Check if we should exclude model from payload (when model is in URL path)
523            let should_exclude_model = if let Some(ref config) = self.provider_config {
524                config.chat_path.contains("{model}")
525            } else {
526                self.chat_path.contains("{model}")
527            };
528
529            if should_exclude_model {
530                // Use ChatRequestWithoutModel for providers that specify model in URL
531                let request_without_model = ChatRequestWithoutModel::from(request);
532                req.json(&request_without_model).send().await?
533            } else {
534                req.json(request).send().await?
535            }
536        };
537
538        if !response.status().is_success() {
539            let status = response.status();
540            let text = response.text().await.unwrap_or_default();
541            anyhow::bail!("API request failed with status {}: {}", status, text);
542        }
543
544        // Get the response text first to handle different formats
545        let response_text = response.text().await?;
546
547        // Check if we have a response template for this provider/model/endpoint
548        if let Some(ref config) = &self.provider_config {
549            if let Some(ref processor) = &self.template_processor {
550                // Get response template for chat endpoint
551                let template = config.get_endpoint_response_template("chat", &request.model);
552
553                if let Some(template_str) = template {
554                    // Parse response as JSON
555                    if let Ok(response_json) = serde_json::from_str::<serde_json::Value>(&response_text) {
556                        // Clone the processor to avoid mutable borrow issues
557                        let mut processor_clone = processor.clone();
558                        // Use template to extract content
559                        match processor_clone.process_response(&response_json, &template_str) {
560                            Ok(extracted) => {
561                                // Extract content from the template result
562                                if let Some(content) = extracted.get("content").and_then(|v| v.as_str()) {
563                                    return Ok(content.to_string());
564                                } else if let Some(tool_calls) = extracted.get("tool_calls").and_then(|v| v.as_array()) {
565                                    if !tool_calls.is_empty() {
566                                        let mut response = String::new();
567                                        response.push_str("🔧 **Tool Calls Made:**\n\n");
568                                        response.push_str(&format!("Tool calls: {:?}\n\n", tool_calls));
569                                        response.push_str("*Tool calls detected - execution handled by chat module*\n\n");
570                                        return Ok(response);
571                                    }
572                                }
573                            }
574                            Err(e) => {
575                                eprintln!("Warning: Failed to process response template: {}. Falling back to default parsing.", e);
576                            }
577                        }
578                    }
579                }
580            }
581        }
582
583        // Fall back to existing parsing logic
584        // Try to parse as standard OpenAI format (with "choices" array)
585        if let Ok(chat_response) = serde_json::from_str::<ChatResponse>(&response_text) {
586            if let Some(choice) = chat_response.choices.first() {
587                // Handle tool calls - check if tool_calls exists AND is not empty
588                if let Some(tool_calls) = &choice.message.tool_calls {
589                    if !tool_calls.is_empty() {
590                        let mut response = String::new();
591                        response.push_str("🔧 **Tool Calls Made:**\n\n");
592
593                        for tool_call in tool_calls {
594                            response.push_str(&format!(
595                                "**Function:** `{}`\n",
596                                tool_call.function.name
597                            ));
598                            response.push_str(&format!(
599                                "**Arguments:** `{}`\n\n",
600                                tool_call.function.arguments
601                            ));
602
603                            // Note: Tool execution is handled by the chat module's tool execution loop
604                            response.push_str(
605                                "*Tool calls detected - execution handled by chat module*\n\n",
606                            );
607                        }
608
609                        return Ok(response);
610                    }
611                    // If tool_calls is empty array, fall through to check content
612                }
613
614                // Handle content (either no tool_calls or empty tool_calls array)
615                if let Some(content) = &choice.message.content {
616                    return Ok(content.clone());
617                } else {
618                    anyhow::bail!("No content or tool calls in response");
619                }
620            } else {
621                anyhow::bail!("No response from API");
622            }
623        }
624
625
626
627        // If all fail, return an error with the response text for debugging
628        anyhow::bail!("Failed to parse chat response. Response: {}", response_text);
629    }
630
631    pub async fn list_models(&self) -> Result<Vec<Model>> {
632        let url = format!("{}{}", self.base_url, self.models_path);
633
634        let mut req = self
635            .client
636            .get(&url)
637            .header("Content-Type", "application/json");
638
639        // Add Authorization header only if no custom headers are present
640        if self.custom_headers.is_empty() {
641            req = req.header("Authorization", format!("Bearer {}", self.api_key));
642        }
643
644        // Add custom headers
645        for (name, value) in &self.custom_headers {
646            req = req.header(name, value);
647        }
648
649        let response = req.send().await?;
650
651        if !response.status().is_success() {
652            let status = response.status();
653            let text = response.text().await.unwrap_or_default();
654            anyhow::bail!("API request failed with status {}: {}", status, text);
655        }
656
657        // Get the response text first to handle different formats
658        let response_text = response.text().await?;
659
660        // Try to parse as ModelsResponse first (with "data" field)
661        let models =
662            if let Ok(models_response) = serde_json::from_str::<ModelsResponse>(&response_text) {
663                models_response.data
664            } else if let Ok(parsed_models) = serde_json::from_str::<Vec<Model>>(&response_text) {
665                // If that fails, try to parse as direct array of models
666                parsed_models
667            } else {
668                // If all fail, return an error with the response text for debugging
669                anyhow::bail!(
670                    "Failed to parse models response. Response: {}",
671                    response_text
672                );
673            };
674
675        // Expand models with providers into separate entries
676        let mut expanded_models = Vec::new();
677
678        for model in models {
679            if model.providers.is_empty() {
680                // No providers, add the model as-is
681                expanded_models.push(model);
682            } else {
683                // Has providers, create a model entry for each provider
684                for provider in &model.providers {
685                    let expanded_model = Model {
686                        id: format!("{}:{}", model.id, provider.provider),
687                        object: model.object.clone(),
688                        providers: vec![], // Clear providers for the expanded model
689                    };
690                    expanded_models.push(expanded_model);
691                }
692            }
693        }
694
695        Ok(expanded_models)
696    }
697
698    // New method that returns the full parsed response for tool handling
699    pub async fn chat_with_tools(&self, request: &ChatRequest) -> Result<ChatResponse> {
700        let url = self.get_chat_url(&request.model);
701
702        let mut req = self
703            .client
704            .post(&url)
705            .header("Content-Type", "application/json");
706
707        // Disable compression for streaming requests
708        if request.stream == Some(true) {
709            req = req.header("Accept-Encoding", "identity");
710        }
711
712        // Add Authorization header only if no custom headers are present
713        if self.custom_headers.is_empty() {
714            req = req.header("Authorization", format!("Bearer {}", self.api_key));
715        }
716
717        // Add custom headers
718        for (name, value) in &self.custom_headers {
719            req = req.header(name, value);
720        }
721
722        // Check if we should exclude model from payload (when model is in URL path)
723        let should_exclude_model = if let Some(ref config) = self.provider_config {
724            config.chat_path.contains("{model}")
725        } else {
726            self.chat_path.contains("{model}")
727        };
728
729        let response = if should_exclude_model {
730            // Use ChatRequestWithoutModel for providers that specify model in URL
731            let request_without_model = ChatRequestWithoutModel::from(request);
732            req.json(&request_without_model).send().await?
733        } else {
734            req.json(request).send().await?
735        };
736
737        if !response.status().is_success() {
738            let status = response.status();
739            let text = response.text().await.unwrap_or_default();
740            anyhow::bail!("API request failed with status {}: {}", status, text);
741        }
742
743        // Get the response text first to handle different formats
744        let response_text = response.text().await?;
745
746        // Try to parse as standard OpenAI format (with "choices" array)
747        if let Ok(chat_response) = serde_json::from_str::<ChatResponse>(&response_text) {
748            return Ok(chat_response);
749        }
750
751        // If parsing fails, return an error with the response text for debugging
752        anyhow::bail!("Failed to parse chat response. Response: {}", response_text);
753    }
754
755    pub async fn get_token_from_url(&self, token_url: &str) -> Result<TokenResponse> {
756        let mut req = self
757            .client
758            .get(token_url)
759            .header("Authorization", format!("token {}", self.api_key))
760            .header("Content-Type", "application/json");
761
762        // Add custom headers
763        for (name, value) in &self.custom_headers {
764            req = req.header(name, value);
765        }
766
767        let response = req.send().await?;
768
769        if !response.status().is_success() {
770            let status = response.status();
771            let text = response.text().await.unwrap_or_default();
772            anyhow::bail!("Token request failed with status {}: {}", status, text);
773        }
774
775        let token_response: TokenResponse = response.json().await?;
776        Ok(token_response)
777    }
778
779    pub async fn embeddings(&self, request: &EmbeddingRequest) -> Result<EmbeddingResponse> {
780        let url = format!("{}/embeddings", self.base_url);
781
782        let mut req = self
783            .client
784            .post(&url)
785            .header("Content-Type", "application/json");
786
787        // Add Authorization header only if no custom headers are present
788        if self.custom_headers.is_empty() {
789            req = req.header("Authorization", format!("Bearer {}", self.api_key));
790        }
791
792        // Add custom headers
793        for (name, value) in &self.custom_headers {
794            req = req.header(name, value);
795        }
796
797        let response = req.json(request).send().await?;
798
799        if !response.status().is_success() {
800            let status = response.status();
801            let text = response.text().await.unwrap_or_default();
802            anyhow::bail!(
803                "Embeddings API request failed with status {}: {}",
804                status,
805                text
806            );
807        }
808
809        let embedding_response: EmbeddingResponse = response.json().await?;
810        Ok(embedding_response)
811    }
812
813    pub async fn generate_images(
814        &self,
815        request: &ImageGenerationRequest,
816    ) -> Result<ImageGenerationResponse> {
817        // Use provider config's images path if available, otherwise default
818        let url = if let Some(ref config) = self.provider_config {
819            format!("{}{}", self.base_url, config.images_path.as_deref().unwrap_or("/images/generations"))
820        } else {
821            format!("{}/images/generations", self.base_url)
822        };
823
824        let mut req = self
825            .client
826            .post(&url)
827            .header("Content-Type", "application/json");
828
829        // Add Authorization header only if no custom headers are present
830        if self.custom_headers.is_empty() {
831            req = req.header("Authorization", format!("Bearer {}", self.api_key));
832        }
833
834        // Add custom headers
835        for (name, value) in &self.custom_headers {
836            req = req.header(name, value);
837        }
838
839        // Check if we have a template for this provider/model/endpoint
840        let request_body = if let Some(ref config) = &self.provider_config {
841            if let Some(ref processor) = &self.template_processor {
842                // Get template for images endpoint
843                let model_name = request.model.as_deref().unwrap_or("");
844                let template = config.get_endpoint_template("images", model_name);
845
846                if let Some(template_str) = template {
847                    // Clone the processor to avoid mutable borrow issues
848                    let mut processor_clone = processor.clone();
849                    // Use template to transform request
850                    match processor_clone.process_image_request(request, &template_str, &config.vars) {
851                        Ok(json_value) => Some(json_value),
852                        Err(e) => {
853                            eprintln!("Warning: Failed to process image request template: {}. Falling back to default.", e);
854                            None
855                        }
856                    }
857                } else {
858                    None
859                }
860            } else {
861                None
862            }
863        } else {
864            None
865        };
866
867        // Send request with template-processed body or fall back to default logic
868        let response = if let Some(json_body) = request_body {
869            req.json(&json_body).send().await?
870        } else {
871            req.json(request).send().await?
872        };
873
874        if !response.status().is_success() {
875            let status = response.status();
876            let text = response.text().await.unwrap_or_default();
877            anyhow::bail!(
878                "Image generation API request failed with status {}: {}",
879                status,
880                text
881            );
882        }
883
884        // Get the response text first to handle different formats
885        let response_text = response.text().await?;
886
887        // Check if we have a response template for this provider/model/endpoint
888        if let Some(ref config) = &self.provider_config {
889            if let Some(ref processor) = &self.template_processor {
890                // Get response template for images endpoint
891                let model_name = request.model.as_deref().unwrap_or("");
892                let template = config.get_endpoint_response_template("images", model_name);
893
894                if let Some(template_str) = template {
895                    // Parse response as JSON
896                    if let Ok(response_json) = serde_json::from_str::<serde_json::Value>(&response_text) {
897                        // Clone the processor to avoid mutable borrow issues
898                        let mut processor_clone = processor.clone();
899                        // Use template to transform response
900                        match processor_clone.process_response(&response_json, &template_str) {
901                            Ok(transformed) => {
902                                // Try to parse the transformed response as ImageGenerationResponse
903                                if let Ok(image_response) = serde_json::from_value::<ImageGenerationResponse>(transformed) {
904                                    return Ok(image_response);
905                                }
906                            }
907                            Err(e) => {
908                                eprintln!("Warning: Failed to process image response template: {}. Falling back to default parsing.", e);
909                            }
910                        }
911                    }
912                }
913            }
914        }
915
916        // Fall back to default parsing
917        let image_response: ImageGenerationResponse = serde_json::from_str(&response_text)?;
918        Ok(image_response)
919    }
920
921    pub async fn chat_stream(&self, request: &ChatRequest) -> Result<()> {
922        use std::io::{stdout, Write};
923
924        let url = self.get_chat_url(&request.model);
925
926        // Use the streaming-optimized client for streaming requests
927        let mut req = self
928            .streaming_client
929            .post(&url)
930            .header("Content-Type", "application/json")
931            .header("Accept", "text/event-stream") // Explicitly request SSE format
932            .header("Cache-Control", "no-cache") // Prevent caching for streaming
933            .header("Accept-Encoding", "identity"); // Explicitly request no compression
934
935        // Wrap stdout in BufWriter for efficiency
936        let stdout = stdout();
937        let mut handle = std::io::BufWriter::new(stdout.lock());
938
939        // Add Authorization header only if no custom headers are present
940        if self.custom_headers.is_empty() {
941            req = req.header("Authorization", format!("Bearer {}", self.api_key));
942        }
943
944        // Add custom headers
945        for (name, value) in &self.custom_headers {
946            req = req.header(name, value);
947        }
948
949        // Check if we should exclude model from payload (when model is in URL path)
950        let should_exclude_model = if let Some(ref config) = self.provider_config {
951            config.chat_path.contains("{model}")
952        } else {
953            self.chat_path.contains("{model}")
954        };
955
956        let response = if should_exclude_model {
957            // Use ChatRequestWithoutModel for providers that specify model in URL
958            let request_without_model = ChatRequestWithoutModel::from(request);
959            req.json(&request_without_model).send().await?
960        } else {
961            req.json(request).send().await?
962        };
963
964        if !response.status().is_success() {
965            let status = response.status();
966            let text = response.text().await.unwrap_or_default();
967            anyhow::bail!("API request failed with status {}: {}", status, text);
968        }
969
970        // Check for compression headers (silent check for potential issues)
971        let headers = response.headers();
972        if headers.get("content-encoding").is_some() {
973            // Content encoding detected - may cause buffering delays but continue silently
974        }
975
976        let mut stream = response.bytes_stream();
977
978        let mut buffer = String::new();
979
980        while let Some(chunk) = stream.next().await {
981            let chunk = chunk?;
982
983            let chunk_str = String::from_utf8_lossy(&chunk);
984            buffer.push_str(&chunk_str);
985
986            // Process complete lines from buffer
987            while let Some(newline_pos) = buffer.find('\n') {
988                let line = buffer[..newline_pos].to_string();
989                buffer.drain(..=newline_pos);
990
991                // Handle Server-Sent Events format
992                if line.starts_with("data: ") {
993                    let data = &line[6..]; // Remove "data: " prefix
994
995                    if data.trim() == "[DONE]" {
996                        handle.write_all(b"\n")?;
997                        handle.flush()?;
998                        return Ok(());
999                    }
1000
1001                    if let Ok(json) = serde_json::from_str::<serde_json::Value>(data) {
1002                        // Try direct "response" field format first
1003                        if let Some(response) = json.get("response") {
1004                            if let Some(text) = response.as_str() {
1005                                if !text.is_empty() {
1006                                    handle.write_all(text.as_bytes())?;
1007                                    handle.flush()?;
1008                                }
1009                            }
1010                        }
1011                        // Try standard OpenAI streaming format
1012                        else if let Some(choices) = json.get("choices") {
1013                            if let Some(choice) = choices.get(0) {
1014                                if let Some(delta) = choice.get("delta") {
1015                                    if let Some(content) = delta.get("content") {
1016                                        if let Some(text) = content.as_str() {
1017                                            // Write directly to stdout and flush immediately
1018                                            handle.write_all(text.as_bytes())?;
1019                                            handle.flush()?;
1020                                        }
1021                                    }
1022                                }
1023                            }
1024                        }
1025                    }
1026                } else if line.trim().is_empty() {
1027                    // Skip empty lines in SSE format
1028                    continue;
1029                } else {
1030                    // Handle non-SSE format (direct JSON stream)
1031                    if let Ok(json) = serde_json::from_str::<serde_json::Value>(&line) {
1032                        // Try direct "response" field format first
1033                        if let Some(response) = json.get("response") {
1034                            if let Some(text) = response.as_str() {
1035                                if !text.is_empty() {
1036                                    handle.write_all(text.as_bytes())?;
1037                                    handle.flush()?;
1038                                }
1039                            }
1040                        }
1041                        // Try standard OpenAI streaming format
1042                        else if let Some(choices) = json.get("choices") {
1043                            if let Some(choice) = choices.get(0) {
1044                                if let Some(delta) = choice.get("delta") {
1045                                    if let Some(content) = delta.get("content") {
1046                                        if let Some(text) = content.as_str() {
1047                                            handle.write_all(text.as_bytes())?;
1048                                            handle.flush()?;
1049                                        }
1050                                    }
1051                                }
1052                            }
1053                        }
1054                    }
1055                }
1056            }
1057        }
1058
1059        // Process any remaining data in buffer
1060        if !buffer.trim().is_empty() {
1061            if let Ok(json) = serde_json::from_str::<serde_json::Value>(&buffer) {
1062                // Try direct "response" field format first
1063                if let Some(response) = json.get("response") {
1064                    if let Some(text) = response.as_str() {
1065                        if !text.is_empty() {
1066                            handle.write_all(text.as_bytes())?;
1067                            handle.flush()?;
1068                        }
1069                    }
1070                }
1071                // Try standard OpenAI streaming format
1072                else if let Some(choices) = json.get("choices") {
1073                    if let Some(choice) = choices.get(0) {
1074                        if let Some(delta) = choice.get("delta") {
1075                            if let Some(content) = delta.get("content") {
1076                                if let Some(text) = content.as_str() {
1077                                    handle.write_all(text.as_bytes())?;
1078                                    handle.flush()?;
1079                                }
1080                            }
1081                        }
1082                    }
1083                }
1084            }
1085        }
1086
1087        // Add newline at the end
1088        handle.write_all(b"\n")?;
1089        handle.flush()?;
1090        Ok(())
1091    }
1092}
lc/provider.rs

lc/
provider.rs
