add: main.py — Stream-режим AI-агента

This commit is contained in:
2026-07-02 04:19:06 +00:00
commit a6f0d869e0
+85
View File
@@ -0,0 +1,85 @@
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 - OpenRouter
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,
)
# Simple backend - filesystem + local shell
backend = CompositeBackend(
[
LocalShellBackend(workspace_dir="./workspace"),
FilesystemBackend(),
]
)
# Example tool - echo (replace with real logic if needed)
@tool
def echo(query: str) -> str:
"""Return the received query unchanged."""
return query
# Create the deep agent
agent = create_deep_agent(
model=llm,
tools=[echo],
backend=backend,
system_prompt="You are a helpful assistant that streams its answer token by token.",
)
def format_message(message) -> str:
"""Convert a LangChain message to a printable string."""
if getattr(message, "content", None):
return message.content
# If the message is a tool call, show the call
if getattr(message, "tool_calls", None):
tc = message.tool_calls[0]
return f"{tc['name']}({tc['args']})"
return ""
def format_chunk_message(chunk):
"""Print a token chunk, adding a separator when the step changes."""
message, meta = chunk
global current_step
step = meta.get("langgraph_step", 0)
if step != current_step:
current_step = step
print("\n--- --- ---\n")
if getattr(message, "content", None):
print(message.content, end="", flush=True)
async def main():
# Prepare the input
user_input = "Расскажи, как приготовить борщ, используя инструмент echo для демонстрации."
stream = agent.stream(
{"messages": [HumanMessage(content=user_input)]},
stream_mode=["messages", "updates"],
)
global current_step
current_step = 0
# Iterate over the stream
for chunk_type, chunk_data in stream:
if chunk_type == "messages":
format_chunk_message(chunk_data)
elif chunk_type == "updates":
# When a model update arrives, print the last model message (tool call or final answer)
model_info = chunk_data.get("model")
if model_info and "messages" in model_info:
last_msg = model_info["messages"][-1]
print("\n" + format_message(last_msg))
# Ensure the final newline
print()
if __name__ == "__main__":
asyncio.run(main())