59 lines
2.2 KiB
Python
59 lines
2.2 KiB
Python
from langgraph.graph import StateGraph
|
|
from typing import TypedDict, List, Any
|
|
from rich.console import Console
|
|
|
|
# Define state schema
|
|
class AgentState(TypedDict):
|
|
conversation_history: List[dict]
|
|
last_tool_call_confirmed: bool | None
|
|
|
|
console = Console()
|
|
|
|
# Memory node: just returns the existing history
|
|
async def memory_node(state: AgentState) -> dict:
|
|
return {"memory": state.get("conversation_history", [])}
|
|
|
|
# Confirm tool call node
|
|
async def confirm_tool_node(state: AgentState, tool_call: Any) -> dict:
|
|
try:
|
|
console.print(f"[bold cyan]Tool call:[/bold cyan] {tool_call}")
|
|
confirmation = input("Confirm execution? (y/n): ")
|
|
confirmed = confirmation.lower().startswith("y")
|
|
state["last_tool_call_confirmed"] = confirmed
|
|
except Exception as e:
|
|
console.print(f"[red]Error during confirmation:[/red] {e}")
|
|
state["last_tool_call_confirmed"] = False
|
|
return {"confirmed_tool_call": state["last_tool_call_confirmed"]}
|
|
|
|
# Agent node: placeholder for actual LangChain agent logic
|
|
async def agent_node(state: AgentState, memory: Any, confirmed_tool_call: Any) -> dict:
|
|
# In a real implementation we would invoke the agent here.
|
|
# For demonstration, echo back the conversation history and confirmation flag.
|
|
response = {
|
|
"memory": memory,
|
|
"confirmed": confirmed_tool_call,
|
|
"message": "Agent executed with memory and confirmation."
|
|
}
|
|
return {"response": response}
|
|
|
|
# Build graph following execution flow: MemoryNode -> AgentNode -> ConfirmToolNode -> AgentNode
|
|
builder = StateGraph(AgentState)
|
|
builder.add_node("MemoryNode", memory_node)
|
|
builder.add_node("AgentNode", agent_node)
|
|
builder.add_node("ConfirmToolNode", confirm_tool_node)
|
|
|
|
# Define transitions according to plan execution flow
|
|
builder.set_entry_point("MemoryNode")
|
|
builder.add_edge("MemoryNode", "AgentNode")
|
|
builder.add_edge("AgentNode", "ConfirmToolNode")
|
|
builder.add_edge("ConfirmToolNode", "AgentNode")
|
|
builder.set_finish_node("AgentNode")
|
|
|
|
# Compile graph
|
|
graph = builder.compile()
|
|
|
|
if __name__ == "__main__":
|
|
# Example initial state
|
|
init_state: AgentState = {"conversation_history": [], "last_tool_call_confirmed": None}
|
|
result = graph.invoke(init_state)
|
|
console.print(result) |