Spaces:
Sleeping
Sleeping
File size: 3,760 Bytes
2e818da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """Outcome-classification tests for ChromaDBClient.query_observed (Task 3, Step 4).
These prove the core reason retrieval failures used to be invisible: a
genuinely-empty successful search and a failed search that fell back to ``[]``
now carry distinct statuses, while ``query()`` stays byte-for-byte compatible.
"""
from __future__ import annotations
import os
os.environ.setdefault("CEREBRAS_API_KEY", "test-key")
import chromadb
import pytest
class _FakeEmbedder:
def embed(self, texts: list[str]) -> list[list[float]]:
return [[0.1, 0.0] for _ in texts]
@pytest.fixture
def db(monkeypatch):
# See test_rag_behavior_unchanged.isolated_db for why this cache-clear is
# required: chromadb.EphemeralClient() shares one process-wide backing
# store, so collection names can otherwise leak data across tests/files.
from chromadb.api.client import SharedSystemClient
SharedSystemClient.clear_system_cache()
monkeypatch.setattr(
chromadb, "PersistentClient", lambda *a, **k: chromadb.EphemeralClient()
)
from app.rag.chromadb_client import ChromaDBClient
return ChromaDBClient(embedder=_FakeEmbedder())
def _seed(db, project_id="proj-1"):
db.upsert(
"library",
["alpha doc", "beta doc"],
[
{"source": "paper-a", "document_id": "doc-a", "project_id": project_id},
{"source": "paper-b", "document_id": "doc-b", "project_id": project_id},
],
["a", "b"],
)
def test_successful_retrieval_outcome(db):
_seed(db)
q = db.embedder.embed(["x"])[0]
outcome = db.query_observed("library", q, n_results=2, where={"project_id": "proj-1"})
assert outcome.status == "success"
assert outcome.retrieval_error is False
assert outcome.empty_result is False
assert len(outcome.candidates) == 2
def test_legitimate_empty_is_success_empty_not_error(db):
_seed(db)
q = db.embedder.embed(["x"])[0]
outcome = db.query_observed(
"library", q, n_results=2, where={"project_id": "no-such-project"}
)
assert outcome.status == "success_empty"
assert outcome.empty_result is True
# An empty *successful* search is NOT an error.
assert outcome.retrieval_error is False
assert outcome.candidates == []
def test_missing_collection_is_error_fallback(db):
q = db.embedder.embed(["x"])[0]
outcome = db.query_observed("never-created", q, n_results=2)
assert outcome.status == "error_fallback"
assert outcome.retrieval_error is True
assert outcome.error_type == "collection_unavailable"
assert outcome.candidates == []
def test_vector_search_failure_is_error_fallback(db, monkeypatch):
_seed(db)
q = db.embedder.embed(["x"])[0]
real_get = db._client.get_collection
class _BoomCollection:
def count(self):
return 2
def query(self, **kwargs):
raise RuntimeError("index corrupted")
monkeypatch.setattr(db._client, "get_collection", lambda name: _BoomCollection())
outcome = db.query_observed("library", q, n_results=2)
assert outcome.status == "error_fallback"
assert outcome.retrieval_error is True
assert outcome.error_type == "vector_search_failed"
assert outcome.candidates == []
def test_query_delegates_and_stays_list_compatible(db):
"""The public query() must remain a plain list[dict] for every caller."""
_seed(db)
q = db.embedder.embed(["x"])[0]
result = db.query("library", q, n_results=2, where={"project_id": "proj-1"})
assert isinstance(result, list)
assert all(isinstance(r, dict) and "text" in r for r in result)
# Missing collection still yields a plain [], never raises.
assert db.query("never-created", q, n_results=2) == []
|