From f75b4907a5c22161c1f1a75aa8d5da95df9a99eb Mon Sep 17 00:00:00 2001
From: vbarda <vadym@langchain.dev>
Date: Wed, 16 Oct 2024 19:42:24 -0400
Subject: [PATCH] disable routing

---
 backend/retrieval_graph/graph.py | 43 ++++++++++++--------------------
 backend/tests/evals/test_e2e.py  |  7 +-----
 2 files changed, 17 insertions(+), 33 deletions(-)

diff --git a/backend/retrieval_graph/graph.py b/backend/retrieval_graph/graph.py
index 10aa21fd..6108dc18 100644
--- a/backend/retrieval_graph/graph.py
+++ b/backend/retrieval_graph/graph.py
@@ -203,7 +203,9 @@ async def respond(
     """
     configuration = AgentConfiguration.from_runnable_config(config)
     model = load_chat_model(configuration.response_model)
-    context = format_docs(state.documents)
+    # TODO: add a re-ranker here
+    top_k = 10
+    context = format_docs(state.documents[:top_k])
     prompt = configuration.response_system_prompt.format(context=context)
     messages = [{"role": "system", "content": prompt}] + state.messages
     response = await model.ainvoke(messages)
@@ -213,29 +215,16 @@ async def respond(
 # Define the graph
 
 
-def make_graph(*, input_schema: Type[Any]):
-    builder = StateGraph(
-        AgentState, input=input_schema, config_schema=AgentConfiguration
-    )
-    builder.add_node(analyze_and_route_query)
-    builder.add_node(ask_for_more_info)
-    builder.add_node(respond_to_general_query)
-    builder.add_node(conduct_research)
-    builder.add_node(create_research_plan)
-    builder.add_node(respond)
-
-    builder.add_edge(START, "analyze_and_route_query")
-    builder.add_conditional_edges("analyze_and_route_query", route_query)
-    builder.add_edge("create_research_plan", "conduct_research")
-    builder.add_conditional_edges("conduct_research", check_finished)
-    builder.add_edge("ask_for_more_info", END)
-    builder.add_edge("respond_to_general_query", END)
-    builder.add_edge("respond", END)
-
-    # Compile into a graph object that you can invoke and deploy.
-    graph = builder.compile()
-    graph.name = "RetrievalGraph"
-    return graph
-
-
-graph = make_graph(input_schema=InputState)
+builder = StateGraph(AgentState, input=InputState, config_schema=AgentConfiguration)
+builder.add_node(create_research_plan)
+builder.add_node(conduct_research)
+builder.add_node(respond)
+
+builder.add_edge(START, "create_research_plan")
+builder.add_edge("create_research_plan", "conduct_research")
+builder.add_conditional_edges("conduct_research", check_finished)
+builder.add_edge("respond", END)
+
+# Compile into a graph object that you can invoke and deploy.
+graph = builder.compile()
+graph.name = "RetrievalGraph"
diff --git a/backend/tests/evals/test_e2e.py b/backend/tests/evals/test_e2e.py
index c2517412..2f26cfe5 100644
--- a/backend/tests/evals/test_e2e.py
+++ b/backend/tests/evals/test_e2e.py
@@ -10,8 +10,7 @@
 from langsmith.evaluation import EvaluationResults, aevaluate
 from langsmith.schemas import Example, Run
 
-from backend.retrieval_graph.graph import make_graph
-from backend.retrieval_graph.state import AgentState, Router
+from backend.retrieval_graph.graph import graph
 from backend.utils import format_docs
 
 DATASET_NAME = "chat-langchain-qa"
@@ -143,15 +142,11 @@ def evaluate_qa_context(run: Run, example: Example) -> dict:
 
 # Run evaluation
 
-# TODO: this is a hack to allow for skipping the router for testing. Add testing for individual components.
-graph = make_graph(input_schema=AgentState)
-
 
 async def run_graph(inputs: dict[str, Any]) -> dict[str, Any]:
     results = await graph.ainvoke(
         {
             "messages": [("human", inputs["question"])],
-            "router": Router(type="langchain", logic="The question is about LangChain"),
         }
     )
     return results