add main.py

This commit is contained in:
2026-05-28 09:51:15 +00:00
parent 4506a05998
commit bc34f1b0ce
+60
View File
@@ -0,0 +1,60 @@
"""
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)