clawmates/crates/tc-llm/src/openai_compat.rs

//! OpenAI-compatible chat/completions provider — the air-gapped inference
//! path (vLLM, Ollama, llama.cpp all speak this protocol).

use async_stream::try_stream;
use eventsource_stream::Eventsource;
use futures::StreamExt;
use serde_json::{json, Value};

use crate::provider::{
    ChatRequest, ChatRole, ContentPart, EventStream, LlmError, LlmEvent, LlmProvider, StopReason,
};

pub struct OpenAiCompatProvider {
    client: reqwest::Client,
    base_url: String,
    api_key: Option<String>,
}

impl OpenAiCompatProvider {
    /// `base_url` includes the version prefix, e.g. `http://local-llm:8000/v1`.
    pub fn new(base_url: String, api_key: Option<String>) -> OpenAiCompatProvider {
        OpenAiCompatProvider {
            client: reqwest::Client::new(),
            base_url,
            api_key,
        }
    }

    /// Maps provider-neutral messages to the OpenAI wire shape: tool calls
    /// ride on assistant messages, tool results become `role: "tool"`.
    fn wire_messages(request: &ChatRequest) -> Vec<Value> {
        let mut wire = vec![json!({"role": "system", "content": request.system})];
        for message in &request.messages {
            let mut text = String::new();
            let mut tool_calls: Vec<Value> = Vec::new();
            for part in &message.parts {
                match part {
                    ContentPart::Text { text: t } => text.push_str(t),
                    ContentPart::ToolUse { id, name, input } => tool_calls.push(json!({
                        "id": id,
                        "type": "function",
                        "function": {"name": name, "arguments": input.to_string()},
                    })),
                    ContentPart::ToolResult {
                        tool_use_id,
                        content,
                    } => wire.push(json!({
                        "role": "tool",
                        "tool_call_id": tool_use_id,
                        "content": content.to_string(),
                    })),
                }
            }
            if !text.is_empty() || !tool_calls.is_empty() {
                let role = match message.role {
                    ChatRole::User => "user",
                    ChatRole::Assistant => "assistant",
                };
                let mut entry = json!({"role": role, "content": text});
                if !tool_calls.is_empty() {
                    entry["tool_calls"] = Value::Array(tool_calls);
                }
                wire.push(entry);
            }
        }
        wire
    }
}

fn stop_reason(wire: &str) -> StopReason {
    match wire {
        "tool_calls" => StopReason::ToolUse,
        "length" => StopReason::MaxTokens,
        _ => StopReason::EndTurn,
    }
}

#[async_trait::async_trait]
impl LlmProvider for OpenAiCompatProvider {
    async fn stream(&self, request: ChatRequest) -> Result<EventStream, LlmError> {
        let tools: Vec<Value> = request
            .tools
            .iter()
            .map(|t| {
                json!({
                    "type": "function",
                    "function": {
                        "name": t.name,
                        "description": t.description,
                        "parameters": t.input_schema,
                    },
                })
            })
            .collect();
        let mut body = json!({
            "model": request.model,
            "max_tokens": request.max_tokens,
            "messages": OpenAiCompatProvider::wire_messages(&request),
            "stream": true,
            "stream_options": {"include_usage": true},
        });
        if !tools.is_empty() {
            body["tools"] = Value::Array(tools);
        }

        let mut http = self
            .client
            .post(format!("{}/chat/completions", self.base_url))
            .json(&body);
        if let Some(key) = &self.api_key {
            http = http.bearer_auth(key);
        }
        let response = http
            .send()
            .await
            .map_err(|e| LlmError::Transport(e.to_string()))?;
        if !response.status().is_success() {
            let status = response.status();
            let detail = response.text().await.unwrap_or_default();
            return Err(LlmError::Api(format!("{status}: {detail}")));
        }

        let mut sse = response.bytes_stream().eventsource();
        let stream = try_stream! {
            // Tool-call arguments stream as partial JSON keyed by index.
            let mut pending: Vec<(String, String, String)> = Vec::new(); // id, name, args
            let mut finish: Option<StopReason> = None;
            while let Some(event) = sse.next().await {
                let event = event.map_err(|e| LlmError::Transport(e.to_string()))?;
                if event.data.trim() == "[DONE]" {
                    break;
                }
                let data: Value = serde_json::from_str(&event.data)
                    .map_err(|e| LlmError::Wire(format!("{e}: {}", event.data)))?;
                let choice = &data["choices"][0];
                let delta = &choice["delta"];
                if let Some(text) = delta["content"].as_str() {
                    if !text.is_empty() {
                        yield LlmEvent::TextDelta(text.to_owned());
                    }
                }
                if let Some(calls) = delta["tool_calls"].as_array() {
                    for call in calls {
                        let index = call["index"].as_u64().unwrap_or(0) as usize;
                        while pending.len() <= index {
                            pending.push((String::new(), String::new(), String::new()));
                        }
                        let slot = &mut pending[index];
                        if let Some(id) = call["id"].as_str() {
                            slot.0 = id.to_owned();
                        }
                        if let Some(name) = call["function"]["name"].as_str() {
                            slot.1.push_str(name);
                        }
                        if let Some(args) = call["function"]["arguments"].as_str() {
                            slot.2.push_str(args);
                        }
                    }
                }
                if let Some(reason) = choice["finish_reason"].as_str() {
                    finish = Some(stop_reason(reason));
                }
                if let Some(usage) = data["usage"].as_object() {
                    yield LlmEvent::Usage {
                        input_tokens: usage["prompt_tokens"].as_u64().unwrap_or(0) as u32,
                        output_tokens: usage["completion_tokens"].as_u64().unwrap_or(0) as u32,
                    };
                }
            }
            for (id, name, args) in pending.drain(..) {
                if name.is_empty() {
                    continue;
                }
                let input: Value = if args.trim().is_empty() {
                    json!({})
                } else {
                    serde_json::from_str(&args)
                        .map_err(|e| LlmError::Wire(format!("tool arguments: {e}")))?
                };
                yield LlmEvent::ToolUse { id, name, input };
            }
            yield LlmEvent::Stop(finish.unwrap_or(StopReason::EndTurn));
        };
        Ok(Box::pin(stream))
    }
}