add: main.py

This commit is contained in:
+75
View File
@@ -0,0 +1,75 @@
import os
import asyncio
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
# ---------- LLM ----------
llm = ChatOpenAI(
model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1",
api_key=os.getenv("OPENAI_API_KEY"),
temperature=0.7,
)
# ---------- Backend ----------
backend = CompositeBackend([
LocalShellBackend(workspace_dir="./workspace"),
FilesystemBackend(),
])
# ---------- Subagent for price generation ----------
# The subagent simply asks the LLM to produce a realistic price table.
# It is wrapped in a tool so that the main agent can call it.
@tool
def get_price(product: str, city: str) -> str:
"""Return a realistic price for a product in a given city.
The response must be a Markdown table with columns: Продукт, Цена (руб.), Магазин.
"""
# Create a tiny agent that only generates the table.
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage
system_prompt = (
"You are a market price generator. "
"Given a product and a city, produce a realistic price table in Markdown. "
"Use plausible Russian store names and prices."
)
sub_agent = create_agent(
model=llm,
tools=[],
system_prompt=system_prompt,
)
prompt = f"Product: {product}\nCity: {city}"
result = sub_agent.invoke({"messages": [HumanMessage(content=prompt)]})
# The subagent returns a dict with 'messages'; take the last content.
return result["messages"][-1].content
# ---------- Main agent ----------
agent = create_deep_agent(
model=llm,
tools=[get_price],
backend=backend,
system_prompt="Ты помощник по планированию покупок.",
)
# ---------- Run ----------
async def main():
user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
result = await agent.ainvoke(
{"messages": [HumanMessage(content=user_query)]},
{"configurable": {"thread_id": "session-1"}},
)
# Print all messages in order
for msg in result["messages"]:
if msg.content:
print(msg.content)
elif msg.tool_calls:
for call in msg.tool_calls:
print(f"{call['name']}({call['args']})")
if __name__ == "__main__":
asyncio.run(main())