hello-agents/Co-creation-projects/tino-chen-HelloClaw/src/agent/enhanced_llm.py

"""增强版 HelloAgentsLLM - 支持流式工具调用"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional, List, Dict, Union, Any, AsyncIterator

from hello_agents.core.llm import HelloAgentsLLM
from hello_agents.core.exceptions import HelloAgentsException


# ==================== 流式工具调用数据结构 ====================

class StreamToolEventType(Enum):
    """流式工具调用事件类型"""
    CONTENT = "content"  # 文本内容增量
    TOOL_CALL_START = "tool_call_start"  # 工具调用开始（收到ID和名称）
    TOOL_CALL_DELTA = "tool_call_delta"  # 工具调用参数增量
    FINISH = "finish"  # 流结束


@dataclass
class StreamToolEvent:
    """流式工具调用事件

    封装流式响应中的不同类型数据，统一处理文本内容和工具调用。
    """
    event_type: StreamToolEventType
    # 文本内容
    content: Optional[str] = None
    # 工具调用
    tool_call_index: Optional[int] = None  # 工具调用索引（用于增量累积）
    tool_call_id: Optional[str] = None  # 工具调用ID
    tool_name: Optional[str] = None  # 工具名称
    tool_arguments_delta: Optional[str] = None  # 参数增量
    # 结束信息
    finish_reason: Optional[str] = None

    @property
    def is_content(self) -> bool:
        """是否为文本内容事件"""
        return self.event_type == StreamToolEventType.CONTENT

    @property
    def is_tool_call(self) -> bool:
        """是否为工具调用事件"""
        return self.event_type in (
            StreamToolEventType.TOOL_CALL_START,
            StreamToolEventType.TOOL_CALL_DELTA
        )

    @property
    def is_finish(self) -> bool:
        """是否为结束事件"""
        return self.event_type == StreamToolEventType.FINISH


@dataclass
class StreamToolCallResult:
    """流式工具调用完成后的结果

    包含累积的文本内容和工具调用列表。
    """
    content: str = ""
    tool_calls: List[Dict[str, Any]] = field(default_factory=list)
    finish_reason: Optional[str] = None

    def add_content(self, delta: str):
        """添加文本内容"""
        self.content += delta

    def add_tool_call_start(self, index: int, tool_id: str, tool_name: str):
        """添加工具调用开始"""
        # 确保列表足够长
        while len(self.tool_calls) <= index:
            self.tool_calls.append({"id": "", "name": "", "arguments": ""})
        self.tool_calls[index]["id"] = tool_id
        self.tool_calls[index]["name"] = tool_name

    def add_tool_call_delta(self, index: int, arguments_delta: str):
        """添加工具调用参数增量"""
        while len(self.tool_calls) <= index:
            self.tool_calls.append({"id": "", "name": "", "arguments": ""})
        self.tool_calls[index]["arguments"] += arguments_delta

    def get_complete_tool_calls(self) -> List[Dict[str, Any]]:
        """获取完整的工具调用列表（过滤不完整的）"""
        return [
            tc for tc in self.tool_calls
            if tc["id"] and tc["name"]
        ]

    def to_assistant_message(self) -> Dict[str, Any]:
        """转换为助手消息格式（用于追加到消息历史）"""
        message: Dict[str, Any] = {"role": "assistant", "content": self.content or None}
        if self.tool_calls:
            message["tool_calls"] = [
                {
                    "id": tc["id"],
                    "type": "function",
                    "function": {
                        "name": tc["name"],
                        "arguments": tc["arguments"]
                    }
                }
                for tc in self.get_complete_tool_calls()
            ]
        return message


# ==================== 增强版 LLM 类 ====================

class EnhancedHelloAgentsLLM(HelloAgentsLLM):
    """
    增强版 HelloAgentsLLM - 添加流式工具调用支持

    继承自 HelloAgentsLLM，新增以下方法：
    - astream_invoke_with_tools: 异步流式工具调用
    - get_last_stream_tool_result: 获取最后一次流式工具调用的累积结果
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._last_stream_tool_result: Optional[StreamToolCallResult] = None

    async def astream_invoke_with_tools(
        self,
        messages: List[Dict],
        tools: List[Dict],
        tool_choice: Union[str, Dict] = "auto",
        **kwargs
    ) -> AsyncIterator[StreamToolEvent]:
        """
        异步流式调用 LLM 并支持工具调用（Function Calling）

        这是最优雅的流式工具调用方法，封装了所有流式处理的复杂逻辑。

        Args:
            messages: 消息列表
            tools: 工具 schema 列表
            tool_choice: 工具选择策略
            **kwargs: 其他参数（temperature, max_tokens 等）

        Yields:
            StreamToolEvent: 流式事件，可能是文本内容或工具调用增量

        Example:
            async for event in llm.astream_invoke_with_tools(messages, tools):
                if event.is_content:
                    print(event.content, end="")
                elif event.event_type == StreamToolEventType.TOOL_CALL_START:
                    print(f"\\n调用工具: {event.tool_name}")

            # 获取累积结果
            result = llm.get_last_stream_tool_result()
        """
        from openai import AsyncOpenAI

        # 创建异步客户端
        client = AsyncOpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
            timeout=self.timeout
        )

        # 构建请求参数
        request_params: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "tools": tools,
            "tool_choice": tool_choice,
            "stream": True,
        }
        if kwargs.get("temperature") is not None:
            request_params["temperature"] = kwargs["temperature"]
        if self.max_tokens:
            request_params["max_tokens"] = self.max_tokens

        # 初始化累积结果
        result = StreamToolCallResult()

        try:
            response = await client.chat.completions.create(**request_params)

            async for chunk in response:
                if not chunk.choices:
                    continue

                choice = chunk.choices[0]
                delta = choice.delta

                # 处理文本内容
                if delta.content:
                    result.add_content(delta.content)
                    yield StreamToolEvent(
                        event_type=StreamToolEventType.CONTENT,
                        content=delta.content
                    )

                # 处理工具调用增量
                if delta.tool_calls:
                    for tc_delta in delta.tool_calls:
                        idx = tc_delta.index

                        # 工具调用开始（收到 ID 或名称）
                        if tc_delta.id or (tc_delta.function and tc_delta.function.name):
                            tool_id = tc_delta.id or ""
                            tool_name = tc_delta.function.name if tc_delta.function else ""
                            if tool_id or tool_name:
                                result.add_tool_call_start(idx, tool_id, tool_name)
                                yield StreamToolEvent(
                                    event_type=StreamToolEventType.TOOL_CALL_START,
                                    tool_call_index=idx,
                                    tool_call_id=tool_id,
                                    tool_name=tool_name
                                )

                        # 工具调用参数增量
                        if tc_delta.function and tc_delta.function.arguments:
                            args_delta = tc_delta.function.arguments
                            result.add_tool_call_delta(idx, args_delta)
                            yield StreamToolEvent(
                                event_type=StreamToolEventType.TOOL_CALL_DELTA,
                                tool_call_index=idx,
                                tool_arguments_delta=args_delta
                            )

                # 处理结束原因
                if choice.finish_reason:
                    result.finish_reason = choice.finish_reason
                    yield StreamToolEvent(
                        event_type=StreamToolEventType.FINISH,
                        finish_reason=choice.finish_reason
                    )

        except Exception as e:
            raise HelloAgentsException(f"流式工具调用失败: {str(e)}")

        # 保存累积结果供后续使用
        self._last_stream_tool_result = result

    def get_last_stream_tool_result(self) -> Optional[StreamToolCallResult]:
        """
        获取最后一次流式工具调用的累积结果

        Returns:
            StreamToolCallResult 或 None
        """
        return self._last_stream_tool_result