ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/stores/bigquery.py

@@ -0,0 +1,1377 @@
+from typing import Any, Dict, List, Optional, Union, Callable
+import uuid
+import time
+import asyncio
+from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances
+from google.cloud import bigquery as bq
+from google.oauth2 import service_account
+from google.cloud.exceptions import NotFound, Conflict
+import numpy as np
+import pandas as pd
+from navconfig.logging import logging
+from .abstract import AbstractStore
+from ..conf import (
+    BIGQUERY_CREDENTIALS,
+    BIGQUERY_PROJECT_ID,
+    BIGQUERY_DATASET
+)
+from .models import SearchResult, Document, DistanceStrategy
+from .utils.chunking import LateChunkingProcessor
+from ..exceptions import DriverError
+
+
+
+class BigQueryStore(AbstractStore):
+    """
+    A BigQuery vector store implementation for storing and searching embeddings.
+    This store provides vector similarity search capabilities using BigQuery's ML functions.
+    """
+
+    def __init__(
+        self,
+        table: str = None,
+        dataset: str = None,
+        project_id: str = None,
+        credentials: str = None,
+        id_column: str = 'id',
+        embedding_column: str = 'embedding',
+        document_column: str = 'document',
+        text_column: str = 'text',
+        metadata_column: str = 'metadata',
+        embedding_model: Union[dict, str] = "sentence-transformers/all-mpnet-base-v2",
+        embedding: Optional[Callable] = None,
+        distance_strategy: DistanceStrategy = DistanceStrategy.COSINE,
+        auto_initialize: bool = True,
+        **kwargs
+    ):
+        """Initialize the BigQueryStore with the specified parameters."""
+        self.table_name = table
+        self.dataset = dataset or BIGQUERY_DATASET
+        self._project_id = project_id or BIGQUERY_PROJECT_ID
+        self._credentials = credentials or BIGQUERY_CREDENTIALS
+
+        # Column definitions
+        self._id_column: str = id_column
+        self._embedding_column: str = embedding_column
+        self._document_column: str = document_column
+        self._text_column: str = text_column
+        self._metadata_column: str = metadata_column
+
+        # Configuration
+        self.distance_strategy = distance_strategy
+        self._auto_initialize: bool = auto_initialize
+        self._collection_store_cache: Dict[str, Any] = {}
+
+        # Initialize parent class
+        super().__init__(
+            embedding_model=embedding_model,
+            embedding=embedding,
+            **kwargs
+        )
+
+        # BigQuery client and session management
+        self._connection: Optional[bq.Client] = None
+        self._connected: bool = False
+        self.credentials = None
+        self._account = None
+
+        # Initialize logger
+        self.logger = logging.getLogger("BigQueryStore")
+
+    def get_vector(self, metric_type: str = None, **kwargs):
+        raise NotImplementedError("This method is part of the old implementation.")
+
+    def _execute_query(
+        self,
+        query: str,
+        job_config: Optional[bq.QueryJobConfig] = None
+    ) -> List[Dict[str, Any]]:
+        query_job = self._connection.query(query, job_config=job_config)
+        return list(query_job.result())
+
+    async def _thread_func(self, func, *args, **kwargs):
+        """
+        Run a synchronous function in an async context.
+        Helper for running blocking calls in a non-blocking way.
+        """
+        return await asyncio.to_thread(func, *args, **kwargs)
+
+    async def connection(self):
+        """Initialize BigQuery client.
+        Assuming that authentication is handled outside
+        (via environment variables or similar)
+        """
+        try:
+            if self._credentials:  # usage of explicit credentials
+                self.credentials = service_account.Credentials.from_service_account_file(
+                    self._credentials
+                )
+                if not self._project_id:
+                    self._project_id = self.credentials.project_id
+                self._connection = bq.Client(credentials=self.credentials, project=self._project_id)
+                self._connected = True
+            else:
+                self.credentials = self._account
+                self._connection = bq.Client(project=self._project_id)
+                self._connected = True
+
+            if self._auto_initialize:
+                await self.initialize_database()
+
+            self.logger.debug("Successfully connected to BigQuery.")
+
+        except Exception as e:
+            self._connected = False
+            raise DriverError(f"BigQuery: Error initializing client: {e}")
+        return self
+
+    async def initialize_database(self):
+        """Initialize BigQuery dataset and any required setup."""
+        try:
+            # Ensure dataset exists
+            dataset_id = f"{self._project_id}.{self.dataset}"
+            try:
+                self._connection.get_dataset(dataset_id)
+                self.logger.info(f"Dataset {dataset_id} already exists")
+            except NotFound:
+                # Create dataset if it doesn't exist
+                dataset = bq.Dataset(dataset_id)
+                dataset.location = "US"
+                dataset = self._connection.create_dataset(dataset, timeout=30)
+                self.logger.info(f"Created dataset {dataset_id}")
+
+        except Exception as e:
+            self.logger.warning(
+                f"⚠️ Database auto-initialization failed: {e}"
+            )
+
+    def _define_collection_store(
+        self,
+        table: str,
+        dataset: str,
+        dimension: int = 384,
+        id_column: str = 'id',
+        embedding_column: str = 'embedding',
+        document_column: str = 'document',
+        metadata_column: str = 'metadata',
+        text_column: str = 'text'
+    ) -> str:
+        """Define a collection store table name for BigQuery.
+
+        Returns:
+            str: Fully qualified table name in format project.dataset.table
+        """
+        full_table_name = f"{self._project_id}.{dataset}.{table}"
+
+        if full_table_name in self._collection_store_cache:
+            return self._collection_store_cache[full_table_name]
+
+        # Cache the table reference
+        self._collection_store_cache[full_table_name] = {
+            'table_name': full_table_name,
+            'dimension': dimension,
+            'columns': {
+                'id': id_column,
+                'embedding': embedding_column,
+                'document': document_column,
+                'metadata': metadata_column,
+                'text': text_column
+            }
+        }
+
+        self.logger.debug(
+            f"Defined collection store: {full_table_name}"
+        )
+        return full_table_name
+
+    async def dataset_exists(self, dataset: str = None) -> bool:
+        """Check if a dataset exists in BigQuery."""
+        if not self._connected:
+            await self.connection()
+
+        dataset = dataset or self.dataset
+        dataset_id = f"{self._project_id}.{dataset}"
+
+        try:
+            self._connection.get_dataset(dataset_id)
+            return True
+        except NotFound:
+            return False
+
+    async def create_dataset(self, dataset: str = None, location: str = "US") -> Any:
+        """Create a new dataset in BigQuery."""
+        if not self._connected:
+            await self.connection()
+
+        dataset = dataset or self.dataset
+
+        try:
+            dataset_ref = bq.DatasetReference(self._project_id, dataset)
+            dataset_obj = bq.Dataset(dataset_ref)
+            dataset_obj.location = location
+            dataset_obj = self._connection.create_dataset(dataset_obj)
+            self.logger.debug(f"Created dataset {self._project_id}.{dataset}")
+            return dataset_obj
+        except Conflict:
+            self.logger.warning(f"Dataset {self._project_id}.{dataset} already exists")
+            # Get the existing dataset to return it
+            return self._connection.get_dataset(f"{self._project_id}.{dataset}")
+        except Exception as exc:
+            self.logger.error(f"Error creating Dataset: {exc}")
+            raise DriverError(
+                f"Error creating Dataset: {exc}"
+            ) from exc
+
+    async def collection_exists(self, table: str, dataset: str = None) -> bool:
+        """Check if a collection (table) exists in BigQuery."""
+        if not self._connected:
+            await self.connection()
+
+        dataset = dataset or self.dataset
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        try:
+            await self._thread_func(self._connection.get_table, table_id)
+            return True
+        except NotFound:
+            return False
+
+    async def create_collection(
+        self,
+        table: str,
+        dataset: str = None,
+        dimension: int = 768,
+        id_column: str = None,
+        embedding_column: str = None,
+        document_column: str = None,
+        metadata_column: str = None,
+        **kwargs
+    ) -> None:
+        """Create a new collection (table) in BigQuery."""
+        if not self._connected:
+            await self.connection()
+
+        dataset = dataset or self.dataset
+        id_column = id_column or self._id_column
+        embedding_column = embedding_column or self._embedding_column
+        document_column = document_column or self._document_column
+        metadata_column = metadata_column or self._metadata_column
+
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        try:
+            # Check if table already exists
+            if await self.collection_exists(table, dataset):
+                self.logger.info(f"Collection {table_id} already exists")
+                return
+
+            # Define table schema - use FLOAT64 REPEATED instead of STRUCT
+            schema = [
+                bq.SchemaField(id_column, "STRING", mode="REQUIRED"),
+                bq.SchemaField(embedding_column, "FLOAT64", mode="REPEATED"),
+                bq.SchemaField(document_column, "STRING", mode="NULLABLE"),
+                bq.SchemaField(metadata_column, "JSON", mode="NULLABLE"),
+                bq.SchemaField(self._text_column, "STRING", mode="NULLABLE"),
+                bq.SchemaField("collection_id", "STRING", mode="NULLABLE"),
+            ]
+
+            table_ref = bq.Table(table_id, schema=schema)
+            table_ref = await self._thread_func(self._connection.create_table, table_ref)
+
+            self.logger.debug(f"Created collection {table_id}")
+
+            # Cache the collection store
+            self._define_collection_store(
+                table=table,
+                dataset=dataset,
+                dimension=dimension,
+                id_column=id_column,
+                embedding_column=embedding_column,
+                document_column=document_column,
+                metadata_column=metadata_column
+            )
+
+        except Exception as e:
+            self.logger.error(f"Error creating collection: {e}")
+            raise RuntimeError(
+                f"Failed to create collection: {e}"
+            ) from e
+
+    async def drop_collection(self, table: str, dataset: str = None) -> None:
+        """
+        Drops the specified table in the given dataset.
+
+        Args:
+            table: The name of the table to drop.
+            dataset: The dataset where the table resides (optional, uses default if not provided).
+        """
+        if not self._connected:
+            await self.connection()
+
+        dataset = dataset or self.dataset
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        try:
+            # Check if table exists first
+            if not await self.collection_exists(table, dataset):
+                self.logger.warning(f"Table '{table_id}' does not exist, nothing to drop")
+                return
+
+            # Drop the table
+            self._connection.delete_table(table_id, not_found_ok=True)
+
+            # Remove from cache if it exists
+            if table_id in self._collection_store_cache:
+                del self._collection_store_cache[table_id]
+
+            self.logger.debug(f"Table '{table_id}' dropped successfully")
+
+        except Exception as e:
+            self.logger.error(
+                f"Error dropping table '{table_id}': {e}"
+            )
+            raise RuntimeError(
+                f"Failed to drop table '{table_id}': {e}"
+            ) from e
+
+    async def prepare_embedding_table(
+        self,
+        table: str,
+        dataset: str = None,
+        dimension: int = 768,
+        id_column: str = 'id',
+        embedding_column: str = 'embedding',
+        document_column: str = 'document',
+        metadata_column: str = 'metadata',
+        **kwargs
+    ) -> bool:
+        """Prepare an existing BigQuery table for embedding storage."""
+        if not self._connected:
+            await self.connection()
+
+        dataset = dataset or self.dataset
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        try:
+            # Get existing table
+            table_ref = self._connection.get_table(table_id)
+            current_schema = table_ref.schema
+
+            # Check if embedding columns already exist
+            existing_fields = {field.name for field in current_schema}
+            new_fields = []
+
+            if embedding_column not in existing_fields:
+                new_fields.append(
+                    bq.SchemaField(
+                        embedding_column,
+                        "REPEATED",
+                        mode="NULLABLE",
+                        fields=[bq.SchemaField("value", "FLOAT64")]
+                    )
+                )
+
+            if metadata_column not in existing_fields:
+                new_fields.append(
+                    bq.SchemaField(metadata_column, "JSON", mode="NULLABLE")
+                )
+
+            if "collection_id" not in existing_fields:
+                new_fields.append(
+                    bq.SchemaField("collection_id", "STRING", mode="NULLABLE")
+                )
+
+            # Add new fields if any
+            if new_fields:
+                new_schema = list(current_schema) + new_fields
+                table_ref.schema = new_schema
+                table_ref = self._connection.update_table(table_ref, ["schema"])
+                self.logger.info(f"Updated table {table_id} schema with embedding columns")
+
+            # Cache the collection store
+            self._define_collection_store(
+                table=table,
+                dataset=dataset,
+                dimension=dimension,
+                id_column=id_column,
+                embedding_column=embedding_column,
+                document_column=document_column,
+                metadata_column=metadata_column
+            )
+
+            return True
+
+        except Exception as e:
+            self.logger.error(f"Error preparing embedding table: {e}")
+            raise RuntimeError(f"Failed to prepare embedding table: {e}") from e
+
+    async def _wait_for_table_insert_ready(self, table_id: str, max_wait_seconds: int = 30, poll_interval: float = 0.5) -> bool:
+        """
+        Wait for a table to be ready for insert operations specifically.
+
+        Args:
+            table_id: Fully qualified table ID (project.dataset.table)
+            max_wait_seconds: Maximum time to wait in seconds
+            poll_interval: Time between polling attempts in seconds
+
+        Returns:
+            bool: True if table is ready for inserts, False if timeout
+        """
+        start_time = time.time()
+
+        while time.time() - start_time < max_wait_seconds:
+            try:
+                # First check if table exists via get_table
+                table_ref = await self._thread_func(self._connection.get_table, table_id)
+
+                # Then test with a simple query to see if the table is accessible for operations
+                # This is a lightweight way to test table readiness
+                test_query = f"SELECT COUNT(*) as row_count FROM `{table_id}` LIMIT 1"
+                query_job = await self._thread_func(self._connection.query, test_query)
+                await self._thread_func(query_job.result)
+
+                # If we get here without exception, table is ready for operations
+                self.logger.debug(f"Table {table_id} is ready for insert operations")
+                return True
+
+            except NotFound:
+                # Table not ready yet, wait and retry
+                self.logger.debug(f"Table {table_id} not ready yet, waiting {poll_interval}s...")
+                await asyncio.sleep(poll_interval)
+            except Exception as e:
+                # For other errors, we might want to continue trying for a bit
+                # as BigQuery can have temporary inconsistencies
+                self.logger.debug(f"Checking table readiness, got error (will retry): {e}")
+                await asyncio.sleep(poll_interval)
+
+        self.logger.warning(f"Table {table_id} not ready after {max_wait_seconds} seconds")
+        return False
+
+    async def add_documents(
+        self,
+        documents: List[Document],
+        table: str = None,
+        dataset: str = None,
+        embedding_column: str = 'embedding',
+        content_column: str = 'document',
+        metadata_column: str = 'metadata',
+        **kwargs
+    ) -> None:
+        """Add documents to BigQuery table with embeddings."""
+        if not self._connected:
+            await self.connection()
+
+        table = table or self.table_name
+        dataset = dataset or self.dataset
+
+        if not table:
+            raise ValueError("Table name must be provided")
+
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        # Ensure collection exists
+        if not await self.collection_exists(table, dataset):
+            await self.create_collection(
+                table=table,
+                dataset=dataset,
+                dimension=self.dimension
+            )
+
+        # If we just created the table, wait for it to be ready for inserts
+        is_ready = await self._wait_for_table_insert_ready(table_id, max_wait_seconds=60)
+        if not is_ready:
+            raise RuntimeError(
+                f"Table {table_id} was created but not ready for insert operations within timeout"
+            )
+
+        # Process documents
+        texts = [doc.page_content for doc in documents]
+        # Thread the embedding generation as it can be slow
+        embeddings = await self._thread_func(self._embed_.embed_documents, texts)
+        metadatas = [doc.metadata for doc in documents]
+
+        # Prepare data for BigQuery (this is fast, no threading needed)
+        rows_to_insert = []
+        for i, doc in enumerate(documents):
+            embedding_vector = embeddings[i]
+            if isinstance(embedding_vector, np.ndarray):
+                embedding_vector = embedding_vector.tolist()
+
+            embedding_array = [float(val) for val in embedding_vector]
+
+            metadata_value = metadatas[i] or {}
+            metadata_json = self._json.dumps(metadata_value) if metadata_value else self._json.dumps({})
+
+            row = {
+                self._id_column: str(uuid.uuid4()),
+                embedding_column: embedding_array,
+                content_column: texts[i],
+                metadata_column: metadata_json,
+                "collection_id": str(uuid.uuid4())
+            }
+            rows_to_insert.append(row)
+
+        # Insert data with retry logic and longer delays
+        max_retries = 5
+        retry_delay = 2.0
+
+        for attempt in range(max_retries):
+            try:
+                # Always get a fresh table reference
+                table_ref = await self._thread_func(self._connection.get_table, table_id)
+
+                # Add a small delay even on first attempt if table was just created
+                if attempt == 0:
+                    await asyncio.sleep(1.0)
+
+                errors = await self._thread_func(
+                    self._connection.insert_rows_json, table_ref, rows_to_insert
+                )
+
+                if errors:
+                    self.logger.error(f"Errors inserting rows: {errors}")
+                    raise RuntimeError(f"Failed to insert documents: {errors}")
+
+                self.logger.info(f"Successfully added {len(documents)} documents to {table_id}")
+                return  # Success, exit the retry loop
+
+            except NotFound as e:
+                if attempt < max_retries - 1:
+                    self.logger.warning(
+                        f"Table {table_id} not found for insert (attempt {attempt + 1}/{max_retries}), retrying in {retry_delay}s..."
+                    )
+                    await asyncio.sleep(retry_delay)
+                    retry_delay = min(retry_delay * 1.5, 10.0)  # Cap at 10 seconds
+                else:
+                    self.logger.error(
+                        f"Table {table_id} still not found for insert after {max_retries} attempts"
+                    )
+                    raise
+            except Exception as e:
+                if attempt < max_retries - 1 and "not found" in str(e).lower():
+                    # Treat any "not found" error as retryable
+                    self.logger.warning(
+                        f"Insert failed with 'not found' error (attempt {attempt + 1}/{max_retries}), retrying in {retry_delay}s..."
+                    )
+                    await asyncio.sleep(retry_delay)
+                    retry_delay = min(retry_delay * 1.5, 10.0)
+                else:
+                    self.logger.error(f"Error adding documents: {e}")
+                    raise
+
+    def _get_distance_function(self, metric: str = None) -> str:
+        """Get BigQuery ML distance function based on strategy."""
+        strategy = metric or self.distance_strategy
+
+        if isinstance(strategy, str):
+            metric_mapping = {
+                'COSINE': DistanceStrategy.COSINE,
+                'L2': DistanceStrategy.EUCLIDEAN_DISTANCE,
+                'EUCLIDEAN': DistanceStrategy.EUCLIDEAN_DISTANCE,
+                'DOT': DistanceStrategy.DOT_PRODUCT,
+            }
+            strategy = metric_mapping.get(strategy.upper(), DistanceStrategy.COSINE)
+
+        if strategy == DistanceStrategy.COSINE:
+            return "ML.DISTANCE"  # Cosine distance in BigQuery ML
+        elif strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:
+            return "ML.EUCLIDEAN_DISTANCE"
+        elif strategy == DistanceStrategy.DOT_PRODUCT:
+            return "ML.DOT_PRODUCT"
+        else:
+            return "ML.DISTANCE"  # Default to cosine
+
+    async def similarity_search(
+        self,
+        query: str,
+        table: str = None,
+        dataset: str = None,
+        k: Optional[int] = None,
+        limit: int = None,
+        metadata_filters: Optional[Dict[str, Any]] = None,
+        score_threshold: Optional[float] = None,
+        metric: str = None,
+        embedding_column: str = 'embedding',
+        content_column: str = 'document',
+        metadata_column: str = 'metadata',
+        id_column: str = 'id',
+        **kwargs
+    ) -> List[SearchResult]:
+        """Perform similarity search using BigQuery ML functions."""
+        if not self._connected:
+            await self.connection()
+
+        table = table or self.table_name
+        dataset = dataset or self.dataset
+
+        if k and not limit:
+            limit = k
+        if not limit:
+            limit = 10
+
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        # Get query embedding
+        # query_embedding = self._embed_.embed_query(query)
+        query_embedding = await self._thread_func(self._embed_.embed_query, query)
+        if isinstance(query_embedding, np.ndarray):
+            query_embedding = query_embedding.tolist()
+
+        # Convert embedding to BigQuery array literal format
+        embedding_literal = "[" + ",".join([str(float(val)) for val in query_embedding]) + "]"
+
+        # Build the SQL query
+        distance_func = self._get_distance_function(metric)
+
+        # Create the SQL query with embedded array literal
+        sql_query = f"""
+        SELECT
+            {id_column},
+            {content_column},
+            {metadata_column},
+            {distance_func}({embedding_column}, {embedding_literal}) as distance
+        FROM `{table_id}`
+        WHERE {embedding_column} IS NOT NULL
+        """
+
+        # Add metadata filters
+        filter_params = []
+        if metadata_filters:
+            filter_conditions = []
+            for key, value in metadata_filters.items():
+                if isinstance(value, str):
+                    filter_conditions.append(f"JSON_EXTRACT_SCALAR({metadata_column}, '$.{key}') = @filter_{key}")
+                    filter_params.append(bq.ScalarQueryParameter(f"filter_{key}", "STRING", value))
+                else:
+                    filter_conditions.append(f"JSON_EXTRACT_SCALAR({metadata_column}, '$.{key}') = @filter_{key}")
+                    filter_params.append(bq.ScalarQueryParameter(f"filter_{key}", "STRING", str(value)))
+
+            if filter_conditions:
+                sql_query += " AND " + " AND ".join(filter_conditions)
+
+        # Add score threshold
+        if score_threshold is not None:
+            sql_query += f" AND {distance_func}({embedding_column}, {embedding_literal}) <= {score_threshold}"
+
+        # Order and limit
+        sql_query += f" ORDER BY distance ASC"
+        if limit:
+            sql_query += f" LIMIT {limit}"
+
+        # Configure query parameters
+        job_config = None
+        if filter_params:
+            job_config = bq.QueryJobConfig(
+                query_parameters=filter_params
+            )
+
+        try:
+            # Execute query
+            query_job = await self._thread_func(
+                self._connection.query, sql_query, job_config=job_config
+            )
+            results = await self._thread_func(query_job.result)
+
+            # Process results
+            search_results = []
+            for row in results:
+                metadata_str = row[metadata_column]
+                if isinstance(metadata_str, str):
+                    # Ensure metadata is a JSON string
+                    metadata_str = metadata_str.strip()
+                    metadata = self._json.loads(metadata_str)
+                else:
+                    metadata = dict(metadata_str) if metadata_str else {}
+
+                search_result = SearchResult(
+                    id=row[id_column],
+                    content=row[content_column],
+                    metadata=metadata,
+                    score=float(row.distance)
+                )
+                search_results.append(search_result)
+
+            self.logger.debug(
+                f"Similarity search returned {len(search_results)} results"
+            )
+            return search_results
+
+        except Exception as e:
+            self.logger.error(f"Error during similarity search: {e}")
+            raise
+
+    async def mmr_search(
+        self,
+        query: str,
+        table: str = None,
+        dataset: str = None,
+        k: int = 10,
+        fetch_k: int = None,
+        lambda_mult: float = 0.5,
+        metadata_filters: Optional[Dict[str, Any]] = None,
+        score_threshold: Optional[float] = None,
+        metric: str = None,
+        embedding_column: str = 'embedding',
+        content_column: str = 'document',
+        metadata_column: str = 'metadata',
+        id_column: str = 'id',
+        **kwargs
+    ) -> List[SearchResult]:
+        """
+        Perform Maximal Marginal Relevance (MMR) search.
+
+        Since BigQuery doesn't have native MMR support, we fetch more candidates
+        and perform MMR selection in Python.
+        """
+        if not self._connected:
+            await self.connection()
+
+        # Default to fetching 3x more candidates than final results
+        if fetch_k is None:
+            fetch_k = max(k * 3, 20)
+
+        # Step 1: Get initial candidates using similarity search
+        candidates = await self.similarity_search(
+            query=query,
+            table=table,
+            dataset=dataset,
+            limit=fetch_k,
+            metadata_filters=metadata_filters,
+            score_threshold=score_threshold,
+            metric=metric,
+            embedding_column=embedding_column,
+            content_column=content_column,
+            metadata_column=metadata_column,
+            id_column=id_column,
+            **kwargs
+        )
+
+        if len(candidates) <= k:
+            return candidates
+
+        # Step 2: Fetch embeddings for MMR computation
+        candidate_embeddings = await self._fetch_embeddings_for_mmr(
+            candidate_ids=[result.id for result in candidates],
+            table=table,
+            dataset=dataset,
+            embedding_column=embedding_column,
+            id_column=id_column
+        )
+
+        # Step 3: Get query embedding
+        query_embedding = self._embed_.embed_query(query)
+
+        # Step 4: Run MMR algorithm
+        selected_results = self._mmr_algorithm(
+            query_embedding=query_embedding,
+            candidates=candidates,
+            candidate_embeddings=candidate_embeddings,
+            k=k,
+            lambda_mult=lambda_mult,
+            metric=metric or self.distance_strategy
+        )
+
+        self.logger.info(
+            f"MMR search selected {len(selected_results)} results from {len(candidates)} candidates"
+        )
+
+        return selected_results
+
+    async def _fetch_embeddings_for_mmr(
+        self,
+        candidate_ids: List[str],
+        table: str,
+        dataset: str,
+        embedding_column: str,
+        id_column: str
+    ) -> Dict[str, np.ndarray]:
+        """Fetch embedding vectors for candidate documents from BigQuery."""
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        # Create placeholders for the IDs
+        id_placeholders = ', '.join([f"'{id_}'" for id_ in candidate_ids])
+
+        sql_query = f"""
+        SELECT {id_column}, {embedding_column}
+        FROM `{table_id}`
+        WHERE {id_column} IN ({id_placeholders})
+        """
+
+        try:
+            query_job = await self._thread_func(self._connection.query, sql_query)
+            results = await self._thread_func(query_job.result)
+
+            embeddings_dict = {}
+            for row in results:
+                doc_id = row[id_column]
+                embedding_data = row[embedding_column]
+
+                # Convert BigQuery array format back to numpy array
+                if embedding_data:
+                    embedding_values = [item['value'] for item in embedding_data]
+                    embedding_values = embedding_data if isinstance(embedding_data, list) else embedding_data.tolist()
+                    embeddings_dict[doc_id] = np.array(embedding_values, dtype=np.float32)
+
+            return embeddings_dict
+
+        except Exception as e:
+            self.logger.error(f"Error fetching embeddings for MMR: {e}")
+            raise
+
+    def _mmr_algorithm(
+        self,
+        query_embedding: np.ndarray,
+        candidates: List[SearchResult],
+        candidate_embeddings: Dict[str, np.ndarray],
+        k: int,
+        lambda_mult: float,
+        metric: str
+    ) -> List[SearchResult]:
+        """Core MMR algorithm implementation (same as PgVectorStore)."""
+        if len(candidates) <= k:
+            return candidates
+
+        # Convert query embedding to numpy array
+        if not isinstance(query_embedding, np.ndarray):
+            query_embedding = np.array(query_embedding, dtype=np.float32)
+
+        # Prepare data structures
+        selected_indices = []
+        remaining_indices = list(range(len(candidates)))
+
+        # Step 1: Select the most relevant document first
+        query_similarities = []
+        for candidate in candidates:
+            doc_embedding = candidate_embeddings.get(candidate.id)
+            if doc_embedding is not None:
+                similarity = self._compute_similarity(query_embedding, doc_embedding, metric)
+                query_similarities.append(similarity)
+            else:
+                # Fallback to distance score if embedding not available
+                query_similarities.append(1.0 / (1.0 + candidate.score))
+
+        # Select the most similar document first
+        best_idx = np.argmax(query_similarities)
+        selected_indices.append(best_idx)
+        remaining_indices.remove(best_idx)
+
+        # Step 2: Iteratively select remaining documents using MMR
+        for _ in range(min(k - 1, len(remaining_indices))):
+            mmr_scores = []
+
+            for idx in remaining_indices:
+                candidate = candidates[idx]
+                doc_embedding = candidate_embeddings.get(candidate.id)
+
+                if doc_embedding is None:
+                    # Fallback scoring if embedding not available
+                    mmr_score = lambda_mult * query_similarities[idx]
+                    mmr_scores.append(mmr_score)
+                    continue
+
+                # Relevance: similarity to query
+                relevance = query_similarities[idx]
+
+                # Diversity: maximum similarity to already selected documents
+                max_similarity_to_selected = 0.0
+                for selected_idx in selected_indices:
+                    selected_candidate = candidates[selected_idx]
+                    selected_embedding = candidate_embeddings.get(selected_candidate.id)
+
+                    if selected_embedding is not None:
+                        similarity = self._compute_similarity(doc_embedding, selected_embedding, metric)
+                        max_similarity_to_selected = max(max_similarity_to_selected, similarity)
+
+                # MMR formula: λ * relevance - (1-λ) * max_similarity_to_selected
+                mmr_score = (
+                    lambda_mult * relevance -
+                    (1.0 - lambda_mult) * max_similarity_to_selected
+                )
+                mmr_scores.append(mmr_score)
+
+            # Select document with highest MMR score
+            if mmr_scores:
+                best_remaining_idx = np.argmax(mmr_scores)
+                best_idx = remaining_indices[best_remaining_idx]
+                selected_indices.append(best_idx)
+                remaining_indices.remove(best_idx)
+
+        # Step 3: Return selected results with MMR scores in metadata
+        selected_results = []
+        for i, idx in enumerate(selected_indices):
+            result = candidates[idx]
+            # Add MMR ranking to metadata
+            enhanced_metadata = dict(result.metadata)
+            enhanced_metadata['mmr_rank'] = i + 1
+            enhanced_metadata['mmr_lambda'] = lambda_mult
+            enhanced_metadata['original_rank'] = idx + 1
+
+            enhanced_result = SearchResult(
+                id=result.id,
+                content=result.content,
+                metadata=enhanced_metadata,
+                score=result.score
+            )
+            selected_results.append(enhanced_result)
+
+        return selected_results
+
+    def _compute_similarity(
+        self,
+        embedding1: np.ndarray,
+        embedding2: np.ndarray,
+        metric: Union[str, Any]
+    ) -> float:
+        """Compute similarity between two embeddings (same as PgVectorStore)."""
+        # Convert to numpy arrays if needed
+        if isinstance(embedding1, list):
+            embedding1 = np.array(embedding1, dtype=np.float32)
+        if isinstance(embedding2, list):
+            embedding2 = np.array(embedding2, dtype=np.float32)
+
+        # Ensure embeddings are numpy arrays
+        if not isinstance(embedding1, np.ndarray):
+            embedding1 = np.array(embedding1, dtype=np.float32)
+        if not isinstance(embedding2, np.ndarray):
+            embedding2 = np.array(embedding2, dtype=np.float32)
+
+        # Ensure embeddings are 2D arrays for sklearn
+        emb1 = embedding1.reshape(1, -1)
+        emb2 = embedding2.reshape(1, -1)
+
+        # Convert string metrics to DistanceStrategy enum if needed
+        if isinstance(metric, str):
+            metric_mapping = {
+                'COSINE': DistanceStrategy.COSINE,
+                'L2': DistanceStrategy.EUCLIDEAN_DISTANCE,
+                'EUCLIDEAN': DistanceStrategy.EUCLIDEAN_DISTANCE,
+                'DOT': DistanceStrategy.DOT_PRODUCT,
+            }
+            strategy = metric_mapping.get(metric.upper(), DistanceStrategy.COSINE)
+        else:
+            strategy = metric
+
+        if strategy == DistanceStrategy.COSINE:
+            # Cosine similarity
+            similarity = cosine_similarity(emb1, emb2)[0, 0]
+            return float(similarity)
+        elif strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:
+            # Convert Euclidean distance to similarity
+            distance = euclidean_distances(emb1, emb2)[0, 0]
+            similarity = 1.0 / (1.0 + distance)
+            return float(similarity)
+        elif strategy == DistanceStrategy.DOT_PRODUCT:
+            # Dot product
+            similarity = np.dot(embedding1.flatten(), embedding2.flatten())
+            return float(similarity)
+        else:
+            # Default to cosine similarity
+            similarity = cosine_similarity(emb1, emb2)[0, 0]
+            return float(similarity)
+
+    async def delete_documents(
+        self,
+        documents: Optional[List[Document]] = None,
+        pk: str = 'source_type',
+        values: Optional[Union[str, List[str]]] = None,
+        table: Optional[str] = None,
+        dataset: Optional[str] = None,
+        metadata_column: Optional[str] = None,
+        **kwargs
+    ) -> int:
+        """Delete documents from BigQuery table based on metadata field values."""
+        if not self._connected:
+            await self.connection()
+
+        table = table or self.table_name
+        dataset = dataset or self.dataset
+        metadata_column = metadata_column or self._metadata_column
+
+        if not table:
+            raise ValueError("Table name must be provided")
+
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        # Extract values to delete
+        delete_values = []
+        if values is not None:
+            if isinstance(values, str):
+                delete_values = [values]
+            else:
+                delete_values = list(values)
+        elif documents:
+            for doc in documents:
+                if hasattr(doc, 'metadata') and doc.metadata and pk in doc.metadata:
+                    value = doc.metadata[pk]
+                    if value and value not in delete_values:
+                        delete_values.append(value)
+        else:
+            raise ValueError("Either 'documents' or 'values' parameter must be provided")
+
+        if not delete_values:
+            self.logger.warning(f"No values found for field '{pk}' to delete")
+            return 0
+
+        deleted_count = 0
+
+        try:
+            for value in delete_values:
+                # Create delete query using JSON extraction
+                delete_query = f"""
+                DELETE FROM `{table_id}`
+                WHERE JSON_EXTRACT_SCALAR({metadata_column}, '$.{pk}') = @value
+                """
+
+                job_config = bq.QueryJobConfig(
+                    query_parameters=[
+                        bq.ScalarQueryParameter("value", "STRING", str(value))
+                    ]
+                )
+
+                query_job = await self._thread_func(
+                    self._connection.query,
+                    delete_query,
+                    job_config=job_config
+                )
+                await self._thread_func(query_job.result)
+
+                rows_deleted = query_job.num_dml_affected_rows or 0
+                deleted_count += rows_deleted
+
+                self.logger.info(
+                    f"Deleted {rows_deleted} documents with {pk}='{value}' from {table_id}"
+                )
+
+            self.logger.info(f"Total deleted: {deleted_count} documents")
+            return deleted_count
+
+        except Exception as e:
+            self.logger.error(f"Error deleting documents: {e}")
+            raise RuntimeError(f"Failed to delete documents: {e}") from e
+
+    async def delete_documents_by_filter(
+        self,
+        filter_dict: Dict[str, Union[str, List[str]]],
+        table: Optional[str] = None,
+        dataset: Optional[str] = None,
+    ) -> int:
+        """Deletes documents based on multiple metadata field conditions."""
+        if not self._connected: await self.connection()
+        if not filter_dict: raise ValueError("filter_dict cannot be empty")
+
+        table = table or self.table_name
+        dataset = dataset or self.dataset
+        table_id = f"`{self._project_id}.{dataset}.{table}`"
+
+        where_conditions = []
+        query_params = []
+        for i, (field, values) in enumerate(filter_dict.items()):
+            safe_field = field.replace("'", "\\'")
+            if isinstance(values, (list, tuple)):
+                param_name = f"val_{i}"
+                where_conditions.append(f"JSON_VALUE({self._metadata_column}, '$.{safe_field}') IN UNNEST(@{param_name})")
+                query_params.append(bq.ArrayQueryParameter(param_name, "STRING", [str(v) for v in values]))
+            else:
+                param_name = f"val_{i}"
+                where_conditions.append(f"JSON_VALUE({self._metadata_column}, '$.{safe_field}') = @{param_name}")
+                query_params.append(bq.ScalarQueryParameter(param_name, "STRING", str(values)))
+
+        where_clause = " AND ".join(where_conditions)
+        delete_query = f"DELETE FROM {table_id} WHERE {where_clause}"
+        job_config = bq.QueryJobConfig(query_parameters=query_params)
+
+        try:
+            query_job = await self._thread_func(
+                self._connection.query, delete_query, job_config=job_config
+            )
+            await self._thread_func(query_job.result)
+            deleted_count = query_job.num_dml_affected_rows or 0
+            self.logger.debug(
+                f"Deleted {deleted_count} documents from {table_id} with filter: {filter_dict}"
+            )
+            return deleted_count
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to delete documents by filter: {e}"
+            ) from e
+
+
+    async def delete_documents_by_ids(
+        self,
+        document_ids: List[str],
+        table: Optional[str] = None,
+        dataset: Optional[str] = None,
+        id_column: Optional[str] = None,
+        **kwargs
+    ) -> int:
+        """Delete documents by their IDs."""
+        if not self._connected:
+            await self.connection()
+
+        if not document_ids:
+            self.logger.warning("No document IDs provided for deletion")
+            return 0
+
+        table = table or self.table_name
+        dataset = dataset or self.dataset
+        id_column = id_column or self._id_column
+
+        if not table:
+            raise ValueError("Table name must be provided")
+
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        # Build parameterized query for multiple IDs
+        query_parameters = []
+        value_params = []
+        for i, doc_id in enumerate(document_ids):
+            param_name = f"id_{i}"
+            value_params.append(f"@{param_name}")
+            query_parameters.append(
+                bq.ScalarQueryParameter(param_name, "STRING", str(doc_id))
+            )
+
+        delete_query = f"""
+        DELETE FROM `{table_id}`
+        WHERE {id_column} IN ({', '.join(value_params)})
+        """
+
+        try:
+            job_config = bq.QueryJobConfig(query_parameters=query_parameters)
+            query_job = self._connection.query(delete_query, job_config=job_config)
+            query_job.result()  # Wait for completion
+
+            deleted_count = query_job.num_dml_affected_rows or 0
+
+            self.logger.info(
+                f"Deleted {deleted_count} documents by IDs from {table_id}"
+            )
+
+            return deleted_count
+
+        except Exception as e:
+            self.logger.error(f"Error deleting documents by IDs: {e}")
+            raise RuntimeError(
+                f"Failed to delete documents by IDs: {e}"
+            ) from e
+
+    async def delete_all_documents(
+        self,
+        table: Optional[str] = None,
+        dataset: Optional[str] = None,
+        confirm: bool = False,
+        **kwargs
+    ) -> int:
+        """Delete ALL documents from the BigQuery table."""
+        if not confirm:
+            raise ValueError(
+                "This operation will delete ALL documents. "
+                "Set confirm=True to proceed."
+            )
+
+        if not self._connected:
+            await self.connection()
+
+        table = table or self.table_name
+        dataset = dataset or self.dataset
+
+        if not table:
+            raise ValueError("Table name must be provided")
+
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        try:
+            # First count existing documents
+            count_query = f"SELECT COUNT(*) as total FROM `{table_id}`"
+            count_job = self._connection.query(count_query)
+            count_result = list(count_job.result())[0]
+            total_docs = count_result.total
+
+            if total_docs == 0:
+                self.logger.info(f"No documents to delete from {table_id}")
+                return 0
+
+            # Delete all documents
+            delete_query = f"DELETE FROM `{table_id}` WHERE TRUE"
+            query_job = self._connection.query(delete_query)
+            query_job.result()  # Wait for completion
+
+            deleted_count = query_job.num_dml_affected_rows or 0
+
+            self.logger.warning(
+                f"DELETED ALL {deleted_count} documents from {table_id}"
+            )
+
+            return deleted_count
+
+        except Exception as e:
+            self.logger.error(f"Error deleting all documents: {e}")
+            raise RuntimeError(f"Failed to delete all documents: {e}") from e
+
+    async def count_documents_by_filter(
+        self,
+        filter_dict: Dict[str, Union[str, List[str]]],
+        table: Optional[str] = None,
+        dataset: Optional[str] = None,
+        metadata_column: Optional[str] = None,
+        **kwargs
+    ) -> int:
+        """Count documents that would be affected by a filter."""
+        if not self._connected:
+            await self.connection()
+
+        if not filter_dict:
+            return 0
+
+        table = table or self.table_name
+        dataset = dataset or self.dataset
+        metadata_column = metadata_column or self._metadata_column
+
+        if not table:
+            raise ValueError("Table name must be provided")
+
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        # Build WHERE conditions (same logic as delete_documents_by_filter)
+        where_conditions = []
+        query_parameters = []
+
+        for field, values in filter_dict.items():
+            if isinstance(values, (list, tuple)):
+                value_params = []
+                for i, value in enumerate(values):
+                    param_name = f"{field}_{i}"
+                    value_params.append(f"@{param_name}")
+                    query_parameters.append(
+                        bq.ScalarQueryParameter(param_name, "STRING", str(value))
+                    )
+
+                condition = f"JSON_EXTRACT_SCALAR({metadata_column}, '$.{field}') IN ({', '.join(value_params)})"
+                where_conditions.append(condition)
+            else:
+                param_name = f"{field}_single"
+                where_conditions.append(f"JSON_EXTRACT_SCALAR({metadata_column}, '$.{field}') = @{param_name}")
+                query_parameters.append(
+                    bq.ScalarQueryParameter(param_name, "STRING", str(values))
+                )
+
+        where_clause = " AND ".join(where_conditions)
+        count_query = f"""
+        SELECT COUNT(*) as total FROM `{table_id}`
+        WHERE {where_clause}
+        """
+
+        try:
+            job_config = bq.QueryJobConfig(query_parameters=query_parameters)
+            query_job = self._connection.query(count_query, job_config=job_config)
+            result = list(query_job.result())[0]
+            count = result.total
+
+            self.logger.info(
+                f"Found {count} documents matching filter: {filter_dict}"
+            )
+
+            return count
+
+        except Exception as e:
+            self.logger.error(f"Error counting documents: {e}")
+            raise RuntimeError(f"Failed to count documents: {e}") from e
+
+    async def delete_collection(
+        self,
+        table: str,
+        dataset: str = None
+    ) -> None:
+        """Delete a collection (table) from BigQuery."""
+        if not self._connected:
+            await self.connection()
+
+        dataset = dataset or self.dataset
+        table_id = f"{self._project_id}.{dataset}.{table}"
+
+        if not await self.collection_exists(table, dataset):
+            raise RuntimeError(f"Collection {table_id} does not exist")
+
+        try:
+            self._connection.delete_table(table_id)
+            self.logger.info(f"Collection {table_id} deleted successfully")
+
+            # Remove from cache
+            if table_id in self._collection_store_cache:
+                del self._collection_store_cache[table_id]
+
+        except Exception as e:
+            self.logger.error(f"Error deleting collection: {e}")
+            raise RuntimeError(f"Failed to delete collection: {e}") from e
+
+    async def from_documents(
+        self,
+        documents: List[Document],
+        table: str = None,
+        dataset: str = None,
+        embedding_column: str = 'embedding',
+        content_column: str = 'document',
+        metadata_column: str = 'metadata',
+        chunk_size: int = 8192,
+        chunk_overlap: int = 200,
+        store_full_document: bool = True,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """Add documents using late chunking strategy (if available)."""
+        if not self._connected:
+            await self.connection()
+
+        table = table or self.table_name
+        dataset = dataset or self.dataset
+
+        if not table:
+            raise ValueError("Table name must be provided")
+
+        # For BigQuery, we'll implement a simpler version without late chunking
+        # since LateChunkingProcessor might not be available
+        await self.add_documents(
+            documents=documents,
+            table=table,
+            dataset=dataset,
+            embedding_column=embedding_column,
+            content_column=content_column,
+            metadata_column=metadata_column,
+            **kwargs
+        )
+
+        stats = {
+            'documents_processed': len(documents),
+            'chunks_created': 0,  # Not implementing chunking in this version
+            'full_documents_stored': len(documents)
+        }
+
+        return stats
+
+    # Context manager support
+    async def __aenter__(self):
+        """Context manager entry."""
+        if not self._connected:
+            await self.connection()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        # BigQuery client doesn't need explicit cleanup
+        pass
+
+    async def disconnect(self) -> None:
+        """Disconnect from BigQuery (cleanup resources)."""
+        if self._connection:
+            self._connection.close()
+            self._connection = None
+            self._connected = False
+            self.logger.info("BigQuery client disconnected")
+
+    def __str__(self) -> str:
+        return f"BigQueryStore(project={self._project_id}, dataset={self.dataset})"
+
+    def __repr__(self) -> str:
+        return (
+            f"<BigQueryStore(project='{self._project_id}', "
+            f"dataset='{self.dataset}', table='{self.table_name}')>"
+        )