INNER CODE UNIT · Python

route_to_retriever

mfmezger/conversational-agent-langchain · src/agent/backend/graph.py:51

    def route_to_retriever(
        self,
        state: AgentState,
    ) -> Literal["retriever", "retriever_with_chat_history"]:
        """Route to the appropriate retriever based on the state."""
        if len(state["messages"]) == 1:
            return "retriever"
        else:
            return "retriever_with_chat_history"

    def build_graph(self) -> StateGraph:
        """Build the graph for the agent."""
        workflow = StateGraph(state_schema=AgentState)

        # define nodes
        workflow.add_node("retriever", functools.partial(retrieve_documents, cfg=self.cfg))
        workflow.add_node("retriever_with_chat_history", functools.partial(retrieve_documents_with_chat_history, cfg=self.cfg, llm=self.llm))
        workflow.add_node("reranker", functools.partial(rerank_documents, reranker=self.reranker))

View source record →