RAG with CrewAI using Couchbase Hyperscale and Composite Vector Index

%pip install --quiet datasets==4.1.0 langchain-couchbase==0.5.0 langchain-openai==0.3.33 crewai==0.186.1 python-dotenv==1.1.1

Note: you may need to restart the kernel to use updated packages.

import getpass
import json
import logging
import os
import time
from datetime import timedelta
from uuid import uuid4

from couchbase.auth import PasswordAuthenticator
from couchbase.cluster import Cluster
from couchbase.diagnostics import PingState, ServiceType
from couchbase.exceptions import (InternalServerFailureException,
                                  QueryIndexAlreadyExistsException,
                                  ServiceUnavailableException,
                                  CouchbaseException)
from couchbase.management.buckets import CreateBucketSettings
from couchbase.options import ClusterOptions
from datasets import load_dataset
from dotenv import load_dotenv
from crewai.tools import tool
from langchain_couchbase.vectorstores import CouchbaseQueryVectorStore
from langchain_couchbase.vectorstores import DistanceStrategy, IndexType
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

from crewai import Agent, Crew, Process, Task

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

# Suppress httpx logging
logging.getLogger('httpx').setLevel(logging.CRITICAL)

# Load environment variables
load_dotenv("./.env")

# Configuration
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY') or input("Enter your OpenAI API key: ")
if not OPENAI_API_KEY:
    raise ValueError("OPENAI_API_KEY is not set")

CB_HOST = os.getenv('CB_HOST') or 'couchbase://localhost'
CB_USERNAME = os.getenv('CB_USERNAME') or 'Administrator'
CB_PASSWORD = os.getenv('CB_PASSWORD') or 'password'
CB_BUCKET_NAME = os.getenv('CB_BUCKET_NAME') or 'vector-search-testing'
SCOPE_NAME = os.getenv('SCOPE_NAME') or 'shared'
COLLECTION_NAME = os.getenv('COLLECTION_NAME') or 'crew'

print("Configuration loaded successfully")

Configuration loaded successfully

# Connect to Couchbase
try:
    auth = PasswordAuthenticator(CB_USERNAME, CB_PASSWORD)
    options = ClusterOptions(auth)
    cluster = Cluster(CB_HOST, options)
    cluster.wait_until_ready(timedelta(seconds=5))
    print("Successfully connected to Couchbase")
except Exception as e:
    print(f"Failed to connect to Couchbase: {str(e)}")
    raise

Successfully connected to Couchbase

def setup_collection(cluster, bucket_name, scope_name, collection_name):
    try:
        # Check if bucket exists, create if it doesn't
        try:
            bucket = cluster.bucket(bucket_name)
            logging.info(f"Bucket '{bucket_name}' exists.")
        except Exception as e:
            logging.info(f"Bucket '{bucket_name}' does not exist. Creating it...")
            bucket_settings = CreateBucketSettings(
                name=bucket_name,
                bucket_type='couchbase',
                ram_quota_mb=1024,
                flush_enabled=True,
                num_replicas=0
            )
            cluster.buckets().create_bucket(bucket_settings)
            time.sleep(2)  # Wait for bucket creation to complete and become available
            bucket = cluster.bucket(bucket_name)
            logging.info(f"Bucket '{bucket_name}' created successfully.")

        bucket_manager = bucket.collections()

        # Check if scope exists, create if it doesn't
        scopes = bucket_manager.get_all_scopes()
        scope_exists = any(scope.name == scope_name for scope in scopes)
        
        if not scope_exists and scope_name != "_default":
            logging.info(f"Scope '{scope_name}' does not exist. Creating it...")
            bucket_manager.create_scope(scope_name)
            logging.info(f"Scope '{scope_name}' created successfully.")

        # Check if collection exists, create if it doesn't
        collections = bucket_manager.get_all_scopes()
        collection_exists = any(
            scope.name == scope_name and collection_name in [col.name for col in scope.collections]
            for scope in collections
        )

        if not collection_exists:
            logging.info(f"Collection '{collection_name}' does not exist. Creating it...")
            bucket_manager.create_collection(scope_name, collection_name)
            logging.info(f"Collection '{collection_name}' created successfully.")
        else:
            logging.info(f"Collection '{collection_name}' already exists. Skipping creation.")

        # Wait for collection to be ready
        collection = bucket.scope(scope_name).collection(collection_name)
        time.sleep(2)  # Give the collection time to be ready for queries

        # Clear all documents in the collection
        try:
            query = f"DELETE FROM `{bucket_name}`.`{scope_name}`.`{collection_name}`"
            cluster.query(query).execute()
            logging.info("All documents cleared from the collection.")
        except Exception as e:
            logging.warning(f"Error while clearing documents: {str(e)}. The collection might be empty.")

        return collection
    except Exception as e:
        raise RuntimeError(f"Error setting up collection: {str(e)}")
    
