LazyBoy2/crates/grokboy-core/src/model.rs

use crate::config::Config;
use anyhow::{anyhow, Context, Result};
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};

#[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<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

impl ChatMessage {
    pub fn system(content: impl Into<String>) -> Self {
        Self {
            role: Role::System,
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: None,
            name: None,
        }
    }

    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: Role::User,
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: None,
            name: None,
        }
    }

    pub fn assistant(content: impl Into<String>) -> 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<String>, tool_calls: Vec<ToolCall>) -> 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<String>, content: impl Into<String>) -> 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<Vec<StreamChoice>>,
    error: Option<ApiErrorBody>,
}

#[derive(Debug, Deserialize)]
struct StreamChoice {
    delta: Option<Delta>,
}

#[derive(Debug, Deserialize)]
struct Delta {
    content: Option<String>,
}

#[derive(Debug, Deserialize)]
struct ApiErrorBody {
    message: Option<String>,
}

#[derive(Debug, Deserialize)]
struct CompletionResponse {
    choices: Option<Vec<CompletionChoice>>,
    error: Option<ApiErrorBody>,
}

#[derive(Debug, Deserialize)]
struct CompletionChoice {
    message: Option<ChatMessage>,
    finish_reason: Option<String>,
}

/// 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<String> {
    let client = reqwest::Client::builder()
        .connect_timeout(std::time::Duration::from_secs(15))
        .timeout(std::time::Duration::from_secs(180))
        .build()?;
    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<ChatMessage> {
    let client = reqwest::Client::builder()
        .connect_timeout(std::time::Duration::from_secs(15))
        .timeout(std::time::Duration::from_secs(180))
        .build()?;
    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"))?;

    validate_completion(choice)
}

fn validate_completion(choice: CompletionChoice) -> Result<ChatMessage> {
    let message = choice
        .message
        .ok_or_else(|| anyhow!("empty message in completion choice"))?;
    if message.role != Role::Assistant {
        return Err(anyhow!("completion message must have assistant role"));
    }
    let calls = message.tool_calls.as_deref().unwrap_or_default();
    match choice.finish_reason.as_deref() {
        Some("stop") if calls.is_empty() => {}
        Some("tool_calls") if !calls.is_empty() => {}
        reason => {
            return Err(anyhow!(
                "incomplete or invalid completion finish_reason: {reason:?}; no tools executed"
            ))
        }
    }
    let mut ids = std::collections::HashSet::new();
    for call in calls {
        if call.id.trim().is_empty()
            || !ids.insert(&call.id)
            || call.kind != "function"
            || call.function.name.trim().is_empty()
        {
            return Err(anyhow!("invalid or duplicate tool call; no tools executed"));
        }
    }
    Ok(message)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn truncated_or_filtered_responses_cannot_finish_or_execute_tools() {
        for reason in ["length", "content_filter", "unknown"] {
            assert!(validate_completion(CompletionChoice {
                message: Some(ChatMessage::assistant("partial answer")),
                finish_reason: Some(reason.into()),
            })
            .is_err());
        }
        assert!(validate_completion(CompletionChoice {
            message: Some(ChatMessage::assistant("answer")),
            finish_reason: None,
        })
        .is_err());
    }

    #[test]
    fn validates_finish_reason_against_tool_calls() {
        let call = ToolCall {
            id: "x".into(),
            kind: "function".into(),
            function: FunctionCall {
                name: "shell".into(),
                arguments: "{}".into(),
            },
        };
        for (reason, calls, valid) in [
            ("stop", vec![call.clone()], false),
            ("tool_calls", vec![], false),
            ("length", vec![call.clone()], false),
            ("tool_calls", vec![call.clone(), call.clone()], false),
            ("tool_calls", vec![call], true),
        ] {
            assert_eq!(
                validate_completion(CompletionChoice {
                    message: Some(ChatMessage::assistant_tool_calls(None, calls)),
                    finish_reason: Some(reason.into()),
                })
                .is_ok(),
                valid
            );
        }
    }

    #[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());
    }
}