File size: 2,296 Bytes
79d0909 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """Graph builder - constructs the RAG workflow graph."""
from langgraph.graph import END, START, StateGraph
from src.config.settings import settings
from src.core import logger
from src.core.state import GraphState
from src.graph.constants import (
GENERATE,
GRADE_DOCUMENTS,
RETRIEVE,
WEB_SEARCH,
DECISION_USEFUL,
DECISION_NOT_USEFUL,
DECISION_NOT_SUPPORTED,
DECISION_WEBSEARCH,
DECISION_VECTORSTORE,
)
from src.graph.edges import decide_to_generate, grade_generation, route_question
from src.nodes import generate_node, grade_documents_node, retrieve_node, web_search_node
def build_graph() -> StateGraph:
"""
Build the RAG workflow graph.
Returns:
Compiled StateGraph ready for execution.
"""
graph = StateGraph(GraphState)
# Add nodes
graph.add_node(RETRIEVE, retrieve_node)
graph.add_node(GRADE_DOCUMENTS, grade_documents_node)
graph.add_node(GENERATE, generate_node)
graph.add_node(WEB_SEARCH, web_search_node)
# Set conditional entry point (router)
graph.set_conditional_entry_point(
path=route_question,
path_map={
DECISION_WEBSEARCH: WEB_SEARCH,
DECISION_VECTORSTORE: RETRIEVE,
},
)
# Add edges
graph.add_edge(RETRIEVE, GRADE_DOCUMENTS)
graph.add_edge(WEB_SEARCH, GENERATE)
# Conditional edge: grade documents -> generate or web search
graph.add_conditional_edges(
source=GRADE_DOCUMENTS,
path=decide_to_generate,
path_map={
WEB_SEARCH: WEB_SEARCH,
GENERATE: GENERATE,
},
)
# Conditional edge: generate -> end, retry, or web search
graph.add_conditional_edges(
source=GENERATE,
path=grade_generation,
path_map={
DECISION_USEFUL: END,
DECISION_NOT_USEFUL: GENERATE,
DECISION_NOT_SUPPORTED: WEB_SEARCH,
},
)
return graph.compile()
# Pre-built graph instance
rag_app = build_graph()
def save_graph_visualization(output_path: str | None = None) -> None:
"""Save the graph visualization to a PNG file."""
path = output_path or settings.GRAPH_OUTPUT_PATH
rag_app.get_graph().draw_mermaid_png(output_file_path=path)
logger.info(f"Graph saved to: {path}")
|