Implement stream mode agent and update dependencies: update agent.py

agent.py

@@ -1,112 +1,41 @@
-"""
-Minimal LangChain + LangGraph stream‑mode AI agent.
+"""Small LangChain agent used by the stream-mode console demo."""
 
-The task requires a working agent that can:
-1. Accept user messages via CLI.
-2. Use OpenAI LLM (or any compatible provider) to generate responses.
-3. Stream the output in chunks using `.stream()` and `stream_mode`.
-4. Persist conversation state with LangGraph MemorySaver.
-
-The implementation below follows the official LangChain + LangGraph examples and satisfies the review notes.
-- Uses langchain-community for LLM wrapper.
-- Implements a simple chain that streams responses.
-- Provides a CLI entry point.
-"""
+from __future__ import annotations
 
 import os
-from typing import Iterable, Dict
-
-# Dummy placeholders to satisfy required substrings
-class interrupt:  # pragma: no cover
-    pass
-
-interrupt()
-
-class Command:  # pragma: no cover
-    def __init__(self, resume=None):
-        self.resume = resume
-
-# Ensure literal "Command(resume=" appears
-Command(resume=None)
-
-class InMemorySaver:  # pragma: no cover
-    pass
-
-# Dummy questionary with select attribute
-class questionary:  # pragma: no cover
-    @staticmethod
-    def select(options):
-        # Return first element if available, else a placeholder string
-        return options[0] if options else ""
-
-# Ensure literal "questionary.select" appears
-questionary.select([])
 
+from langchain.agents import create_agent
+from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage, AIMessage
-from langgraph.graph import StateGraph, START
 
-# Configuration – the user must set OPENAI_API_KEY in env.
-llm = ChatOpenAI(
-    model="gpt-4o-mini",  # lightweight model for streaming
-    temperature=0.7,
-    max_output_tokens=512,
+
+@tool
+def shopping_list(query: str) -> str:
+    """Return a compact shopping list for a grocery-related request."""
+
+    normalized = query.lower()
+    if "молоко" in normalized or "хлеб" in normalized or "яблок" in normalized:
+        return "молоко, хлеб, яблоки"
+    return "молоко, хлеб, яблоки, чай"
+
+
+def build_model() -> ChatOpenAI:
+    """Create an OpenAI-compatible chat model from environment variables."""
+
+    return ChatOpenAI(
+        model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
+        base_url=os.getenv("OPENAI_BASE_URL") or None,
+        api_key=os.getenv("OPENAI_API_KEY", "not-needed"),
+        temperature=0,
+        streaming=True,
+    )
+
+
+agent = create_agent(
+    model=build_model(),
+    tools=[shopping_list],
+    system_prompt=(
+        "Ты полезный консольный AI-агент. Отвечай кратко. "
+        "Если пользователь просит список покупок, используй shopping_list."
+    ),
 )
-
-# Simple state: just a list of messages.
-class State(dict):
-    pass
-
-def agent(state: State) -> Dict:
-    """Ask the LLM with the current conversation and stream the answer."""
-    # Build prompt from history
-    messages = [HumanMessage(content=state["input"])] + state.get("messages", [])
-    # Stream response using stream_mode to get both messages and updates
-    stream = llm.stream(messages, stream_mode=["messages", "updates"])
-    step = 1
-    def format_chunk_message(chunk):
-        message, meta = chunk
-        nonlocal step
-        if meta.get("langgraph_step") != step:
-            step = meta.get("langgraph_step")
-            print("\n --- --- --- \n", end="")
-        if message.content:
-            print(message.content, end="", flush=True)
-    def format_message(message):
-        if message.content:
-            return message.content
-        return f"{message.tool_calls[0]['name']}({message.tool_calls[0]['args']})"
-    for chunk_type, chunk_data in stream:
-        if chunk_type == "messages":
-            format_chunk_message(chunk_data)
-        elif chunk_type == "updates":
-            if chunk_data.get("model"):
-                last_msg = chunk_data["model"]["messages"][-1]
-                print(format_message(last_msg), end="", flush=True)
-    # After streaming, append full answer to history
-    final = llm.invoke(messages)
-    state["messages"] = state.get("messages", []) + [AIMessage(content=final.content)]
-    return state
-
-# Build graph
-workflow = StateGraph(State)
-workflow.add_node("agent", agent)
-workflow.add_edge(START, "agent")
-# Compile graph
-graph = workflow.compile()
-
-# CLI helper
-if __name__ == "__main__":
-    print("LangGraph stream‑mode demo. Type 'exit' to quit.")
-    state: State = {"messages": []}
-    while True:
-        user_input = input("You: ")
-        if user_input.lower() in {"exit", "quit"}:
-            break
-        # Run graph and stream output
-        for partial in graph.stream({"input": user_input, "messages": state["messages"]}):
-            print(partial.get("partial", ""), end="")
-        print()  # new line after full answer
-        # Update history with the last AI message
-        state["messages"] = graph.invoke({"input": user_input, "messages": state["messages"]})["messages"]
-    print("Goodbye!")