code cleanup v2

Author: shyam-cb

examples/query/indexing_pipeline.py

@@ -8,6 +8,7 @@
 
 import requests
 from couchbase.n1ql import QueryScanConsistency
+from couchbase.options import KnownConfigProfiles, QueryOptions
 from haystack import Pipeline
 from haystack.components.converters import TextFileToDocument
 from haystack.components.embedders import SentenceTransformersDocumentEmbedder
@@ -16,13 +17,12 @@
 from haystack.utils import Secret
 
 from couchbase_haystack import (
+    CouchbaseClusterOptions,
     CouchbasePasswordAuthenticator,
     CouchbaseQueryDocumentStore,
     CouchbaseQueryOptions,
     QueryVectorSearchType,
-    CouchbaseClusterOptions,
 )
-from couchbase.options import KnownConfigProfiles, QueryOptions
 
 logger = logging.getLogger(__name__)
 
@@ -59,7 +59,9 @@ def fetch_archive_from_http(url: str, output_dir: str):
 
 document_store = CouchbaseQueryDocumentStore(
     cluster_connection_string=Secret.from_env_var("CONNECTION_STRING"),
-    authenticator=CouchbasePasswordAuthenticator(username=Secret.from_env_var("USER_NAME"), password=Secret.from_env_var("PASSWORD")),
+    authenticator=CouchbasePasswordAuthenticator(
+        username=Secret.from_env_var("USER_NAME"), password=Secret.from_env_var("PASSWORD")
+    ),
     cluster_options=CouchbaseClusterOptions(profile=KnownConfigProfiles.WanDevelopment),
     bucket=bucket_name,
     scope=scope_name,
@@ -102,6 +104,9 @@ def fetch_archive_from_http(url: str, output_dir: str):
     "description": "IVF,PQ32x8",
     "similarity": "L2",
 }
-document_store.scope.query(f"Create Index {index_name} ON {collection_name} (embedding vector) USING GSI WITH {json.dumps(cfg)}", QueryOptions(timeout=timedelta(seconds=300))).execute()
+document_store.scope.query(
+    f"Create Index {index_name} ON {collection_name} (embedding vector) USING GSI WITH {json.dumps(cfg)}",
+    QueryOptions(timeout=timedelta(seconds=300)),
+).execute()
 
-logger.info(f"Index created: {index_name}")
+logger.info(f"Index created: {index_name}")

examples/query/rag_pipeline.py

@@ -1,20 +1,22 @@
 import os
+
+from couchbase.options import KnownConfigProfiles
 from haystack import GeneratedAnswer, Pipeline
 from haystack.components.builders.answer_builder import AnswerBuilder
 from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
-from haystack.dataclasses import ChatMessage
 from haystack.components.embedders import SentenceTransformersTextEmbedder
 from haystack.components.generators.chat import HuggingFaceAPIChatGenerator
-from haystack.utils.hf import HFGenerationAPIType
+from haystack.dataclasses import ChatMessage
 from haystack.utils import Secret
+from haystack.utils.hf import HFGenerationAPIType
+
 from couchbase_haystack import (
     CouchbaseClusterOptions,
     CouchbasePasswordAuthenticator,
     CouchbaseQueryDocumentStore,
     CouchbaseQueryEmbeddingRetriever,
     QueryVectorSearchType,
 )
-from couchbase.options import KnownConfigProfiles
 
 # Load HF Token from environment variables.
 HF_TOKEN = Secret.from_env_var("HF_API_TOKEN")
