132 lines
4.7 KiB
Python
132 lines
4.7 KiB
Python
"""
|
||
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
|