Clean repository: delete stream_agent.py
This commit is contained in:
@@ -1,80 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.tools import tool
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
|
||||
|
||||
@tool
|
||||
def get_demo_price(product: str, city: str = "Казань") -> str:
|
||||
"""Return a deterministic demo price for the requested product."""
|
||||
prices = {
|
||||
"молоко": "89 рублей",
|
||||
"хлеб": "54 рубля",
|
||||
"сыр": "219 рублей",
|
||||
}
|
||||
return f"{product} в городе {city}: {prices.get(product.lower(), 'цена не найдена')}"
|
||||
|
||||
|
||||
def build_agent():
|
||||
llm = ChatOpenAI(
|
||||
model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"),
|
||||
base_url=os.getenv("OPENAI_BASE_URL") or None,
|
||||
api_key=os.getenv("OPENAI_API_KEY", "not-needed"),
|
||||
temperature=0,
|
||||
streaming=True,
|
||||
)
|
||||
return create_react_agent(llm, tools=[get_demo_price])
|
||||
|
||||
|
||||
def format_message(message: AIMessage) -> str:
|
||||
if message.content:
|
||||
return str(message.content)
|
||||
if message.tool_calls:
|
||||
call = message.tool_calls[0]
|
||||
return f"{call['name']}({call['args']})"
|
||||
return ""
|
||||
|
||||
|
||||
def format_chunk_message(chunk: tuple[Any, dict[str, Any]], state: dict[str, int]) -> None:
|
||||
message, meta = chunk
|
||||
current_step = int(meta.get("langgraph_step", state["step"]))
|
||||
if current_step != state["step"]:
|
||||
state["step"] = current_step
|
||||
print("\n --- --- --- \n")
|
||||
if getattr(message, "content", None):
|
||||
print(message.content, end="", flush=True)
|
||||
|
||||
|
||||
def run_stream(question: str) -> None:
|
||||
agent = build_agent()
|
||||
stream = agent.stream(
|
||||
{"messages": [{"role": "human", "content": question}]},
|
||||
stream_mode=["messages", "updates"],
|
||||
)
|
||||
state = {"step": 1}
|
||||
for chunk_type, chunk_data in stream:
|
||||
if chunk_type == "messages":
|
||||
format_chunk_message(chunk_data, state)
|
||||
elif chunk_type == "updates" and chunk_data.get("model"):
|
||||
last_message = chunk_data["model"]["messages"][-1]
|
||||
formatted = format_message(last_message)
|
||||
if formatted:
|
||||
print(f"\n{formatted}")
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Streaming AI-agent demo")
|
||||
parser.add_argument("question", nargs="*", help="User question")
|
||||
args = parser.parse_args()
|
||||
|
||||
question = " ".join(args.question).strip() or "Сколько стоит молоко и хлеб в Казани?"
|
||||
run_stream(question)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user