Skip to content

Commit

Permalink
disable routing
Browse files Browse the repository at this point in the history
  • Loading branch information
vbarda committed Oct 16, 2024
1 parent 8c547f5 commit f75b490
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 33 deletions.
43 changes: 16 additions & 27 deletions backend/retrieval_graph/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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"
7 changes: 1 addition & 6 deletions backend/tests/evals/test_e2e.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit f75b490

Please sign in to comment.