from openai import OpenAI
import json

client = OpenAI()

tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取指定城市的当前天气信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "城市名称，如：北京、上海、Tokyo"
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "温度单位，默认摄氏度"
                    }
                },
                "required": ["city"]
            }
        }
    }
]

# 完整的工具定义示例
{
    "type": "function",
    "function": {
        "name": "search_products",           # 函数名（snake_case）
        "description": "搜索商品目录中的产品",  # 清晰的描述，LLM 据此判断何时调用
        "strict": True,                       # 启用严格模式（推荐）
        "parameters": {
            "type": "object",
            "properties": {
                "keyword": {
                    "type": "string",
                    "description": "搜索关键词"
                },
                "category": {
                    "type": "string",
                    "enum": ["electronics", "clothing", "food", "books"],
                    "description": "商品分类"
                },
                "min_price": {
                    "type": "number",
                    "description": "最低价格"
                },
                "max_price": {
                    "type": "number",
                    "description": "最高价格"
                },
                "in_stock": {
                    "type": "boolean",
                    "description": "是否只显示有库存商品"
                }
            },
            "required": ["keyword"],
            "additionalProperties": False
        }
    }
}

import json

def get_weather(city: str, unit: str = "celsius") -> str:
    """模拟天气查询（实际应调用真实 API）"""
    weather_data = {
        "北京": {"temp": 22, "condition": "晴", "humidity": 45},
        "上海": {"temp": 26, "condition": "多云", "humidity": 70},
    }
    data = weather_data.get(city, {"temp": 20, "condition": "未知", "humidity": 50})
    temp = data["temp"] if unit == "celsius" else data["temp"] * 9/5 + 32
    unit_str = "°C" if unit == "celsius" else "°F"
    return json.dumps({"city": city, "temp": f"{temp}{unit_str}", "condition": data["condition"]}, ensure_ascii=False)

# 第一步：发送带工具定义的请求
messages = [{"role": "user", "content": "北京今天天气怎么样？"}]

response = client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    tools=tools,
    tool_choice="auto",  # auto | none | required | 具体函数
)

# 第二步：检查是否有工具调用
message = response.choices[0].message

if message.tool_calls:
    # 第三步：执行工具调用
    messages.append(message)  # 添加 assistant 消息

    for tool_call in message.tool_calls:
        function_name = tool_call.function.name
        arguments = json.loads(tool_call.function.arguments)

        if function_name == "get_weather":
            result = get_weather(**arguments)

        # 第四步：将结果返回给 LLM
        messages.append({
            "role": "tool",
            "tool_call_id": tool_call.id,
            "content": result,
        })

    # 第五步：获取最终回复
    final_response = client.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        tools=tools,
    )
    print(final_response.choices[0].message.content)
else:
    # LLM 认为不需要调用工具，直接回复
    print(message.content)

# 强制调用特定函数
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "今天天气好热"}],
    tools=tools,
    tool_choice={"type": "function", "function": {"name": "get_weather"}},
)

tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取天气",
            "parameters": {
                "type": "object",
                "properties": {"city": {"type": "string"}},
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_news",
            "description": "获取新闻",
            "parameters": {
                "type": "object",
                "properties": {"topic": {"type": "string"}},
                "required": ["topic"]
            }
        }
    }
]

messages = [{"role": "user", "content": "北京今天天气如何？另外有什么科技新闻？"}]

response = client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    tools=tools,
)

message = response.choices[0].message

# 并行调用：返回多个 tool_calls
if message.tool_calls:
    messages.append(message)

    for tool_call in message.tool_calls:
        print(f"调用: {tool_call.function.name}({tool_call.function.arguments})")
        # 这里可以并行执行（使用 asyncio.gather 等）
        result = "{}"  # 模拟结果
        messages.append({
            "role": "tool",
            "tool_call_id": tool_call.id,
            "content": result,
        })

    final = client.chat.completions.create(model="gpt-4o", messages=messages, tools=tools)
    print(final.choices[0].message.content)

from typing import Callable
import json

class ToolRegistry:
    """工具注册与调度管理器"""

    def __init__(self):
        self.tools: list[dict] = []
        self.handlers: dict[str, Callable] = {}

    def register(self, name: str, description: str, parameters: dict, handler: Callable):
        """注册工具"""
        self.tools.append({
            "type": "function",
            "function": {
                "name": name,
                "description": description,
                "parameters": parameters,
                "strict": True,
            }
        })
        self.handlers[name] = handler

    def execute(self, tool_call) -> str:
        """执行工具调用"""
        name = tool_call.function.name
        args = json.loads(tool_call.function.arguments)
        if name not in self.handlers:
            return json.dumps({"error": f"Unknown tool: {name}"})
        try:
            return self.handlers[name](**args)
        except Exception as e:
            return json.dumps({"error": str(e)})

    def get_tools(self) -> list[dict]:
        return self.tools


