RAG with Couchbase and smolagents using Couchbase Hyperscale and Composite Vector Index

%pip install --quiet datasets==4.1.1 langchain-couchbase==0.5.0 langchain-openai==0.3.33 python-dotenv==1.1.1 smolagents==1.21.3

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

from couchbase.auth import PasswordAuthenticator
from couchbase.cluster import Cluster
from couchbase.exceptions import (CouchbaseException,
                                  InternalServerFailureException,
                                  QueryIndexAlreadyExistsException)
from couchbase.management.buckets import CreateBucketSettings
from couchbase.options import ClusterOptions
from datasets import load_dataset
from dotenv import load_dotenv
from langchain_couchbase.vectorstores import CouchbaseQueryVectorStore
from langchain_couchbase.vectorstores import DistanceStrategy
from langchain_couchbase.vectorstores import IndexType
from langchain_openai import OpenAIEmbeddings

from smolagents import Tool, OpenAIServerModel, ToolCallingAgent

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', force=True)

# Disable all logging except critical to prevent OpenAI API request logs
logging.getLogger("httpx").setLevel(logging.CRITICAL)

load_dotenv()

OPENAI_API_KEY = os.getenv('OPENAI_API_KEY') or getpass.getpass('Enter your OpenAI API Key: ')

CB_HOST = os.getenv('CB_HOST') or input('Enter your Couchbase host (default: couchbase://localhost): ') or 'couchbase://localhost'
CB_USERNAME = os.getenv('CB_USERNAME') or input('Enter your Couchbase username (default: Administrator): ') or 'Administrator'
CB_PASSWORD = os.getenv('CB_PASSWORD') or getpass.getpass('Enter your Couchbase password (default: password): ') or 'password'
CB_BUCKET_NAME = os.getenv('CB_BUCKET_NAME') or input('Enter your Couchbase bucket name (default: query-vector-search-testing): ') or 'query-vector-search-testing'
SCOPE_NAME = os.getenv('SCOPE_NAME') or input('Enter your scope name (default: shared): ') or 'shared'
COLLECTION_NAME = os.getenv('COLLECTION_NAME') or input('Enter your collection name (default: smolagents): ') or 'smolagents'

# Check if the variables are correctly loaded
if not OPENAI_API_KEY:
    raise ValueError("Missing OpenAI API Key")

if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = OPENAI_API_KEY

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

2025-11-07 16:44:51,506 - INFO - 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-11-07 16:44:53,519 - INFO - Bucket 'travel-sample' exists.
2025-11-07 16:44:53,527 - INFO - Collection 'smolagents' does not exist. Creating it...
2025-11-07 16:44:53,575 - INFO - Collection 'smolagents' created successfully.
2025-11-07 16:44:55,731 - INFO - All documents cleared from the collection.





<couchbase.collection.Collection at 0x17bdb9be0>

try:
    embeddings = OpenAIEmbeddings(
        model="text-embedding-3-small",
        api_key=OPENAI_API_KEY,
    )
    logging.info("Successfully created OpenAIEmbeddings")
except Exception as e:
    raise ValueError(f"Error creating OpenAIEmbeddings: {str(e)}")

2025-11-07 16:44:58,634 - INFO - Successfully created OpenAIEmbeddings

# Alternative: Create a Composite index for filtered searches
vector_store.create_index(
    index_type=IndexType.COMPOSITE,
    index_description="IVF,SQ8",
    distance_metric=DistanceStrategy.COSINE,
    index_name="smolagents_composite_index",
)

try:
    vector_store = CouchbaseQueryVectorStore(
        cluster=cluster,
        bucket_name=CB_BUCKET_NAME,
        scope_name=SCOPE_NAME,
        collection_name=COLLECTION_NAME,
        embedding=embeddings,
        distance_metric=DistanceStrategy.COSINE
    )
    logging.info("Successfully created vector store")
except Exception as e:
    raise ValueError(f"Failed to create vector store: {str(e)}")

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)}")

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 = 100

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-11-07 16:46:18,967 - INFO - Document ingestion completed successfully.

# Phase 1: Baseline Performance (Without Hyperscale vector index)
print("="*80)
print("PHASE 1: BASELINE PERFORMANCE (NO Hyperscale vector index)")
print("="*80)

