tests/test_graph_routing.py

from app.pipeline.graph import route_retrieve_result


def test_attempt_three_portfolio_empty_rewrites() -> None:
    state = {
        "retrieval_attempts": 3,
        "reranked_chunks": [],
        "query": "Could you tell me about his tech stack?",
    }
    assert route_retrieve_result(state) == "rewrite"


def test_attempt_three_portfolio_low_confidence_rewrites() -> None:
    state = {
        "retrieval_attempts": 3,
        "reranked_chunks": [{"text": "x", "metadata": {}}],
        "top_rerank_score": -2.0,
        "query": "Could you tell me about his tech stack?",
    }
    assert route_retrieve_result(state) == "rewrite"


def test_attempt_three_unrelated_low_confidence_generates() -> None:
    state = {
        "retrieval_attempts": 3,
        "reranked_chunks": [{"text": "x", "metadata": {}}],
        "top_rerank_score": -2.0,
        "query": "what is the weather in london",
    }
    assert route_retrieve_result(state) == "generate"