endee-llamaindex 0.1.3__py3-none-any.whl → 0.1.5a1__py3-none-any.whl

endee_llamaindex/constants.py

@@ -0,0 +1,70 @@
+"""
+Constants for EndeeVectorStore.
+
+This module contains all constants used by the Endee LlamaIndex integration.
+"""
+
+from llama_index.core.vector_stores.types import FilterOperator
+
+# Endee default constants
+# These may be overridden by importing from endee.constants if available
+MAX_VECTORS_PER_BATCH = 1000
+DEFAULT_EF_SEARCH = 128
+MAX_TOP_K_ALLOWED = 512
+MAX_EF_SEARCH_ALLOWED = 1024
+MAX_DIMENSION_ALLOWED = 10000
+MAX_INDEX_NAME_LENGTH_ALLOWED = 48
+
+# Try to import constants from endee package to stay in sync
+try:
+    from endee.constants import (
+        DEFAULT_EF_SEARCH as _DEFAULT_EF_SEARCH,
+        MAX_DIMENSION_ALLOWED as _MAX_DIMENSION_ALLOWED,
+        MAX_EF_SEARCH_ALLOWED as _MAX_EF_SEARCH_ALLOWED,
+        MAX_INDEX_NAME_LENGTH_ALLOWED as _MAX_INDEX_NAME_LENGTH_ALLOWED,
+        MAX_TOP_K_ALLOWED as _MAX_TOP_K_ALLOWED,
+        MAX_VECTORS_PER_BATCH as _MAX_VECTORS_PER_BATCH,
+    )
+    # Override defaults with values from endee package
+    DEFAULT_EF_SEARCH = _DEFAULT_EF_SEARCH
+    MAX_DIMENSION_ALLOWED = _MAX_DIMENSION_ALLOWED
+    MAX_EF_SEARCH_ALLOWED = _MAX_EF_SEARCH_ALLOWED
+    MAX_INDEX_NAME_LENGTH_ALLOWED = _MAX_INDEX_NAME_LENGTH_ALLOWED
+    MAX_TOP_K_ALLOWED = _MAX_TOP_K_ALLOWED
+    MAX_VECTORS_PER_BATCH = _MAX_VECTORS_PER_BATCH
+except ImportError:
+    pass
+
+# Space types and precision types for index creation
+SPACE_TYPES_VALID = ("cosine", "l2", "ip")
+PRECISION_VALID = ("binary", "float16", "float32", "int16d", "int8d")
+
+# Space type mapping (aliases)
+SPACE_TYPE_MAP = {
+    "cosine": "cosine",
+    "l2": "l2",
+    "ip": "ip",
+    "euclidean": "l2",
+    "inner_product": "ip",
+}
+
+# Vector store keys
+ID_KEY = "id"
+VECTOR_KEY = "values"
+SPARSE_VECTOR_KEY = "sparse_values"
+METADATA_KEY = "metadata"
+
+# Batch size for add(); capped by MAX_VECTORS_PER_BATCH
+DEFAULT_BATCH_SIZE = 100
+
+# Supported filter operations: currently only EQ and IN.
+# Map FilterOperator -> endee/backend filter symbol.
+SUPPORTED_FILTER_OPERATORS = (
+    FilterOperator.EQ,   # eq  -> $eq
+    FilterOperator.IN,   # in  -> $in
+)
+
+REVERSE_OPERATOR_MAP = {
+    FilterOperator.EQ: "$eq",
+    FilterOperator.IN: "$in",
+}

endee_llamaindex/utils.py

@@ -1,29 +1,8 @@
 import logging
-from collections import Counter
-from typing import Any, Callable, Dict, List, Optional, cast
-import json
-
-from llama_index.core.bridge.pydantic import PrivateAttr
-from llama_index.core.schema import BaseNode, TextNode
-from llama_index.core.vector_stores.types import (
-    BasePydanticVectorStore,
-    MetadataFilters,
-    VectorStoreQuery,
-    VectorStoreQueryMode,
-    VectorStoreQueryResult,
-)
-from llama_index.core.vector_stores.utils import (
-    DEFAULT_TEXT_KEY,
-    legacy_metadata_dict_to_node,
-    metadata_dict_to_node,
-    node_to_metadata_dict,
-)
-from llama_index.core.vector_stores.types import MetadataFilter, FilterOperator
+from typing import Any, Callable, List, Optional
 
 _logger = logging.getLogger(__name__)
 
