src/graphs/actionGraph.py

from langgraph.graph import StateGraph, START, END
from src.states.actionState import ExecutorState, ExecutorOutputState

from src.nodes.actionNode import ExecutorNode
from src.llms.groqllm import GroqLLM
from src.utils.prompts import execution_agent_prompt, compress_execution_human_message, compress_execution_system_prompt

# Import the MongoDB tools
from src.utils.utils import tools

class ExecutorGraphBuilder:
    def __init__(self, llm):
        self.llm = llm 
        self.graph = StateGraph(ExecutorState, output=ExecutorOutputState)
        self.tools = tools  # Use MongoDB tools
        self.execution_agent_prompt = execution_agent_prompt
        self.compress_execution_system_prompt = compress_execution_system_prompt
        self.compress_execution_human_message = compress_execution_human_message
    
    def build_executor_graph(self):
        """Build a graph to build the executor with MongoDB tools"""
        self.executor_node_obj = ExecutorNode(self.llm)

        self.graph.add_node("llm_call", self.executor_node_obj.llm_call)
        self.graph.add_node("tool_node", self.executor_node_obj.tool_node)
        self.graph.add_node("compress_execution", self.executor_node_obj.compress_execution)

        # Flow
        self.graph.add_edge(START, "llm_call")
        self.graph.add_conditional_edges(
            "llm_call",
            self.executor_node_obj.guard_llm,
            {
                "tool_node": "tool_node",
                "compress_execution": "compress_execution",
            },
        )
        self.graph.add_edge("tool_node", "llm_call")
        self.graph.add_edge("compress_execution", END)

        return self.graph
    
    def setup_graph(self):
        return self.graph.compile()

# Initialize with MongoDB-enabled tools
llm = GroqLLM().get_llm()

## Creating the graph with MongoDB tools
graph_builder = ExecutorGraphBuilder(llm)
graph = graph_builder.build_executor_graph().compile()

print("MongoDB-enabled Action Graph created successfully")
print(f"Available tools: {[tool.name for tool in tools]}")