71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
import json
|
|
from langchain_openai import ChatOpenAI
|
|
from pydantic import SecretStr
|
|
from langchain.tools import tool
|
|
from langchain.agents import create_agent
|
|
|
|
# Configure the local LLM
|
|
llm = ChatOpenAI(
|
|
model="gpt-4o-mini", # Replace with your LM Studio model name
|
|
base_url="http://localhost:1234/v1",
|
|
api_key=SecretStr("fake"),
|
|
temperature=0.7,
|
|
)
|
|
|
|
@tool
|
|
def get_price(product: str, city: str) -> str:
|
|
"""Get realistic price for a product in a city. Returns a markdown table."""
|
|
# Create a sub-agent that generates a price table
|
|
sub_agent = create_agent(
|
|
model=llm,
|
|
tools=[],
|
|
system_prompt=f"You are a price estimator. Provide a realistic price for {product} in {city}. Return a markdown table with columns: Product, Price (rub.), Store.",
|
|
)
|
|
# Invoke the sub-agent
|
|
result = sub_agent.invoke(
|
|
{
|
|
"messages": [
|
|
{"role": "user", "content": f"Provide price for {product} in {city}."}
|
|
]
|
|
}
|
|
)
|
|
# Extract the assistant message content
|
|
messages = result.get("messages", [])
|
|
for msg in messages:
|
|
if msg.get("role") == "assistant" and msg.get("content"):
|
|
return msg["content"]
|
|
return "No price data available."
|
|
|
|
def main():
|
|
# Main agent that uses the get_price tool
|
|
agent = create_agent(
|
|
model=llm,
|
|
tools=[get_price],
|
|
system_prompt="You are a shopping list planner. Use the get_price tool to find prices for items.",
|
|
)
|
|
# Sample user query
|
|
user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
|
# Invoke the agent
|
|
result = agent.invoke(
|
|
{
|
|
"messages": [
|
|
{"role": "human", "content": user_query}
|
|
]
|
|
}
|
|
)
|
|
# Print all messages, including tool calls and final answer
|
|
for msg in result.get("messages", []):
|
|
role = msg.get("role")
|
|
content = msg.get("content")
|
|
tool_calls = msg.get("tool_calls")
|
|
if content:
|
|
print(f"{role}: {content}")
|
|
elif tool_calls:
|
|
for call in tool_calls:
|
|
print(f"{role} calls {call.get('name')} with args {call.get('arguments')}")
|
|
else:
|
|
print(f"{role}: (no content)")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|