feat: enable weaviate usage in codegate

Closes: #63

Author: yrobla

src/codegate/storage/__init__.py

@@ -0,0 +1,3 @@
+from codegate.storage.storage_engine import StorageEngine
+
+__all__ = [StorageEngine]

src/codegate/storage/storage_engine.py

@@ -0,0 +1,76 @@
+from codegate.codegate_logging import setup_logging
+from codegate.inference.inference_engine import LlamaCppInferenceEngine
+from weaviate.classes.config import DataType, Property
+from weaviate.classes.query import MetadataQuery
+import weaviate
+
+
+schema_config = [
+    {
+        "name": "Package",
+        "properties": [
+            {"name": "name", "data_type": DataType.TEXT},
+            {"name": "type", "data_type": DataType.TEXT},
+            {"name": "status", "data_type": DataType.TEXT},
+            {"name": "description", "data_type": DataType.TEXT},
+        ]
+    },
+]
+
+
+class StorageEngine:
+    def __init__(self, data_path='./weaviate_data'):
+        self.client = weaviate.WeaviateClient(
+            embedded_options=weaviate.EmbeddedOptions(
+                persistence_data_path=data_path
+            ),
+        )
+        self.__logger = setup_logging()
+        self.inference_engine = LlamaCppInferenceEngine()
+        self.model_path = "./models/all-minilm-L6-v2-q5_k_m.gguf"
+        self.schema_config = schema_config
+        self.connect()
+        self.setup_schema()
+
+    def connect(self):
+        self.client.connect()
+        if self.client.is_ready():
+            self.__logger.info("Weaviate connection established and client is ready.")
+        else:
+            raise Exception("Weaviate client is not ready.")
+
+    def setup_schema(self):
+        for class_config in self.schema_config:
+            if not self.client.collections.exists(class_config['name']):
+                self.client.collections.create(class_config['name'], properties=class_config['properties'])
+                self.__logger.info(f"Weaviate schema for class {class_config['name']} setup complete.")
+
+    async def search(self, query, limit=5, distance=0.1):
+        """
+        Search the 'Package' collection based on a query string.
+
+        Args:
+            query (str): The text query for which to search.
+            limit (int): The number of results to return.
+
+        Returns:
+            list: A list of matching results with their properties and distances.
+        """
+        # Generate the vector for the query
+        query_vector = await self.inference_engine.embed(self.model_path, [query])
+
+        # Perform the vector search
+        try:
+            collection = self.client.collections.get("Package")
+            response = collection.query.near_vector(query_vector, limit=limit, distance=distance, return_metadata=MetadataQuery(distance=True))
+            if not response:
+                return []
+            return response.objects
+
+        except Exception as e:
+            self.__logger.error(f"Error during search: {str(e)}")
+            return []
+
+    def close(self):
+        self.client.close()
+        self.__logger.info("Connection closed.")

tests/test_storage.py

@@ -0,0 +1,56 @@
+import pytest
+from unittest.mock import Mock, AsyncMock
+from codegate.storage.storage_engine import StorageEngine  # Adjust the import according to your project structure
+
+
+@pytest.fixture
+def mock_client():
+    client = Mock()
+    client.connect = Mock()
+    client.is_ready = Mock(return_value=True)
+    client.schema.contains = Mock(return_value=False)
+    client.schema.create_class = Mock()
+    client.collections.get = Mock()
+    client.close = Mock()
+    return client
+
+
+@pytest.fixture
+def mock_logger():
+    logger = Mock()
+    return logger
+
+
+@pytest.fixture
+def mock_inference_engine():
+    inference_engine = AsyncMock()
+    inference_engine.embed = AsyncMock(return_value=[0.1, 0.2, 0.3])  # Adjust based on expected vector dimensions
+    return inference_engine
+
+
+@pytest.fixture
+def storage_engine(mock_client, mock_logger, mock_inference_engine):
+    engine = StorageEngine(data_path='./weaviate_data')
+    engine.client = mock_client
+    engine.__logger = mock_logger
+    engine.inference_engine = mock_inference_engine
+    return engine
+
+
+def test_connect(storage_engine, mock_client):
+    storage_engine.connect()
+    mock_client.connect.assert_called_once()
+    mock_client.is_ready.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_search(storage_engine, mock_client):
+    query = "test query"
+    results = await storage_engine.search(query)
+    storage_engine.inference_engine.embed.assert_called_once_with("./models/all-minilm-L6-v2-q5_k_m.gguf", [query])
+    assert results is not None  # Further asserts can be based on your application logic
+
+
+def test_close(storage_engine, mock_client):
+    storage_engine.close()
+    mock_client.close.assert_called_once()