projects/repobot/v4/agent.py

"""RepoBot v4 的智能体：和 v3 的逻辑相同，但每次调用模型都用流式，
一边接收一边把事件交给调用方：正在调用哪个工具、回答的每一段文字、最后的统计。
"""
import json

import llm
from tools import TOOLS

SYSTEM = """你是 RepoBot，Python HTTP 客户端库 httpx 的答疑助手。你可以查 httpx 的官方文档和源码。

做法：
- 先用 search_docs 查文档。文档里有答案，就根据文档回答。
- 文档里没有答案（比如默认值、内部逻辑、某个异常什么时候抛出），再用 grep_source 和 read_source 查源码。
- 回答里注明依据：文档写成 [文档 文件名]，源码写成 [源码 文件路径:行号]。
- 只根据查到的内容回答。查了还是找不到，就如实说没有找到，不要猜。
- 用中文回答，简洁，代码保持原样。"""

MAX_STEPS = 10
MAX_OBSERVATION = 4000


def run_stream(question, history, tracer):
    """生成器。依次产出事件字典：{"type": "tool", ...}、{"type": "token", "text": ...}、{"type": "done", ...}"""
    messages = [{"role": "system", "content": SYSTEM}, *history, {"role": "user", "content": question}]
    schemas = [t["schema"] for t in TOOLS.values()]
    stats = {"steps": 0, "tool_calls": 0, "cost": 0.0}

    for step in range(1, MAX_STEPS + 1):
        stats["steps"] = step
        with tracer.span("llm", llm.MODEL, step=step) as span:
            stream = llm.with_retry(lambda: llm.client.chat.completions.create(
                model=llm.MODEL, messages=messages, tools=schemas, stream=True,
                stream_options={"include_usage": True}, extra_body={"thinking": {"type": "disabled"}}))
            content, calls, usage = [], {}, None
            for chunk in stream:
                if chunk.usage:
                    usage = chunk.usage
                if not chunk.choices:
                    continue
                delta = chunk.choices[0].delta
                if delta.content:
                    content.append(delta.content)
                    yield {"type": "token", "text": delta.content}
                # 流式时，工具调用也是分成很多块发来的：第一块带 id 和函数名，后面的块陆续补上参数。
                # 用 index 区分同一轮里的不同调用，把碎片拼起来
                for piece in delta.tool_calls or []:
                    call = calls.setdefault(piece.index, {"id": "", "name": "", "arguments": ""})
                    call["id"] = piece.id or call["id"]
                    if piece.function and piece.function.name:
                        call["name"] += piece.function.name
                    if piece.function and piece.function.arguments:
                        call["arguments"] += piece.function.arguments
            cost = llm.cost_usd(usage)
            stats["cost"] += cost
            span.update(prompt_tokens=usage.prompt_tokens if usage else None, cost=round(cost, 6),
                        tool_calls=[c["name"] for c in calls.values()])

        if not calls:
            yield {"type": "done", **stats}
            return

        messages.append({"role": "assistant", "content": "".join(content) or None, "tool_calls": [
            {"id": c["id"], "type": "function", "function": {"name": c["name"], "arguments": c["arguments"]}}
            for c in calls.values()]})
        for call in calls.values():
            stats["tool_calls"] += 1
            yield {"type": "tool", "name": call["name"], "args": call["arguments"]}
            with tracer.span("tool", call["name"], args=call["arguments"]) as span:
                try:
                    result = TOOLS[call["name"]]["fn"](**json.loads(call["arguments"] or "{}"))
                except KeyError:
                    result = f"错误：没有叫 {call['name']} 的工具"
                except Exception as e:
                    result = f"错误：{type(e).__name__}: {e}"
                if result.startswith("错误"):
                    span["status"] = "tool_error"
                span["result_chars"] = len(result)
            if len(result) > MAX_OBSERVATION:
                result = result[:MAX_OBSERVATION] + f"\n……（已截断，共 {len(result)} 字符）"
            messages.append({"role": "tool", "tool_call_id": call["id"], "content": result})

    yield {"type": "token", "text": "\n这个问题查了很多步还没有得出结论，我先停下来。可以换个问法，或者问得更具体一些。"}
    yield {"type": "done", **stats}