智能体（Agent）入门系列1

┌──────────┐
│ 用户目标  │
└────┬─────┘
     ▼
┌────────────────────┐
│ Agent Runtime      │
│ 上下文、循环、状态管理 │
└────┬───────────────┘
     ▼
┌──────────────┐
│ 大模型 / 策略 │
└────┬─────────┘
     │
     ├── 已经能够回答 ──────────────┐
     │                              ▼
     │                         ┌────────┐
     │                         │最终答案│
     │                         └────────┘
     │
     └── 需要外部信息或执行动作
                    │
                    ▼
             ┌──────────────┐
             │ 工具执行层    │
             │ API/数据库/代码│
             └──────┬───────┘
                    │ 执行结果
                    └──────────────▶ 返回模型继续决策

"""
最小可运行的 DeepSeek 智能体。

运行环境：
- Python 3.10+
- openai Python SDK

功能：
- 模型自主决定是否调用工具
- 程序执行工具
- 工具结果返回给模型
- 循环运行直到得到最终答案
"""

import ast
import json
import operator
import os
import sys
from datetime import datetime
from zoneinfo import ZoneInfo

from openai import OpenAI


# ============================================================
# 1. 基础配置
# ============================================================

# DeepSeek 提供 OpenAI 兼容接口。
BASE_URL = "https://api.deepseek.com"

# 可通过环境变量覆盖模型名称。
MODEL = os.getenv("DEEPSEEK_MODEL", "deepseek-flash")

# 防止模型无限调用工具。
MAX_STEPS = 8


def create_client() -> OpenAI:
    """创建 DeepSeek API 客户端。"""

    api_key = os.getenv("DEEPSEEK_API_KEY")

    if not api_key:
        raise RuntimeError(
            "没有找到 DEEPSEEK_API_KEY。\n"
            "请先执行：export DEEPSEEK_API_KEY='你的 API Key'"
        )

    return OpenAI(
        api_key=api_key,
        base_url=BASE_URL,
        timeout=30.0,
        max_retries=2,
    )


# ============================================================
# 2. 工具实现
# ============================================================

# 允许计算器使用的运算符。
BINARY_OPERATORS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.FloorDiv: operator.floordiv,
    ast.Mod: operator.mod,
    ast.Pow: operator.pow,
}

UNARY_OPERATORS = {
    ast.UAdd: operator.pos,
    ast.USub: operator.neg,
}


def calculate(expression: str) -> str:
    """
    安全计算数学表达式。

    不能直接使用 eval(expression)，因为 expression 是模型生成的，
    直接执行可能带来代码执行风险。
    """

    if not isinstance(expression, str):
        raise ValueError("expression 必须是字符串")

    if len(expression) > 100:
        raise ValueError("表达式过长")

    tree = ast.parse(expression, mode="eval")

    def evaluate(node):
        # 处理整个表达式。
        if isinstance(node, ast.Expression):
            return evaluate(node.body)

        # 只允许整数和浮点数。
        if isinstance(node, ast.Constant):
            if (
                isinstance(node.value, bool)
                or not isinstance(node.value, (int, float))
            ):
                raise ValueError("只允许使用数字")

            if abs(node.value) > 1_000_000_000_000:
                raise ValueError("数字过大")

            return node.value

        # 处理加、减、乘、除等二元运算。
        if isinstance(node, ast.BinOp):
            operation = BINARY_OPERATORS.get(type(node.op))

            if operation is None:
                raise ValueError("包含不支持的运算符")

            left = evaluate(node.left)
            right = evaluate(node.right)

            # 防止生成过于庞大的指数运算。
            if isinstance(node.op, ast.Pow) and abs(right) > 10:
                raise ValueError("指数不能超过 10")

            return operation(left, right)

        # 处理正数和负数。
        if isinstance(node, ast.UnaryOp):
            operation = UNARY_OPERATORS.get(type(node.op))

            if operation is None:
                raise ValueError("包含不支持的一元运算符")

            return operation(evaluate(node.operand))

        raise ValueError("表达式包含不支持的语法")

    return str(evaluate(tree))


# 工具允许访问的时区白名单。
ALLOWED_TIMEZONES = {
    "Asia/Shanghai",
    "Asia/Tokyo",
    "America/New_York",
    "Europe/London",
    "UTC",
}


def get_current_time(timezone: str) -> str:
    """获取指定时区的当前时间。"""

    if timezone not in ALLOWED_TIMEZONES:
        raise ValueError(f"不支持的时区：{timezone}")

    current_time = datetime.now(ZoneInfo(timezone))

    return current_time.isoformat(timespec="seconds")


# 工具白名单。
# 只有这里注册过的函数才允许被模型调用。
TOOL_FUNCTIONS = {
    "calculate": calculate,
    "get_current_time": get_current_time,
}


