From 9fbaf8deb9d4a185163076e48d096016b4ceefd9 Mon Sep 17 00:00:00 2001 From: Albert Yosef Date: Fri, 12 Sep 2025 18:08:22 +0000 Subject: [PATCH] Refactor VectorDBClient to use settings for configuration; update tests for index creation and vector operations --- services/vector_db_service.py | 49 ++++++++----- tests/test_vector_db_service.py | 119 +++++++++----------------------- 2 files changed, 64 insertions(+), 104 deletions(-) diff --git a/services/vector_db_service.py b/services/vector_db_service.py index 0989847..b1c8ee5 100644 --- a/services/vector_db_service.py +++ b/services/vector_db_service.py @@ -1,39 +1,38 @@ + """ Service for storing and retrieving vectors using Pinecone (scaffold). """ # To use: pip install pinecone-client -import os from typing import List +from config.settings import settings class VectorDBClient: - def __init__(self, api_key: str, dimension: int = None, cloud: str = "aws", region: str = "us-east-1", environment: str = None): + def __init__(self, dimension: int = None, index_name: str = None): """ - Initialize Pinecone client for serverless (cloud/region) or legacy (environment). - :param api_key: Pinecone API key - :param dimension: Vector dimension (required for index ops) - :param cloud: Serverless cloud provider (default 'aws') - :param region: Serverless region (default 'us-east-1') - :param environment: Legacy environment (optional, unused for serverless) + Initialize Pinecone client for serverless (cloud/region). + Loads API key and config from settings. """ from pinecone import Pinecone - self.pc = Pinecone(api_key=api_key) # For serverless, environment is ignored - self.cloud = cloud - self.region = region - self.environment = environment + self.pc = Pinecone(api_key=settings.pinecone_api_key.get_secret_value()) + self.cloud = settings.pinecone_cloud + self.region = settings.pinecone_region + self.index_name = index_name or settings.pinecone_index_name self.index = None - self.dimension = dimension # Must be set for upsert/query; can be set at index creation + self.dimension = dimension or settings.embedding_dimension - def create_index(self, index_name: str, dimension: int): + def create_index(self, index_name: str = None, dimension: int = None): from pinecone import ServerlessSpec + idx_name = index_name or self.index_name + dim = dimension or self.dimension existing = [idx.name for idx in self.pc.list_indexes()] - if index_name not in existing: + if idx_name not in existing: self.pc.create_index( - name=index_name, - dimension=dimension, + name=idx_name, + dimension=dim, spec=ServerlessSpec(cloud=self.cloud, region=self.region), ) - self.index = self.pc.Index(index_name) + self.index = self.pc.Index(idx_name) def upsert_vectors(self, vectors: List[List[float]], ids: List[str]): # Input validation @@ -47,7 +46,6 @@ def upsert_vectors(self, vectors: List[List[float]], ids: List[str]): raise ValueError("vector dimensionality mismatch with index.") if not self.index: raise RuntimeError("Index is not initialized. Call create_index first.") - # Upsert format may vary by pinecone-client version. self.index.upsert(vectors=[(id, vec) for id, vec in zip(ids, vectors)]) def query(self, vector: List[float], top_k: int = 5): @@ -56,3 +54,16 @@ def query(self, vector: List[float], top_k: int = 5): if self.dimension is not None and len(vector) != self.dimension: raise ValueError("query vector dimensionality mismatch with index.") return self.index.query(vector=vector, top_k=top_k) + +# Interactive test block +if __name__ == "__main__": + vdb = VectorDBClient() + vdb.create_index() + print("Index created successfully!") + # Example upsert and query (uncomment to use): + # ids = ["id1", "id2"] + # vectors = [[0.0]*vdb.dimension, [1.0]*vdb.dimension] + # vdb.upsert_vectors(vectors, ids) + # print("Upserted vectors.") + # result = vdb.query([0.0]*vdb.dimension) + # print("Query result:", result) diff --git a/tests/test_vector_db_service.py b/tests/test_vector_db_service.py index e33142a..01d5ed9 100644 --- a/tests/test_vector_db_service.py +++ b/tests/test_vector_db_service.py @@ -1,97 +1,46 @@ -# Patch config.settings and env vars before any other imports -import sys -import types -import os -from pydantic import SecretStr -# Set required env vars and patch config.settings before any other imports -os.environ["INTERNAL_API_KEY"] = "test-key" - -class DummySettings: - pinecone_api_key = SecretStr("test-key") - pinecone_cloud = "aws" - pinecone_region = "us-east-1" - pinecone_index_name = "authormaton-core" - embedding_model = "test-model" - embedding_dimension = 16 - embed_batch_size = 64 - max_upload_mb = 25 - -dummy_config = types.ModuleType("config.settings") -dummy_config.settings = DummySettings() -sys.modules["config.settings"] = dummy_config - - -# Patch Pinecone client for tests to avoid real API calls -from unittest.mock import MagicMock -class MockIndex: - def __init__(self, dimension): - self.dimension = dimension - self.upserted = [] - def upsert(self, items, namespace=None): - self.upserted.extend(items) - def query(self, vector, top_k=8, namespace=None, filter=None): - return {'matches': [{'id': 'test', 'score': 0.99}]} - -mock_pc = MagicMock() -mock_pc.list_indexes.return_value = [] -mock_pc.create_index.return_value = None -mock_pc.Index.side_effect = lambda name: MockIndex(16) -mock_pc.describe_index.return_value = {"dimension": 16} - -mock_pinecone = MagicMock() -mock_pinecone.Pinecone = MagicMock(return_value=mock_pc) -mock_pinecone.ServerlessSpec = MagicMock() -sys.modules["pinecone"] = mock_pinecone - -# Now import everything else import pytest -from services.vector_db_service import VectorDBService +from services.vector_db_service import VectorDBClient -class MockIndex: +class DummyIndex: def __init__(self, dimension): self.dimension = dimension self.upserted = [] - def upsert(self, items, namespace=None): - self.upserted.extend(items) - def query(self, vector, top_k=8, namespace=None, filter=None): - return {'matches': [{'id': 'test', 'score': 0.99}]} - -def test_ensure_index_idempotent(monkeypatch): - svc = VectorDBService() - # Simulate: first call creates the index; second call sees it and skips creation. - mock_pc.create_index.reset_mock() - import types - mock_pc.list_indexes.side_effect = [ - [], - [types.SimpleNamespace(name=dummy_config.settings.pinecone_index_name)], - ] - - svc.ensure_index(svc.embedding_dimension) - svc.ensure_index(svc.embedding_dimension) - # Should only ever create the index once - assert mock_pc.create_index.call_count == 1 - assert svc.index.dimension == svc.embedding_dimension - -def test_upsert_dimension_guard(monkeypatch): - svc = VectorDBService() - monkeypatch.setattr(svc, 'index', MockIndex(svc.embedding_dimension)) - ids = ['a', 'b'] - vectors = [[0.0]*svc.embedding_dimension, [0.0]*svc.embedding_dimension] - metadata = [{}, {}] - count = svc.upsert(namespace='proj', ids=ids, vectors=vectors, metadata=metadata) - assert count == 2 + def upsert(self, vectors=None): + if vectors: + self.upserted.extend(vectors) + def query(self, vector, top_k=5): + return {'matches': [{'id': 'id1', 'score': 0.99}]} + +@pytest.fixture +def vdb(monkeypatch): + svc = VectorDBClient(dimension=8, index_name="test-index") + monkeypatch.setattr(svc, 'index', DummyIndex(svc.dimension)) + return svc + +def test_create_index(monkeypatch): + svc = VectorDBClient(dimension=8, index_name="test-index") + monkeypatch.setattr(svc.pc, 'list_indexes', lambda: []) + monkeypatch.setattr(svc.pc, 'create_index', lambda **kwargs: None) + monkeypatch.setattr(svc.pc, 'Index', lambda name: DummyIndex(8)) + svc.create_index() + assert svc.index.dimension == 8 + +def test_upsert_vectors(vdb): + ids = ["id1", "id2"] + vectors = [[0.0]*8, [1.0]*8] + vdb.upsert_vectors(vectors, ids) + assert len(vdb.index.upserted) == 2 # Wrong dimension with pytest.raises(ValueError): - svc.upsert(namespace='proj', ids=ids, vectors=[[0.0]*10, [0.0]*10], metadata=metadata) + vdb.upsert_vectors([[0.0]*5, [1.0]*5], ids) -def test_query_dimension_guard(monkeypatch): - svc = VectorDBService() - monkeypatch.setattr(svc, 'index', MockIndex(svc.embedding_dimension)) - vector = [0.0]*svc.embedding_dimension - matches = svc.query(namespace='proj', vector=vector) - assert matches[0]['id'] == 'test' +def test_query(vdb): + vector = [0.0]*8 + result = vdb.query(vector) + assert result['matches'][0]['id'] == 'id1' # Wrong dimension with pytest.raises(ValueError): - svc.query(namespace='proj', vector=[0.0]*10) + vdb.query([0.0]*5) +