test_client.py

"""opencode 风格测试客户端：验证 ima 上游能否还原 tool call
用法：
  1. 起 server:  uvicorn server:app --port 8300        （真实链路）
     或 mock:    MOCK_UPSTREAM=1 uvicorn server:app --port 8300 （无token验链路）
  2. 跑本脚本:   python test_client.py "帮我看下 demo.txt 里写了什么"

循环逻辑（标准 agent loop）：
  messages + tools → /v1/chat/completions
  → 若返回 tool_calls：本地执行工具 → role=tool 回传 → 继续
  → 否则输出最终回答，结束
"""
from __future__ import annotations

import json
import os
import sys

import httpx

BASE = os.getenv("OPENAI_BASE", "http://127.0.0.1:8300/v1")
MODEL = os.getenv("MODEL", "ima-copilot")
MAX_ROUNDS = 6

from tools_def import TOOLS_SCHEMA, execute_tool, AGENT_ROOT


def chat(messages: list[dict], stream: bool = False) -> dict:
    """OpenAI 风格调用（httpx 手写，等价 openai sdk）"""
    r = httpx.post(f"{BASE}/chat/completions",
                   json={"model": MODEL, "messages": messages, "tools": TOOLS_SCHEMA, "stream": stream},
                   timeout=600)
    r.raise_for_status()
    if not stream:
        return r.json()
    # SSE 解析
    content, tool_calls, finish = "", [], None
    for line in r.iter_lines():
        if not line.startswith("data:"):
            continue
        data = line[5:].strip()
        if data == "[DONE]":
            break
        obj = json.loads(data)
        if "error" in obj:
            raise RuntimeError(obj["error"]["message"])
        for ch in obj.get("choices", []):
            d = ch.get("delta") or {}
            if d.get("content"):
                content += d["content"]
            for tc in d.get("tool_calls") or []:
                idx = tc.get("index", 0)
                while len(tool_calls) <= idx:
                    tool_calls.append({"id": "", "type": "function",
                                       "function": {"name": "", "arguments": ""}})
                acc = tool_calls[idx]
                acc["id"] = tc.get("id") or acc["id"]
                fn = tc.get("function") or {}
                acc["function"]["name"] += fn.get("name", "")
                acc["function"]["arguments"] += fn.get("arguments", "")
            if ch.get("finish_reason"):
                finish = ch["finish_reason"]
    msg: dict = {"role": "assistant", "content": content or None}
    if tool_calls and any(tc["function"]["name"] for tc in tool_calls):
        msg["tool_calls"] = tool_calls
    return {"choices": [{"message": msg, "finish_reason": finish or "stop"}]}


def run(task: str, stream: bool = False):
    print(f"=== 任务: {task}")
    print(f"=== 端点: {BASE}  模型: {MODEL}  流式: {stream}\n")
    messages = [{"role": "user", "content": task}]

    for rnd in range(1, MAX_ROUNDS + 1):
        resp = chat(messages, stream=stream)
        choice = resp["choices"][0]
        msg = choice["message"]

        if msg.get("tool_calls"):
            print(f"--- 第{rnd}轮：收到 {len(msg['tool_calls'])} 个 tool_call")
            messages.append({"role": "assistant", "content": msg.get("content"),
                             "tool_calls": msg["tool_calls"]})
            for tc in msg["tool_calls"]:
                fn = tc["function"]
                print(f"    ▶ {fn['name']}({fn['arguments'][:120]})")
                result = execute_tool(fn["name"], fn["arguments"])
                print(f"    ◀ {result[:150].replace(chr(10), ' ⏎ ')}{'…' if len(result) > 150 else ''}")
                messages.append({"role": "tool", "tool_call_id": tc["id"],
                                 "name": fn["name"], "content": result})
            continue

        print(f"--- 第{rnd}轮：最终回答\n{msg.get('content', '')}")
        return msg.get("content", "")

    print("!! 达到最大轮次仍未结束")
    return None


if __name__ == "__main__":
    task = " ".join(sys.argv[1:]) or "帮我看下 demo.txt 里写了什么，然后用 bash 列出目录文件"
    os.makedirs(AGENT_ROOT, exist_ok=True)
    run(task, stream=os.getenv("STREAM", "") == "1")