Perform Maximal Marginal Relevance search with LangChain on Bigtable

This page describes how to perform a Maximal Marginal Relevance (MMR) search using the BigtableVectorStore integration for LangChain in Bigtable and Vertex AI as the embedding service.

MMR is a search technique used in information retrieval to return a set ofUf results that are both relevant to the query and diverse, avoiding redundancy. While standard vector similarity search (for example, a search that uses the k-nearest neighbor method) might return many similar items, MMR provides a more varied set of top results. This is useful when you have potentially overlapping or duplicative data in your vector store.

For example, in an ecommerce application, if a user searches for "red tomatoes", then a vector similarity search might return multiple listings of the same type of fresh red tomato. An MMR search would aim to return a more diverse set, such as "fresh red tomatoes", "canned diced red tomatoes", "organic cherry tomatoes", and perhaps even "tomato salad recipe".

Before you read this page, it's important that you know the following concepts:

  • Relevance: A measure of how closely the document matches the query.
  • Diversity: A measure of how different a document is from the documents already selected in the result set.
  • Lambda Multiplier: A factor between 0 and 1 that balances relevance and diversity. A value closer to 1 prioritizes relevance, while a value closer to 0 prioritizes diversity.

The BigtableVectorStore class for LangChain implements MMR as a re-ranking algorithm. This algorithm first fetches a larger set of documents relevant to the query, and selects the documents that match the query in a way that's balanced for relevance and diversity.

Before you begin

This guide uses Vertex AI as the embedding service. Make sure you have the Vertex AI API enabled in your project.

Enable the Vertex AI API

Required roles

To use Bigtable with LangChain, you need the following IAM roles:

Set up your environment

  1. Install the required LangChain packages:

    pip install --upgrade --quiet langchain-google-bigtable langchain-google-vertexai

  2. Authenticate to Google Cloud with your user account.

    gcloud auth application-default login

  3. Set your project ID, Bigtable instance ID, and table ID:

    PROJECT_ID = "your-project-id"
INSTANCE_ID = "your-instance-id"
TABLE_ID = "your-table-id"

Initialize the embedding service and create a table

To use the Bigtable vector store, you need to provide embeddings generated by an AI model. In this guide, you use the text embedding gemini-embedding-004 model on Vertex AI.

from langchain_google_vertexai import VertexAIEmbeddings
from langchain_google_bigtable.vector_store import init_vector_store_table, BigtableVectorStore, ColumnConfig

# Initialize an embedding service
embedding_service = VertexAIEmbeddings(model_name="text-embedding-004", project=PROJECT_ID)

# Define column families
DATA_COLUMN_FAMILY = "product_data"

# Initialize the table (if it doesn't exist)
try:
    init_vector_store_table(
        project_id=PROJECT_ID,
        instance_id=INSTANCE_ID,
        table_id=TABLE_ID,
        content_column_family=DATA_COLUMN_FAMILY,
        embedding_column_family=DATA_COLUMN_FAMILY,
    )
    print(f"Table {TABLE_ID} created successfully.")
except ValueError as e:
    print(e) # Table likely already exists

Instantiate BigtableVectorStore

Create the store instance by passing the embedding service and Bigtable table identifiers.

# Configure columns
content_column = ColumnConfig(
    column_family=DATA_COLUMN_FAMILY, column_qualifier="product_description"
)
embedding_column = ColumnConfig(
    column_family=DATA_COLUMN_FAMILY, column_qualifier="embedding"
)

# Create the vector store instance
vector_store = BigtableVectorStore.create_sync(
    project_id=PROJECT_ID,
    instance_id=INSTANCE_ID,
    table_id=TABLE_ID,
    embedding_service=embedding_service,
    collection="ecommerce_products",
    content_column=content_column,
    embedding_column=embedding_column,
)
print("BigtableVectorStore instantiated.")

Populate the vector store

In this guide, we use a sample scenario of a fictional ecommerce service where users want to search for items related to red tomatoes. First, we need to add some tomato-related products and descriptions to the vector store.

from langchain_core.documents import Document

