use crate::config::Config; use anyhow::{Context, Result, anyhow}; use futures_util::StreamExt; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum Role { System, User, Assistant, Tool, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct FunctionCall { pub name: String, /// JSON-encoded argument object as a string (OpenAI-compatible). pub arguments: String, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolCall { pub id: String, #[serde(rename = "type")] pub kind: String, pub function: FunctionCall, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ChatMessage { pub role: Role, #[serde(default, skip_serializing_if = "Option::is_none")] pub content: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_calls: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_call_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub name: Option, } impl ChatMessage { pub fn system(content: impl Into) -> Self { Self { role: Role::System, content: Some(content.into()), tool_calls: None, tool_call_id: None, name: None, } } pub fn user(content: impl Into) -> Self { Self { role: Role::User, content: Some(content.into()), tool_calls: None, tool_call_id: None, name: None, } } pub fn assistant(content: impl Into) -> Self { Self { role: Role::Assistant, content: Some(content.into()), tool_calls: None, tool_call_id: None, name: None, } } pub fn assistant_tool_calls(content: Option, tool_calls: Vec) -> Self { Self { role: Role::Assistant, content, tool_calls: Some(tool_calls), tool_call_id: None, name: None, } } pub fn tool(tool_call_id: impl Into, content: impl Into) -> Self { Self { role: Role::Tool, content: Some(content.into()), tool_calls: None, tool_call_id: Some(tool_call_id.into()), name: None, } } pub fn text(&self) -> &str { self.content.as_deref().unwrap_or("") } } #[derive(Debug, Deserialize)] struct StreamChunk { choices: Option>, error: Option, } #[derive(Debug, Deserialize)] struct StreamChoice { delta: Option, } #[derive(Debug, Deserialize)] struct Delta { content: Option, } #[derive(Debug, Deserialize)] struct ApiErrorBody { message: Option, } #[derive(Debug, Deserialize)] struct CompletionResponse { choices: Option>, error: Option, } #[derive(Debug, Deserialize)] struct CompletionChoice { message: Option, finish_reason: Option, } /// Stream a chat completion; invoke `on_delta` for each text piece. /// Returns the full assistant text. (No tools — used by interactive `chat`.) pub async fn stream_chat( config: &Config, messages: &[ChatMessage], mut on_delta: impl FnMut(&str), ) -> Result { let client = reqwest::Client::new(); let url = format!("{}/chat/completions", config.base_url); let body = json!({ "model": config.model, "messages": messages, "stream": true, }); let response = client .post(&url) .bearer_auth(&config.api_key) .header("content-type", "application/json") .json(&body) .send() .await .context("request to chat completions failed")?; if !response.status().is_success() { let status = response.status(); let text = response.text().await.unwrap_or_default(); return Err(anyhow!("chat completions HTTP {status}: {text}")); } let mut full = String::new(); let mut stream = response.bytes_stream(); let mut buffer = String::new(); while let Some(item) = stream.next().await { let chunk = item.context("reading SSE stream")?; buffer.push_str(&String::from_utf8_lossy(&chunk)); while let Some(pos) = buffer.find('\n') { let mut line = buffer[..pos].to_string(); buffer.drain(..=pos); if line.ends_with('\r') { line.pop(); } let line = line.trim(); if line.is_empty() { continue; } if !line.starts_with("data:") { continue; } let data = line[5..].trim(); if data == "[DONE]" { return Ok(full); } let parsed: StreamChunk = match serde_json::from_str(data) { Ok(v) => v, Err(_) => continue, }; if let Some(err) = parsed.error { return Err(anyhow!( "provider error: {}", err.message.unwrap_or_else(|| data.to_string()) )); } if let Some(choices) = parsed.choices { for choice in choices { if let Some(delta) = choice.delta.and_then(|d| d.content) { if !delta.is_empty() { on_delta(&delta); full.push_str(&delta); } } } } } } Ok(full) } /// Non-streaming chat completion, optionally with tool definitions. /// Returns the assistant message (may include `tool_calls`). pub async fn chat_completion( config: &Config, messages: &[ChatMessage], tools: Option<&Value>, ) -> Result { let client = reqwest::Client::new(); let url = format!("{}/chat/completions", config.base_url); let mut body = json!({ "model": config.model, "messages": messages, "stream": false, }); if let Some(tools) = tools { body["tools"] = tools.clone(); body["tool_choice"] = json!("auto"); } let response = client .post(&url) .bearer_auth(&config.api_key) .header("content-type", "application/json") .json(&body) .send() .await .context("request to chat completions failed")?; if !response.status().is_success() { let status = response.status(); let text = response.text().await.unwrap_or_default(); return Err(anyhow!("chat completions HTTP {status}: {text}")); } let parsed: CompletionResponse = response .json() .await .context("parsing chat completions JSON")?; if let Some(err) = parsed.error { return Err(anyhow!( "provider error: {}", err.message.unwrap_or_else(|| "unknown".into()) )); } let choice = parsed .choices .and_then(|mut c| c.pop()) .ok_or_else(|| anyhow!("no choices in completion response"))?; let mut message = choice .message .ok_or_else(|| anyhow!("empty message in completion choice"))?; // Normalize: ensure assistant role even if provider omits it. if message.role != Role::Assistant && message.role != Role::Tool { message.role = Role::Assistant; } let _ = choice.finish_reason; Ok(message) } #[cfg(test)] mod tests { use super::*; #[test] fn messages_serialize_roles() { let msg = ChatMessage::user("hi"); let v = serde_json::to_value(&msg).unwrap(); assert_eq!(v["role"], "user"); assert_eq!(v["content"], "hi"); assert!(v.get("tool_calls").is_none()); } #[test] fn tool_message_serializes() { let msg = ChatMessage::tool("call_1", r#"{"ok":true}"#); let v = serde_json::to_value(&msg).unwrap(); assert_eq!(v["role"], "tool"); assert_eq!(v["tool_call_id"], "call_1"); assert_eq!(v["content"], r#"{"ok":true}"#); } #[test] fn assistant_tool_calls_serialize() { let msg = ChatMessage::assistant_tool_calls( None, vec![ToolCall { id: "c1".into(), kind: "function".into(), function: FunctionCall { name: "list_dir".into(), arguments: r#"{"path":"."}"#.into(), }, }], ); let v = serde_json::to_value(&msg).unwrap(); assert_eq!(v["role"], "assistant"); assert_eq!(v["tool_calls"][0]["function"]["name"], "list_dir"); assert!(v.get("content").is_none() || v["content"].is_null()); } }