mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 20:13:04 +00:00
261 lines
9.6 KiB
Python
261 lines
9.6 KiB
Python
"""QdrantStore tests, run against a real in-process Qdrant.
|
|
|
|
qdrant-client is a first-party dependency now that the store no longer goes
|
|
through langchain, so these exercise the real client in ``:memory:`` mode
|
|
rather than asserting against mocks.
|
|
"""
|
|
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
|
|
class _FakeEmbeddings:
|
|
"""Deterministic 3-dim embeddings: distinct texts get distinct directions."""
|
|
|
|
dimension = 3
|
|
|
|
_VECTORS = {
|
|
"paris": [1.0, 0.0, 0.0],
|
|
"database": [0.0, 1.0, 0.0],
|
|
"celery": [0.0, 0.0, 1.0],
|
|
}
|
|
|
|
def _vector(self, text):
|
|
lowered = (text or "").lower()
|
|
for keyword, vector in self._VECTORS.items():
|
|
if keyword in lowered:
|
|
return vector
|
|
return [0.577, 0.577, 0.577]
|
|
|
|
def embed_query(self, query):
|
|
return self._vector(query)
|
|
|
|
def embed_documents(self, documents):
|
|
return [self._vector(d) for d in documents]
|
|
|
|
|
|
def _settings(mock_settings, collection="test_collection"):
|
|
mock_settings.EMBEDDINGS_NAME = "test_model"
|
|
mock_settings.QDRANT_COLLECTION_NAME = collection
|
|
mock_settings.QDRANT_LOCATION = ":memory:"
|
|
mock_settings.QDRANT_DISTANCE_FUNC = "Cosine"
|
|
mock_settings.QDRANT_PREFER_GRPC = False
|
|
mock_settings.QDRANT_GRPC_PORT = 6334
|
|
for unset in (
|
|
"QDRANT_URL", "QDRANT_HOST", "QDRANT_PORT", "QDRANT_HTTPS",
|
|
"QDRANT_API_KEY", "QDRANT_PREFIX", "QDRANT_TIMEOUT", "QDRANT_PATH",
|
|
):
|
|
setattr(mock_settings, unset, None)
|
|
|
|
|
|
@pytest.fixture
|
|
def store():
|
|
from docsgpt.vectorstore.qdrant import QdrantStore
|
|
|
|
with patch(
|
|
"docsgpt.vectorstore.base.BaseVectorStore._get_embeddings",
|
|
return_value=_FakeEmbeddings(),
|
|
), patch("docsgpt.vectorstore.qdrant.settings") as mock_settings:
|
|
_settings(mock_settings)
|
|
yield QdrantStore(source_id="src-A", embeddings_key="k")
|
|
|
|
|
|
@pytest.fixture
|
|
def populated(store):
|
|
store.add_texts(
|
|
["The capital of France is Paris.",
|
|
"Postgres is a relational database.",
|
|
"Celery runs background tasks."],
|
|
[{"source": "geo.txt"}, {"source": "db.txt"}, {"source": "queue.txt"}],
|
|
)
|
|
return store
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestQdrantStore:
|
|
def test_source_id_strips_index_prefix(self, store):
|
|
assert store._source_id == "src-A"
|
|
|
|
def test_add_texts_returns_one_id_per_text(self, store):
|
|
ids = store.add_texts(["a", "b"], [{}, {}])
|
|
assert len(ids) == 2 and len(set(ids)) == 2
|
|
|
|
def test_add_texts_empty_is_noop(self, store):
|
|
assert store.add_texts([], []) == []
|
|
|
|
def test_search_ranks_by_similarity(self, populated):
|
|
hits = populated.search("Tell me about Paris", k=2)
|
|
assert "Paris" in str(hits[0])
|
|
assert hits[0].metadata["source"] == "geo.txt"
|
|
|
|
def test_search_with_scores_reports_cosine(self, populated):
|
|
scored = populated.search_with_scores("Tell me about Paris", k=3)
|
|
assert populated.score_kind == "cosine_similarity"
|
|
assert scored[0][1] == pytest.approx(1.0, abs=1e-3)
|
|
# Scores must come back in descending rank order.
|
|
assert [s for _, s in scored] == sorted(
|
|
(s for _, s in scored), reverse=True
|
|
)
|
|
|
|
def test_search_honours_score_threshold(self, populated):
|
|
assert populated.search_with_scores("Paris", k=3, score_threshold=0.99)
|
|
assert not populated.search_with_scores("Paris", k=3, score_threshold=1.01)
|
|
|
|
def test_add_texts_stamps_source_id(self, populated):
|
|
assert all(c["metadata"]["source_id"] == "src-A" for c in populated.get_chunks())
|
|
|
|
def test_get_chunks_returns_all(self, populated):
|
|
chunks = populated.get_chunks()
|
|
assert len(chunks) == 3
|
|
assert {c["metadata"]["source"] for c in chunks} == {
|
|
"geo.txt", "db.txt", "queue.txt"
|
|
}
|
|
|
|
def test_add_and_delete_chunk(self, populated):
|
|
chunk_id = populated.add_chunk("Redis caches things.", {"source": "cache.txt"})
|
|
assert len(populated.get_chunks()) == 4
|
|
assert populated.delete_chunk(chunk_id) is True
|
|
assert len(populated.get_chunks()) == 3
|
|
|
|
def test_delete_chunk_returns_false_when_client_raises(self, populated):
|
|
with patch.object(
|
|
populated._client, "delete", side_effect=RuntimeError("qdrant down")
|
|
):
|
|
assert populated.delete_chunk("some-id") is False
|
|
|
|
def test_get_chunks_returns_empty_when_client_raises(self, populated):
|
|
with patch.object(
|
|
populated._client, "scroll", side_effect=RuntimeError("qdrant down")
|
|
):
|
|
assert populated.get_chunks() == []
|
|
|
|
def test_delete_chunks_by_source_path(self, populated):
|
|
assert populated.delete_chunks_by_source_path("db.txt") == 1
|
|
assert len(populated.get_chunks()) == 2
|
|
|
|
def test_delete_index_removes_only_this_source(self, populated):
|
|
from docsgpt.vectorstore.qdrant import QdrantStore
|
|
|
|
with patch(
|
|
"docsgpt.vectorstore.base.BaseVectorStore._get_embeddings",
|
|
return_value=_FakeEmbeddings(),
|
|
), patch("docsgpt.vectorstore.qdrant.settings") as mock_settings:
|
|
_settings(mock_settings)
|
|
other = QdrantStore(source_id="src-B", embeddings_key="k")
|
|
# Share the in-memory backend so both sources live in one collection.
|
|
other._client = populated._client
|
|
other.add_texts(["Another source entirely."], [{"source": "other.txt"}])
|
|
|
|
populated.delete_index()
|
|
|
|
assert populated.get_chunks() == []
|
|
assert len(other.get_chunks()) == 1
|
|
|
|
def test_save_local_is_noop(self, store):
|
|
assert store.save_local() is None
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestQdrantClientKwargs:
|
|
def test_unset_settings_are_omitted(self):
|
|
from docsgpt.vectorstore.qdrant import QdrantStore
|
|
|
|
with patch("docsgpt.vectorstore.qdrant.settings") as mock_settings:
|
|
_settings(mock_settings)
|
|
kwargs = QdrantStore._client_kwargs()
|
|
# location/url/path are mutually exclusive in qdrant-client, so only
|
|
# the configured one may be forwarded.
|
|
assert kwargs["location"] == ":memory:"
|
|
assert "url" not in kwargs and "path" not in kwargs and "host" not in kwargs
|
|
|
|
def test_configured_settings_are_forwarded(self):
|
|
from docsgpt.vectorstore.qdrant import QdrantStore
|
|
|
|
with patch("docsgpt.vectorstore.qdrant.settings") as mock_settings:
|
|
_settings(mock_settings)
|
|
mock_settings.QDRANT_LOCATION = None
|
|
mock_settings.QDRANT_URL = "http://qdrant:6333"
|
|
mock_settings.QDRANT_API_KEY = "secret"
|
|
kwargs = QdrantStore._client_kwargs()
|
|
assert kwargs["url"] == "http://qdrant:6333"
|
|
assert kwargs["api_key"] == "secret"
|
|
assert "location" not in kwargs
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestQdrantUpdateChunk:
|
|
def test_keeps_id_and_position(self, populated):
|
|
before = populated.get_chunks()
|
|
target = before[1]["doc_id"]
|
|
|
|
returned = populated.update_chunk(target, "Redis caches things.", {"source": "cache.txt"})
|
|
|
|
assert returned == target
|
|
after = populated.get_chunks()
|
|
assert [c["doc_id"] for c in after] == [c["doc_id"] for c in before]
|
|
assert after[1]["text"] == "Redis caches things."
|
|
assert after[1]["metadata"] == {"source": "cache.txt", "source_id": "src-A"}
|
|
|
|
def test_replaces_the_vector(self, populated):
|
|
target = next(
|
|
c["doc_id"] for c in populated.get_chunks() if c["metadata"]["source"] == "geo.txt"
|
|
)
|
|
|
|
populated.update_chunk(target, "Celery workers are busy.", {"source": "geo.txt"})
|
|
|
|
hits = populated.search_with_scores("celery", k=3)
|
|
assert sorted(doc.page_content for doc, score in hits if score > 0.99) == [
|
|
"Celery runs background tasks.",
|
|
"Celery workers are busy.",
|
|
]
|
|
assert not any("Paris" in doc.page_content for doc, _ in hits)
|
|
|
|
def test_payload_shape_matches_add_texts(self, populated):
|
|
target = populated.get_chunks()[0]["doc_id"]
|
|
|
|
populated.update_chunk(target, "Redis caches things.", {"source": "cache.txt"})
|
|
|
|
record = populated._client.retrieve(
|
|
populated._collection, ids=[target], with_payload=True
|
|
)[0]
|
|
assert record.payload == {
|
|
"page_content": "Redis caches things.",
|
|
"metadata": {"source": "cache.txt", "source_id": "src-A"},
|
|
}
|
|
|
|
def test_embedding_failure_leaves_the_point_untouched(self, populated):
|
|
before = populated.get_chunks()
|
|
target = before[0]["doc_id"]
|
|
|
|
with patch.object(
|
|
populated._embeddings, "embed_documents", side_effect=RuntimeError("embed down")
|
|
):
|
|
with pytest.raises(RuntimeError):
|
|
populated.update_chunk(target, "Redis caches things.", {})
|
|
|
|
assert populated.get_chunks() == before
|
|
|
|
def test_unknown_id_raises(self, populated):
|
|
import uuid
|
|
|
|
with pytest.raises(KeyError):
|
|
populated.update_chunk(str(uuid.uuid4()), "text", {})
|
|
assert len(populated.get_chunks()) == 3
|
|
|
|
def test_other_sources_point_is_refused(self, populated):
|
|
from docsgpt.vectorstore.qdrant import QdrantStore
|
|
|
|
with patch(
|
|
"docsgpt.vectorstore.base.BaseVectorStore._get_embeddings",
|
|
return_value=_FakeEmbeddings(),
|
|
), patch("docsgpt.vectorstore.qdrant.settings") as mock_settings:
|
|
_settings(mock_settings)
|
|
other = QdrantStore(source_id="src-B", embeddings_key="k")
|
|
other._client = populated._client
|
|
target = populated.get_chunks()[0]["doc_id"]
|
|
|
|
with pytest.raises(KeyError):
|
|
other.update_chunk(target, "hijack", {})
|
|
assert populated.get_chunks()[0]["text"] != "hijack"
|