tiny-agent 0.1.0

// src/agent/providers/openai_compatible.rs
use super::LlmProvider;
mod utils;
use crate::{
    providers::{FinishReason, LLMRequest, ProviderError, ProviderStreamChunk, TokenUsage},
    shared::ToolDefinition,
};
use async_trait::async_trait;
use futures_util::{StreamExt, stream::BoxStream};
use reqwest::{self, StatusCode, header::HeaderMap};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use std::time::Duration;
use utils::transform_messages;

pub struct OpenAiCompatibleProvider {
    pub api_key: String,
    pub base_url: String,
    client: reqwest::Client,
}

impl OpenAiCompatibleProvider {
    pub fn new(api_key: impl Into<String>, base_url: impl Into<String>) -> Self {
        Self::with_timeout(api_key, base_url, Duration::from_secs(120))
    }

    pub fn with_timeout(
        api_key: impl Into<String>,
        base_url: impl Into<String>,
        timeout: Duration,
    ) -> Self {
        let client = reqwest::Client::builder()
            .connect_timeout(Duration::from_secs(10))
            .timeout(timeout)
            .build()
            .expect("reqwest client configuration should be valid");
        Self {
            api_key: api_key.into(),
            base_url: base_url.into(),
            client,
        }
    }
}

#[derive(Serialize)]
struct ChatReq {
    model: String,
    messages: Vec<ResponseItem>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    tools: Vec<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    stream: bool,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct OpenAiToolCall {
    pub index: usize,
    pub id: Option<String>,
    pub r#type: Option<String>,
    pub function: Option<OpenAiFunc>,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct OpenAiFunc {
    pub name: Option<String>,
    pub arguments: Option<String>,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ResponseItem {
    pub role: String,
    pub content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<OpenAiToolCall>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}
#[derive(Deserialize, Debug)]
struct ChatStreamChunk {
    choices: Vec<ChatStreamChoice>,
    usage: Option<OpenAiUsage>,
}

#[derive(Deserialize, Debug)]
struct ChatStreamChoice {
    delta: ChatStreamDelta,
    finish_reason: Option<String>,
}

#[derive(Deserialize, Debug)]
struct ChatStreamDelta {
    content: Option<String>,
    #[serde(default, alias = "reasoning_content", alias = "thinking")]
    thinking: Option<String>,
    tool_calls: Option<Vec<StreamToolCall>>,
}

#[derive(Deserialize, Debug)]
struct StreamToolCall {
    index: usize,
    id: Option<String>,
    function: Option<StreamFunc>,
}

#[derive(Deserialize, Debug)]
struct StreamFunc {
    name: Option<String>,
    arguments: Option<String>,
}

#[derive(Deserialize, Debug)]
struct OpenAiUsage {
    prompt_tokens: Option<u32>,
    completion_tokens: Option<u32>,
    total_tokens: Option<u32>,
}

impl From<OpenAiUsage> for TokenUsage {
    fn from(value: OpenAiUsage) -> Self {
        Self {
            input_tokens: value.prompt_tokens,
            output_tokens: value.completion_tokens,
            total_tokens: value.total_tokens,
        }
    }
}

fn map_finish_reason(reason: &str) -> FinishReason {
    match reason {
        "stop" => FinishReason::Stop,
        "tool_calls" | "function_call" => FinishReason::ToolCalls,
        "length" => FinishReason::Length,
        "content_filter" => FinishReason::ContentFilter,
        other => FinishReason::Unknown(other.to_string()),
    }
}

#[async_trait]
impl LlmProvider for OpenAiCompatibleProvider {
    fn name(&self) -> &str {
        "openai_compatible"
    }

    fn request_snapshot(&self, req: &LLMRequest) -> Value {
        serde_json::to_value(build_chat_payload(req.clone())).unwrap_or(Value::Null)
    }

    async fn chat_completion(
        &self,
        req: LLMRequest,
    ) -> Result<BoxStream<'static, Result<ProviderStreamChunk, ProviderError>>, ProviderError> {
        let payload = build_chat_payload(req);

        let res = self
            .client
            .post(format!(
                "{}/chat/completions",
                self.base_url.trim_end_matches('/')
            ))
            .header("Authorization", format!("Bearer {}", self.api_key))
            .json(&payload)
            .send()
            .await
            .map_err(map_transport_error)?;

        let status = res.status();

        if !status.is_success() {
            let headers = res.headers().clone();
            let err_text = res.text().await.unwrap_or_default();
            return Err(map_status_error(status, &headers, err_text));
        }
        let mut res = res;

        let output_stream = async_stream::stream! {
            let mut emitted_done = false;
            let mut pending_usage: Option<TokenUsage> = None;
            let mut byte_buffer = Vec::<u8>::new();
            loop {
                let chunk = match res.chunk().await {
                    Ok(Some(chunk)) => chunk,
                    // 流自然结束。是否合法(收到过 finish_reason)在循环外判定。
                    Ok(None) => break,
                    // 传输中途断开:必须显式 yield 错误,不能静默吞掉当成正常结束。
                    Err(e) => {
                        yield Err(map_transport_error(e));
                        return;
                    }
                };
                let lines = match take_sse_lines(&mut byte_buffer, &chunk) {
                    Ok(lines) => lines,
                    Err(e) => {
                        yield Err(e);
                        return;
                    }
                };

                for line in lines {
                    let line = line.trim();
                    if line.is_empty() { continue; }
                    if line == "data: [DONE]" {
                        if !emitted_done {
                            emitted_done = true;
                            yield Ok(ProviderStreamChunk::Done {
                                finish_reason: FinishReason::Stop,
                                usage: pending_usage.take(),
                            });
                        }
                        break;
                    }

                    if let Some(raw_json) = line.strip_prefix("data: ") {
                        match serde_json::from_str::<ChatStreamChunk>(raw_json) {
                            Ok(chunk_data) => {
                                let usage = chunk_data.usage.map(Into::into);
                                if usage.is_some() {
                                    pending_usage = usage.clone();
                                }
                                if let Some(choice) = chunk_data.choices.into_iter().next() {
                                    let delta = choice.delta;
                                    let finish_reason = choice.finish_reason;

                                    if let Some(thinking) = delta.thinking {
                                        if !thinking.is_empty() {
                                            yield Ok(ProviderStreamChunk::ThinkingDelta(thinking));
                                        }
                                    }

                                    if let Some(text) = delta.content {
                                        if !text.is_empty() {
                                            yield Ok(ProviderStreamChunk::ContentDelta(text));
                                        }
                                    }

                                    if let Some(tool_calls) = delta.tool_calls {
                                        for tool in tool_calls {
                                            let (name, args) = match tool.function {
                                                Some(f) => (f.name, f.arguments.unwrap_or_default()),
                                                None => (None, String::new()),
                                            };

                                            yield Ok(ProviderStreamChunk::ToolCallDelta {
                                                index: tool.index,
                                                id: tool.id,
                                                name,
                                                arguments_delta: args,
                                            });
                                        }
                                    }

                                    if let Some(reason) = finish_reason {
                                        emitted_done = true;
                                        yield Ok(ProviderStreamChunk::Done {
                                            finish_reason: map_finish_reason(&reason),
                                            usage,
                                        });
                                    }
                                }
                            }
                            Err(e) => {
                                yield Err(ProviderError::transient(format!(
                                    "failed to parse SSE JSON line: {e}"
                                )));
                                return;
                            }
                        }
                    }
                }
            }

            // 流结束了却从未见过 [DONE] / finish_reason:多半是被截断的半截响应,
            // 视为错误而非成功,避免上层把残缺内容当正常结果落库。
            if !emitted_done {
                yield Err(ProviderError::transient("Stream ended before a finish reason was received"));
            }
        };

        Ok(output_stream.boxed())
    }
}

fn build_chat_payload(req: LLMRequest) -> ChatReq {
    let mut messages = req.messages;
    if let Some(system) = req.system {
        messages.insert(
            0,
            crate::shared::Message::text(crate::shared::Role::System, system),
        );
    }
    let response_items: Vec<ResponseItem> = transform_messages(messages);
    let tools_schema: Vec<serde_json::Value> = req.tools.iter().map(build_tool_schema).collect();

    ChatReq {
        model: req.model,
        messages: response_items,
        tools: tools_schema,
        temperature: req.temperature,
        max_tokens: req.max_tokens,
        stream: true,
    }
}

fn build_tool_schema(tool: &ToolDefinition) -> Value {
    serde_json::json!({
        "type": "function",
        "function": {
            "name": tool.name,
            "description": tool.desc,
            "parameters": normalize_openai_tool_parameters(tool.arguments.clone())
        }
    })
}

fn normalize_openai_tool_parameters(mut schema: Value) -> Value {
    let definitions = schema
        .get("definitions")
        .or_else(|| schema.get("$defs"))
        .cloned()
        .unwrap_or_else(|| Value::Object(Map::new()));
    inline_local_refs(&mut schema, &definitions);
    simplify_openai_schema(&mut schema);
    schema
}

fn inline_local_refs(value: &mut Value, definitions: &Value) {
    match value {
        Value::Object(map) => {
            if let Some(ref_value) = map.get("$ref").and_then(Value::as_str) {
                if let Some(mut resolved) = resolve_local_ref(ref_value, definitions) {
                    inline_local_refs(&mut resolved, definitions);
                    if let Value::Object(resolved_map) = &mut resolved {
                        for (key, value) in map.iter() {
                            if key != "$ref" {
                                resolved_map
                                    .entry(key.clone())
                                    .or_insert_with(|| value.clone());
                            }
                        }
                    }
                    *value = resolved;
                    return;
                }
            }

            for child in map.values_mut() {
                inline_local_refs(child, definitions);
            }
        }
        Value::Array(items) => {
            for item in items {
                inline_local_refs(item, definitions);
            }
        }
        _ => {}
    }
}

fn resolve_local_ref(ref_value: &str, definitions: &Value) -> Option<Value> {
    let key = ref_value
        .strip_prefix("#/definitions/")
        .or_else(|| ref_value.strip_prefix("#/$defs/"))?;
    definitions.get(key).cloned()
}

fn simplify_openai_schema(value: &mut Value) {
    match value {
        Value::Object(map) => {
            for child in map.values_mut() {
                simplify_openai_schema(child);
            }

            if let Some(compressed) = compress_enum_union(map) {
                *value = compressed;
                return;
            }

            map.remove("$schema");
            map.remove("$defs");
            map.remove("definitions");
            map.remove("title");
            map.remove("format");

            if is_object_schema(map) && !map.contains_key("additionalProperties") {
                map.insert("additionalProperties".to_string(), Value::Bool(false));
            }
        }
        Value::Array(items) => {
            for item in items {
                simplify_openai_schema(item);
            }
        }
        _ => {}
    }
}

fn compress_enum_union(map: &Map<String, Value>) -> Option<Value> {
    let variants = map
        .get("oneOf")
        .or_else(|| map.get("anyOf"))
        .and_then(Value::as_array)?;

    let mut enum_values = Vec::<Value>::new();
    let mut has_null = false;
    for variant in variants {
        let variant = variant.as_object()?;
        if is_null_schema(variant) {
            has_null = true;
            continue;
        }

        let values = variant.get("enum")?.as_array()?;
        if values.is_empty() || !values.iter().all(|value| value.is_string()) {
            return None;
        }
        enum_values.extend(values.iter().cloned());
    }
    if enum_values.is_empty() {
        return None;
    }
    if has_null {
        enum_values.push(Value::Null);
    }

    let mut compressed = Map::new();
    if let Some(description) = map.get("description") {
        compressed.insert("description".to_string(), description.clone());
    }
    compressed.insert(
        "type".to_string(),
        if has_null {
            Value::Array(vec![
                Value::String("string".to_string()),
                Value::String("null".to_string()),
            ])
        } else {
            Value::String("string".to_string())
        },
    );
    compressed.insert("enum".to_string(), Value::Array(enum_values));
    Some(Value::Object(compressed))
}

fn is_null_schema(map: &Map<String, Value>) -> bool {
    matches!(map.get("type").and_then(Value::as_str), Some("null"))
        || matches!(
            map.get("enum").and_then(Value::as_array),
            Some(values) if values.len() == 1 && values[0].is_null()
        )
}

fn is_object_schema(map: &Map<String, Value>) -> bool {
    matches!(map.get("type").and_then(Value::as_str), Some("object"))
        || map.contains_key("properties")
}

fn map_transport_error(error: reqwest::Error) -> ProviderError {
    ProviderError::transient(format!("Network transport failed: {}", error))
}

fn map_status_error(status: StatusCode, headers: &HeaderMap, body: String) -> ProviderError {
    let message = format!("OpenAI Compatible API error ({}): {}", status, body);
    if status == StatusCode::TOO_MANY_REQUESTS {
        return ProviderError::rate_limited(message, retry_after_secs(headers));
    }
    if status.is_client_error() {
        return ProviderError::permanent(message);
    }
    ProviderError::transient(message)
}

fn retry_after_secs(headers: &HeaderMap) -> Option<u64> {
    headers
        .get(reqwest::header::RETRY_AFTER)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.parse::<u64>().ok())
}

fn take_sse_lines(buffer: &mut Vec<u8>, chunk: &[u8]) -> Result<Vec<String>, ProviderError> {
    buffer.extend_from_slice(chunk);
    let mut lines = Vec::new();
    while let Some(pos) = buffer.iter().position(|byte| *byte == b'\n') {
        let mut line = buffer.drain(..=pos).collect::<Vec<_>>();
        if line.last() == Some(&b'\n') {
            line.pop();
        }
        if line.last() == Some(&b'\r') {
            line.pop();
        }
        lines
            .push(String::from_utf8(line).map_err(|e| {
                ProviderError::transient(format!("invalid UTF-8 in SSE line: {e}"))
            })?);
    }
    Ok(lines)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::providers::ProviderStreamChunk;
    use crate::shared::{Message, Role, ToolDefinition};
    use futures_util::StreamExt;
    use std::env;
    use std::io::Write;

    #[tokio::test]
    async fn test_openai_compatible_stream_completion() {
        // 凭证从环境读取,绝不硬编码进源码。未配置则跳过这个联网测试。
        let Ok(api_key) = env::var("COMPATIBLE_API_KEY") else {
            eprintln!("跳过: 未设置 COMPATIBLE_API_KEY");
            return;
        };
        let base_url = env::var("COMPATIBLE_BASE_URL")
            .unwrap_or_else(|_| "https://api.deepseek.com/v1".to_string());

        let model_name =
            env::var("COMPATIBLE_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string());

        let provider = OpenAiCompatibleProvider::new(api_key, base_url);

        let req = LLMRequest {
            model: model_name,
            system: Some("你是 kb，一个文学家".to_string()),
            messages: vec![Message::text(
                Role::User,
                "你是谁？用三个字回答。并给我写一首中文诗歌",
            )],
            tools: vec![],
            temperature: Some(0.3),
            max_tokens: None,
        };

        let stream_result = provider.chat_completion(req).await;

        match stream_result {
            Ok(mut stream) => {
                println!("\n=========== [Compatible Stream 开启流式接收] ===========");

                let mut full_content = String::new();

                while let Some(chunk_res) = stream.next().await {
                    match chunk_res {
                        Ok(ProviderStreamChunk::ContentDelta(text)) => {
                            print!("{}", text);
                            std::io::stdout().flush().ok();
                            full_content.push_str(&text);
                        }
                        Ok(ProviderStreamChunk::ThinkingDelta(_)) => {}
                        Ok(ProviderStreamChunk::ToolCallDelta { .. }) => {}
                        Ok(ProviderStreamChunk::Done { .. }) => {}
                        Err(stream_err) => {
                            panic!("流式传输中途发生网络断裂或解析崩溃: {}", stream_err);
                        }
                    }
                }

                println!("\n=====================================================\n");

                assert!(!full_content.is_empty(), "模型最终回复的内容不能为空！");
            }
            Err(e) => {
                panic!("stream failed: {}", e);
            }
        }
    }

    #[test]
    fn sse_line_decoder_preserves_utf8_split_across_chunks() {
        let line =
            "data: {\"choices\":[{\"delta\":{\"content\":\"你好\"},\"finish_reason\":null}]}\n";
        let bytes = line.as_bytes();
        let split = bytes
            .windows(3)
            .position(|window| window == "你".as_bytes())
            .expect("test line contains the target character")
            + 1;

        let mut buffer = Vec::new();
        let first = take_sse_lines(&mut buffer, &bytes[..split]).unwrap();
        assert!(first.is_empty());
        let second = take_sse_lines(&mut buffer, &bytes[split..]).unwrap();

        assert_eq!(second, vec![line.trim_end().to_string()]);
    }

    #[test]
    fn openai_tool_schema_is_flattened_for_wire_payload() {
        let tool = ToolDefinition {
            name: "write_file".to_string(),
            desc: "edit files".to_string(),
            arguments: serde_json::json!({
                "$schema": "http://json-schema.org/draft-07/schema#",
                "title": "WriteFileArgs",
                "type": "object",
                "definitions": {
                    "WriteFileMode": {
                        "oneOf": [
                            {
                                "description": "write file",
                                "enum": ["write"],
                                "type": "string"
                            },
                            {
                                "description": "replace text",
                                "enum": ["replace"],
                                "type": "string"
                            }
                        ]
                    }
                },
                "properties": {
                    "path": {
                        "title": "Path",
                        "type": "string"
                    },
                    "mode": {
                        "description": "edit mode",
                        "anyOf": [
                            { "$ref": "#/definitions/WriteFileMode" },
                            { "type": "null" }
                        ]
                    }
                },
                "required": ["path"]
            }),
        };

        let wire = build_tool_schema(&tool);
        let parameters = &wire["function"]["parameters"];
        assert!(!contains_key_recursive(parameters, "$schema"));
        assert!(!contains_key_recursive(parameters, "definitions"));
        assert!(!contains_key_recursive(parameters, "$defs"));
        assert!(!contains_key_recursive(parameters, "$ref"));
        assert!(!contains_key_recursive(parameters, "title"));
        assert_eq!(parameters["additionalProperties"], false);
        assert_eq!(
            parameters["properties"]["mode"],
            serde_json::json!({
                "description": "edit mode",
                "type": ["string", "null"],
                "enum": ["write", "replace", null]
            })
        );
    }

    fn contains_key_recursive(value: &Value, key: &str) -> bool {
        match value {
            Value::Object(map) => {
                map.contains_key(key)
                    || map.values().any(|value| contains_key_recursive(value, key))
            }
            Value::Array(items) => items.iter().any(|value| contains_key_recursive(value, key)),
            _ => false,
        }
    }
}