Spaces:
Running
Running
| from app.rag import retriever | |
| def test_transform_query_includes_original_and_dedupes(monkeypatch): | |
| monkeypatch.setattr( | |
| retriever, | |
| "_generate_query_variants", | |
| lambda _query: [ | |
| "How do taxes work?", | |
| "how do taxes work?", | |
| "How does healthcare work?", | |
| "healthcare overview", | |
| ], | |
| ) | |
| queries = retriever.transform_query("How do taxes and healthcare work?") | |
| assert queries == [ | |
| "How do taxes and healthcare work?", | |
| "How do taxes work?", | |
| "How does healthcare work?", | |
| "healthcare overview", | |
| ] | |
| def test_retrieve_fans_out_transformed_queries_and_merges_duplicates(monkeypatch): | |
| from unittest.mock import MagicMock | |
| searched_queries = [] | |
| mock_db = MagicMock() | |
| mock_db.__enter__.return_value = mock_db | |
| mock_query = MagicMock() | |
| mock_db.query.return_value = mock_query | |
| mock_query.filter.return_value.all.return_value = [("doc-1",)] | |
| monkeypatch.setattr("app.database.SessionLocal", lambda: mock_db) | |
| monkeypatch.setattr(retriever, "transform_query", lambda _query: ["taxes", "healthcare"]) | |
| monkeypatch.setattr(retriever, "embed_query", lambda query: f"embedding:{query}") | |
| monkeypatch.setattr(retriever, "get_reranker", lambda: None) | |
| def fake_query_chunks(query_embedding, user_id, document_id=None, document_ids=None, top_k=10): | |
| searched_queries.append(query_embedding) | |
| if query_embedding == "embedding:taxes": | |
| return [ | |
| { | |
| "id": "shared", | |
| "text": "Shared chunk", | |
| "filename": "policy.pdf", | |
| "page": 1, | |
| "score": 0.2, | |
| }, | |
| { | |
| "id": "taxes", | |
| "text": "Tax chunk", | |
| "filename": "policy.pdf", | |
| "page": 2, | |
| "score": 0.7, | |
| }, | |
| ] | |
| return [ | |
| { | |
| "id": "shared", | |
| "text": "Shared chunk", | |
| "filename": "policy.pdf", | |
| "page": 1, | |
| "score": 0.9, | |
| }, | |
| { | |
| "id": "healthcare", | |
| "text": "Healthcare chunk", | |
| "filename": "policy.pdf", | |
| "page": 3, | |
| "score": 0.8, | |
| }, | |
| ] | |
| monkeypatch.setattr(retriever, "query_chunks", fake_query_chunks) | |
| chunks = retriever.retrieve("How do taxes and healthcare work?", user_id="user-1") | |
| assert searched_queries == ["embedding:taxes", "embedding:healthcare"] | |
| assert [chunk["id"] for chunk in chunks] == ["shared", "taxes", "healthcare"] | |
| assert chunks[0]["score"] > 0 # RRF score, not raw similarity | |
| assert chunks[0]["id"] == "shared" # highest RRF score — appears in both query results | |
| assert chunks[0]["confidence"] == 100.0 | |
| def test_retrieve_includes_active_but_excludes_deleted(monkeypatch): | |
| from unittest.mock import MagicMock | |
| mock_db = MagicMock() | |
| mock_db.__enter__.return_value = mock_db | |
| mock_query = MagicMock() | |
| mock_db.query.return_value = mock_query | |
| mock_query.filter.return_value.all.return_value = [("doc-active",)] | |
| monkeypatch.setattr("app.database.SessionLocal", lambda: mock_db) | |
| monkeypatch.setattr(retriever, "transform_query", lambda _query: ["test"]) | |
| monkeypatch.setattr(retriever, "embed_query", lambda query: f"embedding:{query}") | |
| monkeypatch.setattr(retriever, "get_reranker", lambda: None) | |
| captured_ids = [] | |
| def fake_query_chunks(query_embedding, user_id, document_id=None, document_ids=None, top_k=10): | |
| captured_ids.append(document_ids) | |
| return [{"id": "chunk-1", "text": "sample text", "filename": "policy.pdf", "page": 1, "score": 0.5}] | |
| monkeypatch.setattr(retriever, "query_chunks", fake_query_chunks) | |
| retriever.retrieve("hello", user_id="user-1") | |
| assert captured_ids == [None] # retrieve() passes document_id=None, not document_ids | |