81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
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()
|