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

vector_sdk/namespaces/search.py

@@ -0,0 +1,258 @@
+"""
+Search namespace for vector similarity search operations.
+"""
+
+import json
+import uuid
+from datetime import datetime
+from typing import Optional
+
+from vector_sdk.namespaces.base import BaseNamespace
+from vector_sdk.types import (
+    EmbeddingConfigOverride,
+    QueryConfig,
+    QueryRequest,
+    QueryResult,
+    get_query_stream_for_priority,
+    validate_model,
+)
+
+
+class SearchNamespace(BaseNamespace):
+    """
+    Namespace for vector search operations.
+
+    Example:
+        ```python
+        client = VectorClient("redis://localhost:6379")
+
+        # Search for similar vectors
+        result = client.search.query_and_wait(
+            query_text="What is machine learning?",
+            database="turbopuffer",
+            namespace="topic_vectors",
+            top_k=10,
+        )
+
+        for match in result.matches:
+            print(f"{match.id}: {match.score}")
+        ```
+    """
+
+    def query(
+        self,
+        query_text: str,
+        database: str,
+        top_k: int = 10,
+        min_score: Optional[float] = None,
+        filters: Optional[dict[str, str]] = None,
+        namespace: Optional[str] = None,
+        collection: Optional[str] = None,
+        database_name: Optional[str] = None,
+        include_vectors: bool = False,
+        include_metadata: bool = True,
+        embedding_model: Optional[str] = None,
+        embedding_dimensions: Optional[int] = None,
+        priority: str = "normal",
+        metadata: Optional[dict[str, str]] = None,
+        request_id: Optional[str] = None,
+    ) -> str:
+        """
+        Submit a vector search query.
+
+        This method embeds the query text and searches for similar vectors in
+        the specified database. Returns a request ID - use `wait_for()` to get
+        the results, or use `query_and_wait()` for a combined operation.
+
+        Args:
+            query_text: The text to search for (will be embedded)
+            database: Which vector database to search ("mongodb", "turbopuffer", "pinecone")
+            top_k: Number of results to return (default: 10)
+            min_score: Minimum similarity score threshold (0.0 to 1.0)
+            filters: Metadata filters for filtering results
+            namespace: Namespace for Pinecone/TurboPuffer
+            collection: Collection name for MongoDB
+            database_name: Database name for MongoDB
+            include_vectors: Whether to include vector values in response
+            include_metadata: Whether to include metadata in response
+            embedding_model: Optional embedding model override
+            embedding_dimensions: Optional embedding dimensions override
+            priority: Queue priority (default: "normal")
+            metadata: Optional key-value pairs for tracking
+            request_id: Optional custom request ID
+
+        Returns:
+            The request ID for tracking the query
+
+        Raises:
+            ValueError: If query_text is empty
+            ModelValidationError: If embedding model is not supported
+        """
+        if not query_text or query_text.strip() == "":
+            raise ValueError("query_text cannot be empty")
+
+        # Validate embedding model if specified
+        if embedding_model:
+            validate_model(embedding_model, embedding_dimensions)
+
+        if request_id is None:
+            request_id = str(uuid.uuid4())
+
+        query_config = QueryConfig(
+            top_k=top_k,
+            min_score=min_score,
+            filters=filters,
+            namespace=namespace,
+            collection=collection,
+            database=database_name,
+            include_vectors=include_vectors,
+            include_metadata=include_metadata,
+        )
+
+        embedding_config = None
+        if embedding_model or embedding_dimensions:
+            embedding_config = EmbeddingConfigOverride(
+                model=embedding_model,
+                dimensions=embedding_dimensions,
+            )
+
+        request = QueryRequest(
+            request_id=request_id,
+            query_text=query_text,
+            database=database,
+            priority=priority,
+            query_config=query_config,
+            embedding_config=embedding_config,
+            metadata=metadata or {},
+            created_at=datetime.utcnow(),
+        )
+
+        stream = get_query_stream_for_priority(priority)
+        payload = json.dumps(request.to_dict())
+        self._redis.xadd(stream, {"data": payload})
+
+        return request_id
+
+    def wait_for(
+        self,
+        request_id: str,
+        timeout: int = 30,
+    ) -> QueryResult:
+        """
+        Wait for a search query to complete.
+
+        Args:
+            request_id: The request ID to wait for
+            timeout: Maximum time to wait in seconds (default: 30)
+
+        Returns:
+            The query result
+
+        Raises:
+            TimeoutError: If no result is received within the timeout
+        """
+        channel = f"query: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 QueryResult.from_dict(data)
+
+                elapsed = (datetime.utcnow() - start_time).total_seconds()
+                if elapsed >= timeout:
+                    raise TimeoutError(
+                        f"No query result received for {request_id} within {timeout}s"
+                    )
+        finally:
+            pubsub.unsubscribe(channel)
+            pubsub.close()
+
+    def query_and_wait(
+        self,
+        query_text: str,
+        database: str,
+        top_k: int = 10,
+        min_score: Optional[float] = None,
+        filters: Optional[dict[str, str]] = None,
+        namespace: Optional[str] = None,
+        collection: Optional[str] = None,
+        database_name: Optional[str] = None,
+        include_vectors: bool = False,
+        include_metadata: bool = True,
+        embedding_model: Optional[str] = None,
+        embedding_dimensions: Optional[int] = None,
+        priority: str = "normal",
+        metadata: Optional[dict[str, str]] = None,
+        timeout: int = 30,
+    ) -> QueryResult:
+        """
+        Submit a search query 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:
+            query_text: The text to search for
+            database: Which vector database to search
+            top_k: Number of results to return
+            min_score: Minimum similarity score threshold
+            filters: Metadata filters
+            namespace: Namespace for Pinecone/TurboPuffer
+            collection: Collection name for MongoDB
+            database_name: Database name for MongoDB
+            include_vectors: Include vectors in response
+            include_metadata: Include metadata in response
+            embedding_model: Optional embedding model override
+            embedding_dimensions: Optional embedding dimensions override
+            priority: Queue priority
+            metadata: Optional metadata for tracking
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            The query result
+        """
+        request_id = str(uuid.uuid4())
+        channel = f"query:result:{request_id}"
+
+        pubsub = self._redis.pubsub()
+        pubsub.subscribe(channel)
+
+        try:
+            self.query(
+                query_text=query_text,
+                database=database,
+                top_k=top_k,
+                min_score=min_score,
+                filters=filters,
+                namespace=namespace,
+                collection=collection,
+                database_name=database_name,
+                include_vectors=include_vectors,
+                include_metadata=include_metadata,
+                embedding_model=embedding_model,
+                embedding_dimensions=embedding_dimensions,
+                priority=priority,
+                metadata=metadata,
+                request_id=request_id,
+            )
+
+            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 QueryResult.from_dict(data)
+
+                elapsed = (datetime.utcnow() - start_time).total_seconds()
+                if elapsed >= timeout:
+                    raise TimeoutError(
+                        f"No query result received for {request_id} within {timeout}s"
+                    )
+        finally:
+            pubsub.unsubscribe(channel)
+            pubsub.close()