# 注册工具
registry = ToolRegistry()

registry.register(
    name="calculate",
    description="执行数学计算",
    parameters={
        "type": "object",
        "properties": {
            "expression": {"type": "string", "description": "数学表达式，如 2+3*4"}
        },
        "required": ["expression"],
        "additionalProperties": False
    },
    handler=lambda expression: json.dumps({"result": eval(expression)})
)

registry.register(
    name="get_current_time",
    description="获取当前时间",
    parameters={"type": "object", "properties": {}, "additionalProperties": False},
    handler=lambda: json.dumps({"time": "2025-06-01 12:00:00"})
)

# 使用工具注册器
client = OpenAI()
messages = [{"role": "user", "content": "计算 (15 + 27) * 3 等于多少？"}]

response = client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    tools=registry.get_tools(),
)

message = response.choices[0].message
if message.tool_calls:
    messages.append(message)
    for tc in message.tool_calls:
        result = registry.execute(tc)
        messages.append({"role": "tool", "tool_call_id": tc.id, "content": result})

    final = client.chat.completions.create(
        model="gpt-4o", messages=messages, tools=registry.get_tools()
    )
    print(final.choices[0].message.content)

from pydantic import BaseModel

class MovieReview(BaseModel):
    title: str
    year: int
    rating: float
    summary: str
    pros: list[str]
    cons: list[str]

response = client.beta.chat.completions.parse(
    model="gpt-4o",
    messages=[
        {"role": "user", "content": "评价电影《盗梦空间》"}
    ],
    response_format=MovieReview,
)

review = response.choices[0].message.parsed
print(f"电影: {review.title} ({review.year})")
print(f"评分: {review.rating}/10")
print(f"优点: {', '.join(review.pros)}")

structured_tool = [{
    "type": "function",
    "function": {
        "name": "output_review",
        "description": "输出电影评价",
        "strict": True,
        "parameters": {
            "type": "object",
            "properties": {
                "title": {"type": "string", "description": "电影名称"},
                "rating": {"type": "number", "description": "评分 1-10"},
                "summary": {"type": "string", "description": "一句话评价"},
                "tags": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "电影标签"
                }
            },
            "required": ["title", "rating", "summary", "tags"],
            "additionalProperties": False
        }
    }
}]

response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "评价《盗梦空间》"}],
    tools=structured_tool,
    tool_choice={"type": "function", "function": {"name": "output_review"}},
)

args = json.loads(response.choices[0].message.tool_calls[0].function.arguments)
print(json.dumps(args, ensure_ascii=False, indent=2))

def stream_with_tools(messages, tools):
    stream = client.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        tools=tools,
        stream=True,
    )

    tool_calls = {}
    for chunk in stream:
        delta = chunk.choices[0].delta

        if delta.tool_calls:
            for tc in delta.tool_calls:
                idx = tc.index
                if idx not in tool_calls:
                    tool_calls[idx] = {"id": "", "name": "", "arguments": ""}
                if tc.id:
                    tool_calls[idx]["id"] = tc.id
                if tc.function:
                    if tc.function.name:
                        tool_calls[idx]["name"] = tc.function.name
                    if tc.function.arguments:
                        tool_calls[idx]["arguments"] += tc.function.arguments

        if delta.content:
            print(delta.content, end="", flush=True)

    return list(tool_calls.values())

ecommerce_tools = [
    {"type": "function", "function": {"name": "search_products", "description": "搜索商品", ...}},
    {"type": "function", "function": {"name": "get_product_detail", "description": "商品详情", ...}},
    {"type": "function", "function": {"name": "add_to_cart", "description": "加入购物车", ...}},
    {"type": "function", "function": {"name": "check_inventory", "description": "查询库存", ...}},
    {"type": "function", "function": {"name": "create_order", "description": "创建订单", ...}},
]

data_tools = [
    {"type": "function", "function": {"name": "query_sql", "description": "执行SQL查询", ...}},
    {"type": "function", "function": {"name": "create_chart", "description": "生成图表", ...}},
    {"type": "function", "function": {"name": "export_csv", "description": "导出CSV", ...}},
]