model-gateway-rs 0.2.6

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use toolcraft_request::{ByteStream, HeaderMap, Request, response::Response};

use crate::{
    error::{Error, Result},
    llm::Llm,
    model::{
        llm::{
            ChatContentPart, ChatImageUrl, ChatMessage, ChatMessageContent, LlmInput, LlmOutput,
        },
        role::Role,
    },
};

const OLLAMA_CHAT_ENDPOINT: &str = "api/chat";

/// Native Ollama chat client.
///
/// This client uses Ollama's `/api/chat` endpoint instead of the
/// OpenAI-compatible `/v1/chat/completions` endpoint. Use this when you need
/// Ollama-native request fields such as `think: false`.
pub struct OllamaLlm {
    request: Request,
    model: String,
    chat_endpoint: String,
    think: Option<OllamaThink>,
    options: Map<String, Value>,
    extra_body: Map<String, Value>,
}

impl OllamaLlm {
    pub fn new(base_url: &str, model: &str) -> Result<Self> {
        let mut request = Request::new()?;
        request.set_base_url(base_url)?;

        let mut headers = HeaderMap::new();
        headers.insert("Content-Type", "application/json".to_string())?;
        request.set_default_headers(headers);

        Ok(Self {
            request,
            model: model.to_string(),
            chat_endpoint: default_ollama_chat_endpoint(base_url).to_string(),
            think: None,
            options: Map::new(),
            extra_body: Map::new(),
        })
    }

    pub fn without_thinking(mut self) -> Self {
        self.think = Some(OllamaThink::disabled());
        self
    }

    pub fn with_think(mut self, think: Option<OllamaThink>) -> Self {
        self.think = think;
        self
    }

    pub fn with_temperature(mut self, temperature: Option<f32>) -> Result<Self> {
        self.set_option("temperature", temperature)?;
        Ok(self)
    }

    pub fn with_max_tokens(mut self, max_tokens: Option<u32>) -> Result<Self> {
        self.set_option("num_predict", max_tokens)?;
        Ok(self)
    }

    pub fn with_option(mut self, key: impl Into<String>, value: impl Serialize) -> Result<Self> {
        self.options
            .insert(key.into(), serde_json::to_value(value)?);
        Ok(self)
    }

    pub fn with_extra_body_param(
        mut self,
        key: impl Into<String>,
        value: impl Serialize,
    ) -> Result<Self> {
        self.extra_body
            .insert(key.into(), serde_json::to_value(value)?);
        Ok(self)
    }

    pub fn with_chat_endpoint(mut self, endpoint: impl Into<String>) -> Self {
        self.chat_endpoint = endpoint.into().trim_start_matches('/').to_string();
        self
    }

    fn set_option(&mut self, key: impl Into<String>, value: impl Serialize) -> Result<()> {
        let key = key.into();
        let value = serde_json::to_value(value)?;

        if value.is_null() {
            self.options.remove(&key);
        } else {
            self.options.insert(key, value);
        }

        Ok(())
    }
}

#[async_trait]
impl Llm for OllamaLlm {
    async fn chat_once(&self, input: LlmInput) -> Result<LlmOutput> {
        let body = OllamaChatRequest {
            model: self.model.clone(),
            messages: input
                .messages
                .into_iter()
                .map(OllamaMessage::from)
                .collect(),
            stream: Some(false),
            think: self.think.clone(),
            options: non_empty_map(self.options.clone()),
            extra_body: self.extra_body.clone(),
        };
        let payload = serde_json::to_value(body)?;
        let response = self
            .request
            .post(&self.chat_endpoint, &payload, None)
            .await?;
        parse_chat_response(response).await
    }

    async fn chat_stream(&self, input: LlmInput) -> Result<ByteStream> {
        let body = OllamaChatRequest {
            model: self.model.clone(),
            messages: input
                .messages
                .into_iter()
                .map(OllamaMessage::from)
                .collect(),
            stream: Some(true),
            think: self.think.clone(),
            options: non_empty_map(self.options.clone()),
            extra_body: self.extra_body.clone(),
        };
        let payload = serde_json::to_value(body)?;
        self.request
            .post_stream(&self.chat_endpoint, &payload, None)
            .await
            .map_err(Into::into)
    }
}

async fn parse_chat_response(response: Response) -> Result<LlmOutput> {
    let status = response.status();
    let body = response.text().await?;

    if !status.is_success() {
        return Err(Error::ApiError(format_error_body(status.as_u16(), &body)));
    }

    if let Ok(error) = serde_json::from_str::<OllamaErrorResponse>(&body) {
        return Err(Error::ApiError(error.error));
    }

    let json: OllamaChatResponse = serde_json::from_str(&body)?;
    Ok(json.into())
}

fn format_error_body(status: u16, body: &str) -> String {
    match serde_json::from_str::<OllamaErrorResponse>(body) {
        Ok(error) => format!("status={status}, message={}", error.error),
        Err(_) => format!("status={status}, body={body}"),
    }
}

fn default_ollama_chat_endpoint(base_url: &str) -> &'static str {
    let base_url = base_url.trim_end_matches('/');

    if base_url.ends_with("/api") {
        "chat"
    } else {
        OLLAMA_CHAT_ENDPOINT
    }
}

