sf-vector-sdk 0.2.0__py3-none-any.whl

vector_sdk/namespaces/db.py

@@ -0,0 +1,230 @@
+"""
+Database namespace for direct database operations (no embedding required).
+"""
+
+from typing import Any, Optional
+
+import requests
+
+from vector_sdk.namespaces.base import BaseNamespace
+from vector_sdk.types import (
+    CloneResult,
+    DeleteFromNamespaceResult,
+    LookupResult,
+)
+
+
+class DBNamespace(BaseNamespace):
+    """
+    Namespace for direct database operations.
+
+    These operations call the query-gateway HTTP API directly, bypassing
+    the Redis Streams queue. They do not require embedding the query.
+
+    Example:
+        ```python
+        client = VectorClient(
+            redis_url="redis://localhost:6379",
+            http_url="http://localhost:8080",
+        )
+
+        # Lookup documents by ID
+        result = client.db.get_by_ids(
+            ids=["doc1", "doc2"],
+            database="turbopuffer",
+            namespace="my_namespace",
+        )
+
+        # Find by metadata
+        result = client.db.find_by_metadata(
+            filters={"userId": "user123"},
+            database="mongodb",
+            collection="vectors",
+            database_name="mydb",
+        )
+        ```
+    """
+
+    def get_by_ids(
+        self,
+        ids: list[str],
+        database: str,
+        namespace: Optional[str] = None,
+        collection: Optional[str] = None,
+        database_name: Optional[str] = None,
+        include_vectors: bool = False,
+        include_metadata: bool = True,
+    ) -> LookupResult:
+        """
+        Look up documents by their IDs.
+
+        Args:
+            ids: List of document/vector IDs to retrieve
+            database: Which vector database to query ("mongodb", "turbopuffer", "pinecone")
+            namespace: Namespace for TurboPuffer/Pinecone
+            collection: Collection name for MongoDB (also used as index name for Pinecone)
+            database_name: Database name for MongoDB
+            include_vectors: Whether to include vector values in response
+            include_metadata: Whether to include metadata in response (default: True)
+
+        Returns:
+            LookupResult containing retrieved documents
+
+        Raises:
+            ValueError: If http_url is not configured or ids is empty
+            requests.HTTPError: If the request fails
+        """
+        http_url = self._require_http_url("get_by_ids")
+
+        if not ids:
+            raise ValueError("ids list cannot be empty")
+
+        if len(ids) > 100:
+            raise ValueError("Maximum 100 IDs per request")
+
+        url = f"{http_url}/v1/lookup/{database}"
+        body = {
+            "ids": ids,
+            "namespace": namespace,
+            "collection": collection,
+            "database": database_name,
+            "includeVectors": include_vectors,
+            "includeMetadata": include_metadata,
+        }
+
+        response = requests.post(url, json=body, timeout=30)
+        response.raise_for_status()
+
+        return LookupResult.from_dict(response.json())
+
+    def find_by_metadata(
+        self,
+        filters: dict[str, Any],
+        database: str,
+        namespace: Optional[str] = None,
+        collection: Optional[str] = None,
+        database_name: Optional[str] = None,
+        limit: int = 100,
+        include_vectors: bool = False,
+    ) -> LookupResult:
+        """
+        Search for documents by metadata filters.
+
+        Args:
+            filters: Metadata key-value pairs to match
+            database: Which vector database to query ("mongodb", "turbopuffer", "pinecone")
+            namespace: Namespace for TurboPuffer/Pinecone
+            collection: Collection name for MongoDB (also used as index name for Pinecone)
+            database_name: Database name for MongoDB
+            limit: Maximum number of results (default: 100, max: 1000)
+            include_vectors: Whether to include vector values in response
+
+        Returns:
+            LookupResult containing matched documents
+
+        Raises:
+            ValueError: If http_url is not configured or filters is empty
+            requests.HTTPError: If the request fails
+        """
+        http_url = self._require_http_url("find_by_metadata")
+
+        if not filters:
+            raise ValueError("filters dict cannot be empty")
+
+        url = f"{http_url}/v1/search/{database}"
+        body = {
+            "filters": filters,
+            "namespace": namespace,
+            "collection": collection,
+            "database": database_name,
+            "limit": min(limit, 1000),
+            "includeVectors": include_vectors,
+        }
+
+        response = requests.post(url, json=body, timeout=30)
+        response.raise_for_status()
+
+        return LookupResult.from_dict(response.json())
+
+    def clone(
+        self,
+        id: str,
+        source_namespace: str,
+        destination_namespace: str,
+    ) -> CloneResult:
+        """
+        Clone a document from one TurboPuffer namespace to another.
+
+        This method fetches a document by ID from the source namespace (including
+        its vector and metadata) and writes it to the destination namespace.
+        Vectors are stored as f16 in the destination regardless of source format.
+
+        Args:
+            id: Document ID to clone
+            source_namespace: Namespace to clone from
+            destination_namespace: Namespace to clone to
+
+        Returns:
+            CloneResult containing success status and timing
+
+        Raises:
+            ValueError: If http_url is not configured or required params are missing
+            requests.HTTPError: If the request fails
+        """
+        http_url = self._require_http_url("clone")
+
+        if not id:
+            raise ValueError("id is required")
+        if not source_namespace:
+            raise ValueError("source_namespace is required")
+        if not destination_namespace:
+            raise ValueError("destination_namespace is required")
+
+        url = f"{http_url}/v1/clone/turbopuffer"
+        body = {
+            "id": id,
+            "sourceNamespace": source_namespace,
+            "destinationNamespace": destination_namespace,
+        }
+
+        response = requests.post(url, json=body, timeout=30)
+        response.raise_for_status()
+
+        return CloneResult.from_dict(response.json())
+
+    def delete(
+        self,
+        id: str,
+        namespace: str,
+    ) -> DeleteFromNamespaceResult:
+        """
+        Delete a document from a TurboPuffer namespace.
+
+        Args:
+            id: Document ID to delete
+            namespace: Namespace to delete from
+
+        Returns:
+            DeleteFromNamespaceResult containing success status and timing
+
+        Raises:
+            ValueError: If http_url is not configured or required params are missing
+            requests.HTTPError: If the request fails
+        """
+        http_url = self._require_http_url("delete")
+
+        if not id:
+            raise ValueError("id is required")
+        if not namespace:
+            raise ValueError("namespace is required")
+
+        url = f"{http_url}/v1/delete/turbopuffer"
+        body = {
+            "id": id,
+            "namespace": namespace,
+        }
+
+        response = requests.post(url, json=body, timeout=30)
+        response.raise_for_status()
+
+        return DeleteFromNamespaceResult.from_dict(response.json())

