Add main.py

stream-agent/main.py

@@ -0,0 +1,57 @@
+import os
+from langchain_openai import ChatOpenAI
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import PydanticOutputParser
+from langchain_core.messages import HumanMessage
+from langchain_core.tools import tool
+from typing import List
+
+# ---
+# Пример задачи: вывести таблицу цен на продукты в Казани
+# ---
+
+# 1. Определяем модель
+llm = ChatOpenAI(temperature=0, model="gpt-4o")
+
+# 2. Создаём инструмент (псевдо‑вывод таблицы)
+@tool
+def get_price(product: str, city: str = "Казань") -> str:
+    """Возвращает строку с ценами продукта в указанном городе."""
+    # Для примера фиксированные данные
+    prices = {
+        "молоко": "89",
+        "хлеб": "30",
+        "сахар": "70",
+    }
+    return f"{product} в {city}: {prices.get(product, 'не найден')} руб."
+
+# 3. Создаём агента
+from langchain.agents import create_openai_functions_agent, AgentExecutor
+
+agent = create_openai_functions_agent(llm=llm, tools=[get_price])
+
+# 4. Выполняем в режиме stream
+executor = AgentExecutor(agent=agent, tools=[get_price], verbose=False)
+
+# Запускаем
+input_message = "Какая цена на молоко и хлеб в Казани?"
+
+# Stream
+stream = executor.stream(
+    {"messages": [HumanMessage(content=input_message)]},
+    stream_mode=["messages"],
+)
+
+step = 1
+for chunk in stream:
+    chunk_type, chunk_data = chunk
+    if chunk_type == "messages":
+        message, meta = chunk_data
+        if meta["langgraph_step"] != step:
+            step = meta["langgraph_step"]
+            print("\n---------\n")
+        if message.content:
+            print(message.content, end="", flush=True)
+
+# В конце выводим результат
+print("\n\n---\n" + executor.run({"messages": [HumanMessage(content=input_message)]})["output"])