97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
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 FilesystemBackend, LocalShellBackend, CompositeBackend
|
||
|
||
# --- LLM и backend -----------------------------------------------------------
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
temperature=0.0,
|
||
)
|
||
|
||
backend = CompositeBackend([
|
||
LocalShellBackend(workspace_dir="./workspace"),
|
||
FilesystemBackend(),
|
||
])
|
||
|
||
# --- Tool with sub‑agent -----------------------------------------------------
|
||
@tool
|
||
def get_price(query: str) -> str:
|
||
"""Return a fake price table for a product in a city.
|
||
Internally a sub‑agent is used to generate the answer.
|
||
"""
|
||
# Создаём суб‑агента, который просто отвечает на запрос
|
||
sub_agent = create_deep_agent(
|
||
model=llm,
|
||
tools=[],
|
||
backend=backend,
|
||
system_prompt="You are a price lookup assistant.",
|
||
)
|
||
# Запускаем суб‑агента в отдельном цикле событий
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
result = loop.run_until_complete(
|
||
sub_agent.ainvoke(
|
||
{"messages": [HumanMessage(content=f"Provide price for {query}")]},
|
||
{"configurable": {"thread_id": f"price-{query}"}},
|
||
)
|
||
)
|
||
finally:
|
||
loop.close()
|
||
return result["messages"][-1].content
|
||
|
||
# --- Главный агент ----------------------------------------------------------
|
||
agent = create_deep_agent(
|
||
model=llm,
|
||
tools=[get_price],
|
||
backend=backend,
|
||
system_prompt="You are a helpful agent. Use tools when necessary.",
|
||
)
|
||
|
||
# --- Stream ----------------------------------------------
|
||
stream = agent.stream(
|
||
{
|
||
"messages": [HumanMessage(content="What is the price of milk in Kazan?")],
|
||
},
|
||
stream_mode=["messages", "updates"],
|
||
)
|
||
|
||
# --- Handlers ------------------------------------------
|
||
step = 1
|
||
|
||
def format_chunk_message(chunk):
|
||
global step
|
||
message, meta = chunk
|
||
if meta.get("langgraph_step") != step:
|
||
step = meta.get("langgraph_step")
|
||
print("\n--- --- ---\n")
|
||
if message.content:
|
||
print(message.content, end="", flush=True)
|
||
|
||
|
||
def format_message(message):
|
||
if message.content:
|
||
return message.content
|
||
# tool call representation
|
||
if message.tool_calls:
|
||
call = message.tool_calls[0]
|
||
return f"{call['name']}({call['args']})"
|
||
return ""
|
||
|
||
for chunk in stream:
|
||
chunk_type, chunk_data = chunk
|
||
if chunk_type == "messages":
|
||
format_chunk_message(chunk_data)
|
||
elif chunk_type == "updates":
|
||
if chunk_data.get("model"):
|
||
last_msg = chunk_data["model"]["messages"][-1]
|
||
print(format_message(last_msg), end="", flush=True)
|
||
|
||
# --- End of script ----------------------------------------------------------
|