Add main.py
This commit is contained in:
@@ -0,0 +1,63 @@
|
||||
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
|
||||
|
||||
# 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 langchain.agents import create_openai_functions_agent
|
||||
agent = create_openai_functions_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")
|
||||
|
||||
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
|
||||
asyncio.run(run_chat())
|
||||
Reference in New Issue
Block a user