Spaces:
Sleeping
Sleeping
File size: 2,067 Bytes
2fb0cb4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
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]}") |