Add agent.py
This commit is contained in:
@@ -0,0 +1,59 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user