bytedocs_rs/ai/

client.rs

use crate::ai::types::{register_client_factory, AIConfig, Client, ChatRequest, ChatResponse};
use anyhow::Result;
use serde_json::json;

pub struct OpenAIClient {
    config: AIConfig,
}

pub struct GeminiClient {
    config: AIConfig,
}

pub struct OpenRouterClient {
    config: AIConfig,
}

impl OpenAIClient {
    pub fn new(config: &AIConfig) -> Result<Box<dyn Client>> {
        Ok(Box::new(Self {
            config: config.clone(),
        }))
    }

    // Optimize API context to ultra-compact format like Laravel
    fn optimize_api_context(&self, context: &str, user_message: &str) -> String {
        // Parse context to extract API info
        let (api_info, endpoints) = self.parse_context(context);

        // Find relevant endpoints based on user question
        let relevant_endpoints = self.find_relevant_endpoints(&endpoints, user_message);

        // Format in ultra-compact way
        let mut result = format!("API: {} | {}", api_info.0, api_info.1); // title | base_url

        // Add compact endpoints (max 5 to save tokens)
        let max_endpoints = 5;
        for (i, endpoint) in relevant_endpoints.iter().take(max_endpoints).enumerate() {
            if i == 0 {
                result.push('\n');
            }
            result.push_str(&self.format_compact_endpoint(endpoint));
        }

        // Add response format note
        result.push_str("\nRESP: {success:bool,message:str,data:obj}");
        result.push_str("\nRULES: Use only listed endpoints. No invention.");
        result
    }

    // Parse context to extract API info and endpoints
    fn parse_context(&self, context: &str) -> ((String, String), Vec<CompactEndpoint>) {
        let mut api_title = "API".to_string();
        let mut base_url = "localhost:8000".to_string();
        let mut endpoints = Vec::new();

        // Extract API title and base URL
        if let Some(title_start) = context.find("API Title: ") {
            if let Some(title_end) = context[title_start + 11..].find('\n') {
                api_title = context[title_start + 11..title_start + 11 + title_end].trim().to_string();
            }
        }

        if let Some(url_start) = context.find("Base URLs: ") {
            if let Some(url_end) = context[url_start..].find(',') {
                let url_part = &context[url_start..url_start + url_end];
                if let Some(first_url) = url_part.split("Production: ").nth(1) {
                    base_url = first_url.trim().to_string();
                }
            }
        }

        // Parse endpoints from OpenAPI paths
        if let Some(paths_start) = context.find("\"paths\": {") {
            endpoints = self.extract_endpoints_from_openapi(&context[paths_start..]);
        }

        ((api_title, base_url), endpoints)
    }

    // Extract endpoints from OpenAPI JSON in context
    fn extract_endpoints_from_openapi(&self, paths_section: &str) -> Vec<CompactEndpoint> {
        let mut endpoints = Vec::new();

        // Enhanced parsing with request/response structures
        let endpoint_definitions = [
            ("/api/health", "GET", "Health check", vec![], vec![]),
            ("/api/users", "GET", "List users", vec!["page", "limit", "status", "search"], vec![]),
            ("/api/users", "POST", "Create user", vec![], vec!["name*", "email*", "age"]),
            ("/api/users/{id}", "GET", "Get user", vec!["id*"], vec![]),
            ("/api/users/{id}", "PUT", "Update user", vec!["id*"], vec!["name", "email", "age", "status"]),
            ("/api/users/{id}", "DELETE", "Delete user", vec!["id*"], vec![]),
            ("/api/products", "GET", "List products", vec![], vec![]),
            ("/api/products", "POST", "Create product", vec![], vec!["name*", "price*", "category*"]),
        ];

        for (path, method, desc, params, body_fields) in endpoint_definitions.iter() {
            if paths_section.contains(&format!("\"{path}\"")) && paths_section.contains(&format!("\"{}\": {{", method.to_lowercase())) {
                endpoints.push(CompactEndpoint {
                    path: path.to_string(),
                    method: method.to_string(),
                    description: desc.to_string(),
                    params: params.iter().map(|s| s.to_string()).collect(),
                    body_fields: body_fields.iter().map(|s| s.to_string()).collect(),
                });
            }
        }

        endpoints
    }

