procyon 0.1.2

//! The `/chat/completions` wire format: request, response and stream chunk shapes.
//!
//! One dialect covers OpenAI, DeepSeek, Groq, OpenRouter, Together, Cerebras, Fireworks, xAI,
//! Ollama and LM Studio. Two shape differences from the neutral vocabulary drive most of this: the
//! system prompt is a message rather than a top-level field, and a tool result is its own
//! `tool`-role message rather than a block inside the following user turn.
//!
//! Separated from the client and the stream reassembly because these shapes answer to the
//! provider's API, while `ContentPart` — what they translate to — is what the session log
//! persists and cannot change with a vendor's next revision.

use serde::{Deserialize, Serialize};

use crate::agent::{tool_input, ContentPart, Message, TokenUsage, ToolDefinition};

#[derive(Debug, Serialize)]
pub struct ChatRequest {
    model: String,
    max_tokens: u32,
    messages: Vec<WireMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<Vec<WireTool>>,
    stream: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream_options: Option<StreamOptions>,
}

impl ChatRequest {
    pub fn build(
        model: String,
        max_tokens: u32,
        history: &[Message],
        tools: Option<&[ToolDefinition]>,
        system: Option<&str>,
        stream: bool,
    ) -> Self {
        Self {
            model,
            max_tokens,
            messages: to_wire_messages(history, system),
            tools: tools.map(|t| t.iter().map(WireTool::from_definition).collect()),
            stream,
            // Without this the streaming path reports no usage at all, leaving the budget
            // estimator with no anchor to correct itself against.
            stream_options: stream.then_some(StreamOptions {
                include_usage: true,
            }),
        }
    }
}

#[derive(Debug, Serialize)]
struct StreamOptions {
    include_usage: bool,
}