setup_collection(cluster, CB_BUCKET_NAME, SCOPE_NAME, COLLECTION_NAME)

2025-10-06 10:17:53 [INFO] Bucket 'vector-search-testing' exists.
2025-10-06 10:17:53 [INFO] Collection 'crew' already exists. Skipping creation.
2025-10-06 10:17:55 [INFO] All documents cleared from the collection.





<couchbase.collection.Collection at 0x307407a10>

# Hyperscale and Composite Vector Index Configuration
# Unlike Search Vector Index, Hyperscale and Composite vector indexes are created programmatically through the vector store
# We'll configure the parameters that will be used for index creation

# Vector configuration
DISTANCE_STRATEGY = DistanceStrategy.COSINE  # Cosine similarity
INDEX_TYPE = IndexType.HYPERSCALE  # Using HYPERSCALE for high-performance vector search
INDEX_DESCRIPTION = "IVF,SQ8"  # Auto-selected centroids with 8-bit scalar quantization

# To create a Composite Index instead, use the following:
# INDEX_TYPE = IndexType.COMPOSITE  # Combines vector search with scalar filtering

print("Hyperscale and Composite vector index configuration prepared")

Hyperscale and Composite vector index configuration prepared

# Alternative configuration for Composite index
INDEX_TYPE = IndexType.COMPOSITE  # Instead of IndexType.HYPERSCALE
INDEX_DESCRIPTION = "IVF,SQ8"     # Same quantization settings
DISTANCE_STRATEGY = DistanceStrategy.COSINE  # Same distance metric

# The rest of the setup remains identical

# Initialize OpenAI components
embeddings = OpenAIEmbeddings(
    openai_api_key=OPENAI_API_KEY,
    model="text-embedding-3-small"
)

llm = ChatOpenAI(
    openai_api_key=OPENAI_API_KEY,
    model="gpt-4o",
    temperature=0.2
)

print("OpenAI components initialized")

OpenAI components initialized

# Setup Hyperscale vector store with OpenAI embeddings
try:
    vector_store = CouchbaseQueryVectorStore(
        cluster=cluster,
        bucket_name=CB_BUCKET_NAME,
        scope_name=SCOPE_NAME,
        collection_name=COLLECTION_NAME,
        embedding=embeddings,
        distance_metric=DISTANCE_STRATEGY
    )
    print("Hyperscale Vector store initialized successfully")
    logging.info("Hyperscale Vector store setup completed")
except Exception as e:
    logging.error(f"Failed to initialize Hyperscale vector store: {str(e)}")
    raise RuntimeError(f"Hyperscale Vector store initialization failed: {str(e)}")

2025-10-06 10:18:05 [INFO] Hyperscale Vector store setup completed


Hyperscale Vector store initialized successfully

try:
    news_dataset = load_dataset(
        "RealTimeData/bbc_news_alltime", "2024-12", split="train"
    )
    print(f"Loaded the BBC News dataset with {len(news_dataset)} rows")
    logging.info(f"Successfully loaded the BBC News dataset with {len(news_dataset)} rows.")
except Exception as e:
    raise ValueError(f"Error loading the BBC News dataset: {str(e)}")

2025-10-06 10:18:13 [INFO] Successfully loaded the BBC News dataset with 2687 rows.


Loaded the BBC News dataset with 2687 rows

