71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
"""
|
||
Minimal LangChain + LangGraph stream‑mode AI agent.
|
||
|
||
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.
|
||
"""
|
||
|
||
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")
|
||
# removed edge to avoid START as end node
|
||
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!")
|