#[derive(Debug, Serialize)]
pub struct WireMessage {
    pub role: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<WireToolCall>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

#[derive(Debug, Serialize)]
pub struct WireToolCall {
    pub id: String,
    #[serde(rename = "type")]
    kind: String,
    pub function: WireFunctionCall,
}

#[derive(Debug, Serialize)]
pub struct WireFunctionCall {
    pub name: String,
    pub arguments: String,
}

#[derive(Debug, Serialize)]
struct WireTool {
    #[serde(rename = "type")]
    kind: String,
    function: WireToolSchema,
}

impl WireTool {
    fn from_definition(definition: &ToolDefinition) -> Self {
        Self {
            kind: "function".to_string(),
            function: WireToolSchema {
                name: definition.name.clone(),
                description: definition.description.clone(),
                parameters: definition.input_schema.clone(),
            },
        }
    }
}

#[derive(Debug, Serialize)]
struct WireToolSchema {
    name: String,
    description: String,
    parameters: serde_json::Value,
}

/// Translates the neutral history into wire messages.
///
/// One neutral message may expand into several: a user turn carrying three tool results becomes
/// three `tool` messages, because a provider ties each result to the call it answers.
fn to_wire_messages(history: &[Message], system: Option<&str>) -> Vec<WireMessage> {
    let mut out = Vec::new();

    if let Some(system) = system {
        out.push(WireMessage {
            role: "system".to_string(),
            content: Some(system.to_string()),
            tool_calls: None,
            tool_call_id: None,
        });
    }

    for message in history {
        let mut text = String::new();
        let mut tool_calls = Vec::new();
        let mut tool_results = Vec::new();

        for part in &message.content {
            match part {
                ContentPart::Text { text: t } => text.push_str(t),
                ContentPart::ToolUse { id, name, input } => tool_calls.push(WireToolCall {
                    id: id.clone(),
                    kind: "function".to_string(),
                    function: WireFunctionCall {
                        name: name.clone(),
                        arguments: input.to_string(),
                    },
                }),
                ContentPart::ToolResult {
                    tool_use_id,
                    content,
                } => tool_results.push((tool_use_id.clone(), content.clone())),
            }
        }

        // Results first: they answer the tool calls of the message just emitted, and a provider
        // rejects a `tool` message that does not directly follow its call.
        for (tool_call_id, content) in tool_results {
            out.push(WireMessage {
                role: "tool".to_string(),
                content: Some(content),
                tool_calls: None,
                tool_call_id: Some(tool_call_id),
            });
        }

        if text.is_empty() && tool_calls.is_empty() {
            continue;
        }

        out.push(WireMessage {
            role: message.role.to_string(),
            // An assistant turn that only called tools has no text, and some providers reject an
            // empty string where they accept null.
            content: (!text.is_empty()).then_some(text),
            tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
            tool_call_id: None,
        });
    }

    out
}

#[derive(Debug, Deserialize)]
pub struct ChatResponse {
    pub choices: Vec<ResponseChoice>,
}

#[derive(Debug, Deserialize)]
pub struct ResponseChoice {
    pub message: ResponseMessage,
}

#[derive(Debug, Deserialize)]
pub struct ResponseMessage {
    pub content: Option<String>,
    pub tool_calls: Option<Vec<ResponseToolCall>>,
}

#[derive(Debug, Deserialize)]
pub struct ResponseToolCall {
    id: String,
    function: ResponseFunctionCall,
}

#[derive(Debug, Deserialize)]
struct ResponseFunctionCall {
    name: String,
    arguments: String,
}

/// Turns a completed (non-streamed) message into neutral content blocks.
pub fn blocks_from_message(message: ResponseMessage) -> Vec<ContentPart> {
    let mut blocks = Vec::new();

    if let Some(text) = message.content.filter(|t| !t.is_empty()) {
        blocks.push(ContentPart::Text { text });
    }

    for call in message.tool_calls.into_iter().flatten() {
        blocks.push(tool_block(
            call.id,
            call.function.name,
            &call.function.arguments,
        ));
    }

    blocks
}

/// A call, with its arguments resolved by the domain's rule for unparseable ones.
pub fn tool_block(id: String, name: String, arguments: &str) -> ContentPart {
    let input = tool_input(&name, arguments);
    ContentPart::ToolUse { id, name, input }
}

// `stop_reason` is carried for diagnostics rather than control flow, but reporting it in the
// Anthropic spelling keeps one vocabulary in the session log across providers.
pub fn to_neutral_stop_reason(reason: &str) -> String {
    match reason {
        "tool_calls" | "function_call" => "tool_use",
        "length" => "max_tokens",
        "stop" => "end_turn",
        other => other,
    }
    .to_string()
}

#[derive(Debug, Deserialize)]
pub struct StreamChunk {
    #[serde(default)]
    pub choices: Vec<StreamChoice>,
    #[serde(default)]
    pub usage: Option<WireUsage>,
    // A provider may fail after the 200 header is already sent and report it as a chunk. Local
    // servers do this routinely — a model that is not loaded, or that ran out of memory mid
    // response. Without this field the chunk parses as one carrying no choices and the failure is
    // indistinguishable from an empty answer.
    #[serde(default)]
    pub error: Option<WireStreamError>,
}

/// The error object's shape varies across this dialect, so only the message is relied on.
#[derive(Debug, Deserialize)]
#[serde(untagged)]
pub enum WireStreamError {
    Message { message: String },
    Bare(String),
    // Anything else still marks the response as failed, even if the reason cannot be quoted.
    Other(serde_json::Value),
}

impl std::fmt::Display for WireStreamError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Message { message } => write!(f, "{}", message),
            Self::Bare(message) => write!(f, "{}", message),
            Self::Other(value) => write!(f, "{}", value),
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct StreamChoice {
    #[serde(default)]
    pub delta: StreamDelta,
    #[serde(default)]
    pub finish_reason: Option<String>,
}

#[derive(Debug, Default, Deserialize)]
pub struct StreamDelta {
    #[serde(default)]
    pub content: Option<String>,
    #[serde(default)]
    pub tool_calls: Option<Vec<StreamToolCall>>,
}

#[derive(Debug, Deserialize)]
pub struct StreamToolCall {
    // Defaulted because the providers that never emit parallel calls — Ollama and LM Studio among
    // them — omit it. Without a default the whole chunk fails to deserialize and the call is
    // dropped, so the one field meant to keep two calls apart was silently losing single ones.
    #[serde(default)]
    pub index: usize,
    #[serde(default)]
    pub id: Option<String>,
    #[serde(default)]
    pub function: Option<StreamFunctionCall>,
}

#[derive(Debug, Deserialize)]
pub struct StreamFunctionCall {
    #[serde(default)]
    pub name: Option<String>,
    #[serde(default)]
    pub arguments: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct WireUsage {
    #[serde(default)]
    prompt_tokens: Option<usize>,
    #[serde(default)]
    completion_tokens: Option<usize>,
    // OpenAI reports cache hits nested; DeepSeek reports them flat. Reading both keeps the
    // budget anchor honest on either.
    #[serde(default)]
    prompt_tokens_details: Option<PromptTokensDetails>,
    #[serde(default)]
    prompt_cache_hit_tokens: Option<usize>,
}

#[derive(Debug, Deserialize)]
struct PromptTokensDetails {
    #[serde(default)]
    cached_tokens: Option<usize>,
}

impl From<WireUsage> for TokenUsage {
    fn from(usage: WireUsage) -> Self {
        let cache_read = usage
            .prompt_tokens_details
            .and_then(|d| d.cached_tokens)
            .or(usage.prompt_cache_hit_tokens)
            .unwrap_or(0);

        // `prompt_tokens` is the whole prompt including cache hits, while the neutral shape
        // treats the two as disjoint and sums them; subtracting avoids double counting.
        let input = usage.prompt_tokens.unwrap_or(0).saturating_sub(cache_read);

        Self {
            input,
            cache_read,
            // No OpenAI-dialect provider bills cache writes separately.
            cache_write: 0,
            output: usage.completion_tokens.unwrap_or(0),
        }
    }
}

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