fn non_empty_map(map: Map<String, Value>) -> Option<Map<String, Value>> {
    if map.is_empty() { None } else { Some(map) }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum OllamaThink {
    Bool(bool),
    Level(String),
}

impl OllamaThink {
    pub fn enabled() -> Self {
        Self::Bool(true)
    }

    pub fn disabled() -> Self {
        Self::Bool(false)
    }

    pub fn level(level: impl Into<String>) -> Self {
        Self::Level(level.into())
    }
}

impl From<bool> for OllamaThink {
    fn from(value: bool) -> Self {
        Self::Bool(value)
    }
}

impl From<&str> for OllamaThink {
    fn from(value: &str) -> Self {
        Self::Level(value.to_string())
    }
}

#[derive(Debug, Clone, Serialize)]
struct OllamaChatRequest {
    model: String,
    messages: Vec<OllamaMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    think: Option<OllamaThink>,
    #[serde(skip_serializing_if = "Option::is_none")]
    options: Option<Map<String, Value>>,
    #[serde(flatten)]
    extra_body: Map<String, Value>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct OllamaMessage {
    #[serde(default = "assistant_role")]
    role: Role,
    content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    images: Option<Vec<String>>,
}

impl From<ChatMessage> for OllamaMessage {
    fn from(message: ChatMessage) -> Self {
        let (content, images) = match message.content {
            ChatMessageContent::Text(text) => (text, None),
            ChatMessageContent::Parts(parts) => {
                let mut text_parts = Vec::new();
                let mut images = Vec::new();

                for part in parts {
                    match part {
                        ChatContentPart::Text { text } => text_parts.push(text),
                        ChatContentPart::ImageUrl { image_url } => {
                            images.push(ollama_image_value(image_url));
                        }
                    }
                }

                let images = if images.is_empty() {
                    None
                } else {
                    Some(images)
                };

                (text_parts.join("\n"), images)
            }
        };

        Self {
            role: message.role,
            content,
            images,
        }
    }
}

fn ollama_image_value(image_url: ChatImageUrl) -> String {
    image_url
        .url
        .strip_prefix("data:")
        .and_then(|data| data.split_once(',').map(|(_, base64)| base64.to_string()))
        .unwrap_or(image_url.url)
}

fn assistant_role() -> Role {
    Role::Assistant
}

#[derive(Debug, Deserialize)]
struct OllamaChatResponse {
    message: Option<OllamaMessage>,
    prompt_eval_count: Option<u32>,
    eval_count: Option<u32>,
}

#[derive(Debug, Deserialize)]
struct OllamaErrorResponse {
    error: String,
}

impl From<OllamaChatResponse> for LlmOutput {
    fn from(response: OllamaChatResponse) -> Self {
        let message = response.message.map(|message| ChatMessage {
            role: message.role,
            content: message.content.into(),
        });

        let usage = match (response.prompt_eval_count, response.eval_count) {
            (Some(prompt), Some(eval)) => Some(prompt + eval),
            (Some(prompt), None) => Some(prompt),
            (None, Some(eval)) => Some(eval),
            (None, None) => None,
        };

        LlmOutput { message, usage }
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;
    use crate::model::llm::{ChatContentPart, ChatMessage};

    #[test]
    fn root_base_url_uses_ollama_api_chat_path() {
        assert_eq!(
            default_ollama_chat_endpoint("http://127.0.0.1:11434"),
            "api/chat"
        );
    }

    #[test]
    fn api_base_url_uses_chat_path() {
        assert_eq!(
            default_ollama_chat_endpoint("http://127.0.0.1:11434/api"),
            "chat"
        );
    }

    #[test]
    fn chat_request_serializes_think_false() {
        let request = OllamaChatRequest {
            model: "gpt-oss:20b".to_string(),
            messages: vec![OllamaMessage {
                role: Role::User,
                content: "hi".to_string(),
                images: None,
            }],
            stream: Some(false),
            think: Some(OllamaThink::disabled()),
            options: None,
            extra_body: Map::new(),
        };

        let value = serde_json::to_value(request).unwrap();

        assert_eq!(value["think"], false);
        assert!(value.get("reasoning_effort").is_none());
    }

    #[test]
    fn chat_request_serializes_options() {
        let mut options = Map::new();
        options.insert("temperature".to_string(), json!(0.7));
        options.insert("num_predict".to_string(), json!(100));

        let request = OllamaChatRequest {
            model: "gemma4:26b".to_string(),
            messages: vec![OllamaMessage {
                role: Role::User,
                content: "hi".to_string(),
                images: None,
            }],
            stream: Some(false),
            think: None,
            options: Some(options),
            extra_body: Map::new(),
        };

        let value = serde_json::to_value(request).unwrap();

        assert_eq!(value["options"]["temperature"], 0.7);
        assert_eq!(value["options"]["num_predict"], 100);
    }

    #[test]
    fn openai_content_parts_convert_to_ollama_text_content() {
        let message = ChatMessage::user_with_parts(vec![
            ChatContentPart::text("line one"),
            ChatContentPart::image_url("data:image/png;base64,aW1hZ2U="),
            ChatContentPart::text("line two"),
        ]);

        let message = OllamaMessage::from(message);

        assert_eq!(message.content, "line one\nline two");
        assert_eq!(message.images, Some(vec!["aW1hZ2U=".to_string()]));
    }
}