# ============================================================
# 3. 提供给模型的工具定义
# ============================================================

# 这些定义只是告诉模型有哪些工具。
# 模型不能直接执行对应的 Python 函数。
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "准确计算数学表达式。遇到数学计算时使用。",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "数学表达式，例如：(23 + 19) * 7",
                    }
                },
                "required": ["expression"],
                "additionalProperties": False,
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "获取指定时区的当前日期和时间。",
            "parameters": {
                "type": "object",
                "properties": {
                    "timezone": {
                        "type": "string",
                        "enum": sorted(ALLOWED_TIMEZONES),
                        "description": "IANA 时区名称",
                    }
                },
                "required": ["timezone"],
                "additionalProperties": False,
            },
        },
    },
]


SYSTEM_PROMPT = """
你是一个可靠、简洁的中文智能体。

规则：
1. 遇到数学计算时，必须调用 calculate。
2. 遇到当前时间问题时，必须调用 get_current_time。
3. 不要编造工具执行结果。
4. 工具失败时，可以修正参数后重新调用。
5. 获得足够信息后，直接给出最终答案。
"""


# ============================================================
# 4. 工具执行器
# ============================================================

def execute_tool(tool_name: str, raw_arguments: str) -> str:
    """
    校验并执行模型请求的工具。

    工具结果需要以字符串形式返回给模型。
    """

    tool_function = TOOL_FUNCTIONS.get(tool_name)

    # 拒绝执行没有注册的函数。
    if tool_function is None:
        return json.dumps(
            {
                "ok": False,
                "error": f"未知或未授权的工具：{tool_name}",
            },
            ensure_ascii=False,
        )

    try:
        # 模型返回的 arguments 是 JSON 字符串。
        arguments = json.loads(raw_arguments)

        if not isinstance(arguments, dict):
            raise ValueError("工具参数必须是 JSON 对象")

        # 执行真正的 Python 工具函数。
        result = tool_function(**arguments)

        payload = {
            "ok": True,
            "result": result,
        }

    except json.JSONDecodeError:
        payload = {
            "ok": False,
            "error": "工具参数不是合法 JSON",
        }

    except TypeError as error:
        payload = {
            "ok": False,
            "error": f"工具参数错误：{error}",
        }

    except Exception as error:
        payload = {
            "ok": False,
            "error": str(error),
        }

    return json.dumps(payload, ensure_ascii=False)


# ============================================================
# 5. Agent Loop
# ============================================================

def run_agent(user_input: str) -> str:
    """运行一次完整的智能体任务。"""

    if not user_input.strip():
        raise ValueError("用户输入不能为空")

    client = create_client()

    # messages 保存当前任务的上下文和执行历史。
    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        },
        {
            "role": "user",
            "content": user_input,
        },
    ]

    for step in range(1, MAX_STEPS + 1):
        # 把当前上下文和工具列表发送给 DeepSeek。
        response = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            tools=TOOLS,
            stream=False,

            # 为了保持最小示例简单，关闭思考模式。
            extra_body={
                "thinking": {
                    "type": "disabled",
                }
            },
        )

        assistant_message = response.choices[0].message

        # 将模型消息保存到上下文。
        # 如果其中包含工具调用，后续请求也必须携带该消息。
        messages.append(assistant_message)

        tool_calls = assistant_message.tool_calls or []

        # 模型没有调用工具，说明它已经给出了最终答案。
        if not tool_calls:
            return assistant_message.content or "模型没有返回文本。"

        # 执行模型请求的所有工具。
        for tool_call in tool_calls:
            tool_name = tool_call.function.name
            raw_arguments = tool_call.function.arguments

            print(
                f"[步骤 {step}] 调用工具：{tool_name}\n"
                f"参数：{raw_arguments}"
            )

            tool_result = execute_tool(
                tool_name=tool_name,
                raw_arguments=raw_arguments,
            )

            print(f"结果：{tool_result}\n")

            # 将工具执行结果加入上下文。
            # tool_call_id 用于关联具体的工具调用。
            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": tool_result,
                }
            )

    raise RuntimeError(
        f"超过最大执行步数：{MAX_STEPS}"
    )


# ============================================================
# 6. 命令行入口
# ============================================================

def main() -> None:
    """读取命令行参数并启动智能体。"""

    user_input = " ".join(sys.argv[1:]).strip()

    # 没有传入参数时，使用默认测试问题。
    if not user_input:
        user_input = (
            "请计算 (23 + 19) * 7，"
            "然后告诉我上海现在几点。"
        )

    try:
        answer = run_agent(user_input)

        print("智能体最终回答：")
        print(answer)

    except KeyboardInterrupt:
        print("\n任务已取消。")
        raise SystemExit(130)

    except Exception as error:
        print(f"运行失败：{error}", file=sys.stderr)
        raise SystemExit(1)


if __name__ == "__main__":
    main()