import sys

from langchain_core.tools import tool, StructuredTool
from pydantic import BaseModel, Field
from sqlalchemy import True_


# 定义函数/tool

## 方式1：使用注解
@tool
def get_weather1(city: str) -> str:
    """查询指定城市的最新天气信息"""
    # 通过api调用天气网站的接口，得到最新的天气信息
    return f"{city}当前的温度为18°C"


class FieldInfo(BaseModel):
    city: str = Field(description="要查询的城市名称")

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


## 方式2：
def get_weather3(city: str) -> str:
    """查询指定城市的最新天气信息"""
    # 通过api调用天气网站的接口，得到最新的天气信息
    return f"{city}当前的温度为18°C"

get_weather3_tool = StructuredTool.from_function(
    func=get_weather3,
    name="get_weather3",
    args_schema=FieldInfo,
    description="第三个返回天气的函数"
)


if __name__ == '__main__':
    # 调用函数，并打印返回结果
    print(get_weather1("北京"))
    print(get_weather1("巴黎"))

    print(f"name={get_weather3_tool.name}")
    print(f"args={get_weather3_tool.args}")
    print(f"description={get_weather3_tool.description}")
    print(f"return_direct={get_weather3_tool.return_direct}")  # 直接返回：如果为false，就是会将返回值给到大模型，让大模型进一步加工后再返回。如果是true，则直接返回给用户。