utils/question_route.py

from loggers.logger import logger
from template.router_prompt import QUESTION_ROUTER

import sys
sys.dont_write_bytecode = True

query_router = QUESTION_ROUTER

def route_question(state):
    """
    Route question to web search or generation.
    Args:
        state (dict): The current graph state
    Returns:
        str: Next node to call
    """

    logger.info("INSIDE ROUTE QUESTION FUNCTION")

    question = state['question']
    
    output = query_router.invoke(
        {
            "question": question
        }
    )

    if output['choice'] == "web_search":
        logger.info("ROUTING QUERY TO WEB SEARCH CONTENT")
        return "websearch"

    elif output['choice'] == "generate":
        logger.info("ROUTING QUERY TO GENERATION CONTENT")
        return "generate"