    // Find endpoints relevant to user question
    fn find_relevant_endpoints<'a>(&self, endpoints: &'a [CompactEndpoint], user_message: &str) -> Vec<&'a CompactEndpoint> {
        let message = user_message.to_lowercase();
        let mut scored_endpoints = Vec::new();

        for endpoint in endpoints {
            let mut score = 0;

            // Score based on path keywords
            let path_lower = endpoint.path.to_lowercase();
            if message.contains("user") && path_lower.contains("user") { score += 10; }
            if message.contains("product") && path_lower.contains("product") { score += 10; }
            if message.contains("health") && path_lower.contains("health") { score += 10; }

            // Score based on method
            if message.contains("create") && endpoint.method == "POST" { score += 8; }
            if message.contains("update") && endpoint.method == "PUT" { score += 8; }
            if message.contains("delete") && endpoint.method == "DELETE" { score += 8; }
            if message.contains("get") || message.contains("list") {
                if endpoint.method == "GET" { score += 8; }
            }

            // Score based on description
            let desc_lower = endpoint.description.to_lowercase();
            for word in message.split_whitespace() {
                if desc_lower.contains(word) { score += 3; }
            }

            scored_endpoints.push((endpoint, score));
        }

        // Sort by score and return top endpoints
        scored_endpoints.sort_by(|a, b| b.1.cmp(&a.1));

        // If no relevant endpoints found, return first few
        if scored_endpoints.is_empty() || scored_endpoints[0].1 == 0 {
            return endpoints.iter().take(5).collect();
        }

        scored_endpoints.iter().map(|(endpoint, _)| *endpoint).collect()
    }

    // Format endpoint in compact but informative format
    fn format_compact_endpoint(&self, endpoint: &CompactEndpoint) -> String {
        let mut result = format!("{} {}", endpoint.method, endpoint.path);

        // Add request body structure for POST/PUT methods
        if !endpoint.body_fields.is_empty() && (endpoint.method == "POST" || endpoint.method == "PUT") {
            result.push_str(" {");
            result.push_str(&endpoint.body_fields.join(","));
            result.push('}');
        }

        // Add parameters for GET methods or path parameters
        if !endpoint.params.is_empty() {
            result.push_str(" ?");
            result.push_str(&endpoint.params.join(","));
        }

        result.push('\n');
        result
    }
}

// Simple struct for compact endpoint representation
#[derive(Debug)]
struct CompactEndpoint {
    path: String,
    method: String,
    description: String,
    params: Vec<String>,
    body_fields: Vec<String>,
}

#[async_trait::async_trait]
impl Client for OpenAIClient {
    async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
        let client = reqwest::Client::new();

        // Build messages for OpenAI API
        let mut messages = Vec::new();

        if let Some(ref context) = request.context {
            if !context.is_empty() {
                // Extract and optimize context in Laravel-style compact format
                let optimized_context = self.optimize_api_context(context, &request.message);

                messages.push(json!({
                    "role": "system",
                    "content": optimized_context
                }));
            } else {
                // Add basic system message if no context
                messages.push(json!({
                    "role": "system",
                    "content": "API assistant. Use only documented endpoints."
                }));
            }
        } else {
            // Add basic system message if no context
            messages.push(json!({
                "role": "system",
                "content": "API assistant. Use only documented endpoints."
            }));
        }

        messages.push(json!({
            "role": "user",
            "content": request.message
        }));

        // Prepare request body
        let mut body = json!({
            "model": self.config.features.model,
            "messages": messages
        });

        // Handle temperature based on model type
        if self.config.features.model.starts_with("gpt-5") {
            // GPT-5 models (nano, mini, standard) only support integer temperature of 1
            body["temperature"] = json!(1);
        } else {
            // Other models use floating point temperature
            body["temperature"] = json!(self.config.features.temperature as f64 / 10.0);
        }