vector_sdk/namespaces/embeddings.py

@@ -0,0 +1,268 @@
+"""
+Embeddings namespace for creating and managing vector embeddings.
+"""
+
+import json
+import uuid
+from datetime import datetime
+from typing import Any, Optional
+
+from vector_sdk.namespaces.base import BaseNamespace
+from vector_sdk.types import (
+    EmbeddingConfigOverride,
+    EmbeddingRequest,
+    EmbeddingResult,
+    StorageConfig,
+    TextInput,
+    get_stream_for_priority,
+    validate_model,
+)
+
+
+class EmbeddingsNamespace(BaseNamespace):
+    """
+    Namespace for embedding generation operations.
+
+    Example:
+        ```python
+        client = VectorClient("redis://localhost:6379")
+
+        # Create embeddings asynchronously
+        request_id = client.embeddings.create(
+            texts=[{"id": "doc1", "text": "Hello world"}],
+            content_type="document",
+        )
+
+        # Wait for the result
+        result = client.embeddings.wait_for(request_id)
+
+        # Or do both in one call
+        result = client.embeddings.create_and_wait(
+            texts=[{"id": "doc1", "text": "Hello world"}],
+            content_type="document",
+        )
+        ```
+    """
+
+    def create(
+        self,
+        texts: list[dict[str, Any]],
+        content_type: str,
+        priority: str = "normal",
+        storage: Optional[StorageConfig] = None,
+        metadata: Optional[dict[str, str]] = None,
+        request_id: Optional[str] = None,
+        embedding_model: Optional[str] = None,
+        embedding_dimensions: Optional[int] = None,
+    ) -> str:
+        """
+        Create embeddings for the given texts.
+
+        This method submits an embedding request to the gateway and returns immediately
+        with a request ID. Use `wait_for()` to get the result, or use `create_and_wait()`
+        for a combined operation.
+
+        Args:
+            texts: List of text inputs. Each item should have:
+                - id: Unique identifier for the text
+                - text: The actual text content to embed
+                - document: (optional) Full document to store with embedding
+            content_type: Type of content being embedded (e.g., "topic", "flashcard")
+            priority: Queue priority - one of "critical", "high", "normal", "low"
+            storage: Configuration for where to store embeddings
+            metadata: Optional key-value pairs for tracking
+            request_id: Optional custom request ID (auto-generated if not provided)
+            embedding_model: Optional embedding model override
+            embedding_dimensions: Optional embedding dimensions override
+
+        Returns:
+            The request ID for tracking the request
+
+        Raises:
+            ValueError: If texts list is empty or invalid
+            ModelValidationError: If embedding model is not supported
+        """
+        if not texts:
+            raise ValueError("texts list cannot be empty")
+
+        # Validate embedding model if specified
+        if embedding_model:
+            validate_model(embedding_model, embedding_dimensions)
+
+        # Generate request ID if not provided
+        if request_id is None:
+            request_id = str(uuid.uuid4())
+
+        # Convert text dicts to TextInput objects
+        text_inputs = []
+        for t in texts:
+            if isinstance(t, TextInput):
+                text_inputs.append(t)
+            elif isinstance(t, dict):
+                text_inputs.append(TextInput(
+                    id=t["id"],
+                    text=t["text"],
+                    document=t.get("document"),
+                ))
+            else:
+                raise ValueError(f"Invalid text input type: {type(t)}")
+
+        # Build embedding config if model or dimensions specified
+        embedding_config = None
+        if embedding_model or embedding_dimensions:
+            embedding_config = EmbeddingConfigOverride(
+                model=embedding_model,
+                dimensions=embedding_dimensions,
+            )
+
+        # Build request
+        request = EmbeddingRequest(
+            request_id=request_id,
+            content_type=content_type,
+            priority=priority,
+            texts=text_inputs,
+            storage=storage,
+            embedding_config=embedding_config,
+            metadata=metadata or {},
+            created_at=datetime.utcnow(),
+        )
+
+        # Get the appropriate stream for this priority
+        stream = get_stream_for_priority(priority)
+
+        # Serialize and publish to Redis Stream
+        payload = json.dumps(request.to_dict())
+        self._redis.xadd(stream, {"payload": payload})
+
+        return request_id
+
+    def wait_for(
+        self,
+        request_id: str,
+        timeout: int = 60,
+    ) -> EmbeddingResult:
+        """
+        Wait for an embedding request to complete.
+
+        Args:
+            request_id: The request ID to wait for
+            timeout: Maximum time to wait in seconds (default: 60)
+
+        Returns:
+            The embedding result
+
+        Raises:
+            TimeoutError: If no result is received within the timeout
+        """
+        channel = f"embedding:result:{request_id}"
+        pubsub = self._redis.pubsub()
+        pubsub.subscribe(channel)
+
+        try:
+            start_time = datetime.utcnow()
+            while True:
+                message = pubsub.get_message(timeout=1.0)
+                if message and message["type"] == "message":
+                    data = json.loads(message["data"])
+                    return EmbeddingResult.from_dict(data)
+
+                elapsed = (datetime.utcnow() - start_time).total_seconds()
+                if elapsed >= timeout:
+                    raise TimeoutError(
+                        f"No result received for {request_id} within {timeout}s"
+                    )
+        finally:
+            pubsub.unsubscribe(channel)
+            pubsub.close()
+
+    def create_and_wait(
+        self,
+        texts: list[dict[str, Any]],
+        content_type: str,
+        priority: str = "normal",
+        storage: Optional[StorageConfig] = None,
+        metadata: Optional[dict[str, str]] = None,
+        embedding_model: Optional[str] = None,
+        embedding_dimensions: Optional[int] = None,
+        timeout: int = 60,
+    ) -> EmbeddingResult:
+        """
+        Create embeddings and wait for the result.
+
+        This method subscribes to the result channel BEFORE submitting the request,
+        ensuring no race condition where the result is published before we're listening.
+
+        Args:
+            texts: List of text inputs
+            content_type: Type of content
+            priority: Queue priority
+            storage: Storage configuration
+            metadata: Optional metadata
+            embedding_model: Optional embedding model override
+            embedding_dimensions: Optional embedding dimensions override
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            The embedding result
+        """
+        # Generate request ID upfront so we can subscribe before submitting
+        request_id = str(uuid.uuid4())
+        channel = f"embedding:result:{request_id}"
+
+        # Subscribe BEFORE submitting to avoid race condition
+        pubsub = self._redis.pubsub()
+        pubsub.subscribe(channel)
+
+        try:
+            # Now submit the request (subscription is already active)
+            self.create(
+                texts=texts,
+                content_type=content_type,
+                priority=priority,
+                storage=storage,
+                metadata=metadata,
+                request_id=request_id,
+                embedding_model=embedding_model,
+                embedding_dimensions=embedding_dimensions,
+            )
+
+            # Wait for message with timeout
+            start_time = datetime.utcnow()
+            while True:
+                message = pubsub.get_message(timeout=1.0)
+                if message and message["type"] == "message":
+                    data = json.loads(message["data"])
+                    return EmbeddingResult.from_dict(data)
+
+                elapsed = (datetime.utcnow() - start_time).total_seconds()
+                if elapsed >= timeout:
+                    raise TimeoutError(
+                        f"No result received for {request_id} within {timeout}s"
+                    )
+        finally:
+            pubsub.unsubscribe(channel)
+            pubsub.close()
+
+    def get_queue_depth(self) -> dict[str, int]:
+        """
+        Get the current queue depth for each priority level.
+
+        Returns:
+            Dictionary mapping stream name to pending message count
+        """
+        streams = [
+            "embedding:critical",
+            "embedding:high",
+            "embedding:normal",
+            "embedding:low",
+        ]
+
+        depths = {}
+        for stream in streams:
+            try:
+                info = self._redis.xinfo_stream(stream)
+                depths[stream] = info.get("length", 0)
+            except Exception:
+                depths[stream] = 0
+
+        return depths