query = "What was manchester city manager pep guardiola's reaction to the team's current form?"

try:
    # Perform the semantic search
    start_time = time.time()
    search_results = vector_store.similarity_search_with_score(query, k=10)
    baseline_time = time.time() - start_time

    logging.info(f"Semantic search completed in {baseline_time:.2f} seconds")

    # Display search results
    print(f"\nSemantic Search Results (completed in {baseline_time:.2f} seconds):")
    print("-" * 80)  # Add separator line
    for doc, distance in search_results:
        print(f"Vector Distance: {distance:.4f}, Text: {doc.page_content}")
        print("-" * 80)  # Add separator between results

except CouchbaseException as e:
    raise RuntimeError(f"Error performing semantic search: {str(e)}")
except Exception as e:
    raise RuntimeError(f"Unexpected error: {str(e)}")

================================================================================
PHASE 1: BASELINE PERFORMANCE
================================================================================


2025-11-07 16:46:24,561 - INFO - Semantic search completed in 1.34 seconds



Semantic Search Results (completed in 1.34 seconds):
--------------------------------------------------------------------------------
Vector Distance: 0.2956, Text: Manchester City boss Pep Guardiola has won 18 trophies since he arrived at the club in 2016

Manchester City boss Pep Guardiola says he is "fine" despite admitting his sleep and diet are being affected by the worst run of results in his entire managerial career. In an interview with former Italy international Luca Toni for Amazon Prime Sport before Wednesday's Champions League defeat by Juventus, Guardiola touched on the personal impact City's sudden downturn in form has had. Guardiola said his state of mind was "ugly", that his sleep was "worse" and he was eating lighter as his digestion had suffered. City go into Sunday's derby against Manchester United at Etihad Stadium having won just one of their past 10 games. The Juventus loss means there is a chance they may not even secure a play-off spot in the Champions League. Asked to elaborate on his comments to Toni, Guardiola said: "I'm fine. "In our jobs we always want to do our best or the best as possible. When that doesn't happen you are more uncomfortable than when the situation is going well, always that happened. "In good moments I am happier but when I get to the next game I am still concerned about what I have to do. There is no human being that makes an activity and it doesn't matter how they do." Guardiola said City have to defend better and "avoid making mistakes at both ends". To emphasise his point, Guardiola referred back to the third game of City's current run, against a Sporting side managed by Ruben Amorim, who will be in the United dugout at the weekend. City dominated the first half in Lisbon, led thanks to Phil Foden's early effort and looked to be cruising. Instead, they conceded three times in 11 minutes either side of half-time as Sporting eventually ran out 4-1 winners. "I would like to play the game like we played in Lisbon on Sunday, believe me," said Guardiola, who is facing the prospect of only having three fit defenders for the derby as Nathan Ake and Manuel Akanji try to overcome injury concerns. If there is solace for City, it comes from the knowledge United are not exactly flying. Their comeback Europa League victory against Viktoria Plzen on Thursday was their third win of Amorim's short reign so far but only one of those successes has come in the Premier League, where United have lost their past two games against Arsenal and Nottingham Forest. Nevertheless, Guardiola can see improvements already on the red side of the city. "It's already there," he said. "You see all the patterns, the movements, the runners and the pace. He will do a good job at United, I'm pretty sure of that."

Guardiola says skipper Kyle Walker has been offered support by the club after the City defender highlighted the racial abuse he had received on social media in the wake of the Juventus trip. "It's unacceptable," he said. "Not because it's Kyle - for any human being. "Unfortunately it happens many times in the real world. It is not necessary to say he has the support of the entire club. It is completely unacceptable and we give our support to him."
--------------------------------------------------------------------------------
Vector Distance: 0.3100, Text: Pep Guardiola has said Manchester City will be his final managerial job in club football before he "maybe" coaches a national team.

--------------------------------------------------------------------------------

vector_store.create_index(index_type=IndexType.BHIVE, index_name="smolagents_bhive_index", index_description="IVF,SQ8")

# Phase 2: Hyperscale-Optimized Performance 
print("\n" + "="*80)
print("PHASE 2: HYPERSCALE-OPTIMIZED PERFORMANCE")
print("="*80)

