File size: 1,593 Bytes
028a330 8d56dc3 |
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 55 56 57 58 59 60 |
"""Graph builder for LangGraph workflow"""
from langgraph.graph import StateGraph, END
from src.state.rag_state import RAGState
from src.node.nodes import RAGNodes
class GraphBuilder:
"""Builds and manages the LangGraph workflow"""
def __init__(self, retriever, llm):
"""
Initialize graph builder
Args:
retriever: Document retriever instance
llm: Language model instance
"""
self.nodes = RAGNodes(retriever, llm)
self.graph = None
def build(self):
"""
Build the RAG workflow graph
Returns:
Compiled graph instance
"""
# Create state graph
builder = StateGraph(RAGState)
# Add nodes
builder.add_node("retriever", self.nodes.retrieve_docs)
builder.add_node("responder", self.nodes.generate_answer)
# Set entry point
builder.set_entry_point("retriever")
# Add edges
builder.add_edge("retriever", "responder")
builder.add_edge("responder", END)
# Compile graph
self.graph = builder.compile()
return self.graph
def run(self, question: str) -> dict:
"""
Run the RAG workflow
Args:
question: User question
Returns:
Final state with answer
"""
if self.graph is None:
self.build()
initial_state = RAGState(question=question)
return self.graph.invoke(initial_state) |