vector_sdk/structured/__init__.py

@@ -0,0 +1,60 @@
+"""
+Structured Embeddings Module.
+
+Provides type-safe methods for embedding known tool types (FlashCard, TestQuestion, etc.)
+with automatic text extraction, content hash computation, and database routing.
+"""
+
+from .router import (
+    DatabaseRoutingError,
+    DatabaseRoutingMode,
+    build_storage_config,
+    get_content_type,
+    get_database_routing_mode,
+    validate_database_routing,
+)
+from .structured_embeddings import (
+    StructuredEmbeddingsNamespace,
+    TestQuestionInput,
+    ToolMetadata,
+)
+from .tool_config import (
+    TOOL_CONFIGS,
+    PineconeToolConfig,
+    QuestionType,
+    ToolConfig,
+    ToolDatabaseConfig,
+    TurboPufferToolConfig,
+    get_flashcard_namespace_suffix,
+    get_pinecone_namespace,
+    get_question_namespace_suffix,
+    get_tool_config,
+    get_turbopuffer_namespace,
+)
+
+__all__ = [
+    # Namespace class
+    "StructuredEmbeddingsNamespace",
+    # Types
+    "ToolMetadata",
+    "TestQuestionInput",
+    # Tool configuration
+    "ToolConfig",
+    "ToolDatabaseConfig",
+    "TurboPufferToolConfig",
+    "PineconeToolConfig",
+    "QuestionType",
+    "TOOL_CONFIGS",
+    "get_tool_config",
+    "get_flashcard_namespace_suffix",
+    "get_question_namespace_suffix",
+    "get_turbopuffer_namespace",
+    "get_pinecone_namespace",
+    # Database router
+    "DatabaseRoutingMode",
+    "DatabaseRoutingError",
+    "get_database_routing_mode",
+    "validate_database_routing",
+    "build_storage_config",
+    "get_content_type",
+]

vector_sdk/structured/router.py

