iAs/src/llm/provider.rs

//! ## LLM 提供商抽象（trait） + SSE 流式解析
//!
//! 定义了 LLM 层的核心抽象接口 `LlmProvider`，以及 SSE 流式响应的解析工具。
//!
//! ### 设计模式
//! - `LlmProvider` trait — 所有 LLM 提供商需实现的接口
//! - `create_provider()` — 工厂函数，根据配置创建对应的提供商实例
//! - `parse_chat_chunk()` — SSE 流式 delta 解析器
//!
//! ### 流式架构
//! ```text
//! Provider.chat_stream() → mpsc::Receiver<StreamChunk>
//!   ├── StreamChunk::Text (delta 文本)
//!   ├── StreamChunk::Reasoning (思考内容)
//!   ├── StreamChunk::Done (完成信号 + 工具调用 + 用量)
//!   └── StreamChunk::Error (错误)
//! ```

use crate::llm::types::{ConversationConfig, Message, StreamChunk};
use async_trait::async_trait;
use serde::Deserialize;
use std::sync::Arc;
use tokio::sync::mpsc;

/// 流式返回：每个 StreamChunk 是 delta 或控制信号
pub type StreamReceiver = mpsc::Receiver<StreamChunk>;
pub type StreamSender = mpsc::Sender<StreamChunk>;

/// ## LLM 提供商抽象（trait）
///
/// 这是 LLM 层的核心抽象接口，定义了"一个 LLM 提供商需要提供什么能力"。
/// 目前唯一的实现是 `DeepSeekProvider`，但通过这个 trait 可以轻松添加
/// 其他提供商（如 OpenAI、Claude 等）。
///
/// ### 关键设计
/// * 使用 `async_trait` 宏支持异步 trait 方法
/// * 返回 `StreamReceiver`（mpsc::Receiver），调用方可以按 chunk 消费流式结果
/// * `Send + Sync` 保证可以在 tokio task 间安全共享
#[async_trait]
pub trait LlmProvider: Send + Sync {
    /// 提供商名称
    fn name(&self) -> &str;

    /// 发起流式对话
    async fn chat_stream(
        &self,
        config: &ConversationConfig,
        messages: &[Message],
    ) -> Result<StreamReceiver, String>;
}

// ─── 内置提供商注册 ───

pub type BoxedProvider = Arc<dyn LlmProvider>;

/// ## `create_provider` — 工厂函数：从配置创建恰当的 LLM 提供商
///
/// 目前只返回 DeepSeek，后续可以在这里根据配置 switch 到不同提供商。
pub fn create_provider(_config: &ConversationConfig) -> Result<BoxedProvider, String> {
    Ok(Arc::new(super::deepseek::DeepSeekProvider::new()?))
}

// ─── 内部：SSE 解析工具 ───

/// ## 流式块解析结果（内部枚举）
///
/// 用于从 SSE 的 JSON delta 中解析出不同种类的块。
/// * `Text` — 普通文本 delta
/// * `Reasoning` — 思考内容（DeepSeek reasoning_content）
/// * `ToolCallDelta` — 工具调用的增量信息（可能需要跨多个 chunk 拼接）
/// * `FinishReason` — 结束原因（stop / tool_calls 等）
/// * `Usage` — 用量统计（通常在最后一个 chunk 中）
#[allow(dead_code)]
pub(crate) enum ParsedChunk {
    Text(String),
    Reasoning(String),
    ToolCallDelta {
        index: i32,
        id: Option<String>,
        name: Option<String>,
        arguments: String,
    },
    FinishReason(String),
    Usage(super::types::Usage),
}

/// ## 从 JSON body 中解析 DeepSeek/OpenAI 流式 delta
///
/// ### 输入
/// 一行 SSE 数据，格式为 `data: {json}`。
///
/// ### 处理逻辑
/// 1. 跳过非 `data: ` 前缀的行和 `[DONE]` 行
/// 2. 解析 JSON 为 `ChunkResponse`（内嵌多个 `choices`）
/// 3. 提取 usage（可能在任意一个 chunk 中）
/// 4. 对每个 choice，按优先级提取：
///    - tool_calls delta → `ToolCallDelta`
///    - reasoning_content → `Reasoning`
///    - content → `Text`
///    - finish_reason → `FinishReason`
pub(crate) fn parse_chat_chunk(line: &str) -> Option<ParsedChunk> {
    if !line.starts_with("data: ") {
        return None;
    }

    let data = line["data: ".len()..].trim();
    if data == "[DONE]" {
        return None;
    }

    #[derive(Deserialize)]
    struct ToolCallDelta {
        #[serde(default)]
        index: Option<i32>,
        #[serde(default)]
        id: Option<String>,
        #[serde(rename = "type", default)]
        #[allow(dead_code)]
        call_type: Option<String>,
        #[serde(default)]
        function: Option<ToolFunctionDelta>,
    }

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

    #[derive(Deserialize)]
    struct Delta {
        #[serde(default)]
        content: Option<String>,
        #[serde(default)]
        reasoning_content: Option<String>,
        #[serde(default)]
        tool_calls: Option<Vec<ToolCallDelta>>,
    }

    #[derive(Deserialize)]
    struct ChunkChoice {
        delta: Delta,
        #[serde(default)]
        finish_reason: Option<String>,
    }

    #[derive(Deserialize)]
    struct ChunkResponse {
        choices: Vec<ChunkChoice>,
        #[serde(default)]
        usage: Option<super::types::Usage>,
    }

    let parsed: ChunkResponse = match serde_json::from_str(data) {
        Ok(p) => p,
        Err(_) => return None,
    };

    // 提取 usage（可能在最后一个 chunk）
    if let Some(ref usage) = parsed.usage {
        if usage.total_tokens > 0 {
            return Some(ParsedChunk::Usage(usage.clone()));
        }
    }

    for choice in parsed.choices {
        // 工具调用 delta
        if let Some(tool_calls) = &choice.delta.tool_calls {
            for tc in tool_calls {
                let idx = tc.index.unwrap_or(0);
                let args = tc
                    .function
                    .as_ref()
                    .and_then(|f| f.arguments.clone())
                    .unwrap_or_default();
                let name = tc.function.as_ref().and_then(|f| f.name.clone());
                return Some(ParsedChunk::ToolCallDelta {
                    index: idx,
                    id: tc.id.clone(),
                    name,
                    arguments: args,
                });
            }
        }

        if let Some(reasoning) = &choice.delta.reasoning_content {
            if !reasoning.is_empty() {
                return Some(ParsedChunk::Reasoning(reasoning.clone()));
            }
        }
        if let Some(content) = &choice.delta.content {
            if !content.is_empty() {
                return Some(ParsedChunk::Text(content.clone()));
            }
        }
        if let Some(reason) = &choice.finish_reason {
            if !reason.is_empty() {
                return Some(ParsedChunk::FinishReason(reason.clone()));
            }
        }
    }

    None
}