import sys

from langchain_core.messages import HumanMessage, ToolMessage
from langchain_core.tools import tool, StructuredTool
from pydantic import BaseModel, Field

import logging

from langchain_openai import ChatOpenAI
import os
import dotenv


# 定义函数/tool

## 方式1：使用注解
class FieldInfo(BaseModel):
    city: str = Field(description="要查询的城市名称")


## 注解中可以通过传参覆盖原有函数的描述等信息
@tool(
    name_or_callable="get_weather",
    args_schema=FieldInfo,
    description="查询某个城市的天气，并返回温度信息",
    return_direct=True
)
def get_weather(city: str) -> str:
    """查询指定城市的最新天气信息"""
    # 通过api调用天气网站的接口，得到最新的天气信息
    return f"{city}当前的温度为18°C"

@tool
def test(city: str) -> str:
    """这个是一个测试函数"""
    # 通过api调用天气网站的接口，得到最新的天气信息
    return f"{city}当前的温度为18°C"
######################################################################

logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)

dotenv.load_dotenv()

## 设置环境变量
# os.environ['OPENAI_API_KEY'] = os.getenv("SILICONFLOW_API_KEY")
# os.environ['OPENAI_BASE_URL'] = os.getenv("SILICONFLOW_BASE_URL")
os.environ['OPENAI_API_KEY'] = os.getenv("OLLAMA_API_KEY")
os.environ['OPENAI_BASE_URL'] = os.getenv("OLLAMA_BASE_URL")
# os.environ['OPENAI_API_KEY'] = os.getenv("MINIMAX_API_KEY")
# os.environ['OPENAI_BASE_URL'] = os.getenv("MINIMAX_BASE_URL")

# 如果要使用函数调用，需要选择支持的模型
# 如果不支持，会报错：Error code: 400 - {'code': 20037, 'message': 'Function call is not supported for this model.', 'data': None}
# llm = ChatOpenAI(model="Qwen/Qwen2.5-7B-Instruct")
llm = ChatOpenAI(model="gemma4:e2b")
# llm = ChatOpenAI(model="MiniMax-M2.7")


## 将模型与工具进行绑定
llm_with_tools = llm.bind_tools([get_weather,test],tool_choice="auto")

messages = [HumanMessage(content="巴黎现在的气温是多少")]

## 使用哪个工具，由大模型自己选择
response = llm_with_tools.invoke(messages)
print(response)

## 模型会返回：需要调用get_weather，并且识别出来传参为 Paris，但是不会进行真正的调用

# 如果需要调用工具，则执行并把结果回传
for call in getattr(response, "tool_calls", []) or []:
    # print(call)
    # {'name': 'get_weather', 'args': {'city': 'Paris'}, 'id': '019a305388b40ad0d961da5696e9fd2f', 'type': 'tool_call'}
    if call["name"] == "get_weather":
        args = call["args"]  # 例如 {"city": "Paris"}
        result = get_weather.invoke(args)
        messages.append(response)  # 把模型消息加入对话
        messages.append(
            ToolMessage(
                content=result,
                tool_call_id=call["id"],  # 必须把这次tool调用的id对上
            )
        )
        # 5) 第二轮：把工具结果交回给模型，让它产出最终回答
        final_msg = llm.invoke(messages)
        print("FINAL:", final_msg.content)
        break
else:
    # 模型未请求调用工具，直接给出回答
    print("FINAL(no-tool):", response.content)