code/03-llm-apps/tool_calling.py

78 行 · 2.9 KB
"""让模型调用一个真实的工具:查询 PyPI 上某个包的最新版本。

    python tool_calling.py            关掉思考
    python tool_calling.py --think    开启思考(要把 reasoning_content 传回去)
"""
import json
import os
import sys

import httpx
from openai import OpenAI

client = OpenAI(
    api_key=os.environ["LLM_API_KEY"],
    base_url=os.environ.get("LLM_BASE_URL", "https://api.deepseek.com"),
)
MODEL = os.environ.get("LLM_MODEL", "deepseek-flash")
THINKING = "--think" in sys.argv


def get_pypi_info(package: str) -> dict:
    """真正干活的函数:调用 PyPI 的公开接口。"""
    r = httpx.get(f"https://pypi.org/pypi/{package}/json", timeout=10)
    if r.status_code == 404:
        return {"error": f"PyPI 上没有叫 {package} 的包"}
    info = r.json()["info"]
    return {"name": info["name"], "version": info["version"], "summary": info["summary"],
            "requires_python": info["requires_python"]}


# 告诉模型有哪些工具可以用。模型只能看到这里的描述,看不到函数的代码
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_pypi_info",
            "description": "查询一个 Python 包在 PyPI 上的最新版本、简介和支持的 Python 版本。",
            "parameters": {
                "type": "object",
                "properties": {
                    "package": {"type": "string", "description": "PyPI 上的包名,例如 httpx"},
                },
                "required": ["package"],
            },
        },
    }
]
FUNCTIONS = {"get_pypi_info": get_pypi_info}

messages = [{"role": "user", "content": "httpx 和 requests 在 PyPI 上的最新版本分别是多少?各自要求什么 Python 版本?"}]

for step in range(1, 6):  # 最多 5 轮,防止意外的死循环
    response = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        tools=TOOLS,
        extra_body={"thinking": {"type": "enabled" if THINKING else "disabled"}},
    )
    message = response.choices[0].message
    print(f"第 {step} 轮:finish_reason={response.choices[0].finish_reason}")

    if not message.tool_calls:
        print("最终回答:", message.content)
        break

    # 把模型的这条消息原样放回历史。开思考时,里面的 reasoning_content 也必须带上
    messages.append(message.model_dump(exclude_none=True))

    for call in message.tool_calls:
        args = json.loads(call.function.arguments)
        print(f"  模型要求调用 {call.function.name}({args})")
        try:
            result = FUNCTIONS[call.function.name](**args)
        except Exception as e:  # 工具出错也要告诉模型,而不是让程序崩掉
            result = {"error": f"{type(e).__name__}: {e}"}
        print(f"  返回:{result}")
        messages.append({"role": "tool", "tool_call_id": call.id, "content": json.dumps(result, ensure_ascii=False)})