diff --git a/main.py b/main.py new file mode 100644 index 0000000..9600cfa --- /dev/null +++ b/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