openrouter-rust 0.1.0

use crate::{
    client::OpenRouterClient,
    error::{OpenRouterError, Result},
    types::{Message, Plugin, ProviderPreferences, ResponseFormat, Tool, ToolChoice, Usage},
};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionRequest {
    pub model: String,
    pub messages: Vec<Message>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<Tool>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<ToolChoice>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ResponseFormat>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub seed: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_k: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub repetition_penalty: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub provider: Option<ProviderPreferences>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub plugins: Option<Vec<Plugin>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub transforms: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub models: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub route: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionResponse {
    pub id: String,
    pub object: String,
    pub created: i64,
    pub model: String,
    pub choices: Vec<Choice>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage: Option<Usage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub system_fingerprint: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Choice {
    pub index: u32,
    pub message: Message,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub finish_reason: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub native_finish_reason: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<ChoiceError>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChoiceError {
    pub code: u16,
    pub message: String,
}

pub struct ChatCompletionBuilder {
    request: ChatCompletionRequest,
}

impl ChatCompletionBuilder {
    pub fn new(model: impl Into<String>) -> Self {
        Self {
            request: ChatCompletionRequest {
                model: model.into(),
                messages: Vec::new(),
                temperature: None,
                top_p: None,
                max_tokens: None,
                stop: None,
                stream: None,
                tools: None,
                tool_choice: None,
                response_format: None,
                seed: None,
                top_k: None,
                frequency_penalty: None,
                presence_penalty: None,
                repetition_penalty: None,
                provider: None,
                plugins: None,
                transforms: None,
                models: None,
                route: None,
            },
        }
    }

    pub fn message(mut self, role: crate::types::Role, content: impl Into<String>) -> Self {
        self.request.messages.push(Message {
            role,
            content: Some(content.into()),
            name: None,
            tool_calls: None,
        });
        self
    }

    pub fn system_message(self, content: impl Into<String>) -> Self {
        self.message(crate::types::Role::System, content)
    }

    pub fn user_message(self, content: impl Into<String>) -> Self {
        self.message(crate::types::Role::User, content)
    }

    pub fn assistant_message(self, content: impl Into<String>) -> Self {
        self.message(crate::types::Role::Assistant, content)
    }

    pub fn temperature(mut self, temp: f32) -> Self {
        self.request.temperature = Some(temp);
        self
    }

    pub fn top_p(mut self, top_p: f32) -> Self {
        self.request.top_p = Some(top_p);
        self
    }

    pub fn max_tokens(mut self, max: u32) -> Self {
        self.request.max_tokens = Some(max);
        self
    }

    pub fn stop(mut self, stop: Vec<String>) -> Self {
        self.request.stop = Some(stop);
        self
    }

    pub fn stream(mut self, stream: bool) -> Self {
        self.request.stream = Some(stream);
        self
    }

    pub fn tools(mut self, tools: Vec<Tool>) -> Self {
        self.request.tools = Some(tools);
        self
    }

    pub fn response_format_json(mut self) -> Self {
        self.request.response_format = Some(ResponseFormat {
            response_type: "json_object".to_string(),
            json_schema: None,
        });
        self
    }

    pub fn build(self) -> ChatCompletionRequest {
        self.request
    }
}

impl OpenRouterClient {
    pub async fn chat_completion(
        &self,
        request: ChatCompletionRequest,
    ) -> Result<ChatCompletionResponse> {
        let url = format!("{}/chat/completions", self.base_url);
        let headers = self.build_headers()?;

        let response = self
            .client
            .post(&url)
            .headers(headers)
            .json(&request)
            .send()
            .await
            .map_err(OpenRouterError::HttpError)?;

        let status = response.status();
        
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(OpenRouterError::ApiError {
                code: status.as_u16(),
                message: error_text,
            });
        }

        let completion = response
            .json::<ChatCompletionResponse>()
            .await
            .map_err(OpenRouterError::HttpError)?;

        Ok(completion)
    }
}