lakehouse/crates/aibridge/src/providers/openrouter.rs

use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::time::Duration;

use crate::client::{GenerateRequest, GenerateResponse, EmbedRequest, EmbedResponse};
use crate::provider::ProviderAdapter;

pub struct OpenRouterAdapter {
    client: Client,
    base_url: String,
    api_key: String,
    default_model: String,
}

#[derive(Serialize)]
struct OpenRouterChatRequest {
    model: String,
    messages: Vec<OpenRouterMessage>,
    temperature: Option<f64>,
    max_tokens: Option<u32>,
}

#[derive(Serialize)]
struct OpenRouterMessage {
    role: String,
    content: String,
}

#[derive(Deserialize)]
struct OpenRouterChatResponse {
    choices: Vec<OpenRouterChoice>,
    usage: OpenRouterUsage,
    model: String,
}

#[derive(Deserialize)]
struct OpenRouterChoice {
    message: OpenRouterMessageOut,
}

#[derive(Deserialize)]
struct OpenRouterMessageOut {
    role: String,
    content: String,
}

#[derive(Deserialize)]
struct OpenRouterUsage {
    prompt_tokens: Option<u32>,
    completion_tokens: Option<u32>,
    total_tokens: Option<u32>,
}

impl OpenRouterAdapter {
    pub fn new(base_url: &str, api_key: String, default_model: &str) -> Self {
        let client = Client::builder()
            .timeout(Duration::from_secs(180))
            .build()
            .expect("failed to build HTTP client");
        Self {
            client,
            base_url: base_url.trim_end_matches('/').to_string(),
            api_key,
            default_model: default_model.to_string(),
        }
    }

    fn chat_model(&self, model: &str) -> String {
        // Strip "openrouter/" prefix if present
        let m = model.trim_start_matches("openrouter/");
        if m.is_empty() || m == model {
            self.default_model.clone()
        } else {
            m.to_string()
        }
    }

    fn to_openrouter_messages(req: &GenerateRequest) -> Vec<OpenRouterMessage> {
        let mut out = vec![];
        if let Some(sys) = &req.system {
            out.push(OpenRouterMessage { role: "system".into(), content: sys.clone() });
        }
        out.push(OpenRouterMessage {
            role: "user".into(),
            content: req.prompt.clone(),
        });
        out
    }
}

#[async_trait]
impl ProviderAdapter for OpenRouterAdapter {
    fn name(&self) -> &str {
        "openrouter"
    }

    async fn chat(&self, req: GenerateRequest) -> Result<GenerateResponse, String> {
        let model = self.chat_model(req.model.as_deref().unwrap_or(""));
        let or_req = OpenRouterChatRequest {
            model: model.clone(),
            messages: OpenRouterAdapter::to_openrouter_messages(&req),
            temperature: req.temperature,
            max_tokens: req.max_tokens,
        };

        let resp = self.client
            .post(format!("{}/chat/completions", self.base_url))
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&or_req)
            .send()
            .await
            .map_err(|e| format!("openrouter request failed: {e}"))?;

        let status = resp.status();
        let body = resp.text().await.unwrap_or_default();

        if !status.is_success() {
            return Err(format!("openrouter error ({}): {}", status, body));
        }

        let or_resp: OpenRouterChatResponse = serde_json::from_str(&body)
            .map_err(|e| format!("openrouter parse error: {e}"))?;

        let choice = or_resp.choices.into_iter().next()
            .ok_or("no completion choice returned")?;
        let usage = or_resp.usage;

        Ok(GenerateResponse {
            text: choice.message.content,
            model: or_resp.model,
            tokens_evaluated: usage.prompt_tokens.map(|n| n as u64),
            tokens_generated: usage.completion_tokens.map(|n| n as u64),
        })
    }

    async fn embed(&self, _req: EmbedRequest) -> Result<EmbedResponse, String> {
        Err("openrouter: embed not implemented".into())
    }

    async fn health(&self) -> Result<serde_json::Value, String> {
        // OpenRouter doesn't have a dedicated health endpoint,
        // so we just return a healthy response if the client works.
        Ok(serde_json::json!({
            "status": "ok",
            "provider": "openrouter",
        }))
    }
}