llmkit-ollama 0.1.0

//! Mapping between llmkit types and the Ollama `/api/chat` wire format.

use llmkit_core::{
    ChatRequest, ChatResponse, FinishReason, LlmError, LlmResult, Message, MessageContent, Role,
    ToolCall,
};

use crate::types::*;

pub(crate) fn build_request(req: &ChatRequest, model: String, stream: bool) -> ChatRequestBody {
    let mut messages = Vec::with_capacity(req.messages.len() + 1);
    if let Some(system) = &req.system {
        messages.push(WireMessage { role: "system".into(), content: system.clone(), tool_calls: None });
    }
    for m in &req.messages {
        messages.push(map_message(m));
    }

    let tools = req.tools.as_ref().map(|ts| {
        ts.iter()
            .map(|t| WireTool {
                kind: "function",
                function: WireFunction {
                    name: t.name.clone(),
                    description: t.description.clone(),
                    parameters: t.input_schema.clone(),
                },
            })
            .collect()
    });

    let options = Options {
        temperature: req.temperature,
        num_predict: req.max_tokens,
        stop: req.stop.clone(),
    };

    ChatRequestBody {
        model,
        messages,
        stream,
        tools,
        options: (!options.is_empty()).then_some(options),
    }
}

fn map_message(m: &Message) -> WireMessage {
    match &m.content {
        MessageContent::ToolResult { content, .. } => {
            WireMessage { role: "tool".into(), content: content.clone(), tool_calls: None }
        }
        MessageContent::ToolUse { name, input, .. } => WireMessage {
            role: "assistant".into(),
            content: String::new(),
            tool_calls: Some(vec![WireToolCall {
                function: WireToolCallFunction { name: name.clone(), arguments: input.clone() },
            }]),
        },
        other => WireMessage {
            role: role_str(m.role).into(),
            content: other.as_text().unwrap_or_default(),
            tool_calls: None,
        },
    }
}

fn role_str(role: Role) -> &'static str {
    match role {
        Role::User => "user",
        Role::Assistant => "assistant",
        Role::System => "system",
        Role::Tool => "tool",
    }
}

pub(crate) fn map_response(resp: ChatResponseBody, latency_ms: u64) -> LlmResult<ChatResponse> {
    let message = resp
        .message
        .ok_or_else(|| LlmError::Provider { status: 200, message: "no message in response".into() })?;

    let mut tool_calls = Vec::new();
    if let Some(calls) = message.tool_calls {
        for (i, c) in calls.into_iter().enumerate() {
            // Ollama does not assign tool-call ids; synthesise stable ones.
            tool_calls.push(ToolCall::new(
                format!("call_{i}"),
                c.function.name,
                c.function.arguments,
            ));
        }
    }

    let finish_reason = if !tool_calls.is_empty() {
        FinishReason::ToolUse
    } else {
        match resp.done_reason.as_deref() {
            Some("length") => FinishReason::MaxTokens,
            Some("stop") | None => FinishReason::Stop,
            Some(other) => FinishReason::Other(other.to_string()),
        }
    };

    let usage = llmkit_core::TokenUsage::new(
        resp.prompt_eval_count.unwrap_or(0),
        resp.eval_count.unwrap_or(0),
    );

    Ok(ChatResponse {
        id: String::new(),
        provider: "ollama".into(),
        model: resp.model,
        message: Message::assistant(message.content),
        finish_reason,
        tool_calls,
        usage,
        cost: None,
        latency_ms,
    })
}