add main.py

main.py

@@ -0,0 +1,131 @@
+"""
+Main entry point for the Stream‑mode LangChain agent.
+
+The script demonstrates how to replace a single ``invoke`` call with a streaming
+``stream`` interface.  The output is printed token by token so that the user can
+see the agent “think” in real time.
+
+Three example invocations are provided:
+1. A simple question that requires no tool calls.
+2. A request that triggers the ``get_price`` tool.
+3. A multi‑step conversation that uses the same tool twice.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, Tuple
+
+# ---------------------------------------------------------------------------
+# LLM configuration – BroJS provider
+# ---------------------------------------------------------------------------
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(
+    model="openai/gpt-oss-20b:free",
+    base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
+    api_key=os.getenv("JOURNAL_MCP_PAT"),
+    temperature=0.0,
+)
+
+# ---------------------------------------------------------------------------
+# Tool definition – a tiny mock price lookup
+# ---------------------------------------------------------------------------
+from langchain.tools import tool
+
+@tool
+def get_price(product: str, city: str) -> str:
+    """Return a fake price for *product* in *city*.
+
+    The function is intentionally simple; it only demonstrates how the agent
+    can call tools while streaming.  In a real project this would query an API
+    or database.
+    """
+    prices = {
+        ("milk", "kazan"): "89",
+        ("bread", "kazan"): "45",
+        ("coffee", "moscow"): "120",
+    }
+    key = (product.lower(), city.lower())
+    price = prices.get(key, "unknown")
+    return f"{price} rubles"
+
+# ---------------------------------------------------------------------------
+# Agent construction – zero‑shot react description
+# ---------------------------------------------------------------------------
+from langchain.agents import create_agent
+from langchain_core.messages import HumanMessage
+
+agent = create_agent(
+    llm=llm,
+    tools=[get_price],
+    system_prompt="You are a helpful assistant that can call the get_price tool.",
+)
+
+# ---------------------------------------------------------------------------
+# Streaming helper functions
+# ---------------------------------------------------------------------------
+def format_message(message: Any) -> str:
+    """Return human‑readable representation of a message.
+
+    If the message contains content we return it directly.  Otherwise we
+    construct a string that shows the tool call.
+    """
+    if getattr(message, "content", None):
+        return message.content
+    # Tool call – ``message.tool_calls`` is a list of dicts
+    calls = [f"{c['name']}({c['args']})" for c in message.tool_calls]
+    return ", ".join(calls)
+
+# ---------------------------------------------------------------------------
+# Core streaming logic
+# ---------------------------------------------------------------------------
+def stream_and_print(prompt: str) -> None:
+    """Invoke the agent with ``stream`` and print tokens as they arrive.
+
+    Parameters
+    ----------
+    prompt: str
+        The user message to send to the agent.
+    """
+    # Start streaming – we want both token‑level messages and state updates
+    stream = agent.stream(
+        {"messages": [HumanMessage(content=prompt)]},
+        stream_mode=["messages", "updates"],
+    )
+
+    current_step: int | None = None
+
+    for chunk_type, chunk_data in stream:
+        if chunk_type == "messages":
+            # ``chunk_data`` is a tuple (message, meta)
+            message, meta = chunk_data  # type: ignore[assignment]
+            step = meta.get("langgraph_step")
+            if current_step != step:
+                current_step = step
+                print("\n--- Step {} ---\n".format(step), end="", flush=True)
+            if message.content:
+                print(message.content, end="", flush=True)
+        elif chunk_type == "updates":
+            # ``chunk_data`` contains the finished state of a step
+            model = chunk_data.get("model")
+            if model and model["messages"]:
+                last_msg = model["messages"][-1]
+                print(format_message(last_msg), end="", flush=True)
+    print("\n--- End ---\n")
+
+# ---------------------------------------------------------------------------
+# Example usage – three distinct scenarios
+# ---------------------------------------------------------------------------
+if __name__ == "__main__":
+    examples = [
+        "What is the capital of France?",
+        "How much does milk cost in Kazan?",
+        "I need prices for bread and coffee in Kazan and Moscow.",
+    ]
+
+    for i, ex in enumerate(examples, 1):
+        print(f"\n=== Example {i} ===")
+        stream_and_print(ex)
+
+# End of file