diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..5122237 --- /dev/null +++ b/agent.py @@ -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) \ No newline at end of file