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