products = [
    Document(page_content="Fresh organic red tomatoes, great for salads.", metadata={"type": "fresh produce", "color": "red", "name": "Organic Vine Tomatoes"}),
    Document(page_content="Ripe red tomatoes on the vine.", metadata={"type": "fresh", "color": "red", "name": "Tomatoes on Vine"}),
    Document(page_content="Sweet cherry tomatoes, red and juicy.", metadata={"type": "fresh", "color": "red", "name": "Cherry Tomatoes"}),
    Document(page_content="Canned diced red tomatoes in juice.", metadata={"type": "canned", "color": "red", "name": "Diced Tomatoes"}),
    Document(page_content="Sun-dried tomatoes in oil.", metadata={"type": "preserved", "color": "red", "name": "Sun-Dried Tomatoes"}),
    Document(page_content="Green tomatoes, perfect for frying.", metadata={"type": "fresh", "color": "green", "name": "Green Tomatoes"}),
    Document(page_content="Tomato paste, concentrated flavor.", metadata={"type": "canned", "color": "red", "name": "Tomato Paste"}),
    Document(page_content="Mixed salad greens with cherry tomatoes.", metadata={"type": "prepared", "color": "mixed", "name": "Salad Mix with Tomatoes"}),
    Document(page_content="Yellow pear tomatoes, mild flavor.", metadata={"type": "fresh", "color": "yellow", "name": "Yellow Pear Tomatoes"}),
    Document(page_content="Heirloom tomatoes, various colors.", metadata={"type": "fresh", "color": "various", "name": "Heirloom Tomatoes"}),
]

vector_store.add_documents(products)
print(f"Added {len(products)} products to the vector store.")

The store contains many tomato-related products, but users want to search for offerings associated with "red tomatoes" only. To get a diverse set of results, use the MMR technique to perform the search.

The key method to use is the max_marginal_relevance_search that takes the following arguments:

  • query (str): The search text.
  • k (int): The final number of search results.
  • fetch_k (int): The initial number of similar products to retrieve before applying the MMR algorithm. We recommend that this number is larger than the k parameter.
  • lambda_mult (float): The diversity tuning parameter. Use 0.0 for maximum diversity, 1.0 for maximum relevance.
user_query = "red tomatoes"
k_results = 4
fetch_k_candidates = 10

print(f"Performing MMR search for: '{user_query}'")

# Example 1: Balanced relevance and diversity
mmr_results_balanced = vector_store.max_marginal_relevance_search(
    user_query, k=k_results, fetch_k=fetch_k_candidates, lambda_mult=0.5
)
print(f"MMR Results (lambda=0.5, k={k_results}, fetch_k={fetch_k_candidates}):")
for doc in mmr_results_balanced:
    print(f"  - {doc.metadata['name']}: {doc.page_content}")

print("\n")

# Example 2: Prioritizing Diversity
mmr_results_diverse = vector_store.max_marginal_relevance_search(
    user_query, k=k_results, fetch_k=fetch_k_candidates, lambda_mult=0.1
)
print(f"MMR Results (lambda=0.1, k={k_results}, fetch_k={fetch_k_candidates}):")
for doc in mmr_results_diverse:
    print(f"  - {doc.metadata['name']}: {doc.page_content}")

print("\n")

# Example 3: Prioritizing Relevance
mmr_results_relevant = vector_store.max_marginal_relevance_search(
    user_query, k=k_results, fetch_k=fetch_k_candidates, lambda_mult=0.9
)
print(f"MMR Results (lambda=0.9, k={k_results}, fetch_k={fetch_k_candidates}):")
for doc in mmr_results_relevant:
    print(f"  - {doc.metadata['name']}: {doc.page_content}")

Different lambda_mult values yield different sets of results, balancing the similarity to "red tomatoes" with the uniqueness of the products shown.

Using MMR with a retriever

To use MMR search, you can also configure a LangChain retriever. The retriever provides a uniform interface that lets you seamlessly integrate specialized search methods, such as MMR, directly into your chains and applications.

retriever = vector_store.as_retriever(
    search_type="mmr",
    search_kwargs={
        "k": 3,
        "fetch_k": 10,
        "lambda_mult": 0.3,
    }
)

retrieved_docs = retriever.invoke(user_query)
print(f"\nRetriever MMR Results for '{user_query}':")
for doc in retrieved_docs:
    print(f"  - {doc.metadata['name']}: {doc.page_content}")

What's next