-DEFAULT_BATCH_SIZE = 100
-
 # Supported sparse embedding models
 SUPPORTED_SPARSE_MODELS = {
     "splade_pp": "prithivida/Splade_PP_en_v1",
@@ -36,23 +15,11 @@ SUPPORTED_SPARSE_MODELS = {
     "xlm_roberta": "xlm-roberta-base",
 }
 
-reverse_operator_map = {
-    FilterOperator.EQ: "$eq",
-    FilterOperator.NE: "$ne",
-    FilterOperator.GT: "$gt",
-    FilterOperator.GTE: "$gte",
-    FilterOperator.LT: "$lt",
-    FilterOperator.LTE: "$lte",
-    FilterOperator.IN: "$in",
-    FilterOperator.NIN: "$nin",
-}
-
 
 def _import_endee() -> Any:
     """Import endee module."""
     try:
         import endee
-        from endee.endee import Endee
     except ImportError as e:
         raise ImportError(
             "Could not import endee python package. "
@@ -61,30 +28,6 @@ def _import_endee() -> Any:
     return endee
 
 
-def build_dict(input_batch: List[List[int]]) -> List[Dict[str, Any]]:
-    """
-    Build a list of sparse dictionaries from a batch of input_ids.
-    """
-    sparse_emb = []
-    for token_ids in input_batch:
-        indices = []
-        values = []
-        d = dict(Counter(token_ids))
-        for idx in d:
-            indices.append(idx)
-            values.append(float(d[idx]))
-        sparse_emb.append({"indices": indices, "values": values})
-    return sparse_emb
-
-
-def generate_sparse_vectors(
-    context_batch: List[str], tokenizer: Callable
-) -> List[Dict[str, Any]]:
-    """Generate sparse vectors from a batch of contexts."""
-    inputs = tokenizer(context_batch)["input_ids"]
-    return build_dict(inputs)
-
-
 def _initialize_sparse_encoder_fastembed(
     model_name: str,
     batch_size: int = 256,
@@ -98,9 +41,12 @@ def _initialize_sparse_encoder_fastembed(
         from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
     except ImportError as e:
         raise ImportError(
-            "Could not import FastEmbed. "
-            "Please install it with `pip install fastembed` or "
-            "`pip install fastembed-gpu` for GPU support."
+            "FastEmbed is required for hybrid search but not installed.\n"
+            "Install options:\n"
+            "  - CPU:  pip install endee-llamaindex[hybrid]\n"
+            "  - GPU:  pip install endee-llamaindex[hybrid-gpu]\n"
+            "  - Or:   pip install fastembed\n"
+            "For dense-only search, create vector store without sparse_dim/model_name."
         ) from e
 
     resolved_model_name = SUPPORTED_SPARSE_MODELS.get(model_name, model_name)
@@ -212,529 +158,3 @@ def get_sparse_encoder(
         )
     else:
         return _initialize_sparse_encoder_transformers(model_name=model_name)
-
-
-class EndeeHybridVectorStore(BasePydanticVectorStore):
-    """
-    Endee Hybrid Vector Store for combined dense and sparse vector search.
-    
-    This class provides hybrid search capabilities using both dense embeddings
-    and sparse vectors (e.g., SPLADE, BM25-style) for improved retrieval.
-    """
-
-    stores_text: bool = True
-    flat_metadata: bool = False
-
-    api_token: Optional[str]
-    index_name: Optional[str]
-    space_type: Optional[str]
-    dimension: Optional[int]
-    vocab_size: int
-    insert_kwargs: Optional[Dict]
-    text_key: str
-    batch_size: int
-    remove_text_from_metadata: bool
-    model_name: Optional[str]
-    use_fastembed: bool
-    alpha: float  # Weight for dense vs sparse (0=sparse only, 1=dense only)
-
-    _endee_index: Any = PrivateAttr()
-    _sparse_encoder: Optional[Callable] = PrivateAttr(default=None)
-
-    def __init__(
-        self,
-        endee_index: Optional[Any] = None,
-        api_token: Optional[str] = None,
-        index_name: Optional[str] = None,
-        space_type: Optional[str] = "cosine",
-        dimension: Optional[int] = None,
-        vocab_size: int = 30522,  # Default BERT vocab size
-        insert_kwargs: Optional[Dict] = None,
-        text_key: str = DEFAULT_TEXT_KEY,
-        batch_size: int = DEFAULT_BATCH_SIZE,
-        remove_text_from_metadata: bool = False,
-        model_name: Optional[str] = "splade_pp",
-        use_fastembed: bool = True,
-        alpha: float = 0.5,
-        **kwargs: Any,
-    ) -> None:
-        insert_kwargs = insert_kwargs or {}
-
-        super().__init__(
-            index_name=index_name,
-            api_token=api_token,
-            space_type=space_type,
-            dimension=dimension,
-            vocab_size=vocab_size,
-            insert_kwargs=insert_kwargs,
-            text_key=text_key,
-            batch_size=batch_size,
-            remove_text_from_metadata=remove_text_from_metadata,
-            model_name=model_name,
-            use_fastembed=use_fastembed,
-            alpha=alpha,
-        )
-
-        # Initialize hybrid index
-        if endee_index is not None:
-            self._endee_index = endee_index
-        else:
-            self._endee_index = self._initialize_hybrid_index(
-                api_token, index_name, dimension, space_type, vocab_size
-            )
-
-        # Initialize sparse encoder
-        if model_name:
-            _logger.info(f"Initializing sparse encoder with model: {model_name}")
-            self._sparse_encoder = get_sparse_encoder(
-                model_name=model_name,
-                use_fastembed=use_fastembed,
-                batch_size=batch_size,
-            )
-        else:
-            self._sparse_encoder = None
-
-    @classmethod
-    def _initialize_hybrid_index(
-        cls,
-        api_token: Optional[str],
-        index_name: Optional[str],
-        dimension: Optional[int] = None,
-        space_type: Optional[str] = "cosine",
-        vocab_size: Optional[int] = None,
-    ) -> Any:
-        """Initialize Endee hybrid index."""
-        _import_endee()
-        from endee.endee import Endee
-
-        nd = Endee(token=api_token)
-
-        try:
-            index = nd.get_hybrid_index(name=index_name)
-            _logger.info(f"Retrieved existing hybrid index: {index_name}")
-            return index
-        except Exception as e:
-            if dimension is None:
-                raise ValueError(
-                    "Must provide dimension when creating a new hybrid index"
-                ) from e
-            if vocab_size is None:
-                raise ValueError(
-                    "Must provide vocab_size when creating a new hybrid index"
-                ) from e
-            
-            _logger.info(f"Creating new hybrid index: {index_name}")
-            nd.create_hybrid_index(
-                name=index_name,
-                dimension=dimension,
-                space_type=space_type,
-                vocab_size=vocab_size,
-            )
-            return nd.get_hybrid_index(name=index_name)
-
-    @classmethod
-    def from_params(
-        cls,
-        api_token: Optional[str] = None,
-        index_name: Optional[str] = None,
-        dimension: Optional[int] = None,
-        space_type: str = "cosine",
-        vocab_size: int = 30522,
-        batch_size: int = DEFAULT_BATCH_SIZE,
-        model_name: Optional[str] = "splade_pp",
-        use_fastembed: bool = True,
-        alpha: float = 0.5,
-    ) -> "EndeeHybridVectorStore":
-        """
-        Create EndeeHybridVectorStore from parameters.
-        
-        Args:
-            api_token: API token for Endee service
-            index_name: Name of the hybrid index
-            dimension: Vector dimension for dense embeddings
-            space_type: Distance metric ("cosine", "l2", or "ip")
-            vocab_size: Vocabulary size for sparse vectors
-            batch_size: Batch size for operations
-            model_name: Model name or alias for sparse embeddings
-                       Supported models:
-                       - 'splade_pp': prithivida/Splade_PP_en_v1
-                       - 'splade_cocondenser': naver/splade-cocondenser-ensembledistil
-                       - 'bert_base': bert-base-uncased
-                       - 'distilbert': distilbert-base-uncased
-                       - 'minilm': sentence-transformers/all-MiniLM-L6-v2
-                       - 'mpnet': sentence-transformers/all-mpnet-base-v2
-                       - 'roberta': roberta-base
-                       - 'xlm_roberta': xlm-roberta-base
-            use_fastembed: Use FastEmbed for sparse encoding (recommended)
-            alpha: Weight for hybrid search (0=sparse only, 1=dense only, 0.5=balanced)
-        """
-        endee_index = cls._initialize_hybrid_index(
-            api_token, index_name, dimension, space_type, vocab_size
-        )
-
-        return cls(
-            endee_index=endee_index,
-            api_token=api_token,
-            index_name=index_name,
-            dimension=dimension,
-            space_type=space_type,
-            vocab_size=vocab_size,
-            batch_size=batch_size,
-            model_name=model_name,
-            use_fastembed=use_fastembed,
-            alpha=alpha,
-        )
-
-    @classmethod
-    def class_name(cls) -> str:
-        return "EndeeHybridVectorStore"
-
-    def _compute_sparse_vectors(self, texts: List[str]) -> tuple:
-        """Compute sparse vectors for a list of texts."""
-        if self._sparse_encoder is None:
-            raise ValueError(
-                "Sparse encoder not initialized. "
-                "Please provide model_name when creating the store."
-            )
-        return self._sparse_encoder(texts)
-
-    def add(
-        self,
-        nodes: List[BaseNode],
-        **add_kwargs: Any,
-    ) -> List[str]:
-        """
-        Add nodes to hybrid index with both dense and sparse vectors.
-
-        Args:
-            nodes: List[BaseNode]: list of nodes with embeddings
-        """
-        ids = []
-        entries = []
-        texts = []
-        
-        # Collect all texts for batch sparse encoding
-        for node in nodes:
-            text = node.get_content()
-            texts.append(text)
-        
-        # Compute sparse vectors in batch
-        if self._sparse_encoder is not None and texts:
-            sparse_indices, sparse_values = self._compute_sparse_vectors(texts)
-        else:
-            sparse_indices = [[] for _ in texts]
-            sparse_values = [[] for _ in texts]
-        
-        for i, node in enumerate(nodes):
-            node_id = node.node_id
-            metadata = node_to_metadata_dict(node)
-            
-            # Filter values for hybrid index
-            filter_data = {}
-            for key in ["file_name", "doc_id", "category", "difficulty", 
-                        "language", "field", "type", "feature"]:
-                if key in metadata:
-                    filter_data[key] = metadata[key]
-            
-            entry = {
-                "id": node_id,
-                "vector": node.get_embedding(),
-                "sparse_indices": sparse_indices[i],
-                "sparse_values": sparse_values[i],
-                "meta": metadata,
-                "filter": filter_data
-            }
-
-            ids.append(node_id)
-            entries.append(entry)
-        
-        # Batch upsert
-        batch_size = self.batch_size
-        for i in range(0, len(entries), batch_size):
-            batch = entries[i : i + batch_size]
-            self._endee_index.upsert(batch)
-        
-        return ids
-
-    def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
-        """
-        Delete nodes using ref_doc_id.
-        
-        Args:
-            ref_doc_id (str): The id of the document to delete.
-        """
-        try:
-            self._endee_index.delete_with_filter({"doc_id": ref_doc_id})
-        except Exception as e:
-            _logger.error(f"Error deleting vectors for doc_id {ref_doc_id}: {e}")
-
-    def delete_by_ids(self, ids: List[str], **delete_kwargs: Any) -> None:
-        """
-        Delete nodes by their IDs.
-        
-        Args:
-            ids: List of node IDs to delete.
-        """
-        try:
-            self._endee_index.delete(ids)
-        except Exception as e:
-            _logger.error(f"Error deleting vectors by IDs: {e}")
-
-    def delete_with_filter(self, filter_dict: Dict[str, Any], **delete_kwargs: Any) -> None:
-        """
-        Delete nodes matching a filter.
-        
-        Args:
-            filter_dict: Filter dictionary for deletion.
-        """
-        try:
-            self._endee_index.delete_with_filter(filter_dict)
-        except Exception as e:
-            _logger.error(f"Error deleting vectors with filter: {e}")
-
-    @property
-    def client(self) -> Any:
-        """Return Endee hybrid index client."""
-        return self._endee_index
-
-    def query(
-        self,
-        query: VectorStoreQuery,
-        sparse_query_text: Optional[str] = None,
-        alpha: Optional[float] = None,
-        **kwargs: Any,
-    ) -> VectorStoreQueryResult:
-        """
-        Query hybrid index for top k most similar nodes.
-
-        Args:
-            query: VectorStoreQuery object containing query parameters
-            sparse_query_text: Optional text to compute sparse vector for query.
-                              If not provided, uses query.query_str if available.
-            alpha: Optional weight override for this query (0=sparse only, 1=dense only)
-        """
-        # Get dimension
-        try:
-            dimension = self._endee_index.describe()["dimension"]
-        except:
-            if query.query_embedding is not None:
-                dimension = len(query.query_embedding)
-            else:
-                raise ValueError("Could not determine vector dimension")
-        
-        query_embedding = [0.0] * dimension
-        filters = {}
-        use_alpha = alpha if alpha is not None else self.alpha
-
-        # Build filters
-        if query.filters is not None:
-            for filter_item in query.filters.filters:
-                if hasattr(filter_item, "key") and hasattr(filter_item, "value") and hasattr(filter_item, "operator"):
-                    op_symbol = reverse_operator_map.get(filter_item.operator)
-                    if not op_symbol:
-                        raise ValueError(f"Unsupported filter operator: {filter_item.operator}")
-                    
-                    if filter_item.key not in filters:
-                        filters[filter_item.key] = {}
-                    filters[filter_item.key][op_symbol] = filter_item.value
-
-                elif isinstance(filter_item, dict):
-                    for key, op_dict in filter_item.items():
-                        if isinstance(op_dict, dict):
-                            for op, val in op_dict.items():
-                                if key not in filters:
-                                    filters[key] = {}
-                                filters[key][op] = val
-                else:
-                    raise ValueError(f"Unsupported filter format: {filter_item}")
-
-        _logger.info(f"Final structured filters: {filters}")
-
-        # Get dense query embedding
-        if query.query_embedding is not None:
-            query_embedding = cast(List[float], query.query_embedding)
-
-        # Compute sparse query vector
-        sparse_indices = []
-        sparse_values = []
-        
-        query_text = sparse_query_text or getattr(query, 'query_str', None)
-        if query_text and self._sparse_encoder is not None:
-            sparse_indices_batch, sparse_values_batch = self._compute_sparse_vectors([query_text])
-            sparse_indices = sparse_indices_batch[0]
-            sparse_values = sparse_values_batch[0]
-
-        # Execute hybrid query
-        try:
-            results = self._endee_index.query(
-                vector=query_embedding,
-                sparse_indices=sparse_indices,
-                sparse_values=sparse_values,
-                top_k=query.similarity_top_k,
-                filter=filters if filters else None,
-                include_vectors=True,
-                alpha=use_alpha,
-            )
-        except Exception as e:
-            _logger.error(f"Error querying Endee hybrid index: {e}")
-            return VectorStoreQueryResult(nodes=[], similarities=[], ids=[])
-
-        # Process results
-        nodes = []
-        similarities = []
-        ids = []
-
-        for result in results:
-            node_id = result["id"]
-            score = result.get("similarity", result.get("score", 0.0))
-            metadata = result.get("meta", {})
-            
-            if self.flat_metadata:
-                node = metadata_dict_to_node(
-                    metadata=metadata,
-                    text=metadata.pop(self.text_key, None),
-                    id_=node_id,
-                )
-            else:
-                metadata_dict, node_info, relationships = legacy_metadata_dict_to_node(
-                    metadata=metadata,
-                    text_key=self.text_key,
-                )
-                
-                _node_content_str = metadata.get("_node_content", "{}")
-                try:
-                    node_content = json.loads(_node_content_str)
-                except json.JSONDecodeError:
-                    node_content = {}
-
-                text = node_content.get(self.text_key, "")
-                node = TextNode(
-                    text=text,
-                    metadata=metadata_dict,
-                    relationships=relationships,
-                    node_id=node_id,
-                )
-                
-                for key, val in node_info.items():
-                    if hasattr(node, key):
-                        setattr(node, key, val)
-            
-            if "vector" in result:
-                node.embedding = result["vector"]
-            
-            nodes.append(node)
-            similarities.append(score)
-            ids.append(node_id)
-
-        return VectorStoreQueryResult(nodes=nodes, similarities=similarities, ids=ids)
-
-    def hybrid_query(
-        self,
-        query_text: str,
-        query_embedding: List[float],
-        top_k: int = 10,
-        alpha: Optional[float] = None,
-        filters: Optional[Dict[str, Any]] = None,
-    ) -> VectorStoreQueryResult:
-        """
-        Direct hybrid query method for convenience.
-        
-        Args:
-            query_text: Text query for sparse vector computation
-            query_embedding: Dense embedding vector
-            top_k: Number of results to return
-            alpha: Weight for hybrid search (0=sparse, 1=dense)
-            filters: Optional filter dictionary
-        
-        Returns:
-            VectorStoreQueryResult with combined results
-        """
-        use_alpha = alpha if alpha is not None else self.alpha
-        
-        # Compute sparse vector
-        sparse_indices = []
-        sparse_values = []
-        if self._sparse_encoder is not None:
-            sparse_indices_batch, sparse_values_batch = self._compute_sparse_vectors([query_text])
-            sparse_indices = sparse_indices_batch[0]
-            sparse_values = sparse_values_batch[0]
-        
-        try:
-            results = self._endee_index.query(
-                vector=query_embedding,
-                sparse_indices=sparse_indices,
-                sparse_values=sparse_values,
-                top_k=top_k,
-                filter=filters,
-                include_vectors=True,
-                alpha=use_alpha,
-            )
-        except Exception as e:
-            _logger.error(f"Error in hybrid query: {e}")
-            return VectorStoreQueryResult(nodes=[], similarities=[], ids=[])
-
-        nodes = []
-        similarities = []
-        ids = []
-
-        for result in results:
-            node_id = result["id"]
-            score = result.get("similarity", result.get("score", 0.0))
-            metadata = result.get("meta", {})
-            
-            metadata_dict, node_info, relationships = legacy_metadata_dict_to_node(
-                metadata=metadata,
-                text_key=self.text_key,
-            )
-            
-            _node_content_str = metadata.get("_node_content", "{}")
-            try:
-                node_content = json.loads(_node_content_str)
-            except json.JSONDecodeError:
-                node_content = {}
-
-            text = node_content.get(self.text_key, "")
-            node = TextNode(
-                text=text,
-                metadata=metadata_dict,
-                relationships=relationships,
-                node_id=node_id,
-            )
-            
-            for key, val in node_info.items():
-                if hasattr(node, key):
-                    setattr(node, key, val)
-            
-            if "vector" in result:
-                node.embedding = result["vector"]
-            
-            nodes.append(node)
-            similarities.append(score)
-            ids.append(node_id)
-
-        return VectorStoreQueryResult(nodes=nodes, similarities=similarities, ids=ids)
-
-    def describe(self) -> Dict[str, Any]:
-        """Get index description/stats."""
-        try:
-            return self._endee_index.describe()
-        except Exception as e:
-            _logger.error(f"Error describing index: {e}")
-            return {}
-
-    def list_ids(self, limit: int = 100) -> List[str]:
-        """List IDs in the index."""
-        try:
-            return self._endee_index.list_ids(limit=limit)
-        except Exception as e:
-            _logger.error(f"Error listing IDs: {e}")
-            return []
-
-    def fetch(self, ids: List[str]) -> List[Dict[str, Any]]:
-        """Fetch vectors by IDs."""
-        try:
-            return self._endee_index.fetch(ids)
-        except Exception as e:
-            _logger.error(f"Error fetching vectors: {e}")
-            return []
-