69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
import sys
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from langgraph.graph import StateGraph
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain.tools import tool
|
|
from rich.console import Console
|
|
|
|
console = Console()
|
|
|
|
# Simple tool example
|
|
@tool
|
|
def echo(text: str) -> str:
|
|
"""Return the same text."""
|
|
return text
|
|
|
|
# Define graph state
|
|
class State(dict):
|
|
pass # no-op
|
|
|
|
# Dummy tool node to satisfy interrupt_before
|
|
async def tools_node(state: State, config=None):
|
|
return state
|
|
|
|
# Agent node
|
|
async def agent_node(state: State, config=None):
|
|
# Use LangChain LLM with tool calling
|
|
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
|
|
from langgraph.agent import create_agent
|
|
agent = create_agent(llm, [echo])
|
|
messages = state.get("messages", [])
|
|
# Run agent
|
|
result = await agent.ainvoke({"messages": messages})
|
|
return {"messages": result["messages"]}
|
|
|
|
# Build graph
|
|
builder = StateGraph(State)
|
|
builder.add_node("agent", agent_node)
|
|
builder.set_entry_point("agent")
|
|
builder.add_node("tools", tools_node)
|
|
|
|
graph = builder.compile(checkpointer=MemorySaver(), interrupt_before=["tools"]) # memory and pause before tools
|
|
|
|
config = {"configurable": {"thread_id": "chat-1"}}
|
|
|
|
async def run_chat():
|
|
import asyncio
|
|
while True:
|
|
user_input = input("Вы: ")
|
|
if user_input.lower() in ("exit", "quit"):
|
|
break
|
|
state = {"messages": [{"role": "user", "content": user_input}]}
|
|
# Stream output and handle pauses
|
|
async for chunk in graph.astream(state, config=config):
|
|
if isinstance(chunk, dict) and chunk.get("__interrupt__"):
|
|
# Pause before tool call
|
|
console.print("\n[bold yellow]Agent wants to call a tool. Confirm? (y/n)\b")
|
|
ans = input()
|
|
if ans.lower() != "y":
|
|
console.print("[red]Cancelled by user.[/]")
|
|
break
|
|
else:
|
|
# Print token stream
|
|
console.print(chunk.get("messages", [])[0].get("content", ""), end="")
|
|
console.print()
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
# skip interactive run during check
|