news_articles = news_dataset["content"]
unique_articles = set()
for article in news_articles:
    if article:
        unique_articles.add(article)
unique_news_articles = list(unique_articles)
print(f"We have {len(unique_news_articles)} unique articles in our database.")

We have 1749 unique articles in our database.

batch_size = 50

# Automatic Batch Processing
articles = [article for article in unique_news_articles if article and len(article) <= 50000]

try:
    vector_store.add_texts(
        texts=articles,
        batch_size=batch_size
    )
    logging.info("Document ingestion completed successfully.")
except Exception as e:
    raise ValueError(f"Failed to save documents to vector store: {str(e)}")

2025-10-06 10:19:43 [INFO] Document ingestion completed successfully.

import time

# Create Hyperscale vector retriever optimized for high-performance searches
retriever = vector_store.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 4}  # Return top 4 most similar documents
)

def test_vector_search_performance(query_text, label="Vector Search"):
    """Test pure vector search performance and return timing metrics"""
    print(f"\n[{label}] Testing vector search performance")
    print(f"[{label}] Query: '{query_text}'")
    
    start_time = time.time()
    
    try:
        # Perform vector search using the retriever
        docs = retriever.invoke(query_text)
        end_time = time.time()
        
        search_time = end_time - start_time
        print(f"[{label}] Vector search completed in {search_time:.4f} seconds")
        print(f"[{label}] Found {len(docs)} relevant documents")
        
        # Show a preview of the first result
        if docs:
            preview = docs[0].page_content[:100] + "..." if len(docs[0].page_content) > 100 else docs[0].page_content
            print(f"[{label}] Top result preview: {preview}")
        
        return search_time
    except Exception as e:
        print(f"[{label}] Vector search failed: {str(e)}")
        return None

# Test baseline vector search performance without Hyperscale vector index
test_query = "What are the latest developments in football transfers?"
print("Testing baseline vector search performance without Hyperscale vector index optimization...")
baseline_time = test_vector_search_performance(test_query, "Baseline Search")
print(f"\nBaseline vector search time (without Hyperscale vector index): {baseline_time:.4f} seconds\n")

Testing baseline vector search performance without Hyperscale vector index optimization...

[Baseline Search] Testing vector search performance
[Baseline Search] Query: 'What are the latest developments in football transfers?'
[Baseline Search] Vector search completed in 1.3999 seconds
[Baseline Search] Found 4 relevant documents
[Baseline Search] Top result preview: The latest updates and analysis from the BBC.

Baseline vector search time (without Hyperscale vector index): 1.3999 seconds

# Create Hyperscale Vector Index for high-performance searches
print("Creating Hyperscale vector index...")
try:
    # Create a Hyperscale index optimized for pure vector searches
    vector_store.create_index(
        index_type=INDEX_TYPE,  # Hyperscale index type
        index_description=INDEX_DESCRIPTION  # IVF,SQ8 for optimized performance
    )
    print(f"Hyperscale Vector index created successfully")
    logging.info(f"Hyperscale index created with description '{INDEX_DESCRIPTION}'")
    
    # Wait a moment for index to be available
    print("Waiting for index to become available...")
    time.sleep(5)
    
except Exception as e:
    # Index might already exist, which is fine
    if "already exists" in str(e).lower():
        print(f"Hyperscale Vector index already exists, proceeding...")
        logging.info(f"Index already exists")
    else:
        logging.error(f"Failed to create Hyperscale vector index: {str(e)}")
        raise RuntimeError(f"Hyperscale vector index creation failed: {str(e)}")

Creating Hyperscale vector index...


2025-10-06 10:20:15 [INFO] Hyperscale index created with description 'IVF,SQ8'


Hyperscale Vector index created successfully
Waiting for index to become available...

# Test vector search performance with Hyperscale vector index
print("Testing vector search performance with Hyperscale vector index optimization...")
hyperscale_search_time = test_vector_search_performance(test_query, "Hyperscale-Optimized Search")

Testing vector search performance with Hyperscale vector index optimization...

