pyseekdb 0.1.0.dev3__py3-none-any.whl

pyseekdb/client/filters.py

@@ -0,0 +1,357 @@
+"""
+Filter builder utilities for metadata and document filtering
+
+Supports:
+- Metadata filters: $eq, $lt, $gt, $lte, $gte, $ne, $in, $nin
+- Logical operators: $or, $and, $not
+- Document filters: $contains, $regex
+"""
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+
+class FilterBuilder:
+    """Build SQL WHERE clauses from filter dictionaries"""
+    
+    # Comparison operators mapping
+    COMPARISON_OPS = {
+        "$eq": "=",
+        "$lt": "<",
+        "$gt": ">",
+        "$lte": "<=",
+        "$gte": ">=",
+        "$ne": "!="
+    }
+    
+    # Logical operators
+    LOGICAL_OPS = ["$and", "$or", "$not"]
+    
+    # Document operators
+    DOCUMENT_OPS = ["$contains", "$regex"]
+    
+    @staticmethod
+    def build_metadata_filter(
+        where: Dict[str, Any],
+        metadata_column: str = "metadata"
+    ) -> Tuple[str, List[Any]]:
+        """
+        Build WHERE clause for metadata filtering
+        
+        Args:
+            where: Filter dictionary with operators like $eq, $lt, $gt, $lte, $gte, $ne, $in, $nin, $and, $or, $not
+            metadata_column: Name of metadata column (default: "metadata")
+            
+        Returns:
+            Tuple of (where_clause, params) for parameterized query
+            
+        Examples:
+            where = {"age": {"$gte": 18}}
+            -> ("JSON_EXTRACT(metadata, '$.age') >= %s", [18])
+            
+            where = {"$and": [{"age": {"$gte": 18}}, {"city": "Beijing"}]}
+            -> ("(JSON_EXTRACT(metadata, '$.age') >= %s AND JSON_EXTRACT(metadata, '$.city') = %s)", [18, "Beijing"])
+        """
+        if not where:
+            return "", []
+        
+        return FilterBuilder._build_condition(where, metadata_column)
+    
+    @staticmethod
+    def build_document_filter(
+        where_document: Dict[str, Any],
+        document_column: str = "document"
+    ) -> Tuple[str, List[Any]]:
+        """
+        Build WHERE clause for document filtering
+        
+        Args:
+            where_document: Filter dictionary with $contains, $regex, $and, $or operators
+            document_column: Name of document column (default: "document")
+            
+        Returns:
+            Tuple of (where_clause, params) for parameterized query
+            
+        Examples:
+            where_document = {"$contains": "python"}
+            -> ("MATCH(document) AGAINST (%s IN NATURAL LANGUAGE MODE)", ["python"])
+            
+            where_document = {"$regex": "^hello.*world$"}
+            -> ("document REGEXP %s", ["^hello.*world$"])
+        """
+        if not where_document:
+            return "", []
+        
+        return FilterBuilder._build_document_condition(where_document, document_column)
+    
+    @staticmethod
+    def _build_condition(
+        condition: Dict[str, Any],
+        metadata_column: str,
+        params: Optional[List[Any]] = None
+    ) -> Tuple[str, List[Any]]:
+        """Recursively build condition from nested dictionary"""
+        if params is None:
+            params = []
+        
+        clauses = []
+        
+        for key, value in condition.items():
+            if key in FilterBuilder.LOGICAL_OPS:
+                # Handle logical operators
+                if key == "$and":
+                    sub_clauses = []
+                    for sub_condition in value:
+                        sub_clause, params = FilterBuilder._build_condition(sub_condition, metadata_column, params)
+                        sub_clauses.append(sub_clause)
+                    clauses.append(f"({' AND '.join(sub_clauses)})")
+                
+                elif key == "$or":
+                    sub_clauses = []
+                    for sub_condition in value:
+                        sub_clause, params = FilterBuilder._build_condition(sub_condition, metadata_column, params)
+                        sub_clauses.append(sub_clause)
+                    clauses.append(f"({' OR '.join(sub_clauses)})")
+                
+                elif key == "$not":
+                    sub_clause, params = FilterBuilder._build_condition(value, metadata_column, params)
+                    clauses.append(f"NOT ({sub_clause})")
+            
+            elif isinstance(value, dict):
+                # Handle comparison operators
+                for op, op_value in value.items():
+                    if op in FilterBuilder.COMPARISON_OPS:
+                        sql_op = FilterBuilder.COMPARISON_OPS[op]
+                        clauses.append(f"JSON_EXTRACT({metadata_column}, '$.{key}') {sql_op} %s")
+                        params.append(op_value)
+                    
+                    elif op == "$in":
+                        placeholders = ", ".join(["%s"] * len(op_value))
+                        clauses.append(f"JSON_EXTRACT({metadata_column}, '$.{key}') IN ({placeholders})")
+                        params.extend(op_value)
+                    
+                    elif op == "$nin":
+                        placeholders = ", ".join(["%s"] * len(op_value))
+                        clauses.append(f"JSON_EXTRACT({metadata_column}, '$.{key}') NOT IN ({placeholders})")
+                        params.extend(op_value)
+            
+            else:
+                # Direct equality comparison
+                clauses.append(f"JSON_EXTRACT({metadata_column}, '$.{key}') = %s")
+                params.append(value)
+        
+        where_clause = " AND ".join(clauses) if clauses else "1=1"
+        return where_clause, params
+    
+    @staticmethod
+    def _build_document_condition(
+        condition: Dict[str, Any],
+        document_column: str,
+        params: Optional[List[Any]] = None
+    ) -> Tuple[str, List[Any]]:
+        """Build document filter condition"""
+        if params is None:
+            params = []
+        
+        clauses = []
+        
+        for key, value in condition.items():
+            if key == "$contains":
+                # Full-text search using MATCH AGAINST
+                clauses.append(f"MATCH({document_column}) AGAINST (%s IN NATURAL LANGUAGE MODE)")
+                params.append(value)
+            
+            elif key == "$regex":
+                # Regular expression matching
+                clauses.append(f"{document_column} REGEXP %s")
+                params.append(value)
+            
+            elif key == "$and":
+                sub_clauses = []
+                for sub_condition in value:
+                    sub_clause, params = FilterBuilder._build_document_condition(sub_condition, document_column, params)
+                    sub_clauses.append(sub_clause)
+                clauses.append(f"({' AND '.join(sub_clauses)})")
+            
+            elif key == "$or":
+                sub_clauses = []
+                for sub_condition in value:
+                    sub_clause, params = FilterBuilder._build_document_condition(sub_condition, document_column, params)
+                    sub_clauses.append(sub_clause)
+                clauses.append(f"({' OR '.join(sub_clauses)})")
+        
+        where_clause = " AND ".join(clauses) if clauses else "1=1"
+        return where_clause, params
+    
+    @staticmethod
+    def combine_filters(
+        metadata_filter: Tuple[str, List[Any]],
+        document_filter: Tuple[str, List[Any]]
+    ) -> Tuple[str, List[Any]]:
+        """
+        Combine metadata and document filters
+        
+        Args:
+            metadata_filter: Tuple of (where_clause, params) for metadata
+            document_filter: Tuple of (where_clause, params) for document
+            
+        Returns:
+            Combined (where_clause, params)
+        """
+        meta_clause, meta_params = metadata_filter
+        doc_clause, doc_params = document_filter
+        
+        clauses = []
+        all_params = []
+        
+        if meta_clause:
+            clauses.append(meta_clause)
+            all_params.extend(meta_params)
+        
+        if doc_clause:
+            clauses.append(doc_clause)
+            all_params.extend(doc_params)
+        
+        if clauses:
+            combined_clause = " AND ".join(clauses)
+            return combined_clause, all_params
+        else:
+            return "", []
+    
+    @staticmethod
+    def build_search_filter(where: Optional[Dict[str, Any]]) -> Optional[List[Dict[str, Any]]]:
+        """
+        Build search_params filter format from where condition for hybrid search
+        
+        Args:
+            where: Filter dictionary with operators like $eq, $lt, $gt, $lte, $gte, $ne, $in, $nin, $and, $or, $not
+            
+        Returns:
+            List of filter conditions in search_params format, or None if where is empty
+            
+        Examples:
+            where = {"category": {"$eq": "science"}}
+            -> [{"term": {"metadata.category": {"value": "science"}}}]
+            
+            where = {"$and": [{"page": {"$gte": 5}}, {"page": {"$lte": 10}}]}
+            -> [{"bool": {"must": [{"range": {"metadata.page": {"gte": 5}}}, {"range": {"metadata.page": {"lte": 10}}}]}}]
+        """
+        if not where:
+            return None
+        
+        filter_condition = FilterBuilder._build_search_filter_condition(where)
+        if filter_condition:
+            return [filter_condition]
+        return None
+    
+    @staticmethod
+    def _build_search_filter_condition(condition: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Recursively build search_params filter condition from nested dictionary"""
+        if not condition:
+            return None
+        
+        # Handle logical operators
+        if "$and" in condition:
+            must_conditions = []
+            for sub_condition in condition["$and"]:
+                sub_filter = FilterBuilder._build_search_filter_condition(sub_condition)
+                if sub_filter:
+                    must_conditions.append(sub_filter)
+            if must_conditions:
+                return {"bool": {"must": must_conditions}}
+            return None
+        
+        if "$or" in condition:
+            should_conditions = []
+            for sub_condition in condition["$or"]:
+                sub_filter = FilterBuilder._build_search_filter_condition(sub_condition)
+                if sub_filter:
+                    should_conditions.append(sub_filter)
+            if should_conditions:
+                return {"bool": {"should": should_conditions}}
+            return None
+        
+        if "$not" in condition:
+            not_filter = FilterBuilder._build_search_filter_condition(condition["$not"])
+            if not_filter:
+                return {"bool": {"must_not": [not_filter]}}
+            return None
+        
+        # Handle field conditions
+        result = {"bool": {"must": [], "should": [], "must_not": []}}
+        has_conditions = False
+        
+        for key, value in condition.items():
+            if key in FilterBuilder.LOGICAL_OPS:
+                continue
+            
+            field_name = f"metadata.{key}"
+            
+            if isinstance(value, dict):
+                # Handle comparison operators
+                range_conditions = {}
+                term_conditions = []
+                in_conditions = []
+                nin_conditions = []
+                
+                for op, op_value in value.items():
+                    if op == "$eq":
+                        term_conditions.append({"term": {field_name: {"value": op_value}}})
+                        has_conditions = True
+                    elif op == "$ne":
+                        result["bool"]["must_not"].append({"term": {field_name: {"value": op_value}}})
+                        has_conditions = True
+                    elif op == "$lt":
+                        range_conditions["lt"] = op_value
+                        has_conditions = True
+                    elif op == "$lte":
+                        range_conditions["lte"] = op_value
+                        has_conditions = True
+                    elif op == "$gt":
+                        range_conditions["gt"] = op_value
+                        has_conditions = True
+                    elif op == "$gte":
+                        range_conditions["gte"] = op_value
+                        has_conditions = True
+                    elif op == "$in":
+                        for val in op_value:
+                            in_conditions.append({"term": {field_name: {"value": val}}})
+                        has_conditions = True
+                    elif op == "$nin":
+                        for val in op_value:
+                            nin_conditions.append({"term": {field_name: {"value": val}}})
+                        has_conditions = True
+                
+                if range_conditions:
+                    result["bool"]["must"].append({"range": {field_name: range_conditions}})
+                if term_conditions:
+                    result["bool"]["must"].extend(term_conditions)
+                if in_conditions:
+                    result["bool"]["should"].extend(in_conditions)
+                if nin_conditions:
+                    result["bool"]["must_not"].extend(nin_conditions)
+            else:
+                # Direct equality
+                result["bool"]["must"].append({"term": {field_name: {"value": value}}})
+                has_conditions = True
+        
+        if not has_conditions:
+            return None
+        
+        # Clean up empty arrays
+        if not result["bool"]["must"]:
+            del result["bool"]["must"]
+        if not result["bool"]["should"]:
+            del result["bool"]["should"]
+        if not result["bool"]["must_not"]:
+            del result["bool"]["must_not"]
+        
+        # If only one type of condition, simplify
+        if len(result["bool"]) == 1:
+            key = list(result["bool"].keys())[0]
+            conditions = result["bool"][key]
+            if len(conditions) == 1:
+                return conditions[0]
+            return {"bool": {key: conditions}}
+        
+        return result
+

pyseekdb/client/meta_info.py

@@ -0,0 +1,15 @@
+"""
+Metadata information for collection fields.
+"""
+class CollectionFieldNames:
+    ID = "_id"
+    DOCUMENT = "document"
+    EMBEDDING = "embedding"
+    METADATA = "metadata"
+
+    ALL_FIELDS = [ID, DOCUMENT, EMBEDDING, METADATA]
+
+class CollectionNames:
+    @staticmethod
+    def table_name(collection_name: str) -> str:
+        return f"c$v1${collection_name}"

pyseekdb/client/query_result.py

@@ -0,0 +1,122 @@
+"""
+Query result wrapper class with JSON serialization support
+"""
+import json
+from typing import Any, Dict, List, Optional
+
+
+class QueryResultItem:
+    """Single query result item"""
+    
+    def __init__(
+        self,
+        id: Any,
+        document: Optional[str] = None,
+        embedding: Optional[List[float]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        distance: Optional[float] = None
+    ):
+        """
+        Initialize a query result item
+        
+        Args:
+            id: Record ID
+            document: Document text (optional)
+            embedding: Vector embedding (optional)
+            metadata: Metadata dictionary (optional)
+            distance: Distance/similarity score (optional)
+        """
+        self._id = id
+        self.document = document
+        self.embedding = embedding
+        self.metadata = metadata if metadata is not None else {}
+        self.distance = distance
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary"""
+        result = {"_id": self._id}
+        
+        if self.document is not None:
+            result["document"] = self.document
+        
+        if self.embedding is not None:
+            result["embedding"] = self.embedding
+        
+        if self.metadata:
+            result["metadata"] = self.metadata
+        
+        if self.distance is not None:
+            result["distance"] = self.distance
+        
+        return result
+    
+    def to_json(self) -> str:
+        """Convert to JSON string"""
+        return json.dumps(self.to_dict(), ensure_ascii=False, indent=2)
+    
+    def __repr__(self) -> str:
+        return f"QueryResultItem(id={self._id}, distance={self.distance})"
+
+
+class QueryResult:
+    """Query result wrapper with multiple items"""
+    
+    def __init__(self, items: Optional[List[QueryResultItem]] = None):
+        """
+        Initialize query result
+        
+        Args:
+            items: List of QueryResultItem objects (optional)
+        """
+        self.items = items if items is not None else []
+    
+    def add_item(
+        self,
+        id: Any,
+        document: Optional[str] = None,
+        embedding: Optional[List[float]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        distance: Optional[float] = None
+    ) -> None:
+        """
+        Add a result item
+        
+        Args:
+            id: Record ID
+            document: Document text (optional)
+            embedding: Vector embedding (optional)
+            metadata: Metadata dictionary (optional)
+            distance: Distance/similarity score (optional)
+        """
+        item = QueryResultItem(
+            id=id,
+            document=document,
+            embedding=embedding,
+            metadata=metadata,
+            distance=distance
+        )
+        self.items.append(item)
+    
+    def to_list(self) -> List[Dict[str, Any]]:
+        """Convert to list of dictionaries"""
+        return [item.to_dict() for item in self.items]
+    
+    def to_json(self) -> str:
+        """Convert to JSON string"""
+        return json.dumps(self.to_list(), ensure_ascii=False, indent=2)
+    
+    def __len__(self) -> int:
+        """Return number of items"""
+        return len(self.items)
+    
+    def __getitem__(self, index: int) -> QueryResultItem:
+        """Get item by index"""
+        return self.items[index]
+    
+    def __iter__(self):
+        """Iterate over items"""
+        return iter(self.items)
+    
+    def __repr__(self) -> str:
+        return f"QueryResult(items={len(self.items)})"
+

pyseekdb/client/sql_utils.py

@@ -0,0 +1,48 @@
+"""
+Utility functions and classes for SQL string generation and escaping in SeekDB client.
+
+Provides helpers to safely stringify values and SQL identifiers for insertion into SQL expressions.
+"""
+
+from typing import Optional, Union
+
+
+def _quote_string(value, quote: str):
+    return quote + str(value) + quote
+
+
+class SqlStringifier:
+    """
+    Translate values into strings in SQL.
+    """
+
+    def __init__(self, *, quote: str = "'", identifier: str = "`"):
+        self._quote = quote
+        self._identifier = identifier
+
+    def stringify_value(self, value: Optional[Union[str, int, float, bytes]]):
+        if value is None:
+            return "NULL"
+        if isinstance(value, bytes):
+            # For varbinary type, convert bytes to hex string and use UNHEX function
+            hex_str = value.hex()
+            return f"UNHEX('{hex_str}')"
+        if isinstance(value, str):
+            # Check if it's a hex string (for varbinary IDs)
+            # If it looks like a hex string (even length, only hex chars), use UNHEX
+            # Otherwise, treat as regular string
+            if len(value) > 0 and len(value) % 2 == 0 and all(c in '0123456789abcdefABCDEF' for c in value):
+                # Likely a hex string for varbinary, use UNHEX
+                return f"UNHEX('{value}')"
+            formatted = value.replace('\\', '\\\\').replace(self._quote, f"\\{self._quote}")
+            return _quote_string(formatted, self._quote)
+        if isinstance(value, (int, float)):
+            return str(value)
+        return _quote_string(str(value), self._quote)
+
+    def stringify_id(self, id_name: str):
+        if id_name is None:
+            raise ValueError("Identifier shouldn't be null")
+        if not isinstance(id_name, str):
+            raise ValueError(f"Identifier should be string type, but got {type(id_name).__name__}")
+        return _quote_string(id_name, self._identifier)