feat: Support for hyperscale and composite index (#17)

Author: shyam-cb

README.md

@@ -0,0 +1,101 @@
+import json
+import logging
+import os
+import zipfile
+from datetime import timedelta
+from io import BytesIO
+from pathlib import Path
+
+import requests
+from couchbase.n1ql import QueryScanConsistency
+from haystack import Pipeline
+from haystack.components.converters import TextFileToDocument
+from haystack.components.embedders import SentenceTransformersDocumentEmbedder
+from haystack.components.preprocessors import DocumentCleaner, DocumentSplitter
+from haystack.components.writers import DocumentWriter
+from haystack.utils import Secret
+
+from couchbase_haystack import (
+    CouchbasePasswordAuthenticator,
+    CouchbaseQueryDocumentStore,
+    CouchbaseQueryOptions,
+    QueryVectorSearchType,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def fetch_archive_from_http(url: str, output_dir: str):
+    if Path(output_dir).is_dir():
+        logger.warn(f"'{output_dir}' directory already exists. Skipping data download")
+        return
+
+    with requests.get(url, timeout=10, stream=True) as response:
+        with zipfile.ZipFile(BytesIO(response.content)) as zip_ref:
+            zip_ref.extractall(output_dir)
+
+
+# Let's first get some files that we want to use
+docs_dir = "data/docs"
+fetch_archive_from_http(
+    url="https://s3.eu-central-1.amazonaws.com/deepset.ai-farm-qa/datasets/documents/wiki_gameofthrones_txt6.zip",
+    output_dir=docs_dir,
+)
+
+# Make sure you have a running couchbase database, e.g. with Docker:
+# docker run \
+#     --restart always \
+#     --publish=8091-8096:8091-8096 --publish=11210:11210 \
+#     --env COUCHBASE_ADMINISTRATOR_USERNAME=admin \
+#     --env COUCHBASE_ADMINISTRATOR_PASSWORD=passw0rd \
+#     couchbase:enterprise-8.0.0
+
+bucket_name = "haystack_bucket_name"
+scope_name = "haystack_scope_name"
+collection_name = "haystack_collection_name"
+index_name = "haystack_index_name"
+
+document_store = CouchbaseQueryDocumentStore(
+    cluster_connection_string=Secret.from_env_var("CONNECTION_STRING"),
+    authenticator=CouchbasePasswordAuthenticator(username=Secret.from_token("username"), password=Secret.from_token("password")),
+    bucket=bucket_name,
+    scope=scope_name,
+    collection=collection_name,
+    search_type=QueryVectorSearchType.ANN,
+    similarity="L2",
+    nprobes=10,
+    query_options=CouchbaseQueryOptions(timeout=timedelta(seconds=300), scan_consistency=QueryScanConsistency.REQUEST_PLUS),
+)
+
+
+# Create components and an indexing pipeline that converts txt to documents, cleans and splits them, and
+# indexes them for dense retrieval.
+p = Pipeline()
+p.add_component("text_file_converter", TextFileToDocument())
+p.add_component("cleaner", DocumentCleaner())
+p.add_component("splitter", DocumentSplitter(split_by="sentence", split_length=250, split_overlap=30))
+p.add_component("embedder", SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"))
+p.add_component("writer", DocumentWriter(document_store=document_store))
+
+p.connect("text_file_converter.documents", "cleaner.documents")
+p.connect("cleaner.documents", "splitter.documents")
+p.connect("splitter.documents", "embedder.documents")
+p.connect("embedder.documents", "writer.documents")
+
+# Take the docs data directory as input and run the pipeline
+file_paths = [docs_dir / Path(name) for name in os.listdir(docs_dir)]
+result = p.run({"text_file_converter": {"sources": file_paths}})
+
+# Assuming you have a Docker container running, navigate to <http://localhost:8091>
+# to open the Couchbase Web Console and explore your data.
+
+
+# currently index needs to be created after some documents are available in the collection for training
+# this is a current limitation of the couchbase gsi vector index
+cfg = {
+    "dimension": 384,
+    "train_list": 500,
+    "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)}")

examples/gsi/indexing_pipeline.py

@@ -0,0 +1,95 @@
+from couchbase.options import KnownConfigProfiles
+from haystack import GeneratedAnswer, Pipeline
+from haystack.components.builders.answer_builder import AnswerBuilder
+from haystack.components.builders.prompt_builder import PromptBuilder
+from haystack.components.embedders import SentenceTransformersTextEmbedder
+from haystack.components.generators import HuggingFaceAPIGenerator
+from haystack.utils import Secret
+
+from couchbase_haystack import (
+    CouchbaseClusterOptions,
+    CouchbasePasswordAuthenticator,
+    CouchbaseQueryDocumentStore,
+    CouchbaseQueryEmbeddingRetriever,
+    QueryVectorSearchType,
+)
+
+# Load HF Token from environment variables.
+HF_TOKEN = Secret.from_env_var("HF_API_TOKEN")
+
+# Make sure you have a running couchbase database, e.g. with Docker:
+# docker run \
+#     --restart always \
+#     --publish=8091-8096:8091-8096 --publish=11210:11210 \
+#     --env COUCHBASE_ADMINISTRATOR_USERNAME=admin \
+#     --env COUCHBASE_ADMINISTRATOR_PASSWORD=passw0rd \
+#     couchbase:enterprise-8.0.0
+
+document_store = CouchbaseQueryDocumentStore(
+    cluster_connection_string=Secret.from_token("localhost"),
+    authenticator=CouchbasePasswordAuthenticator(username=Secret.from_token("username"), password=Secret.from_token("password")),
+    cluster_options=CouchbaseClusterOptions(
+        profile=KnownConfigProfiles.WanDevelopment,
+    ),
+    bucket="haystack_bucket_name",
+    scope="haystack_scope_name",
+    collection="haystack_collection_name",
+    search_type=QueryVectorSearchType.ANN,
+    similarity="L2",
+    nprobes=10,
+)
+
+# Build a RAG pipeline with a Retriever to get relevant documents to the query and a HuggingFaceTGIGenerator
+# interacting with LLMs using a custom prompt.
+prompt_template = """
+Given these documents, answer the question.\nDocuments:
+{% for doc in documents %}
+    {{ doc.content }}
+{% endfor %}
+
+\nQuestion: {{question}}
+\nAnswer:
+"""
+rag_pipeline = Pipeline()
+rag_pipeline.add_component(
+    "query_embedder",
+    SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2", progress_bar=False),
+)
+rag_pipeline.add_component("retriever", CouchbaseQueryEmbeddingRetriever(document_store=document_store))
+rag_pipeline.add_component("prompt_builder", PromptBuilder(template=prompt_template))
+rag_pipeline.add_component(
+    "llm",
+    HuggingFaceAPIGenerator(
+        api_type="serverless_inference_api",
+        api_params={"model": "mistralai/Mistral-7B-v0.1"},
+    ),
+)
+rag_pipeline.add_component("answer_builder", AnswerBuilder())
+
+rag_pipeline.connect("query_embedder", "retriever.query_embedding")
+rag_pipeline.connect("retriever.documents", "prompt_builder.documents")
+rag_pipeline.connect("prompt_builder.prompt", "llm.prompt")
+rag_pipeline.connect("llm.replies", "answer_builder.replies")
+rag_pipeline.connect("llm.meta", "answer_builder.meta")
+rag_pipeline.connect("retriever", "answer_builder.documents")
+
+# Ask a question on the data you just added.
+question = "Who created the Dothraki vocabulary?"
+result = rag_pipeline.run(
+    {
+        "query_embedder": {"text": question},
+        "retriever": {"top_k": 3},
+        "prompt_builder": {"question": question},
+        "answer_builder": {"query": question},
+    }
+)
+
+# For details, like which documents were used to generate the answer, look into the GeneratedAnswer object
+answer: GeneratedAnswer = result["answer_builder"]["answers"][0]
+
+# ruff: noqa: T201
+print("Query: ", answer.query)
+print("Answer: ", answer.data)
+print("== Sources:")
+for doc in answer.documents:
+    print("-> ", doc.meta["file_path"])

examples/gsi/rag_pipeline.py

@@ -8,20 +8,29 @@
 It includes authentication methods, document store implementation, and embedding-based retrieval functionality.
 """
 
-from couchbase_haystack.components.retrievers import CouchbaseSearchEmbeddingRetriever
+from couchbase_haystack.components.retrievers import CouchbaseQueryEmbeddingRetriever, CouchbaseSearchEmbeddingRetriever
 from couchbase_haystack.document_stores import (
     CouchbaseAuthenticator,
     CouchbaseCertificateAuthenticator,
     CouchbaseClusterOptions,
     CouchbasePasswordAuthenticator,
+    CouchbaseQueryDocumentStore,
+    CouchbaseQueryOptions,
     CouchbaseSearchDocumentStore,
+    QueryVectorSearchSimilarity,
+    QueryVectorSearchType,
 )
 
 __all__ = [
     "CouchbaseAuthenticator",
     "CouchbaseCertificateAuthenticator",
     "CouchbaseClusterOptions",
     "CouchbasePasswordAuthenticator",
+    "CouchbaseQueryDocumentStore",
+    "CouchbaseQueryEmbeddingRetriever",
+    "CouchbaseQueryOptions",
     "CouchbaseSearchDocumentStore",
     "CouchbaseSearchEmbeddingRetriever",
+    "QueryVectorSearchSimilarity",
+    "QueryVectorSearchType",
 ]

examples/indexing_pipeline.py examples/search/indexing_pipeline.pyexamples/indexing_pipeline.py renamed to examples/search/indexing_pipeline.py

@@ -1,3 +1,3 @@
-from .embedding_retriever import CouchbaseSearchEmbeddingRetriever
+from .embedding_retriever import CouchbaseQueryEmbeddingRetriever, CouchbaseSearchEmbeddingRetriever
 
-__all__ = ["CouchbaseSearchEmbeddingRetriever"]
+__all__ = ["CouchbaseQueryEmbeddingRetriever", "CouchbaseSearchEmbeddingRetriever"]