File size: 2,977 Bytes
f5e247b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import os
from langchain.agents import AgentType, initialize_agent
from langchain.tools import StructuredTool
from langchain_deepseek import ChatDeepSeek
from pydantic import BaseModel, Field
from typing import Optional
import math
import requests
class CalculatorInput(BaseModel):
operation: str = Field(description="Mathematical operation: add, subtract, multiply, divide, power, sqrt")
a: float = Field(description="First number")
b: Optional[float] = Field(default=None, description="Second number (not required for sqrt)")
def calculator(operation: str, a: float, b: Optional[float] = None) -> str:
"""Perform basic arithmetic operations."""
op = operation.lower()
if op == "add":
return f"{a} + {b} = {a + b}"
elif op == "subtract":
return f"{a} - {b} = {a - b}"
elif op == "multiply":
return f"{a} * {b} = {a * b}"
elif op == "divide":
if b == 0:
return "Error: Division by zero"
return f"{a} / {b} = {a / b}"
elif op == "power":
return f"{a}^{b} = {a ** b}"
elif op == "sqrt":
if a < 0:
return "Error: Cannot take square root of negative number"
return f"sqrt({a}) = {math.sqrt(a)}"
else:
return f"Unknown operation: {op}"
calculator_tool = StructuredTool.from_function(
func=calculator,
name="Calculator",
description="Useful for performing mathematical calculations. Input should include operation (add, subtract, multiply, divide, power, sqrt) and numbers.",
args_schema=CalculatorInput
)
class WeatherInput(BaseModel):
city: str = Field(description="Name of the city to get the weather for")
def get_weather(city: str) -> str:
"""模拟天气信息（不访问真实 API）"""
mock_data = {
"北京": {"description": "晴天", "temp": 26},
"上海": {"description": "小雨", "temp": 22},
"广州": {"description": "多云", "temp": 30},
"杭州": {"description": "雷阵雨", "temp": 24}
}
weather = mock_data.get(city, {"description": "未知", "temp": "N/A"})
return f"{city} 当前天气：{weather['description']}，温度：{weather['temp']}°C"
weather_tool = StructuredTool.from_function(
func=get_weather,
name="WeatherTool",
description="查询指定城市的天气信息（模拟数据）",
args_schema=WeatherInput
)
llm = ChatDeepSeek(api_key=os.getenv("DEEPSEEK_API_KEY"), model="deepseek-chat")
tools = [calculator_tool, weather_tool]
agent = initialize_agent(
tools,
llm,
agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
verbose=True
)
if __name__ == "__main__":
while True:
query = input("\nEnter your question (or 'exit' to quit): ")
if query.lower() == 'exit':
break
result = agent.invoke(query)
print(f"Result: {result['output']}") |