    fn assistant_with_call() -> Message {
        Message::assistant(vec![
            ContentPart::Text {
                text: "checking".to_string(),
            },
            ContentPart::ToolUse {
                id: "call_1".to_string(),
                name: "read_file".to_string(),
                input: serde_json::json!({"path": "a.rs"}),
            },
        ])
    }

    #[test]
    fn system_prompt_becomes_the_first_message() {
        let wire = to_wire_messages(&[Message::user("hi")], Some("be brief"));

        assert_eq!(wire.len(), 2);
        assert_eq!(wire[0].role, "system");
        assert_eq!(wire[0].content.as_deref(), Some("be brief"));
        assert_eq!(wire[1].role, "user");
    }

    #[test]
    fn tool_use_becomes_an_assistant_tool_call() {
        let wire = to_wire_messages(&[assistant_with_call()], None);

        assert_eq!(wire.len(), 1);
        assert_eq!(wire[0].role, "assistant");
        assert_eq!(wire[0].content.as_deref(), Some("checking"));
        let calls = wire[0].tool_calls.as_ref().expect("tool_calls");
        assert_eq!(calls[0].id, "call_1");
        assert_eq!(calls[0].function.name, "read_file");
    }

    // One neutral user message carrying several results has to expand into one `tool` message
    // each, or the provider sees a single result answering many calls.
    #[test]
    fn each_tool_result_becomes_its_own_tool_message() {
        let history = vec![Message::tool_results(vec![
            ("call_1".to_string(), "ok".to_string()),
            ("call_2".to_string(), "also ok".to_string()),
        ])];

        let wire = to_wire_messages(&history, None);

        assert_eq!(wire.len(), 2);
        assert!(wire.iter().all(|m| m.role == "tool"));
        assert_eq!(wire[0].tool_call_id.as_deref(), Some("call_1"));
        assert_eq!(wire[1].tool_call_id.as_deref(), Some("call_2"));
    }