query = "What was manchester city manager pep guardiola's reaction to the team's current form?"

try:
    # Perform the semantic search
    start_time = time.time()
    search_results = vector_store.similarity_search_with_score(query, k=10)
    hyperscale_time = time.time() - start_time

    logging.info(f"Semantic search completed in {hyperscale_time:.2f} seconds")

    # Display search results
    print(f"\nSemantic Search Results (completed in {hyperscale_time:.2f} seconds):")
    print("-" * 80)  # Add separator line
    for doc, distance in search_results:
        print(f"Vector Distance: {distance:.4f}, Text: {doc.page_content}")
        print("-" * 80)  # Add separator between results

except CouchbaseException as e:
    raise RuntimeError(f"Error performing semantic search: {str(e)}")
except Exception as e:
    raise RuntimeError(f"Unexpected error: {str(e)}")

================================================================================
PHASE 2: HYPERSCALE-OPTIMIZED PERFORMANCE
================================================================================


2025-11-07 16:47:01,538 - INFO - Semantic search completed in 0.42 seconds



Semantic Search Results (completed in 0.42 seconds):
--------------------------------------------------------------------------------
Vector Distance: 0.2956, Text: Manchester City boss Pep Guardiola has won 18 trophies since he arrived at the club in 2016

Manchester City boss Pep Guardiola says he is "fine" despite admitting his sleep and diet are being affected by the worst run of results in his entire managerial career. In an interview with former Italy international Luca Toni for Amazon Prime Sport before Wednesday's Champions League defeat by Juventus, Guardiola touched on the personal impact City's sudden downturn in form has had. Guardiola said his state of mind was "ugly", that his sleep was "worse" and he was eating lighter as his digestion had suffered. City go into Sunday's derby against Manchester United at Etihad Stadium having won just one of their past 10 games. The Juventus loss means there is a chance they may not even secure a play-off spot in the Champions League. Asked to elaborate on his comments to Toni, Guardiola said: "I'm fine. "In our jobs we always want to do our best or the best as possible. When that doesn't happen you are more uncomfortable than when the situation is going well, always that happened. "In good moments I am happier but when I get to the next game I am still concerned about what I have to do. There is no human being that makes an activity and it doesn't matter how they do." Guardiola said City have to defend better and "avoid making mistakes at both ends". To emphasise his point, Guardiola referred back to the third game of City's current run, against a Sporting side managed by Ruben Amorim, who will be in the United dugout at the weekend. City dominated the first half in Lisbon, led thanks to Phil Foden's early effort and looked to be cruising. Instead, they conceded three times in 11 minutes either side of half-time as Sporting eventually ran out 4-1 winners. "I would like to play the game like we played in Lisbon on Sunday, believe me," said Guardiola, who is facing the prospect of only having three fit defenders for the derby as Nathan Ake and Manuel Akanji try to overcome injury concerns. If there is solace for City, it comes from the knowledge United are not exactly flying. Their comeback Europa League victory against Viktoria Plzen on Thursday was their third win of Amorim's short reign so far but only one of those successes has come in the Premier League, where United have lost their past two games against Arsenal and Nottingham Forest. Nevertheless, Guardiola can see improvements already on the red side of the city. "It's already there," he said. "You see all the patterns, the movements, the runners and the pace. He will do a good job at United, I'm pretty sure of that."

Guardiola says skipper Kyle Walker has been offered support by the club after the City defender highlighted the racial abuse he had received on social media in the wake of the Juventus trip. "It's unacceptable," he said. "Not because it's Kyle - for any human being. "Unfortunately it happens many times in the real world. It is not necessary to say he has the support of the entire club. It is completely unacceptable and we give our support to him."
--------------------------------------------------------------------------------
Vector Distance: 0.3100, Text: Pep Guardiola has said Manchester City will be his final managerial job in club football before he "maybe" coaches a national team.
--------------------------------------------------------------------------------

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

print(f"\n📊 Performance Comparison:")
print(f"{'Optimization Level':<35} {'Time (seconds)':<20} {'Status'}")
print("-" * 80)
print(f"{'Phase 1 - Baseline (No Hyperscale)':<35} {baseline_time:.4f}{'':16} ⚪ Baseline")
print(f"{'Phase 2 - Hyperscale-Optimized Search':<35} {hyperscale_time:.4f}{'':16} ✅ Optimized")