        // Add token limits with special handling for GPT-5 models
        if self.config.features.model.starts_with("gpt-5") {
            // GPT-5 models need higher token limits to account for reasoning tokens
            let completion_tokens = if self.config.features.max_completion_tokens > 0 {
                std::cmp::max(self.config.features.max_completion_tokens, 2000)
            } else {
                2000
            };
            body["max_completion_tokens"] = json!(completion_tokens);
        } else {
            // Other models use configured limits
            if self.config.features.max_completion_tokens > 0 {
                body["max_completion_tokens"] = json!(self.config.features.max_completion_tokens);
            } else if self.config.features.max_tokens > 0 {
                body["max_tokens"] = json!(self.config.features.max_tokens);
            }
        }

        // Make API call to OpenAI
        let response = client
            .post("https://api.openai.com/v1/chat/completions")
            .header("Authorization", format!("Bearer {}", self.config.api_key))
            .header("Content-Type", "application/json")
            .json(&body)
            .send()
            .await
            .map_err(|e| anyhow::anyhow!("Failed to send request to OpenAI: {}", e))?;

        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Ok(ChatResponse {
                response: "".to_string(),
                provider: "openai".to_string(),
                model: self.config.features.model.clone(),
                tokens_used: 0,
                error: format!("OpenAI API error: {}", error_text),
            });
        }

        // Parse response
        let response_json: serde_json::Value = response
            .json()
            .await
            .map_err(|e| anyhow::anyhow!("Failed to parse OpenAI response: {}", e))?;

        // Extract content from response
        let choices = response_json
            .get("choices")
            .and_then(|c| c.as_array())
            .ok_or_else(|| anyhow::anyhow!("No choices in OpenAI response"))?;

        if choices.is_empty() {
            return Ok(ChatResponse {
                response: "".to_string(),
                provider: "openai".to_string(),
                model: self.config.features.model.clone(),
                tokens_used: 0,
                error: "No response choices returned from OpenAI".to_string(),
            });
        }

        let choice = &choices[0];
        let message = choice.get("message");
        let finish_reason = choice.get("finish_reason");


        let content = message
            .and_then(|m| m.get("content"))
            .and_then(|c| c.as_str())
            .unwrap_or("")
            .to_string();


        // Extract token usage if available (do this before early returns)
        let tokens_used = response_json
            .get("usage")
            .and_then(|u| u.get("total_tokens"))
            .and_then(|t| t.as_i64())
            .unwrap_or(0) as i32;

        // Extract model used
        let model_used = response_json
            .get("model")
            .and_then(|m| m.as_str())
            .unwrap_or(&self.config.features.model)
            .to_string();

        // Check if the response was cut off due to length limit or content filtering
        if let Some(reason) = finish_reason.and_then(|r| r.as_str()) {
            if reason == "length" {
                // For GPT-5 models with empty content due to length, provide a helpful message
                if content.is_empty() && self.config.features.model.starts_with("gpt-5") {
                    return Ok(ChatResponse {
                        response: "I understand your question, but my response was truncated due to token limits. Could you please ask a more specific question about the API?".to_string(),
                        provider: "openai".to_string(),
                        model: model_used,
                        tokens_used,
                        error: String::new(),
                    });
                }
            }
        }

        Ok(ChatResponse {
            response: content,
            provider: "openai".to_string(),
            model: model_used,
            tokens_used,
            error: String::new(),
        })
    }

    fn get_provider(&self) -> &str {
        "openai"
    }

    fn get_model(&self) -> &str {
        &self.config.features.model
    }
}

impl GeminiClient {
    pub fn new(config: &AIConfig) -> Result<Box<dyn Client>> {
        Ok(Box::new(Self {
            config: config.clone(),
        }))
    }
}

#[async_trait::async_trait]
impl Client for GeminiClient {
    async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
        // TODO: Implement actual Gemini API call
        Ok(ChatResponse {
            response: format!("Gemini response to: {}", request.message),
            provider: "gemini".to_string(),
            model: self.config.features.model.clone(),
            tokens_used: 100,
            error: String::new(),
        })
    }

    fn get_provider(&self) -> &str {
        "gemini"
    }

    fn get_model(&self) -> &str {
        &self.config.features.model
    }
}

