File size: 1,897 Bytes
c8fb072 | 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 | from langgraph.graph import StateGraph, START,END
from src.sdlc.states.states import CoderState
from src.sdlc.nodes.code_orchestrator import CodeOrchestratorNode
from src.sdlc.nodes.code_generation_node import CodeGenerationNode
from src.sdlc.nodes.synthesizer_node import SynthesizerNode
from src.sdlc.nodes.code_reviewer_node import CodeReviewerNode
class SubGraphBuilder:
def __init__(self,model):
self.llm=model
self.graph_builder=StateGraph(CoderState)
def build_graph(self):
"""
Builds a subgraph for code generation.
"""
self.code_orchestrator_node=CodeOrchestratorNode(self.llm)
self.code_generation_node=CodeGenerationNode(self.llm)
self.synthesizer_node=SynthesizerNode(self.llm)
self.code_reviewer_node=CodeReviewerNode(self.llm)
self.graph_builder.add_node("code_orchestrator",self.code_orchestrator_node.process)
self.graph_builder.add_node("code_generation_node", self.code_generation_node.process)
self.graph_builder.add_node("code_synthesizer",self.synthesizer_node.process)
self.graph_builder.add_node("code_reviewer",self.code_reviewer_node.process)
self.graph_builder.add_edge(START,"code_orchestrator")
self.graph_builder.add_conditional_edges("code_orchestrator",
self.code_orchestrator_node.assign_workers,
path_map=["code_generation_node"])
self.graph_builder.add_edge("code_generation_node", "code_synthesizer")
self.graph_builder.add_edge("code_synthesizer", "code_reviewer")
self.graph_builder.add_edge("code_reviewer", END)
def setup_graph(self):
"""
Sets up the graph
"""
self.build_graph()
return self.graph_builder.compile()
|