[Hyperscale-Optimized Search] Testing vector search performance
[Hyperscale-Optimized Search] Query: 'What are the latest developments in football transfers?'
[Hyperscale-Optimized Search] Vector search completed in 0.5885 seconds
[Hyperscale-Optimized Search] Found 4 relevant documents
[Hyperscale-Optimized Search] Top result preview: Four key areas for Everton's new owners to address

Everton fans last saw silverware in 1995 when th...

# Test cache benefits with a different query to avoid interference
cache_test_query = "What happened in the latest Premier League matches?"

print("Testing cache benefits with vector search...")
print("First execution (cache miss):")
cache_time_1 = test_vector_search_performance(cache_test_query, "Cache Test - First Run")

print("\nSecond execution (cache hit - should be faster):")
cache_time_2 = test_vector_search_performance(cache_test_query, "Cache Test - Second Run")

Testing cache benefits with vector search...
First execution (cache miss):

[Cache Test - First Run] Testing vector search performance
[Cache Test - First Run] Query: 'What happened in the latest Premier League matches?'
[Cache Test - First Run] Vector search completed in 0.6450 seconds
[Cache Test - First Run] Found 4 relevant documents
[Cache Test - First Run] Top result preview: Who has made Troy's Premier League team of the week?

After every round of Premier League matches th...

Second execution (cache hit - should be faster):

[Cache Test - Second Run] Testing vector search performance
[Cache Test - Second Run] Query: 'What happened in the latest Premier League matches?'
[Cache Test - Second Run] Vector search completed in 0.4306 seconds
[Cache Test - Second Run] Found 4 relevant documents
[Cache Test - Second Run] Top result preview: Who has made Troy's Premier League team of the week?

After every round of Premier League matches th...

print("\n" + "="*80)
print("VECTOR SEARCH PERFORMANCE OPTIMIZATION SUMMARY")
print("="*80)

print(f"Phase 1 - Baseline Search (No Hyperscale):     {baseline_time:.4f} seconds")
print(f"Phase 2 - Hyperscale-Optimized Search:         {hyperscale_search_time:.4f} seconds")
if cache_time_1 and cache_time_2:
    print(f"Phase 3 - Cache Benefits:")
    print(f"  First execution (cache miss):         {cache_time_1:.4f} seconds")
    print(f"  Second execution (cache hit):         {cache_time_2:.4f} seconds")

print("\n" + "-"*80)
print("VECTOR SEARCH OPTIMIZATION IMPACT:")
print("-"*80)

# Hyperscale improvement analysis
if baseline_time and hyperscale_search_time:
    speedup = baseline_time / hyperscale_search_time if hyperscale_search_time > 0 else float('inf')
    time_saved = baseline_time - hyperscale_search_time
    percent_improvement = (time_saved / baseline_time) * 100
    print(f"Hyperscale Index Benefit:      {speedup:.2f}x faster ({percent_improvement:.1f}% improvement)")

# Cache improvement analysis
if cache_time_1 and cache_time_2 and cache_time_2 < cache_time_1:
    cache_speedup = cache_time_1 / cache_time_2
    cache_improvement = ((cache_time_1 - cache_time_2) / cache_time_1) * 100
    print(f"Cache Benefit:          {cache_speedup:.2f}x faster ({cache_improvement:.1f}% improvement)")
else:
    print(f"Cache Benefit:          Variable (depends on query complexity and caching mechanism)")

print(f"\nKey Insights for Vector Search Performance:")
print(f"• Hyperscale indexes provide significant performance improvements for vector similarity search")
print(f"• Performance gains are most dramatic for complex semantic queries")
print(f"• Hyperscale optimization is particularly effective for high-dimensional embeddings")
print(f"• Combined with proper quantization (SQ8), Hyperscale vector indexes deliver production-ready performance")
print(f"• These performance improvements directly benefit any application using the vector store")

================================================================================
VECTOR SEARCH PERFORMANCE OPTIMIZATION SUMMARY
================================================================================
Phase 1 - Baseline Search (No Hyperscale):     1.3999 seconds
Phase 2 - Hyperscale-Optimized Search:         0.5885 seconds
Phase 3 - Cache Benefits:
  First execution (cache miss):         0.6450 seconds
  Second execution (cache hit):         0.4306 seconds

