helper/graphs.py

from langgraph.graph import END, StateGraph
from helper.model_load import GraphState
from utils.search_web_content import web_search
from utils.query_transform import transform_query
from utils.generate_content import generate
from utils.question_route import route_question

import sys
sys.dont_write_bytecode = True

builder = StateGraph(GraphState)

builder.add_node("web_search", web_search)

builder.add_node("transform_query", transform_query)

builder.add_node("generate", generate)

builder.set_conditional_entry_point(
    route_question,
    {
        "websearch": "transform_query",
        "generate": "generate"
    },
)

builder.add_edge("transform_query", "web_search")

builder.add_edge("web_search", "generate")

builder.add_edge("generate", END)

local_agent = builder.compile()