    // A `tool` message the provider cannot tie back to its call is rejected, so results must
    // never be emitted before the assistant turn that requested them.
    #[test]
    fn tool_results_follow_the_assistant_turn_that_called_them() {
        let history = vec![
            assistant_with_call(),
            Message::tool_results(vec![("call_1".to_string(), "ok".to_string())]),
        ];

        let wire = to_wire_messages(&history, None);

        assert_eq!(wire.len(), 2);
        assert_eq!(wire[0].role, "assistant");
        assert_eq!(wire[1].role, "tool");
    }

    #[test]
    fn a_tool_only_assistant_turn_sends_null_content() {
        let history = vec![Message::assistant(vec![ContentPart::ToolUse {
            id: "call_1".to_string(),
            name: "list_dir".to_string(),
            input: serde_json::json!({}),
        }])];

        let wire = to_wire_messages(&history, None);

        assert_eq!(wire[0].content, None);
        assert!(serde_json::to_string(&wire[0])
            .unwrap()
            .contains("tool_calls"));
    }

    // Streaming is opt-in per request, and the usage flag only belongs on the streaming one.
    #[test]
    fn usage_is_only_requested_when_streaming() {
        let history = vec![Message::user("hi")];

        let streamed = serde_json::to_value(ChatRequest::build(
            "m".into(),
            10,
            &history,
            None,
            None,
            true,
        ))
        .unwrap();
        assert_eq!(streamed["stream_options"]["include_usage"], true);

        let plain = serde_json::to_value(ChatRequest::build(
            "m".into(),
            10,
            &history,
            None,
            None,
            false,
        ))
        .unwrap();
        assert!(plain.get("stream_options").is_none(), "{}", plain);
    }

    #[test]
    fn arguments_reach_the_block_parsed() {
        assert_eq!(
            tool_block("c1".to_string(), "read".to_string(), r#"{"a":1}"#),
            ContentPart::ToolUse {
                id: "c1".to_string(),
                name: "read".to_string(),
                input: serde_json::json!({"a": 1}),
            }
        );
    }

    // A malformed call must still be a `tool_use`: a `tool_result` here lands inside an assistant
    // turn, which both wire formats reject. See `agent::tool_input`.
    #[test]
    fn malformed_arguments_still_yield_a_tool_use() {
        let _guard = crate::diag::test_lock();

        assert_eq!(
            tool_block("c1".to_string(), "read".to_string(), "{invalid"),
            ContentPart::ToolUse {
                id: "c1".to_string(),
                name: "read".to_string(),
                input: serde_json::json!({}),
            }
        );
    }

    // `prompt_tokens` already includes the cached prefix, while TokenUsage::total sums its
    // fields; counting the hits in both places would inflate the budget anchor.
    #[test]
    fn cache_hits_are_not_counted_twice() {
        let usage: TokenUsage = WireUsage {
            prompt_tokens: Some(1000),
            completion_tokens: Some(50),
            prompt_tokens_details: None,
            prompt_cache_hit_tokens: Some(800),
        }
        .into();

        assert_eq!(usage.input, 200);
        assert_eq!(usage.cache_read, 800);
        assert_eq!(usage.total(), 1050);
    }

    #[test]
    fn nested_cached_tokens_are_read_too() {
        let usage: TokenUsage = WireUsage {
            prompt_tokens: Some(1000),
            completion_tokens: Some(0),
            prompt_tokens_details: Some(PromptTokensDetails {
                cached_tokens: Some(600),
            }),
            prompt_cache_hit_tokens: None,
        }
        .into();

        assert_eq!(usage.cache_read, 600);
        assert_eq!(usage.input, 400);
    }

    #[test]
    fn finish_reasons_map_to_the_neutral_vocabulary() {
        assert_eq!(to_neutral_stop_reason("tool_calls"), "tool_use");
        assert_eq!(to_neutral_stop_reason("length"), "max_tokens");
        assert_eq!(to_neutral_stop_reason("stop"), "end_turn");
    }
}