--------------------------------------------------------------------------------
VECTOR SEARCH OPTIMIZATION IMPACT:
--------------------------------------------------------------------------------
Hyperscale Index Benefit:      2.38x faster (58.0% improvement)
Cache Benefit:          1.50x faster (33.2% improvement)

Key Insights for Vector Search Performance:
• Hyperscale indexes provide significant performance improvements for vector similarity search
• Performance gains are most dramatic for complex semantic queries
• Hyperscale optimization is particularly effective for high-dimensional embeddings
• Combined with proper quantization (SQ8), Hyperscale delivers production-ready performance
• These performance improvements directly benefit any application using the vector store

# Define the Hyperscale vector search tool using the @tool decorator
@tool("hyperscale_vector_search")
def search_tool(query: str) -> str:
    """Search for relevant documents using Hyperscale vector similarity.
    Input should be a simple text query string.
    Returns a list of relevant document contents from Hyperscale vector search.
    Use this tool to find detailed information about topics using high-performance Hyperscale indexes."""
    
    # Invoke the Hyperscale vector retriever (now optimized with HYPERSCALE index)
    docs = retriever.invoke(query)

    # Format the results with distance information
    formatted_docs = "\n\n".join([
        f"Document {i+1}:\n{'-'*40}\n{doc.page_content}"
        for i, doc in enumerate(docs)
    ])
    return formatted_docs

# Create research agent
researcher = Agent(
    role='Research Expert',
    goal='Find and analyze the most relevant documents to answer user queries accurately',
    backstory="""You are an expert researcher with deep knowledge in information retrieval 
    and analysis. Your expertise lies in finding, evaluating, and synthesizing information 
    from various sources. You have a keen eye for detail and can identify key insights 
    from complex documents. You always verify information across multiple sources and 
    provide comprehensive, accurate analyses.""",
    tools=[search_tool],
    llm=llm,
    verbose=False,
    memory=True,
    allow_delegation=False
)

# Create writer agent
writer = Agent(
    role='Technical Writer',
    goal='Generate clear, accurate, and well-structured responses based on research findings',
    backstory="""You are a skilled technical writer with expertise in making complex 
    information accessible and engaging. You excel at organizing information logically, 
    explaining technical concepts clearly, and creating well-structured documents. You 
    ensure all information is properly cited, accurate, and presented in a user-friendly 
    manner. You have a talent for maintaining the reader's interest while conveying 
    detailed technical information.""",
    llm=llm,
    verbose=False,
    memory=True,
    allow_delegation=False
)

print("CrewAI agents created successfully with optimized Hyperscale vector search")

CrewAI agents created successfully with optimized Hyperscale vector search

def process_interactive_query(query, researcher, writer):
    """Run complete RAG workflow with CrewAI agents using optimized Hyperscale vector search"""
    print(f"\nProcessing Query: {query}")
    print("=" * 80)
    
    # Create tasks
    research_task = Task(
        description=f"Research and analyze information relevant to: {query}",
        agent=researcher,
        expected_output="A detailed analysis with key findings"
    )
    
    writing_task = Task(
        description="Create a comprehensive response",
        agent=writer,
        expected_output="A clear, well-structured answer",
        context=[research_task]
    )
    
    # Execute crew
    crew = Crew(
        agents=[researcher, writer],
        tasks=[research_task, writing_task],
        process=Process.sequential,
        verbose=True,
        cache=True,
        planning=True
    )
    
    try:
        start_time = time.time()
        result = crew.kickoff()
        elapsed_time = time.time() - start_time
        
        print(f"\nCompleted in {elapsed_time:.2f} seconds")
        print("=" * 80)
        print("RESPONSE")
        print("=" * 80)
        print(result)
                
        return elapsed_time
    except Exception as e:
        print(f"Error: {str(e)}")
        return None

# Disable logging for cleaner output
logging.disable(logging.CRITICAL)

# Run demo with a sample query
demo_query = "What are the key details about the FA Cup third round draw?"
final_time = process_interactive_query(demo_query, researcher, writer)

if final_time:
    print(f"\n\n✅ CrewAI agent demo completed successfully in {final_time:.2f} seconds")