Files
task-6997111cd6d3a5544a3deffd/main.py
T
2026-05-28 09:51:15 +00:00

61 lines
2.2 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
Simple hierarchical LangChain agent for shopping list.
"""
from langchain_openai import ChatOpenAI
from pydantic import SecretStr
from langchain.tools import tool
from langchain.agents import create_agent, AgentExecutor
import json
# LLM configuration replace <model_name> with your LM Studio model name
llm = ChatOpenAI(
model="<model_name>",
base_url="http://localhost:1234/v1",
api_key=SecretStr("fake"),
temperature=0.7,
)
# Subagent that generates a price table for one product in a city
@tool(name="get_price", description="Return realistic price for a product in a city as a markdown table.")
def get_price(product: str, city: str) -> str:
"""
Generates a markdown table with columns: Product | Price (руб.) | Store.
The subagent uses the same LLM to produce realistic values.
"""
# Create a tiny agent that only returns the price table
prompt = (
f"You are a local market assistant. Provide a markdown table with columns:\n"
f"| Продукт | Цена (руб.) | Магазин |\n"
f"For product '{product}' in city '{city}'. Use realistic Russian prices and store names.")
sub_agent = create_agent(
model=llm,
tools=[],
system_prompt=prompt,
)
result = sub_agent.invoke({"messages": [{"role": "human", "content": "Generate table"}]})
return result["messages"][-1]["content"]
# Main agent with get_price tool
main_agent = create_agent(
model=llm,
tools=[get_price],
system_prompt="Ты помощник по планированию покупок.",
)
if __name__ == "__main__":
user_query = (
"Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
)
response = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]})
# Print all messages
for msg in response["messages"]:
if msg.get("content"):
print(msg["content"])
elif msg.get("tool_calls"):
call = msg["tool_calls"][0]
print(f"{call['name']}({json.dumps(call['args'])})")
# Final answer
final = response["messages"][-1]["content"]
print("\n---\n", final)