mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/lib/adapters/generic_rag_adapter/vector_stores/weaviate_store.py

@@ -0,0 +1,418 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from typing import Any
+
+from mantisdk.algorithm.gepa.lib.adapters.generic_rag_adapter.vector_store_interface import VectorStoreInterface
+
+
+class WeaviateVectorStore(VectorStoreInterface):
+    """
+    Weaviate implementation of the VectorStoreInterface.
+
+    Weaviate is a cloud-native, modular, real-time vector database
+    with powerful search capabilities including hybrid search.
+    Supports both Weaviate Cloud Services and self-hosted deployments.
+    """
+
+    def __init__(self, client, collection_name: str, embedding_function=None):
+        """
+        Initialize WeaviateVectorStore.
+
+        Args:
+            client: Weaviate client instance (WeaviateClient from weaviate-client)
+            collection_name: Name of the collection/class to use
+            embedding_function: Optional function to compute embeddings for queries
+        """
+        import importlib.util
+
+        if importlib.util.find_spec("weaviate") is None:
+            raise ImportError(
+                "Weaviate client is required for WeaviateVectorStore. Install with: pip install litellm weaviate-client"
+            )
+
+        import weaviate.classes as wvc
+
+        self.client = client
+        self.collection_name = collection_name
+        self.wvc = wvc
+        self.embedding_function = embedding_function
+
+        # Get the collection
+        try:
+            self.collection = self.client.collections.get(collection_name)
+        except Exception as e:
+            raise ValueError(
+                f"Collection '{collection_name}' not found. Please create the collection first. Error: {e!s}"
+            ) from e
+
+    def similarity_search(
+        self,
+        query: str,
+        k: int = 5,
+        filters: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Search for documents similar to the query text using embeddings."""
+        if self.embedding_function is None:
+            raise ValueError("No embedding function provided for similarity search")
+
+        # Compute embeddings for the query
+        try:
+            query_vector = self.embedding_function(query)
+            if hasattr(query_vector, "tolist"):
+                query_vector = query_vector.tolist()
+        except Exception as e:
+            raise RuntimeError(f"Failed to compute embeddings for query: {e!s}") from e
+
+        # Use vector search with computed embeddings
+        return self.vector_search(query_vector, k, filters)
+
+    def vector_search(
+        self,
+        query_vector: list[float],
+        k: int = 5,
+        filters: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Search using a pre-computed query vector."""
+        weaviate_filters = self._convert_filters(filters) if filters else None
+
+        try:
+            # Execute query
+            if weaviate_filters:
+                response = self.collection.query.near_vector(near_vector=query_vector, limit=k).where(weaviate_filters)
+            else:
+                response = self.collection.query.near_vector(near_vector=query_vector, limit=k)
+
+            # Handle GenerativeReturn object - access .objects attribute
+            if hasattr(response, "objects"):
+                results = response.objects
+            else:
+                results = response
+
+            return self._format_results(results)
+
+        except Exception as e:
+            raise RuntimeError(f"Weaviate vector search failed: {e!s}") from e
+
+    def hybrid_search(
+        self,
+        query: str,
+        k: int = 5,
+        alpha: float = 0.5,
+        filters: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        Hybrid semantic + keyword search using Weaviate's native hybrid search.
+
+        Args:
+            query: Text query to search for
+            k: Number of documents to return
+            alpha: Weight for semantic vs keyword search (0.0 = pure keyword, 1.0 = pure semantic)
+            filters: Optional metadata filters
+        """
+        weaviate_filters = self._convert_filters(filters) if filters else None
+
+        try:
+            results = self.collection.query.hybrid(
+                query=query,
+                alpha=alpha,
+                limit=k,
+                where=weaviate_filters,
+                return_metadata=self.wvc.query.MetadataQuery(score=True, explain_score=True),
+                return_properties=["content", "*"],
+            )
+
+            return self._format_results(results)
+
+        except Exception as e:
+            raise RuntimeError(f"Weaviate hybrid search failed: {e!s}") from e
+
+    def get_collection_info(self) -> dict[str, Any]:
+        """Get metadata about the Weaviate collection."""
+        try:
+            # Get collection configuration
+            config = self.collection.config.get()
+
+            # Count objects in collection
+            count_result = self.collection.aggregate.over_all(total_count=True)
+            total_count = count_result.total_count
+
+            # Try to determine vector dimensions
+            dimension = 0
+            if hasattr(config, "vector_config") and config.vector_config:
+                # For collections with vector config
+                if hasattr(config.vector_config, "vector_index_config"):
+                    vector_index = config.vector_index_config
+                    if hasattr(vector_index, "distance"):
+                        # This indicates vectors are configured
+                        # Try to get dimension from a sample object
+                        try:
+                            sample = self.collection.query.fetch_objects(limit=1, include_vector=True)
+                            if sample.objects and hasattr(sample.objects[0], "vector") and sample.objects[0].vector:
+                                # Handle both named vectors and default vector
+                                vector_data = sample.objects[0].vector
+                                if isinstance(vector_data, dict):
+                                    # Named vectors case
+                                    for _vector_name, vector_values in vector_data.items():
+                                        if vector_values:
+                                            dimension = len(vector_values)
+                                            break
+                                elif isinstance(vector_data, list):
+                                    # Default vector case
+                                    dimension = len(vector_data)
+                        except Exception:
+                            # If we can't determine dimension from sample, that's ok
+                            pass
+
+            # Determine if hybrid search is supported
+            supports_hybrid = True  # Weaviate generally supports hybrid search
+            try:
+                # Test if hybrid search works by doing a minimal query
+                self.collection.query.hybrid(query="test", limit=1)
+                supports_hybrid = True
+            except Exception:
+                supports_hybrid = False
+
+            return {
+                "name": self.collection_name,
+                "document_count": total_count,
+                "dimension": dimension,
+                "vector_store_type": "weaviate",
+                "supports_hybrid_search": supports_hybrid,
+                "vectorizer": getattr(config, "vectorizer_config", None),
+                "properties": [prop.name for prop in getattr(config, "properties", [])],
+            }
+
+        except Exception as e:
+            # Fallback info if detailed info fails
+            return {
+                "name": self.collection_name,
+                "document_count": 0,
+                "dimension": 0,
+                "vector_store_type": "weaviate",
+                "error": str(e),
+            }
+
+    def supports_hybrid_search(self) -> bool:
+        """Weaviate supports hybrid search."""
+        try:
+            # Test if hybrid search works
+            self.collection.query.hybrid(query="test", limit=1)
+            return True
+        except Exception:
+            return False
+
+    def supports_metadata_filtering(self) -> bool:
+        """Weaviate supports metadata filtering with where clauses."""
+        return True
+
+    def _convert_filters(self, filters: dict[str, Any]) -> Any:
+        """
+        Convert generic filters to Weaviate where clause format.
+
+        Generic format: {"key": "value", "key2": {"$gt": 5}}
+        Weaviate format uses Filter class
+        """
+        if not filters:
+            return None
+
+        try:
+            from weaviate.collections.classes.filters import Filter
+        except ImportError:
+            # Fallback if Filter class not available
+            return None
+
+        filter_conditions = []
+
+        for key, value in filters.items():
+            if isinstance(value, dict):
+                # Handle operator-based filters
+                for op, op_value in value.items():
+                    if op == "$eq" or op == "equal":
+                        filter_conditions.append(Filter.by_property(key).equal(op_value))
+                    elif op == "$ne" or op == "not_equal":
+                        filter_conditions.append(Filter.by_property(key).not_equal(op_value))
+                    elif op == "$gt" or op == "greater_than":
+                        filter_conditions.append(Filter.by_property(key).greater_than(op_value))
+                    elif op == "$gte" or op == "greater_equal":
+                        filter_conditions.append(Filter.by_property(key).greater_or_equal(op_value))
+                    elif op == "$lt" or op == "less_than":
+                        filter_conditions.append(Filter.by_property(key).less_than(op_value))
+                    elif op == "$lte" or op == "less_equal":
+                        filter_conditions.append(Filter.by_property(key).less_or_equal(op_value))
+                    elif op == "$in" or op == "contains_any":
+                        if isinstance(op_value, list):
+                            filter_conditions.append(Filter.by_property(key).contains_any(op_value))
+                    elif op == "$like" or op == "like":
+                        filter_conditions.append(Filter.by_property(key).like(op_value))
+            else:
+                # Simple equality filter
+                filter_conditions.append(Filter.by_property(key).equal(value))
+
+        # Combine all conditions with AND
+        if len(filter_conditions) == 1:
+            return filter_conditions[0]
+        elif len(filter_conditions) > 1:
+            combined_filter = filter_conditions[0]
+            for condition in filter_conditions[1:]:
+                combined_filter = combined_filter & condition
+            return combined_filter
+
+        return None
+
+    def _format_results(self, results) -> list[dict[str, Any]]:
+        """Convert Weaviate results to standardized format."""
+        documents = []
+
+        # Handle both GenerativeReturn objects and direct lists
+        if hasattr(results, "objects"):
+            objects_list = results.objects
+        else:
+            objects_list = results
+
+        if not objects_list:
+            return documents
+
+        for obj in objects_list:
+            # Get the content - try different property names
+            content = ""
+            properties = obj.properties
+
+            # Common content field names
+            content_fields = ["content", "text", "document", "body", "description"]
+            for field in content_fields:
+                if field in properties:
+                    content = properties[field]
+                    break
+
+            # If no content field found, use all text properties
+            if not content:
+                text_properties = []
+                for key, value in properties.items():
+                    if isinstance(value, str) and value.strip():
+                        text_properties.append(f"{key}: {value}")
+                content = " | ".join(text_properties)
+
+            # Extract metadata (all properties except content)
+            metadata = {}
+            for key, value in properties.items():
+                if key not in ["content", "text", "document", "body"] or not content:
+                    metadata[key] = value
+
+            # Add UUID as doc_id in metadata
+            metadata["doc_id"] = str(obj.uuid)
+
+            # Calculate similarity score from distance or use provided score
+            score = 0.0
+            if hasattr(obj.metadata, "score") and obj.metadata.score is not None:
+                score = float(obj.metadata.score)
+            elif hasattr(obj.metadata, "distance") and obj.metadata.distance is not None:
+                # Convert distance to similarity (assuming cosine distance)
+                # Weaviate cosine distance is between 0 and 2, similarity = 1 - (distance/2)
+                distance = float(obj.metadata.distance)
+                score = max(0.0, 1.0 - (distance / 2.0))
+
+            documents.append({"content": content, "metadata": metadata, "score": score})
+
+        return documents
+
+    @classmethod
+    def create_local(
+        cls,
+        host: str = "localhost",
+        port: int = 8080,
+        grpc_port: int = 50051,
+        collection_name: str = "Documents",
+        headers: dict[str, str] | None = None,
+    ) -> "WeaviateVectorStore":
+        """
+        Create a WeaviateVectorStore connected to local Weaviate instance.
+
+        Args:
+            host: Weaviate host
+            port: HTTP port
+            grpc_port: gRPC port
+            collection_name: Name of the collection
+            headers: Optional headers for authentication
+
+        Returns:
+            WeaviateVectorStore instance
+        """
+        try:
+            import weaviate
+        except ImportError as e:
+            raise ImportError("Weaviate client is required. Install with: pip install litellm weaviate-client") from e
+
+        client = weaviate.connect_to_local(host=host, port=port, grpc_port=grpc_port, headers=headers)
+
+        return cls(client, collection_name)
+
+    @classmethod
+    def create_cloud(
+        cls,
+        cluster_url: str,
+        auth_credentials,
+        collection_name: str = "Documents",
+        headers: dict[str, str] | None = None,
+    ) -> "WeaviateVectorStore":
+        """
+        Create a WeaviateVectorStore connected to Weaviate Cloud Services.
+
+        Args:
+            cluster_url: Weaviate cluster URL
+            auth_credentials: Authentication credentials (API key or other auth)
+            collection_name: Name of the collection
+            headers: Optional headers
+
+        Returns:
+            WeaviateVectorStore instance
+        """
+        try:
+            import weaviate
+        except ImportError as e:
+            raise ImportError("Weaviate client is required. Install with: pip install litellm weaviate-client") from e
+
+        client = weaviate.connect_to_weaviate_cloud(
+            cluster_url=cluster_url, auth_credentials=auth_credentials, headers=headers
+        )
+
+        return cls(client, collection_name)
+
+    @classmethod
+    def create_custom(
+        cls,
+        url: str,
+        auth_credentials=None,
+        collection_name: str = "Documents",
+        headers: dict[str, str] | None = None,
+        grpc_port: int | None = None,
+    ) -> "WeaviateVectorStore":
+        """
+        Create a WeaviateVectorStore with custom connection parameters.
+
+        Args:
+            url: Weaviate instance URL
+            auth_credentials: Authentication credentials
+            collection_name: Name of the collection
+            headers: Optional headers
+            grpc_port: Optional gRPC port
+
+        Returns:
+            WeaviateVectorStore instance
+        """
+        try:
+            import weaviate
+        except ImportError as e:
+            raise ImportError("Weaviate client is required. Install with: pip install litellm weaviate-client") from e
+
+        client = weaviate.connect_to_custom(
+            http_host=url,
+            http_port=443 if "https" in url else 80,
+            http_secure="https" in url,
+            grpc_port=grpc_port,
+            grpc_secure="https" in url,
+            auth_credentials=auth_credentials,
+            headers=headers,
+        )
+
+        return cls(client, collection_name)