@@ -29,7 +31,9 @@
 
 document_store = CouchbaseQueryDocumentStore(
     cluster_connection_string=Secret.from_env_var("CONNECTION_STRING"),
-    authenticator=CouchbasePasswordAuthenticator(username=Secret.from_env_var("USER_NAME"), password=Secret.from_env_var("PASSWORD")),
+    authenticator=CouchbasePasswordAuthenticator(
+        username=Secret.from_env_var("USER_NAME"), password=Secret.from_env_var("PASSWORD")
+    ),
     cluster_options=CouchbaseClusterOptions(
         profile=KnownConfigProfiles.WanDevelopment,
     ),
@@ -45,14 +49,16 @@
 # interacting with LLMs using a custom prompt.
 prompt_messages = [
     ChatMessage.from_system("You are a helpful assistant that answers questions based on the provided documents."),
-    ChatMessage.from_user("""Given these documents, answer the question.
+    ChatMessage.from_user(
+        """Given these documents, answer the question.
 Documents:
 {% for doc in documents %}
     {{ doc.content }}
 {% endfor %}
 
 Question: {{question}}
-Answer:""")
+Answer:"""
+    ),
 ]
 rag_pipeline = Pipeline()
 rag_pipeline.add_component(
@@ -61,10 +67,13 @@
 )
 rag_pipeline.add_component("retriever", CouchbaseQueryEmbeddingRetriever(document_store=document_store))
 rag_pipeline.add_component("prompt_builder", ChatPromptBuilder(template=prompt_messages, required_variables=["question"]))
-rag_pipeline.add_component("llm", HuggingFaceAPIChatGenerator(
-    api_type=HFGenerationAPIType.SERVERLESS_INFERENCE_API, 
-    api_params={"model": "mistralai/Mistral-7B-Instruct-v0.2"},
-))
+rag_pipeline.add_component(
+    "llm",
+    HuggingFaceAPIChatGenerator(
+        api_type=HFGenerationAPIType.SERVERLESS_INFERENCE_API,
+        api_params={"model": "mistralai/Mistral-7B-Instruct-v0.2"},
+    ),
+)
 rag_pipeline.add_component("answer_builder", AnswerBuilder())
 
 rag_pipeline.connect("query_embedder", "retriever.query_embedding")

examples/search/rag_pipeline.py

@@ -1,12 +1,13 @@
 import os
+
 from haystack import GeneratedAnswer, Pipeline
 from haystack.components.builders.answer_builder import AnswerBuilder
 from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
-from haystack.dataclasses import ChatMessage
 from haystack.components.embedders import SentenceTransformersTextEmbedder
 from haystack.components.generators.chat import HuggingFaceAPIChatGenerator
-from haystack.utils.hf import HFGenerationAPIType
+from haystack.dataclasses import ChatMessage
 from haystack.utils import Secret
+from haystack.utils.hf import HFGenerationAPIType
 
 from couchbase_haystack import CouchbasePasswordAuthenticator, CouchbaseSearchDocumentStore, CouchbaseSearchEmbeddingRetriever
 
@@ -23,7 +24,9 @@
 
 document_store = CouchbaseSearchDocumentStore(
     cluster_connection_string=Secret.from_env_var("CONNECTION_STRING"),
-    authenticator=CouchbasePasswordAuthenticator(username=Secret.from_env_var("USER_NAME"), password=Secret.from_env_var("PASSWORD")),
+    authenticator=CouchbasePasswordAuthenticator(
+        username=Secret.from_env_var("USER_NAME"), password=Secret.from_env_var("PASSWORD")
+    ),
     bucket=os.getenv("BUCKET_NAME"),
     scope=os.getenv("SCOPE_NAME"),
     collection=os.getenv("COLLECTION_NAME"),
@@ -34,14 +37,16 @@
 # interacting with LLMs using a custom prompt.
 prompt_messages = [
     ChatMessage.from_system("You are a helpful assistant that answers questions based on the provided documents."),
-    ChatMessage.from_user("""Given these documents, answer the question.
+    ChatMessage.from_user(
+        """Given these documents, answer the question.
 Documents:
 {% for doc in documents %}
     {{ doc.content }}
 {% endfor %}
 
 Question: {{question}}
-Answer:""")
+Answer:"""
+    ),
 ]
 rag_pipeline = Pipeline()
 rag_pipeline.add_component(
@@ -50,10 +55,13 @@
 )
 rag_pipeline.add_component("retriever", CouchbaseSearchEmbeddingRetriever(document_store=document_store))
 rag_pipeline.add_component("prompt_builder", ChatPromptBuilder(template=prompt_messages, required_variables=["question"]))
-rag_pipeline.add_component("llm", HuggingFaceAPIChatGenerator(
-    api_type=HFGenerationAPIType.SERVERLESS_INFERENCE_API, 
-    api_params={"model": "mistralai/Mistral-7B-Instruct-v0.2"},
-))
+rag_pipeline.add_component(
+    "llm",
+    HuggingFaceAPIChatGenerator(
+        api_type=HFGenerationAPIType.SERVERLESS_INFERENCE_API,
+        api_params={"model": "mistralai/Mistral-7B-Instruct-v0.2"},
+    ),
+)
 rag_pipeline.add_component("answer_builder", AnswerBuilder())
 
 rag_pipeline.connect("query_embedder", "retriever.query_embedding")

src/couchbase_haystack/components/retrievers/embedding_retriever.py

@@ -16,7 +16,7 @@
 @component
 class CouchbaseSearchEmbeddingRetriever:
     """Retrieves documents from the CouchbaseSearchDocumentStore by embedding similarity.
-    
+
     Uses Search Vector Index (FTS-based) for hybrid searches combining vector, full-text, and geospatial queries.
     See CouchbaseSearchDocumentStore for more information.
 
@@ -139,7 +139,7 @@ def run(
 @component
 class CouchbaseQueryEmbeddingRetriever:
     """Retrieves documents from the CouchbaseQueryDocumentStore using vector similarity search.
-    
+
     Works with both Hyperscale Vector Index and Composite Vector Index.
     Supports ANN (approximate) and KNN (exact) search with various similarity metrics.
     See CouchbaseQueryDocumentStore for more details.

src/couchbase_haystack/document_stores/document_store.py

@@ -580,18 +580,18 @@ def __get_doc_from_kv(self, response: SearchResult) -> List[Document]:
 
 class CouchbaseQueryDocumentStore(CouchbaseDocumentStore):
     """CouchbaseQueryDocumentStore uses Couchbase Global Secondary Index (GSI) for high-performance vector search.
-    
+
     Supports two types of vector indexes:
-    
+
     - **Hyperscale Vector Indexes**: Optimized for pure vector searches, scales to billions of documents.
       Best for chatbot context (RAG), reverse image search, and anomaly detection.
-    
+
     - **Composite Vector Indexes**: Combines vector and scalar indexing. Applies scalar filters before vector search.
       Best for filtered recommendations, job searches, and supply chain management.
-    
+
     Search types: ANN (fast, approximate) or KNN (exact).
     Similarity metrics: COSINE, DOT, L2/EUCLIDEAN, L2_SQUARED/EUCLIDEAN_SQUARED.
-    
+
     See [Couchbase documentation](https://docs.couchbase.com/server/current/vector-index/use-vector-indexes.html).
     """
 
@@ -802,7 +802,6 @@ def _embedding_retrieval(
         """  # noqa: S608  # query_vector_str is a float array, where_clause is normalized by normalize_sql_filters
 
         try:
-
             query_options = self.query_options.cb_query_options()
             # Execute the query
             result: QueryResult = self.connection.query(

src/couchbase_haystack/telemetry.py

@@ -57,4 +57,4 @@ def send_telemetry() -> None:
     """
     with suppress(Exception):
         t = threading.Thread(target=_send_telemetry, daemon=True)
-        t.start()
+        t.start()

tests/search_document_store/test_document_store.py

@@ -149,7 +149,6 @@ def document_store(self):
         cluster.close()
 
     def assert_documents_are_equal(self, received: List[Document], expected: List[Document]):
-
         for r in received:
             r.score = None
             r.embedding = None
@@ -283,7 +282,6 @@ class TestSearchDocumentStoreUnit:
     @pytest.fixture
     def document_store(self):
         with patch("couchbase_haystack.document_stores.document_store.Cluster") as mock_cb_cluster:
-
             cluster = mock_cb_cluster.return_value
             bucket = cluster.bucket.return_value
             scope = bucket.scope.return_value

tests/search_document_store/test_search_filter.py

@@ -216,7 +216,6 @@ def test_filter_lt_condition_array_of_number(self):
 
 @pytest.mark.unit
 class TestFilterLTE:
-
     # def test_filter_gt_condition_str(self):
     #    _filter = {"field": "meta.years", "operator": "==", "value": "2019"}
     #    normalized_filter = _normalize_filters(_filter)