impl OpenRouterClient {
    pub fn new(config: &AIConfig) -> Result<Box<dyn Client>> {
        if config.api_key.is_empty() {
            return Err(anyhow::anyhow!("OpenRouter API key is required"));
        }
        Ok(Box::new(Self {
            config: config.clone(),
        }))
    }

    fn build_system_prompt(&self, request: &ChatRequest) -> String {
        let mut prompt = String::from("You are a helpful AI assistant that provides information about APIs. ");

        if let Some(ref context) = request.context {
            if !context.is_empty() {
                prompt.push_str("Here's the API documentation context:\n\n");
                prompt.push_str(context);
                prompt.push_str("\n\nPlease help the user understand this API based on the provided documentation.");
            }
        }

        prompt
    }
}

#[async_trait::async_trait]
impl Client for OpenRouterClient {
    async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
        let client = reqwest::Client::new();

        // Build messages for OpenAI-compatible API
        let system_prompt = self.build_system_prompt(&request);
        let messages = vec![
            json!({
                "role": "system",
                "content": system_prompt
            }),
            json!({
                "role": "user",
                "content": request.message
            })
        ];

        // Prepare request body
        let mut body = json!({
            "model": self.config.features.model,
            "messages": messages
        });

        // Add optional parameters
        if self.config.features.max_tokens > 0 {
            body["max_tokens"] = json!(self.config.features.max_tokens);
        }
        if self.config.features.max_completion_tokens > 0 {
            body["max_completion_tokens"] = json!(self.config.features.max_completion_tokens);
        }
        if self.config.features.temperature > 0 {
            body["temperature"] = json!(self.config.features.temperature);
        }

        // Make API call to OpenRouter
        let response = client
            .post("https://openrouter.ai/api/v1/chat/completions")
            .header("Authorization", format!("Bearer {}", self.config.api_key))
            .header("Content-Type", "application/json")
            .header("HTTP-Referer", "https://bytedocs.rs") // Required for OpenRouter
            .header("X-Title", "ByteDocs") // Required for OpenRouter
            .json(&body)
            .send()
            .await
            .map_err(|e| anyhow::anyhow!("Failed to send request to OpenRouter: {}", e))?;

        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Ok(ChatResponse {
                response: "".to_string(),
                provider: "openrouter".to_string(),
                model: self.config.features.model.clone(),
                tokens_used: 0,
                error: format!("OpenRouter API error: {}", error_text),
            });
        }

        // Parse response
        let response_json: serde_json::Value = response
            .json()
            .await
            .map_err(|e| anyhow::anyhow!("Failed to parse OpenRouter response: {}", e))?;

        // Extract content from response
        let choices = response_json
            .get("choices")
            .and_then(|c| c.as_array())
            .ok_or_else(|| anyhow::anyhow!("No choices in OpenRouter response"))?;

        if choices.is_empty() {
            return Ok(ChatResponse {
                response: "".to_string(),
                provider: "openrouter".to_string(),
                model: self.config.features.model.clone(),
                tokens_used: 0,
                error: "No response choices returned from OpenRouter".to_string(),
            });
        }

        let content = choices[0]
            .get("message")
            .and_then(|m| m.get("content"))
            .and_then(|c| c.as_str())
            .unwrap_or("")
            .to_string();

        // Extract token usage if available
        let tokens_used = response_json
            .get("usage")
            .and_then(|u| u.get("total_tokens"))
            .and_then(|t| t.as_i64())
            .unwrap_or(0) as i32;

        // Extract model used
        let model_used = response_json
            .get("model")
            .and_then(|m| m.as_str())
            .unwrap_or(&self.config.features.model)
            .to_string();

        Ok(ChatResponse {
            response: content,
            provider: "openrouter".to_string(),
            model: model_used,
            tokens_used,
            error: String::new(),
        })
    }

    fn get_provider(&self) -> &str {
        "openrouter"
    }

    fn get_model(&self) -> &str {
        &self.config.features.model
    }
}

// Register all client factories
pub fn init_client_factories() {
    register_client_factory("openai", |config| OpenAIClient::new(config));
    register_client_factory("gemini", |config| GeminiClient::new(config));
    register_client_factory("openrouter", |config| OpenRouterClient::new(config));
}