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

vector_sdk/namespaces/search.py

@@ -151,26 +151,21 @@ class SearchNamespace(BaseNamespace):
         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()
+        list_key = f"query:response:{request_id}"
+
+        # BRPOP blocks until result is available or timeout
+        result = self._redis.brpop(list_key, timeout=timeout)
+
+        if result is None:
+            raise TimeoutError(
+                f"No query result received for {request_id} within {timeout}s"
+            )
+
+        # result = (key, value)
+        data = json.loads(result[1])
+        # Cleanup the response list
+        self._redis.delete(list_key)
+        return QueryResult.from_dict(data)
 
     def query_and_wait(
         self,
@@ -193,8 +188,8 @@ class SearchNamespace(BaseNamespace):
         """
         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.
+        Uses BRPOP for efficient blocking wait - no race condition since the result
+        is pushed to a list that persists until consumed.
 
         Args:
             query_text: The text to search for
@@ -217,42 +212,25 @@ class SearchNamespace(BaseNamespace):
             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()
+        # Submit the request first
+        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,
+        )
+
+        # Wait for result via BRPOP
+        return self.wait_for(request_id, timeout)

vector_sdk/structured/__init__.py

@@ -14,9 +14,15 @@ from .router import (
     validate_database_routing,
 )
 from .structured_embeddings import (
+    AudioRecapBatchItem,
+    BatchItem,
+    FlashCardBatchItem,
     StructuredEmbeddingsNamespace,
+    TestQuestionBatchItem,
     TestQuestionInput,
     ToolMetadata,
+    TopicBatchItem,
+    TopicMetadata,
 )
 from .tool_config import (
     TOOL_CONFIGS,
@@ -37,7 +43,14 @@ __all__ = [
     "StructuredEmbeddingsNamespace",
     # Types
     "ToolMetadata",
+    "TopicMetadata",
     "TestQuestionInput",
+    # Batch types
+    "BatchItem",
+    "FlashCardBatchItem",
+    "TestQuestionBatchItem",
+    "AudioRecapBatchItem",
+    "TopicBatchItem",
     # Tool configuration
     "ToolConfig",
     "ToolDatabaseConfig",