# Calculate improvement
if baseline_time > hyperscale_time:
    speedup = baseline_time / hyperscale_time
    improvement = ((baseline_time - hyperscale_time) / baseline_time) * 100
    print(f"\n✨ Hyperscale Index Performance Gain: {speedup:.2f}x faster ({improvement:.1f}% improvement)")
elif gsi_time > baseline_time:
    slowdown_pct = ((gsi_time - baseline_time) / baseline_time) * 100
    print(f"\n⚠️  Note: Hyperscale Index was {slowdown_pct:.1f}% slower than baseline in this run")
    print(f"   This can happen with small datasets. Hyperscale Index benefits emerge with scale.")
else:
    print(f"\n⚖️  Performance: Comparable to baseline")

print("\n" + "-"*80)
print("KEY INSIGHTS:")
print("-"*80)
print("1. 🚀 Hyperscale Index Optimization:")
print("   • Hyperscale indexes excel with large-scale datasets (millions+ vectors)")
print("   • Performance gains increase with dataset size and concurrent queries")
print("   • Optimal for production workloads with sustained traffic patterns")

print("\n2. 📦 Dataset Size Impact:")
print(f"   • Current dataset: ~1,700 articles")
print("   • At this scale, performance differences may be minimal or variable")
print("   • Significant gains typically seen with 10M+ vectors")

print("\n3. 🎯 When to Use Hyperscale Index:")
print("   • Large-scale vector search applications")
print("   • High query-per-second (QPS) requirements")
print("   • Multi-user concurrent access scenarios")
print("   • Production environments requiring scalability")

print("\n" + "="*80)

================================================================================
VECTOR SEARCH PERFORMANCE OPTIMIZATION SUMMARY
================================================================================

📊 Performance Comparison:
Optimization Level                  Time (seconds)       Status
--------------------------------------------------------------------------------
Phase 1 - Baseline (No Index)       1.3410                 ⚪ Baseline
Phase 2 - Hyperscale-Optimized     0.4157                 ✅ Optimized

✨ Hyperscale Index Performance Gain: 3.23x faster (69.0% improvement)

--------------------------------------------------------------------------------
KEY INSIGHTS:
--------------------------------------------------------------------------------
1. 🚀 Hperscale Index Optimization:
   • Hyperscale indexes excel with large-scale datasets (millions+ vectors)
   • Performance gains increase with dataset size and concurrent queries
   • Optimal for production workloads with sustained traffic patterns

2. 📦 Dataset Size Impact:
   • Current dataset: ~1,700 articles
   • At this scale, performance differences may be minimal or variable
   • Significant gains typically seen with 10M+ vectors

3. 🎯 When to Use Hyperscale Indexes:
   • Large-scale vector search applications
   • High query-per-second (QPS) requirements
   • Multi-user concurrent access scenarios
   • Production environments requiring scalability

================================================================================

class RetrieverTool(Tool):
    name = "retriever"
    description = "Uses semantic search to retrieve the parts of news documentation that could be most relevant to answer your query."
    inputs = {
        "query": {
            "type": "string",
            "description": "The query to perform. This should be semantically close to your target documents. Use the affirmative form rather than a question.",
        }
    }
    output_type = "string"

    def __init__(self, vector_store: CouchbaseQueryVectorStore, **kwargs):
        super().__init__(**kwargs)
        self.vector_store = vector_store

    def forward(self, query: str) -> str:
        assert isinstance(query, str), "Query must be a string"

        docs = self.vector_store.similarity_search_with_score(query, k=5)
        return "\n\n".join(
            f"# Documents:\n{doc.page_content}"
            for doc, distance in docs
        )

retriever_tool = RetrieverTool(vector_store)

agent = ToolCallingAgent(
    tools=[retriever_tool],
    model=OpenAIServerModel(
        model_id="gpt-4o-2024-08-06",
        api_key=OPENAI_API_KEY,
    ),
    max_steps=4,
    verbosity_level=2
)

query = "What was manchester city manager pep guardiola's reaction to the team's current form?"

agent_output = agent.run(query)