@@ -0,0 +1,190 @@
+"""
+Database Router for Structured Embeddings.
+
+Routes embedding writes to the appropriate databases based on:
+- STRUCTURED_EMBEDDING_DATABASE_ROUTER environment variable
+- Per-tool enabled/disabled configuration
+"""
+
+import os
+from typing import Any, Literal, Optional
+
+from ..hash import ToolCollection
+from ..types import (
+    PineconeStorageConfig,
+    StorageConfig,
+    TurboPufferStorage,
+)
+from .tool_config import (
+    get_pinecone_namespace,
+    get_tool_config,
+    get_turbopuffer_namespace,
+)
+
+# ============================================================================
+# Types
+# ============================================================================
+
+DatabaseRoutingMode = Literal["dual", "turbopuffer", "pinecone"]
+
+
+class DatabaseRoutingError(Exception):
+    """Error thrown when database routing fails."""
+
+    pass
+
+
+# ============================================================================
+# Environment Variable
+# ============================================================================
+
+ENV_VAR_NAME = "STRUCTURED_EMBEDDING_DATABASE_ROUTER"
+
+
+def get_database_routing_mode() -> DatabaseRoutingMode:
+    """
+    Get the database routing mode from environment variable.
+    Defaults to 'turbopuffer' if not set.
+    """
+    env_value = os.environ.get(ENV_VAR_NAME)
+
+    if not env_value:
+        return "turbopuffer"
+
+    mode = env_value.lower()
+
+    if mode not in ("dual", "turbopuffer", "pinecone"):
+        raise DatabaseRoutingError(
+            f"Invalid {ENV_VAR_NAME} value: '{env_value}'. "
+            "Must be 'dual', 'turbopuffer', or 'pinecone'."
+        )
+
+    return mode  # type: ignore
+
+
+# ============================================================================
+# Validation
+# ============================================================================
+
+
+def validate_database_routing(
+    tool_collection: ToolCollection,
+    mode: DatabaseRoutingMode,
+) -> None:
+    """
+    Validate that the required databases are enabled for the routing mode.
+
+    Args:
+        tool_collection: The tool collection type
+        mode: The database routing mode
+
+    Raises:
+        DatabaseRoutingError: If validation fails
+    """
+    config = get_tool_config(tool_collection)
+
+    if mode == "turbopuffer":
+        if not config.turbopuffer.enabled:
+            raise DatabaseRoutingError(
+                f"Database routing mode 'turbopuffer' requested but "
+                f"turbopuffer.enabled is false for {tool_collection}"
+            )
+
+    elif mode == "pinecone":
+        if not config.pinecone.enabled:
+            raise DatabaseRoutingError(
+                f"Database routing mode 'pinecone' requested but "
+                f"pinecone.enabled is false for {tool_collection}"
+            )
+
+    elif mode == "dual":
+        # In dual mode, at least one database must be enabled
+        if not config.turbopuffer.enabled and not config.pinecone.enabled:
+            raise DatabaseRoutingError(
+                f"No databases enabled for {tool_collection}. "
+                "At least one of turbopuffer or pinecone must be enabled."
+            )
+
+
+# ============================================================================
+# Storage Config Builder
+# ============================================================================
+
+
+def build_storage_config(
+    tool_collection: ToolCollection,
+    sub_type: Optional[str],
+    content_hash: str,
+    document_fields: dict[str, Any],
+) -> StorageConfig:
+    """
+    Build the storage configuration for a structured embedding request.
+
+    This function:
+    1. Gets the routing mode from environment
+    2. Validates that required databases are enabled
+    3. Builds the storage config with appropriate namespaces
+
+    Args:
+        tool_collection: Tool collection type
+        sub_type: Sub-type (FlashCardType or QuestionType)
+        content_hash: Content hash (used as vector ID)
+        document_fields: Additional document fields to store
+
+    Returns:
+        StorageConfig for the embedding request
+
+    Raises:
+        DatabaseRoutingError: If validation fails
+    """
+    mode = get_database_routing_mode()
+    validate_database_routing(tool_collection, mode)
+
+    config = get_tool_config(tool_collection)
+
+    # Build TurboPuffer config if enabled and requested
+    turbopuffer = None
+    include_turbopuffer = config.turbopuffer.enabled and mode in (
+        "turbopuffer",
+        "dual",
+    )
+
+    if include_turbopuffer:
+        namespace = get_turbopuffer_namespace(tool_collection, sub_type)
+        turbopuffer = TurboPufferStorage(
+            namespace=namespace,
+            id_field=config.turbopuffer.id_field,
+            metadata=list(config.turbopuffer.metadata_fields),
+        )
+
+    # Build Pinecone config if enabled and requested
+    pinecone = None
+    include_pinecone = config.pinecone.enabled and mode in ("pinecone", "dual")
+
+    if include_pinecone:
+        namespace = get_pinecone_namespace(tool_collection, sub_type)
+        pinecone = PineconeStorageConfig(
+            index_name=config.pinecone.index_name,
+            namespace=namespace,
+            id_field=config.pinecone.id_field,
+            metadata=list(config.pinecone.metadata_fields),
+        )
+
+    return StorageConfig(
+        turbopuffer=turbopuffer,
+        pinecone=pinecone,
+    )
+
+
+def get_content_type(tool_collection: ToolCollection) -> str:
+    """
+    Get the content type string for the embedding request.
+    Maps tool collections to the contentType used in embedding requests.
+    """
+    mapping: dict[ToolCollection, str] = {
+        "FlashCard": "flashcard",
+        "TestQuestion": "testquestion",
+        "SpacedTestQuestion": "spacedtestquestion",
+        "AudioRecapV2Section": "audiorecap",
+    }
+    return mapping.get(tool_collection, "tool")