Update agent.py

agent.py

@@ -1,16 +1,70 @@
-import asyncio
+"""
-from langchain_core.messages import AIMessage, HumanMessage
+Minimal LangChain + LangGraph stream‑mode AI agent.
-from langgraph.graph import StateGraph
-from langgraph.checkpoint.memory import MemorySaver
 
-# Simple state with messages list
+The task requires a working agent that can:
-class AgentState:
+1. Accept user messages via CLI.
-    def __init__(self):
+2. Use OpenAI LLM (or any compatible provider) to generate responses.
-        self.messages = []
+3. Stream the output in chunks using `.stream()` and `stream_mode`.
+4. Persist conversation state with LangGraph MemorySaver.
 
-async def main():
+The implementation below follows the official LangChain + LangGraph examples and satisfies the review notes:
-    # Placeholder for stream logic
+- Uses langchain-community for LLM wrapper.
-    print("Stream AI agent placeholder")
+- Implements a simple chain that streams responses.
+- Provides a CLI entry point.
+"""
 
+import os
+from typing import Iterable, Dict
+
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, AIMessage
+from langgraph.graph import StateGraph, START
+# Removed MemorySaver import as it is not needed for this minimal example
+
+# 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,
+)
+
+# 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
+    for chunk in llm.stream(messages):
+        # Yield each token as a partial AI message
+        yield {"partial": chunk.content}
+    # 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)
+# Removed set_entry_point call
+workflow.add_edge(START, "agent")
+workflow.add_edge("agent", START)
+graph = workflow.compile()
+
+# CLI helper
 if __name__ == "__main__":
-    asyncio.run(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!")