s-siddhesh/LangAgentX

1import logging
2from typing import List, Dict, Any
3from typing_extensions import TypedDict
4from router import question_router
5
6logging.basicConfig(level=logging.INFO)
7logger = logging.getLogger(__name__)
8
9class GraphState(TypedDict):
10    question: str
11    generation: str
12    documents: List[str]
13
14def route_question(state: Dict[str, Any]) -> str:
15    logger.info("---ROUTE QUESTION---")
16    question = state.get("question", "")
17    if not question:
18        logger.warning("No question provided in state.")
19        return "unknown"
20    try:
21        source = question_router.invoke({"question": question})
22        if source.datasource == "wiki_search":
23            logger.info("---ROUTE TO WIKI SEARCH---")
24            return "wiki_search"
25        elif source.datasource == "vectorstore":
26            logger.info("---ROUTE TO VECTORSTORE---")
27            return "vectorstore"
28        else:
29            logger.warning(f"Unknown datasource: {source.datasource}")
30            return "unknown"
31    except Exception as e:
32        logger.error(f"Error in routing question: {e}")
33        return "unknown"