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

parrot/stores/postgres.py

@@ -0,0 +1,2853 @@
+from typing import Any, Dict, List, Optional, Union, Callable
+import uuid
+from contextlib import asynccontextmanager
+import numpy as np
+import sqlalchemy
+from sqlalchemy import (
+    text,
+    Column,
+    insert,
+    Table,
+    MetaData,
+    select,
+    asc,
+    func,
+    event,
+    JSON,
+    Index
+)
+from sqlalchemy.sql import literal_column
+from sqlalchemy import bindparam
+from sqlalchemy.orm import aliased
+from sqlalchemy.ext.asyncio import (
+    create_async_engine,
+    AsyncSession,
+    AsyncEngine,
+    async_sessionmaker
+)
+from sqlalchemy.sql.expression import cast
+from sqlalchemy.dialects.postgresql import JSONB, ARRAY
+from sqlalchemy.orm import (
+    declarative_base,
+    DeclarativeBase,
+    Mapped,
+    mapped_column
+)
+# PgVector
+from pgvector.sqlalchemy import Vector
+from pgvector.asyncpg import register_vector
+from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances
+# Datamodel
+from datamodel.parsers.json import json_encoder  # pylint: disable=E0611
+from navconfig.logging import logging
+from .abstract import AbstractStore
+from ..conf import default_sqlalchemy_pg
+from .models import SearchResult, Document, DistanceStrategy
+from .utils.chunking import LateChunkingProcessor
+
+
+def vector_distance(embedding_column, vector, op):
+    return text(f"{embedding_column} {op} :query_embedding").label("distance")
+
+class Base(DeclarativeBase):
+    pass
+
+class PgVectorStore(AbstractStore):
+    """
+    A PostgreSQL vector store implementation using pgvector, completely independent of Langchain.
+    This store interacts directly with a specified schema and table for robust data isolation.
+    """
+    def __init__(
+        self,
+        table: str = None,
+        schema: str = 'public',
+        id_column: str = 'id',
+        embedding_column: str = 'embedding',
+        document_column: str = 'document',
+        text_column: str = 'text',
+        embedding_model: Union[dict, str] = "sentence-transformers/all-mpnet-base-v2",
+        embedding: Optional[Callable] = None,
+        distance_strategy: DistanceStrategy = DistanceStrategy.COSINE,
+        use_uuid: bool = False,
+        pool_size: int = 50,
+        auto_initialize: bool = True,
+        **kwargs
+    ):
+        """ Initializes the PgVectorStore with the specified parameters.
+        """
+        self.table_name = table
+        self.schema = schema
+        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.distance_strategy = distance_strategy
+        self._use_uuid: bool = use_uuid
+        self._embedding_store_cache: Dict[str, Any] = {}
+        self._max_size = pool_size or 50
+        self._auto_initialize_db: bool = auto_initialize
+        super().__init__(
+            embedding_model=embedding_model,
+            embedding=embedding,
+            **kwargs
+        )
+        self.dsn = kwargs.get('dsn', default_sqlalchemy_pg)
+        self._connection: Optional[AsyncEngine] = None
+        self._session_factory: Optional[async_sessionmaker] = None
+        self._session: Optional[AsyncSession] = None
+        self.logger = logging.getLogger("PgVectorStore")
+        self.embedding_store = None
+        if table:
+            # create a table definition:
+            self.embedding_store = self._define_collection_store(
+                table=table,
+                schema=schema,
+                dimension=self.dimension,
+                id_column=id_column,
+                embedding_column=embedding_column,
+                document_column=self._document_column,
+                text_column=text_column,
+            )
+
+    def get_id_column(self, use_uuid: bool) -> sqlalchemy.Column:
+        """
+        Return the ID column definition based on whether to use UUID or not.
+        If use_uuid is True, the ID column will be a PostgreSQL UUID type with
+        server-side generation using uuid_generate_v4().
+        If use_uuid is False, the ID column will be a String type with a default
+        value generated by Python's uuid.uuid4() function.
+        """
+        if use_uuid:
+            # DB will auto-generate UUID; SQLAlchemy should not set a default!
+            return sqlalchemy.Column(
+                sqlalchemy.dialects.postgresql.UUID(as_uuid=True),
+                primary_key=True,
+                index=True,
+                unique=True,
+                server_default=sqlalchemy.text('uuid_generate_v4()')
+            )
+        else:
+            # Python generates UUID (as string)
+            return sqlalchemy.Column(
+                sqlalchemy.String,
+                primary_key=True,
+                index=True,
+                unique=True,
+                default=lambda: str(uuid.uuid4())
+            )
+
+    def _define_collection_store(
+        self,
+        table: str,
+        schema: str,
+        dimension: int = 384,
+        id_column: str = 'id',
+        embedding_column: str = 'embedding',
+        document_column: str = 'document',
+        metadata_column: str = 'cmetadata',
+        text_column: str = 'text',
+        store_name: str = 'EmbeddingStore',
+        colbert_dimension: int = 128  # ColBERT token dimension
+    ) -> Any:
+        """Dynamically define a SQLAlchemy Table for pgvector storage.
+
+        Args:
+            table: The name of the table to create.
+            schema: The schema in which to create the table.
+            dimension: The dimensionality of the vector embeddings.
+        """
+        fq_table_name = f"{schema}.{table}"
+        if fq_table_name in self._embedding_store_cache:
+            return self._embedding_store_cache[fq_table_name]
+
+        self.logger.notice(
+            f"Defining dynamic ORM class for table {fq_table_name} with dimension {dimension}"
+        )
+        table_args = {
+            "schema": schema,
+            "extend_existing": True
+        }
+        attrs = {
+            '__tablename__': table,
+            '__table_args__': table_args,
+            # id_column: self.get_id_column(use_uuid=self._use_uuid),
+            id_column: mapped_column(
+                sqlalchemy.String if not self._use_uuid else sqlalchemy.dialects.postgresql.UUID(as_uuid=True),
+                primary_key=True,
+                index=True,
+                unique=True,
+                default=lambda: str(uuid.uuid4()) if not self._use_uuid else None,
+                server_default=sqlalchemy.text('uuid_generate_v4()') if self._use_uuid else None
+            ),
+            embedding_column: mapped_column(Vector(dimension)),
+            text_column: mapped_column(sqlalchemy.String, nullable=True),
+            document_column: mapped_column(sqlalchemy.String, nullable=True),
+            metadata_column: mapped_column(JSONB, nullable=True),
+
+            # embedding_column: Column(Vector(dimension)),
+            # text_column: Column(sqlalchemy.String, nullable=True),
+            # document_column: Column(sqlalchemy.String, nullable=True),
+            # metadata_column: Column(JSONB, nullable=True),
+            # ColBERT columns
+            # 'token_embeddings': Column(ARRAY(Vector(colbert_dimension)), nullable=True),
+            # 'num_tokens': Column(sqlalchemy.Integer, nullable=True),
+            # 'collection_id': Column(
+            #     sqlalchemy.dialects.postgresql.UUID(as_uuid=True),
+            #     index=True,
+            #     unique=True,
+            #     default=uuid.uuid4,
+            #     server_default=sqlalchemy.text('uuid_generate_v4()')
+            # )
+            'token_embeddings': mapped_column(ARRAY(Vector(colbert_dimension)), nullable=True),
+            'num_tokens': mapped_column(sqlalchemy.Integer, nullable=True),
+            'collection_id': mapped_column(
+                sqlalchemy.dialects.postgresql.UUID(as_uuid=True),
+                index=True,
+                unique=True,
+                default=uuid.uuid4,
+                server_default=sqlalchemy.text('uuid_generate_v4()')
+            )
+        }
+
+        # Create dynamic ORM class
+        EmbeddingStore = type(store_name, (Base,), attrs)
+        EmbeddingStore.__name__ = store_name
+        EmbeddingStore.__qualname__ = store_name
+
+        # Cache the store
+        self._embedding_store_cache[fq_table_name] = EmbeddingStore
+        self.logger.debug(
+            f"Created dynamic ORM class {store_name} for table {fq_table_name}"
+        )
+
+        return EmbeddingStore
+
+    def define_collection_table(
+        self,
+        table: str,
+        schema: str,
+        dimension: int = 384,
+        metadata: Optional[MetaData] = None,
+        use_uuid: bool = False,
+        id_column: str = 'id',
+        embedding_column: str = 'embedding'
+    ) -> sqlalchemy.Table:
+        """Dynamically define a SQLAlchemy Table for pgvector storage."""
+        columns = []
+
+        if use_uuid:
+            columns.append(Column(
+                id_column,
+                sqlalchemy.dialects.postgresql.UUID(as_uuid=True),
+                primary_key=True,
+                server_default=sqlalchemy.text("uuid_generate_v4()")
+            ))
+        else:
+            columns.append(Column(
+                id_column,
+                sqlalchemy.String,
+                primary_key=True,
+                default=lambda: str(uuid.uuid4())
+            ))
+
+        columns.extend([
+            Column(embedding_column, Vector(dimension)),
+            Column('text', sqlalchemy.String, nullable=True),
+            Column('document', sqlalchemy.String, nullable=True),
+            Column('cmetadata', JSONB, nullable=True)
+        ])
+
+        return Table(
+            table,
+            metadata,
+            *columns,
+            schema=schema
+        )
+
+    async def connection(self, dsn: str = None) -> AsyncEngine:
+        """Establishes and returns an async database connection."""
+        if self._connection is not None:
+            return self._connection
+        if not dsn:
+            dsn = self.dsn or default_sqlalchemy_pg
+        try:
+            self._connection = create_async_engine(
+                dsn,
+                future=True,
+                pool_size=self._max_size,           # High concurrency support
+                max_overflow=100,       # Burst capacity
+                pool_pre_ping=True,     # Connection health checks
+                pool_recycle=3600,      # Prevent stale connections (1 hour)
+                pool_timeout=30,        # Wait max 30s for connection
+                connect_args={
+                    "server_settings": {
+                        "jit": "off",                    # Disable JIT for vector queries
+                        "random_page_cost": "1.1",       # SSD optimization
+                        "effective_cache_size": "24GB",  # Memory configuration
+                        "work_mem": "256MB"
+                    }
+                }
+            )
+            # @event.listens_for(self._connection.sync_engine, "first_connect")
+            # def connect(dbapi_connection, connection_record):
+            #     dbapi_connection.run_async(register_vector)
+
+            # Create session factory
+            self._session_factory = async_sessionmaker(
+                bind=self._connection,
+                class_=AsyncSession,
+                expire_on_commit=False,
+                autoflush=False,  # Manual control over flushing
+                autocommit=False
+            )
+            if self._auto_initialize_db:
+                await self.initialize_database()
+            self._connected = True
+            self.logger.info(
+                "Successfully connected to PostgreSQL."
+            )
+        except Exception as e:
+            self.logger.error(
+                f"Failed to connect to PostgreSQL: {e}"
+            )
+            self._connected = False
+            raise
+
+    async def get_session(self) -> AsyncSession:
+        """Get a session from the pool. This is the main method for getting connections."""
+        if not self._connection:
+            await self.connection()
+
+        if not self._session_factory:
+            raise RuntimeError("Session factory not initialized")
+
+        return self._session_factory()
+
+    @asynccontextmanager
+    async def session(self):
+        """
+        Context manager for handling database sessions with proper cleanup.
+        This is the recommended way to handle database operations.
+
+        Usage:
+            async with store.session() as session:
+                result = await session.execute(stmt)
+                await session.commit()  # if needed
+        """
+        if not self._connection:
+            await self.connection()
+
+        session = await self.get_session()
+        try:
+            yield session
+            # Auto-commit if no exception occurred
+            if session.in_transaction():
+                await session.commit()
+        except Exception:
+            # Auto-rollback on exception
+            if session.in_transaction():
+                await session.rollback()
+            raise
+        finally:
+            # Always close session (returns connection to pool)
+            await session.close()
+
+    async def initialize_database(self):
+        """Initialize with PgVector 0.8.0+ optimizations"""
+        try:
+            async with self.session() as session:
+                await session.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
+
+                # Enable iterative scanning (breakthrough feature)
+                await session.execute(text("SET hnsw.iterative_scan = 'relaxed_order'"))
+                await session.execute(text("SET hnsw.max_scan_tuples = 20000"))
+                await session.execute(text("SET hnsw.ef_search = 200"))
+                await session.execute(text("SET ivfflat.iterative_scan = 'on'"))
+                await session.execute(text("SET ivfflat.max_probes = 100"))
+
+                # Performance tuning
+                await session.execute(text("SET maintenance_work_mem = '2GB'"))
+                await session.execute(text("SET max_parallel_maintenance_workers = 8"))
+                await session.execute(text("SET enable_seqscan = off"))
+
+                # Create ColBERT MaxSim function
+                await self._create_maxsim_function(session)
+
+                await session.commit()
+        except Exception as e:
+            self.logger.warning(f"⚠️ Database auto-initialization failed: {e}")
+            # Don't raise - let the engine continue to work
+
+    async def _create_maxsim_function(self, session):
+        """Create the MaxSim function for ColBERT late interaction scoring."""
+        maxsim_function = text("""
+            CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[])
+            RETURNS double precision AS $$
+            DECLARE
+                query_vec vector;
+                doc_vec vector;
+                max_similarity double precision;
+                total_score double precision := 0;
+                similarity double precision;
+            BEGIN
+                -- For each query token, find the maximum similarity with any document token
+                FOR i IN 1..array_length(query, 1) LOOP
+                    query_vec := query[i];
+                    max_similarity := -1;
+
+                    -- Find max similarity with all document tokens
+                    FOR j IN 1..array_length(document, 1) LOOP
+                        doc_vec := document[j];
+                        similarity := 1 - (query_vec <=> doc_vec);  -- Convert distance to similarity
+
+                        IF similarity > max_similarity THEN
+                            max_similarity := similarity;
+                        END IF;
+                    END LOOP;
+
+                    -- Add the maximum similarity for this query token
+                    total_score := total_score + max_similarity;
+                END LOOP;
+
+                RETURN total_score;
+            END;
+            $$ LANGUAGE plpgsql IMMUTABLE STRICT;
+        """)
+
+        await session.execute(maxsim_function)
+        self.logger.info("✅ Created ColBERT MaxSim function")
+
+
+    # Async Context Manager - improved pattern
+    async def __aenter__(self):
+        """
+        Context manager entry. Ensures engine is initialized and manages session lifecycle.
+        """
+        if not self._connection:
+            await self.connection()
+
+        # Create a session for this context if we don't have one
+        if self._session is None:
+            self._session = await self.get_session()
+
+        self._context_depth += 1
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """
+        Context manager exit. Properly handles session cleanup.
+        """
+        self._context_depth -= 1
+
+        # Only close session when we exit the outermost context
+        if self._context_depth == 0 and self._session:
+            try:
+                if exc_type is not None:
+                    # Exception occurred, rollback
+                    if self._session.in_transaction():
+                        await self._session.rollback()
+                else:
+                    # No exception, commit if in transaction
+                    if self._session.in_transaction():
+                        await self._session.commit()
+            finally:
+                # Always close the session (returns connection to pool)
+                await self._session.close()
+                self._session = None
+
+    async def _free_resources(self):
+        """Clean up resources but keep the engine/pool available."""
+        if self._embed_:
+            self._embed_.free()
+        self._embed_ = None
+
+        # Close current session if exists
+        if self._session:
+            await self._session.close()
+            self._session = None
+
+    async def disconnect(self) -> None:
+        """
+        Completely dispose of the engine and close all connections.
+        Call this when you're completely done with the store.
+        """
+        # Close current session first
+        if self._session:
+            await self._session.close()
+            self._session = None
+
+        # Dispose of the engine (closes all pooled connections)
+        if self._connection:
+            await self._connection.dispose()
+            self._connection = None
+            self._connected = False
+            self._session_factory = None
+            self.logger.info(
+                "🔌 PostgreSQL engine disposed and all connections closed"
+            )
+
+    async def add_documents(
+        self,
+        documents: List[Document],
+        table: str = None,
+        schema: str = None,
+        embedding_column: str = 'embedding',
+        content_column: str = 'document',
+        metadata_column: str = 'cmetadata',
+        **kwargs
+    ) -> None:
+        """
+        Embeds and adds documents to the specified table.
+
+        Args:
+            documents: A list of Document objects to add.
+            table: The name of the table.
+            schema: The database schema where the table resides.
+            embedding_column: The name of the column to store embeddings.
+            content_column: The name of the column to store the main text content.
+            metadata_column: The name of the JSONB column for metadata.
+        """
+        if not self._connected:
+            await self.connection()
+
+        if not table:
+            table = self.table_name
+        if not schema:
+            schema = self.schema
+
+        texts = [doc.page_content for doc in documents]
+        embeddings = self._embed_.embed_documents(texts)
+        metadatas = [doc.metadata for doc in documents]
+
+        # Step 1: Ensure the ORM table is initialized
+        if self.embedding_store is None:
+            self.embedding_store = self._define_collection_store(
+                table=table,
+                schema=schema,
+                dimension=self.dimension,
+                id_column=self._id_column,
+                embedding_column=embedding_column,
+                document_column=content_column,
+                metadata_column=metadata_column,
+                text_column=self._text_column,
+            )
+
+        # Step 2: Prepare values for bulk insert
+        values = [
+            {
+                self._id_column: str(uuid.uuid4()),
+                embedding_column: embeddings[i].tolist() if isinstance(
+                    embeddings[i], np.ndarray
+                ) else embeddings[i],
+                content_column: texts[i],
+                metadata_column: metadatas[i] or {}
+            }
+            for i in range(len(documents))
+        ]
+
+        # Step 3: Build insert statement using SQLAlchemy's insert()
+        insert_stmt = insert(self.embedding_store)
+
+        # Step 4: Execute using async executemany
+        try:
+            async with self.session() as session:
+                await session.execute(insert_stmt, values)
+            self.logger.info(
+                f"✅ Successfully added {len(documents)} documents to '{schema}.{table}'"
+            )
+        except Exception as e:
+            self.logger.error(f"Error adding documents: {e}")
+            raise
+
+    def get_distance_strategy(
+        self,
+        embedding_column_obj,
+        query_embedding,
+        metric: str = None
+    ) -> Any:
+        """
+        Return the appropriate distance expression based on the metric or configured strategy.
+
+        Args:
+            embedding_column_obj: The SQLAlchemy column object for embeddings
+            query_embedding: The query embedding vector
+            metric: Optional metric string ('COSINE', 'L2', 'IP', 'DOT')
+                - if None, uses self.distance_strategy
+        """
+        # Use provided metric or fall back to instance distance_strategy
+        strategy = metric or self.distance_strategy
+        # self.logger.debug(
+        #     f"PgVector: using distance strategy → {strategy}"
+        # )
+
+        # Convert string metrics to DistanceStrategy enum if needed
+        if isinstance(strategy, str):
+            metric_mapping = {
+                'COSINE': DistanceStrategy.COSINE,
+                'L2': DistanceStrategy.EUCLIDEAN_DISTANCE,
+                'EUCLIDEAN': DistanceStrategy.EUCLIDEAN_DISTANCE,
+                'IP': DistanceStrategy.MAX_INNER_PRODUCT,
+                'DOT': DistanceStrategy.DOT_PRODUCT,
+                'DOT_PRODUCT': DistanceStrategy.DOT_PRODUCT,
+                'MAX_INNER_PRODUCT': DistanceStrategy.MAX_INNER_PRODUCT
+            }
+            strategy = metric_mapping.get(strategy.upper(), DistanceStrategy.COSINE)
+
+        # self.logger.debug(
+        #     f"PgVector: using distance strategy → {strategy}"
+        # )
+
+        # Convert numpy array to list if needed
+        if isinstance(query_embedding, np.ndarray):
+            query_embedding = query_embedding.tolist()
+
+        if strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:
+            return embedding_column_obj.l2_distance(query_embedding)
+        elif strategy == DistanceStrategy.COSINE:
+            return embedding_column_obj.cosine_distance(query_embedding)
+        elif strategy == DistanceStrategy.MAX_INNER_PRODUCT:
+            return embedding_column_obj.max_inner_product(query_embedding)
+        elif strategy == DistanceStrategy.DOT_PRODUCT:
+            # Note: pgvector doesn't have dot_product, using max_inner_product
+            return embedding_column_obj.max_inner_product(query_embedding)
+        else:
+            raise ValueError(
+                f"Got unexpected value for distance: {strategy}. "
+                f"Should be one of {', '.join([ds.value for ds in DistanceStrategy])}."
+            )
+
+    async def similarity_search(
+        self,
+        query: str,
+        table: str = None,
+        schema: 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 = 'cmetadata',
+        id_column: str = 'id',
+        additional_columns: Optional[List[str]] = None
+    ) -> List[SearchResult]:
+        """
+        Perform similarity search with optional threshold filtering.
+
+        Args:
+            query: The search query text
+            table: Table name (optional, uses default if not provided)
+            schema: Schema name (optional, uses default if not provided)
+            limit: Maximum number of results to return
+            score_threshold: Maximum distance threshold
+                results with distance > threshold will be filtered out)
+            metadata_filters: Dictionary of metadata filters to apply
+            metric: Distance metric to use ('COSINE', 'L2', 'IP')
+            embedding_column: Name of the embedding column
+            content_column: Name of the content column
+            metadata_column: Name of the metadata column
+            id_column: Name of the ID column
+            additional_columns: List of additional columns to include in results.
+        Returns:
+            List of SearchResult objects with content, metadata, score, collection_id, and record_id
+        """
+        if not self._connected:
+            await self.connection()
+
+        table = table or self.table_name
+        schema = schema or self.schema
+
+        if k and not limit:
+            limit = k
+        if not limit:
+            limit = 10
+
+        # Step 1: Ensure the ORM class exists
+        if not self.embedding_store:
+            self.embedding_store = self._define_collection_store(
+                table=table,
+                schema=schema,
+                dimension=self.dimension,
+                id_column=self._id_column,
+                embedding_column=embedding_column,
+                document_column=content_column,
+                metadata_column=metadata_column,
+                text_column=self._text_column,
+            )
+
+        # Step 2: Embed the query
+        query_embedding = self._embed_.embed_query(query)
+
+        # Get the actual column objects
+        content_col = getattr(self.embedding_store, content_column)
+        metadata_col = getattr(self.embedding_store, metadata_column)
+        embedding_col = getattr(self.embedding_store, embedding_column)
+        id_col = getattr(self.embedding_store, id_column)
+        collection_id_col = getattr(self.embedding_store, 'collection_id')
+
+        # Get the distance expression using the appropriate method
+        distance_expr = self.get_distance_strategy(
+            embedding_col,
+            query_embedding,
+            metric=metric
+        ).label("distance")
+        # self.logger.debug(f"Compiled distance expr → {distance_expr}")
+
+
+        # Build the select columns list
+        select_columns = [
+            id_col,
+            content_col,
+            metadata_col,
+            distance_expr,
+            collection_id_col,
+        ]
+
+        # Add additional columns dynamically using literal_column (no validation)
+        if additional_columns:
+            for col_name in additional_columns:
+                # Use literal_column to reference any column name without ORM validation
+                additional_col = literal_column(f'"{col_name}"').label(col_name)
+                select_columns.append(additional_col)
+                self.logger.debug(f"Added dynamic column: {col_name}")
+
+        # Step 5: Construct statement
+        stmt = (
+            select(*select_columns)
+            .select_from(self.embedding_store)  # Explicitly specify the table
+            .order_by(asc(distance_expr))
+        )
+
+        # Apply threshold filter if provided
+        if score_threshold is not None:
+            stmt = stmt.where(distance_expr <= score_threshold)
+
+        if limit:
+            stmt = stmt.limit(limit)
+
+        # 6) Apply any JSONB metadata filters
+        if metadata_filters:
+            for key, val in metadata_filters.items():
+                if isinstance(val, bool):
+                    # Handle boolean values properly in JSONB
+                    stmt = stmt.where(
+                        metadata_col[key].astext.cast(sqlalchemy.Boolean) == val
+                    )
+                else:
+                    stmt = stmt.where(
+                        metadata_col[key].astext == str(val)
+                    )
+
+        try:
+            # Execute query
+            async with self.session() as session:
+                result = await session.execute(stmt)
+                rows = result.fetchall()
+                # Create enhanced SearchResult objects
+                results = []
+                for row in rows:
+                    metadata = row[2]
+                    metadata['collection_id'] = row[4]
+                    # Add additional columns as a dictionary (starting from index 5)
+                    if additional_columns:
+                        for i, col_name in enumerate(additional_columns):
+                            metadata[col_name] = row[5 + i]
+                    # Create an enhanced SearchResult with additional fields
+                    search_result = SearchResult(
+                        id=row[0],
+                        content=row[1],        # content_col
+                        metadata=metadata,       # metadata_col
+                        score=row[3]           # distance
+                    )
+                    results.append(search_result)
+
+                return results
+        except Exception as e:
+            self.logger.error(f"Error during similarity search: {e}")
+            raise
+
+    def get_vector(self, metric_type: str = None, **kwargs):
+        raise NotImplementedError("This method is part of the old implementation.")
+
+    async def drop_collection(self, table: str, schema: str = 'public') -> None:
+        """
+        Drops the specified table in the given schema.
+
+        Args:
+            table: The name of the table to drop.
+            schema: The database schema where the table resides.
+        """
+        if not self._connected:
+            await self.connection()
+
+        full_table_name = f"{schema}.{table}"
+        async with self._connection.begin() as conn:
+            await conn.execute(text(f"DROP TABLE IF EXISTS {full_table_name}"))
+        self.logger.info(f"Table '{full_table_name}' dropped successfully.")
+
+    async def prepare_embedding_table(
+        self,
+        table: str,
+        schema: str = 'public',
+        conn: AsyncEngine = None,
+        id_column: str = 'id',
+        embedding_column: str = 'embedding',
+        document_column: str = 'document',
+        metadata_column: str = 'cmetadata',
+        dimension: int = 768,
+        colbert_dimension: int = 128,  # ColBERT token dimension
+        use_jsonb: bool = True,
+        drop_columns: bool = False,
+        create_all_indexes: bool = True,
+        **kwargs
+    ):
+        """
+        Prepare a Postgres Table as an embedding table in PostgreSQL with advanced features.
+        This method prepares a table with the following columns:
+        - id: unique identifier (String)
+        - embedding: the vector column (Vector(dimension) or JSONB)
+        - document: text column containing the document
+        - collection_id: UUID column for collection identification.
+        - metadata: JSONB column for metadata
+        - Additional columns based on the provided `columns` list
+        - Enhanced indexing strategies for efficient querying
+        - Support for multiple distance strategies (COSINE, L2, IP, etc.)
+        Args:
+        - tablename (str): Name of the table to create.
+        - embedding_column (str): Name of the column for storing embeddings.
+        - document_column (str): Name of the column for storing document text.
+        - metadata_column (str): Name of the column for storing metadata.
+        - dimension (int): Dimension of the embedding vector.
+        - id_column (str): Name of the column for storing unique identifiers.
+        - use_jsonb (bool): Whether to use JSONB for metadata storage.
+        - drop_columns (bool): Whether to drop existing columns.
+        - create_all_indexes (bool): Whether to create all distance strategies.
+    """
+        tablename = f"{schema}.{table}"
+        # Drop existing columns if requested
+        if drop_columns:
+            columns_to_drop = [
+                document_column, embedding_column, metadata_column,
+                'token_embeddings', 'num_tokens'
+            ]
+            for column in columns_to_drop:
+                await conn.execute(
+                    sqlalchemy.text(
+                        f'ALTER TABLE {tablename} DROP COLUMN IF EXISTS {column};'
+                    )
+                )
+        # Create metadata column as a jsonb field
+        if use_jsonb:
+            await conn.execute(
+                sqlalchemy.text(
+                    f'ALTER TABLE {tablename} ADD COLUMN IF NOT EXISTS {metadata_column} JSONB;'
+                )
+            )
+        # Use pgvector type for dense embeddings
+        await conn.execute(
+            sqlalchemy.text(
+                f'ALTER TABLE {tablename} ADD COLUMN IF NOT EXISTS {embedding_column} vector({dimension});'
+            )
+        )
+        # Add ColBERT columns for token-level embeddings
+        await conn.execute(
+            sqlalchemy.text(
+                f'ALTER TABLE {tablename} ADD COLUMN IF NOT EXISTS token_embeddings vector({colbert_dimension})[];'
+            )
+        )
+        await conn.execute(
+            sqlalchemy.text(
+                f'ALTER TABLE {tablename} ADD COLUMN IF NOT EXISTS num_tokens INTEGER;'
+            )
+        )
+        # Create the additional columns
+        for col_name, col_type in [
+            (document_column, 'TEXT'),
+            (id_column, 'varchar'),
+        ]:
+            await conn.execute(
+                sqlalchemy.text(
+                    f'ALTER TABLE {tablename} ADD COLUMN IF NOT EXISTS {col_name} {col_type};'
+                )
+            )
+        # Create the Collection Column:
+        await conn.execute(
+            sqlalchemy.text(
+                f"ALTER TABLE {tablename} ADD COLUMN IF NOT EXISTS collection_id UUID;"
+            )
+        )
+        await conn.execute(
+            sqlalchemy.text(
+                f"ALTER TABLE {tablename} ADD COLUMN IF NOT EXISTS collection_id UUID DEFAULT uuid_generate_v4();"
+            )
+        )
+        # Set the value on null values before declaring not null:
+        await conn.execute(
+            sqlalchemy.text(
+                f"UPDATE {tablename} SET collection_id = uuid_generate_v4() WHERE collection_id IS NULL;"
+            )
+        )
+        await conn.execute(
+            sqlalchemy.text(
+                f"ALTER TABLE {tablename} ALTER COLUMN collection_id SET NOT NULL;"
+            )
+        )
+        await conn.execute(
+            sqlalchemy.text(
+                f"CREATE UNIQUE INDEX IF NOT EXISTS idx_{table}_{schema}_collection_id ON {tablename} (collection_id);"
+            )
+        )
+        # ✅ CREATE COMPREHENSIVE INDEXES
+        if create_all_indexes:
+            await self._create_all_indexes(conn, tablename, embedding_column)
+        else:
+            # Create index only for current strategy
+            distance_strategy_ops = {
+                DistanceStrategy.COSINE: "vector_cosine_ops",
+                DistanceStrategy.EUCLIDEAN_DISTANCE: "vector_l2_ops",
+                DistanceStrategy.MAX_INNER_PRODUCT: "vector_ip_ops",
+                DistanceStrategy.DOT_PRODUCT: "vector_ip_ops"
+            }
+
+            ops = distance_strategy_ops.get(self.distance_strategy, "vector_cosine_ops")
+            strategy_name = str(self.distance_strategy).rsplit('.', maxsplit=1)[-1].lower()
+
+            await conn.execute(
+                sqlalchemy.text(
+                    f"CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_{strategy_name} "
+                    f"ON {tablename} USING ivfflat ({embedding_column} {ops});"
+                )
+            )
+            print(f"✅ Created {strategy_name.upper()} index")
+
+        # Create ColBERT-specific indexes
+        await self._create_colbert_indexes(conn, tablename)
+
+        # Create JSONB indexes for better performance
+        await self._create_jsonb_indexes(
+            conn,
+            tablename,
+            metadata_column,
+            id_column
+        )
+        # Ensure the table is ready for embedding operations
+        self.embedding_store = self._define_collection_store(
+            table=table,
+            schema=schema,
+            dimension=dimension,
+            id_column=id_column,
+            embedding_column=embedding_column,
+            document_column=self._document_column
+        )
+        return True
+
+    async def _create_all_indexes(self, conn, tablename: str, embedding_column: str):
+        """Create all standard vector indexes."""
+        print("🔧 Creating indexes for all distance strategies...")
+
+        # COSINE index (most common for text embeddings)
+        await conn.execute(
+            sqlalchemy.text(
+                f"CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_cosine "
+                f"ON {tablename} USING ivfflat ({embedding_column} vector_cosine_ops);"
+            )
+        )
+        print("✅ Created COSINE index")
+
+        # L2/Euclidean index
+        await conn.execute(
+            sqlalchemy.text(
+                f"CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_l2 "
+                f"ON {tablename} USING ivfflat ({embedding_column} vector_l2_ops);"
+            )
+        )
+        print("✅ Created L2 index")
+
+        # Inner Product index
+        try:
+            await conn.execute(
+                sqlalchemy.text(
+                    f"CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_ip "
+                    f"ON {tablename} USING ivfflat ({embedding_column} vector_ip_ops);"
+                )
+            )
+            print("✅ Created Inner Product index")
+        except Exception as e:
+            print(f"⚠️ Inner Product index creation failed: {e}")
+
+        # HNSW indexes for better performance (requires more memory)
+        try:
+            await conn.execute(
+                sqlalchemy.text(
+                    f"CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_hnsw_cosine "
+                    f"ON {tablename} USING hnsw ({embedding_column} vector_cosine_ops);"
+                )
+            )
+            print("✅ Created HNSW COSINE index")
+
+            await conn.execute(
+                sqlalchemy.text(
+                    f"""CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_hnsw_l2
+                        ON {tablename} USING hnsw ({embedding_column} vector_l2_ops) WITH (
+                        m = 16,  -- graph connectivity (higher → better recall, more memory)
+                        ef_construction = 200 -- controls indexing time vs. recall
+                    );"""
+                )
+            )
+            print("✅ Created HNSW EUCLIDEAN index")
+        except Exception as e:
+            print(f"⚠️ HNSW index creation failed (this is optional): {e}")
+
+    async def _create_colbert_indexes(self, conn, tablename: str):
+        """Create ColBERT-specific indexes for token embeddings."""
+        print("🔧 Creating ColBERT indexes...")
+
+        try:
+            # GIN index for array operations on token embeddings
+            await conn.execute(
+                sqlalchemy.text(
+                    f"CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_token_embeddings_gin "
+                    f"ON {tablename} USING gin(token_embeddings);"
+                )
+            )
+            print("✅ Created GIN index for token embeddings")
+
+            # Index on num_tokens for filtering
+            await conn.execute(
+                sqlalchemy.text(
+                    f"CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_num_tokens "
+                    f"ON {tablename} (num_tokens);"
+                )
+            )
+            print("✅ Created index for num_tokens")
+
+            # Partial index for non-null token embeddings
+            await conn.execute(
+                sqlalchemy.text(
+                    f"CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_has_tokens "
+                    f"ON {tablename} (id) WHERE token_embeddings IS NOT NULL;"
+                )
+            )
+            print("✅ Created partial index for documents with token embeddings")
+
+        except Exception as e:
+            print(f"⚠️ ColBERT index creation failed: {e}")
+
+    async def create_embedding_table(
+        self,
+        table: str,
+        columns: List[str],
+        schema: str = 'public',
+        embedding_column: str = 'embedding',
+        document_column: str = 'document',
+        metadata_column: str = 'cmetadata',
+        dimension: int = None,
+        id_column: str = 'id',
+        use_jsonb: bool = False,
+        drop_columns: bool = True,
+        create_all_indexes: bool = True,
+        **kwargs
+    ):
+        """
+        Create an embedding table in PostgreSQL with advanced features.
+        This method creates a table with the following columns:
+        - id: unique identifier (String)
+        - embedding: the vector column (Vector(dimension) or JSONB)
+        - document: text column containing the document
+        - cmetadata: JSONB column for metadata
+        - Additional columns based on the provided `columns` list
+        - Enhanced indexing strategies for efficient querying
+        - Support for multiple distance strategies (COSINE, L2, IP, etc.)
+        Args:
+        - table (str): Name of the table to create.
+        - columns (List[str]): List of column names to include in the table.
+        - schema (str): Database schema where the table will be created.
+        - embedding_column (str): Name of the column for storing embeddings.
+        - document_column (str): Name of the column for storing document text.
+        - metadata_column (str): Name of the column for storing metadata.
+        - dimension (int): Dimension of the embedding vector.
+        - id_column (str): Name of the column for storing unique identifiers.
+        - use_jsonb (bool): Whether to use JSONB for metadata storage.
+        - drop_columns (bool): Whether to drop existing columns.
+        - create_all_indexes (bool): Whether to create all distance strategies.
+
+        Enhanced embedding table creation with JSONB strategy for better semantic search.
+
+        This approach creates multiple document representations:
+        1. Primary search content (emphasizing store ID)
+        2. Location-based content
+        3. Structured metadata for filtering
+        4. Multiple embedding variations
+        """
+        tablename = f'{schema}.{table}'
+        cols = ', '.join(columns)
+        _qry = f'SELECT {cols} FROM {tablename};'
+        dimension = dimension or self.dimension
+
+        # Generate a sample embedding to determine its dimension
+        sample_vector = self._embed_.embedding.embed_query("sample text")
+        vector_dim = len(sample_vector)
+        self.logger.notice(
+            f"USING EMBED {self._embed_} with dimension {vector_dim}"
+        )
+
+        if vector_dim != dimension:
+            raise ValueError(
+                f"Expected embedding dimension {dimension}, but got {vector_dim}"
+            )
+
+        async with self._connection.begin() as conn:
+            result = await conn.execute(sqlalchemy.text(_qry))
+            rows = result.fetchall()
+
+            await self.prepare_embedding_table(
+                table=table,
+                schema=schema,
+                embedding_column=embedding_column,
+                document_column=document_column,
+                metadata_column=metadata_column,
+                dimension=dimension,
+                id_column=id_column,
+                use_jsonb=use_jsonb,
+                drop_columns=drop_columns,
+                create_all_indexes=create_all_indexes,
+                **kwargs
+            )
+
+            # Populate the embedding data
+            for i, row in enumerate(rows):
+                _id = getattr(row, id_column)
+                metadata = {col: getattr(row, col) for col in columns}
+                data = await self._create_metadata_structure(metadata, id_column, _id)
+
+                # Generate embedding
+                searchable_text = data['structured_search']
+                print(f"🔍 Row {i + 1}/{len(rows)} - {_id}")
+                print(f"   Text: {searchable_text[:100]}...")
+
+                vector = self._embed_.embedding.embed_query(searchable_text)
+                vector_str = "[" + ",".join(str(v) for v in vector) + "]"
+
+                await conn.execute(
+                    sqlalchemy.text(f"""
+                        UPDATE {tablename}
+                        SET {embedding_column} = :embeddings,
+                            {document_column} = :document,
+                            {metadata_column} = :metadata
+                        WHERE {id_column} = :id
+                    """),
+                    {
+                        "embeddings": vector_str,
+                        "document": searchable_text,
+                        "metadata": json_encoder(data),
+                        "id": _id
+                    }
+                )
+
+        print("✅ Updated Table embeddings with comprehensive indexes.")
+
+    def _create_natural_searchable_text(
+        self,
+        metadata: dict,
+        id_column: str,
+        record_id: str
+    ) -> str:
+        """
+        Create well-structured, natural language text with proper separation.
+
+        This creates clean, readable text that embedding models can understand better.
+        """
+        # Start with the ID in multiple formats for exact matching
+        text_parts = [
+            f"ID: {record_id}",
+            f"Identifier: {record_id}",
+            id_column + ": " + record_id
+        ]
+
+        # Process each field to create natural language descriptions
+        for key, value in metadata.items():
+            if value is None or value == '':
+                continue
+            clean_value = value.strip() if isinstance(value, str) else str(value)
+            text_parts.append(f"{key}: {clean_value}")
+            # Add the field in natural language format
+            clean_key = key.replace('_', ' ').title()
+            text_parts.append(f"{clean_key}={clean_value}")
+
+        # Join with spaces and clean up
+        searchable_text = ', '.join(text_parts) + '.'
+
+        return searchable_text
+
+    def _create_structured_search_text(self, metadata: dict, id_column: str, record_id: str) -> str:
+        """
+        Create a more structured but still readable search text.
+
+        This emphasizes key-value relationships while staying readable.
+        """
+        # ID section with emphasis
+        kv_sections = [
+            f"ID: {record_id}",
+            f"Identifier: {record_id}",
+            id_column + ": " + record_id
+        ]
+
+        # Key-value sections with clean separation
+        for key, value in metadata.items():
+            if value is None or value == '':
+                continue
+
+            # Clean key-value representation
+            clean_key = key.replace('_', ' ').title()
+            kv_sections.append(f"{clean_key}: {value}")
+            kv_sections.append(f"{key}: {value}")
+
+        # Combine with proper separation
+        return ' | '.join(kv_sections)
+
+    async def _create_metadata_structure(
+        self,
+        metadata: dict,
+        id_column: str,
+        _id: str
+    ):
+        """Create a structured metadata representation for the document."""
+        # Create a structured metadata representation
+        enhanced_metadata = {
+            "id": _id,
+            id_column: _id,
+            "_variants": [
+                _id,
+                _id.lower(),
+                _id.upper()
+            ]
+        }
+        for key, value in metadata.items():
+            enhanced_metadata[key] = value
+            # Create searchable variants for key fields
+            if value and isinstance(value, str):
+                variants = [value, value.lower(), value.upper()]
+                # Add variants without special characters
+                clean_value = ''.join(c for c in str(value) if c.isalnum() or c.isspace())
+                if clean_value != value:
+                    variants.append(clean_value)
+                enhanced_metadata[f"_{key}_variants"] = list(set(variants))
+        # create a full-text search field of searchable content
+        enhanced_metadata['searchable_content'] = self._create_natural_searchable_text(
+            metadata, id_column, _id
+        )
+
+        # Also create a structured search text that emphasizes important fields
+        enhanced_metadata['structured_search'] = self._create_structured_search_text(
+            metadata, id_column, _id
+        )
+
+        return enhanced_metadata
+
+    async def _create_jsonb_indexes(
+        self,
+        conn,
+        tablename: str,
+        metadata_col: str,
+        id_column: str
+    ):
+        """Create optimized JSONB indexes for better search performance."""
+
+        print("🔧 Creating JSONB indexes on Metadata for optimized search...")
+
+        # Index for ID searches
+        await conn.execute(
+            sqlalchemy.text(f"""
+                CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_{id_column}
+                ON {tablename} USING BTREE (({metadata_col}->>'{id_column}'));
+            """)
+        )
+        await conn.execute(
+            sqlalchemy.text(f"""
+                CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_id
+                ON {tablename} USING BTREE (({metadata_col}->>'id'));
+            """)
+        )
+
+        # GIN index for full-text search on searchable content
+        await conn.execute(
+            sqlalchemy.text(f"""
+                CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_fulltext
+                ON {tablename} USING GIN (to_tsvector('english', {metadata_col}->>'searchable_content'));
+            """)
+        )
+
+        # GIN index for JSONB structure searches
+        await conn.execute(
+            sqlalchemy.text(f"""
+                CREATE INDEX IF NOT EXISTS idx_{tablename.replace('.', '_')}_metadata_gin
+                ON {tablename} USING GIN ({metadata_col});
+            """)
+        )
+        print("✅ Created optimized JSONB indexes")
+
+    async def add_colbert_document(
+        self,
+        document_id: str,
+        content: str,
+        token_embeddings: np.ndarray,
+        table: str,
+        schema: str = 'public',
+        metadata: Optional[Dict[str, Any]] = None,
+        document_column: str = 'document',
+        metadata_column: str = 'cmetadata',
+        id_column: str = 'id',
+        **kwargs
+    ) -> None:
+        """
+        Add a document with ColBERT token embeddings to the specified table.
+
+        Args:
+            document_id: Unique identifier for the document
+            content: The document text content
+            token_embeddings: NumPy array of token embeddings (shape: [num_tokens, embedding_dim])
+            table: The name of the table
+            schema: The database schema where the table resides
+            metadata: Optional metadata dictionary
+            document_column: Name of the document content column
+            metadata_column: Name of the metadata column
+            id_column: Name of the ID column
+        """
+        if not self._connected:
+            await self.connection()
+
+        # Ensure the ORM table is initialized
+        if self.embedding_store is None:
+            self.embedding_store = self._define_collection_store(
+                table=table,
+                schema=schema,
+                dimension=self.dimension,
+                id_column=id_column,
+                document_column=document_column,
+                metadata_column=metadata_column,
+                text_column=self._text_column,
+            )
+
+        # Convert numpy array to list format for PostgreSQL
+        if isinstance(token_embeddings, np.ndarray):
+            token_embeddings_list = token_embeddings.tolist()
+        else:
+            token_embeddings_list = token_embeddings
+
+        num_tokens = len(token_embeddings_list)
+
+        # Prepare the insert/upsert data
+        values = {
+            id_column: document_id,
+            document_column: content,
+            'token_embeddings': token_embeddings_list,
+            'num_tokens': num_tokens,
+            metadata_column: metadata or {}
+        }
+
+        # Build insert statement with upsert capability
+        insert_stmt = insert(self.embedding_store).values(values)
+
+        # Create upsert statement (ON CONFLICT DO UPDATE)
+        upsert_stmt = insert_stmt.on_conflict_do_update(
+            index_elements=[id_column],
+            set_={
+                # document_column: insert_stmt.excluded.__getattr__(document_column),
+                document_column: getattr(insert_stmt.excluded, document_column),
+                'token_embeddings': insert_stmt.excluded.token_embeddings,
+                'num_tokens': insert_stmt.excluded.num_tokens,
+                metadata_column: getattr(insert_stmt.excluded, metadata_column),
+            }
+        )
+
+        try:
+            async with self._connection.begin() as conn:
+                await conn.execute(upsert_stmt)
+
+            self.logger.info(
+                f"Successfully added ColBERT document '{document_id}' with {num_tokens} tokens to '{schema}.{table}'"
+            )
+        except Exception as e:
+            self.logger.error(f"Error adding ColBERT document: {e}")
+            raise
+
+    async def colbert_search(
+        self,
+        query_tokens: np.ndarray,
+        table: str,
+        schema: str = 'public',
+        top_k: int = 10,
+        metadata_filters: Optional[Dict[str, Any]] = None,
+        min_tokens: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        id_column: str = 'id',
+        document_column: str = 'document',
+        metadata_column: str = 'cmetadata',
+        additional_columns: Optional[List[str]] = None
+    ) -> List[SearchResult]:
+        """
+        Perform ColBERT search with late interaction using MaxSim scoring.
+
+        Args:
+            query_tokens: NumPy array of query token embeddings (shape: [num_query_tokens, embedding_dim])
+            table: Table name
+            schema: Schema name
+            top_k: Number of results to return
+            metadata_filters: Optional metadata filters
+            min_tokens: Minimum number of tokens in documents to consider
+            max_tokens: Maximum number of tokens in documents to consider
+            id_column: Name of the ID column
+            document_column: Name of the document content column
+            metadata_column: Name of the metadata column
+            additional_columns: Additional columns to include in results
+
+        Returns:
+            List of SearchResult objects ordered by ColBERT score (descending)
+        """
+        if not self._connected:
+            await self.connection()
+
+        # Ensure the ORM table is initialized
+        if self.embedding_store is None:
+            self.embedding_store = self._define_collection_store(
+                table=table,
+                schema=schema,
+                dimension=self.dimension,
+                id_column=id_column,
+                document_column=document_column,
+                metadata_column=metadata_column,
+                text_column=self._text_column,
+            )
+
+        # Convert query tokens to list format
+        if isinstance(query_tokens, np.ndarray):
+            query_tokens_list = query_tokens.tolist()
+        else:
+            query_tokens_list = query_tokens
+
+        # Get column objects
+        id_col = getattr(self.embedding_store, id_column)
+        content_col = getattr(self.embedding_store, document_column)
+        metadata_col = getattr(self.embedding_store, metadata_column)
+        token_embeddings_col = getattr(self.embedding_store, 'token_embeddings')
+        num_tokens_col = getattr(self.embedding_store, 'num_tokens')
+        collection_id_col = getattr(self.embedding_store, 'collection_id')
+
+        # Build select columns
+        select_columns = [
+            id_col,
+            content_col,
+            metadata_col,
+            collection_id_col,
+            func.max_sim(token_embeddings_col, query_tokens_list).label('colbert_score')
+        ]
+
+        # Add additional columns dynamically
+        if additional_columns:
+            for col_name in additional_columns:
+                additional_col = literal_column(f'"{col_name}"').label(col_name)
+                select_columns.append(additional_col)
+
+        # Build the query
+        stmt = (
+            select(*select_columns)
+            .select_from(self.embedding_store)
+            .where(token_embeddings_col.isnot(None))  # Only documents with token embeddings
+            .order_by(func.max_sim(token_embeddings_col, query_tokens_list).desc())
+            .limit(top_k)
+        )
+
+        # Apply token count filters
+        if min_tokens is not None:
+            stmt = stmt.where(num_tokens_col >= min_tokens)
+        if max_tokens is not None:
+            stmt = stmt.where(num_tokens_col <= max_tokens)
+
+        # Apply metadata filters
+        if metadata_filters:
+            for key, value in metadata_filters.items():
+                stmt = stmt.where(metadata_col[key].astext == str(value))
+
+        try:
+            async with self._connection.connect() as conn:
+                result = await conn.execute(stmt)
+                rows = result.fetchall()
+
+                # Create SearchResult objects
+                results = []
+                for row in rows:
+                    # Enhance metadata with additional info
+                    metadata = dict(row[2]) if row[2] else {}
+                    metadata['collection_id'] = row[3]
+                    metadata['colbert_score'] = float(row[4])
+
+                    # Add additional columns to metadata
+                    if additional_columns:
+                        for i, col_name in enumerate(additional_columns):
+                            metadata[col_name] = row[5 + i]
+
+                    search_result = SearchResult(
+                        id=row[0],
+                        content=row[1],
+                        metadata=metadata,
+                        score=float(row[4])  # ColBERT score
+                    )
+                    results.append(search_result)
+
+                self.logger.info(
+                    f"ColBERT search returned {len(results)} results from {schema}.{table}"
+                )
+                return results
+
+        except Exception as e:
+            self.logger.error(f"Error during ColBERT search: {e}")
+            raise
+
+    async def hybrid_search(
+        self,
+        query: str,
+        query_tokens: Optional[np.ndarray] = None,
+        table: str = None,
+        schema: str = None,
+        top_k: int = 10,
+        dense_weight: float = 0.7,
+        colbert_weight: float = 0.3,
+        metadata_filters: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> List[SearchResult]:
+        """
+        Perform hybrid search combining dense embeddings and ColBERT token matching.
+
+        Args:
+            query: Text query
+            query_tokens: Optional pre-computed query token embeddings
+            table: Table name
+            schema: Schema name
+            top_k: Number of final results
+            dense_weight: Weight for dense similarity scores (0-1)
+            colbert_weight: Weight for ColBERT scores (0-1)
+            metadata_filters: Metadata filters to apply
+
+        Returns:
+            List of SearchResult objects with combined scores
+        """
+        if not self._connected:
+            await self.connection()
+
+        table = table or self.table_name
+        schema = schema or self.schema
+
+        # Fetch more candidates for reranking
+        candidate_count = min(top_k * 3, 100)
+
+        # Get dense similarity results
+        dense_results = await self.similarity_search(
+            query=query,
+            table=table,
+            schema=schema,
+            limit=candidate_count,
+            metadata_filters=metadata_filters,
+            **kwargs
+        )
+
+        # Get ColBERT results if token embeddings provided
+        colbert_results = []
+        if query_tokens is not None:
+            colbert_results = await self.colbert_search(
+                query_tokens=query_tokens,
+                table=table,
+                schema=schema,
+                top_k=candidate_count,
+                metadata_filters=metadata_filters
+            )
+
+        # Combine and rerank results
+        combined_results = self._combine_search_results(
+            dense_results=dense_results,
+            colbert_results=colbert_results,
+            dense_weight=dense_weight,
+            colbert_weight=colbert_weight
+        )
+
+        # Return top-k results
+        return combined_results[:top_k]
+
+    def _combine_search_results(
+        self,
+        dense_results: List[SearchResult],
+        colbert_results: List[SearchResult],
+        dense_weight: float,
+        colbert_weight: float
+    ) -> List[SearchResult]:
+        """
+        Combine and rerank results from dense and ColBERT searches.
+        """
+        # Create lookup dictionaries
+        dense_lookup = {result.id: result for result in dense_results}
+        colbert_lookup = {result.id: result for result in colbert_results}
+
+        # Get all unique document IDs
+        all_ids = set(dense_lookup.keys()) | set(colbert_lookup.keys())
+
+        # Normalize scores to 0-1 range
+        if dense_results:
+            dense_scores = [r.score for r in dense_results]
+            dense_min, dense_max = min(dense_scores), max(dense_scores)
+            dense_range = dense_max - dense_min if dense_max != dense_min else 1
+        else:
+            dense_min, dense_range = 0, 1
+
+        if colbert_results:
+            colbert_scores = [r.score for r in colbert_results]
+            colbert_min, colbert_max = min(colbert_scores), max(colbert_scores)
+            colbert_range = colbert_max - colbert_min if colbert_max != colbert_min else 1
+        else:
+            colbert_min, colbert_range = 0, 1
+
+        # Combine results
+        combined_results = []
+        for doc_id in all_ids:
+            dense_result = dense_lookup.get(doc_id)
+            colbert_result = colbert_lookup.get(doc_id)
+
+            # Normalize scores
+            dense_score_norm = 0
+            if dense_result:
+                dense_score_norm = (dense_result.score - dense_min) / dense_range
+
+            colbert_score_norm = 0
+            if colbert_result:
+                colbert_score_norm = (colbert_result.score - colbert_min) / colbert_range
+
+            # Calculate combined score
+            combined_score = (
+                dense_weight * dense_score_norm +
+                colbert_weight * colbert_score_norm
+            )
+
+            # Use the result with more complete information
+            primary_result = dense_result or colbert_result
+
+            # Create combined result
+            combined_result = SearchResult(
+                id=primary_result.id,
+                content=primary_result.content,
+                metadata={
+                    **primary_result.metadata,
+                    'dense_score': dense_result.score if dense_result else 0,
+                    'colbert_score': colbert_result.score if colbert_result else 0,
+                    'combined_score': combined_score
+                },
+                score=combined_score
+            )
+            combined_results.append(combined_result)
+
+        # Sort by combined score (descending)
+        combined_results.sort(key=lambda x: x.score, reverse=True)
+
+        return combined_results
+
+    async def mmr_search(
+        self,
+        query: str,
+        table: str = None,
+        schema: 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 = 'cmetadata',
+        id_column: str = 'id',
+        additional_columns: Optional[List[str]] = None
+    ) -> List[SearchResult]:
+        """
+        Perform Maximal Marginal Relevance (MMR) search to balance relevance and diversity.
+
+        MMR helps avoid redundant results by selecting documents that are relevant to the query
+        but diverse from each other.
+
+        Args:
+            query: The search query text
+            table: Table name (optional, uses default if not provided)
+            schema: Schema name (optional, uses default if not provided)
+            k: Number of final results to return
+            fetch_k: Number of candidate documents to fetch (default: k * 3)
+            lambda_mult: MMR diversity parameter (0-1):
+                        - 1.0 = pure relevance (no diversity)
+                        - 0.0 = pure diversity (no relevance)
+                        - 0.5 = balanced (default)
+            metadata_filters: Dictionary of metadata filters to apply
+            score_threshold: Maximum distance threshold for initial candidates
+            metric: Distance metric to use ('COSINE', 'L2', 'IP')
+            embedding_column: Name of the embedding column
+            content_column: Name of the content column
+            metadata_column: Name of the metadata column
+            id_column: Name of the ID column
+            additional_columns: Additional columns to include in results
+
+        Returns:
+            List of SearchResult objects selected via MMR algorithm
+        """
+        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,
+            schema=schema,
+            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,
+            additional_columns=additional_columns
+        )
+
+        if len(candidates) <= k:
+            # If we have fewer candidates than requested results, return all
+            return candidates
+
+        # Step 2: Get embeddings for MMR computation
+        # We need to fetch the actual embedding vectors for similarity computation
+        candidate_embeddings = await self._fetch_embeddings_for_mmr(
+            candidate_ids=[result.id for result in candidates],
+            table=table,
+            schema=schema,
+            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 "
+            f"(λ={lambda_mult})"
+        )
+
+        return selected_results
+
+    async def _fetch_embeddings_for_mmr(
+        self,
+        candidate_ids: List[str],
+        table: str,
+        schema: str,
+        embedding_column: str,
+        id_column: str
+    ) -> Dict[str, np.ndarray]:
+        """
+        Fetch embedding vectors for candidate documents.
+
+        Args:
+            candidate_ids: List of document IDs to fetch embeddings for
+            table: Table name
+            schema: Schema name
+            embedding_column: Name of the embedding column
+            id_column: Name of the ID column
+
+        Returns:
+            Dictionary mapping document ID to embedding vector
+        """
+        if not self.embedding_store:
+            self.embedding_store = self._define_collection_store(
+                table=table,
+                schema=schema,
+                dimension=self.dimension,
+                id_column=self._id_column,
+                embedding_column=embedding_column,
+                document_column=self._document_column,
+                metadata_column='cmetadata',
+                text_column=self._text_column,
+            )
+
+        # Get column objects
+        id_col = getattr(self.embedding_store, id_column)
+        embedding_col = getattr(self.embedding_store, embedding_column)
+
+        # Build query to fetch embeddings
+        stmt = (
+            select(id_col, embedding_col)
+            .select_from(self.embedding_store)
+            .where(id_col.in_(candidate_ids))
+        )
+
+        embeddings_dict = {}
+        async with self.session() as session:
+            result = await session.execute(stmt)
+            rows = result.fetchall()
+
+            for row in rows:
+                doc_id = row[0]
+                embedding = row[1]
+
+                # Convert to numpy array if needed
+                if isinstance(embedding, (list, tuple)):
+                    embeddings_dict[doc_id] = np.array(embedding, dtype=np.float32)
+                elif hasattr(embedding, '__array__'):
+                    embeddings_dict[doc_id] = np.array(embedding, dtype=np.float32)
+                else:
+                    # Handle pgvector Vector type
+                    embeddings_dict[doc_id] = np.array(embedding, dtype=np.float32)
+
+        return embeddings_dict
+
+    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.
+
+        Args:
+            query_embedding: Query embedding vector
+            candidates: List of candidate SearchResult objects
+            candidate_embeddings: Dictionary mapping doc ID to embedding vector
+            k: Number of results to select
+            lambda_mult: MMR diversity parameter (0-1)
+            metric: Distance metric to use
+
+        Returns:
+            List of selected SearchResult objects ordered by MMR score
+        """
+        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
+                # Convert distance to similarity (lower distance = higher similarity)
+                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 based on the specified metric.
+
+        Args:
+            embedding1: First embedding vector (numpy array or list)
+            embedding2: Second embedding vector (numpy array or list)
+            metric: Distance metric ('COSINE', 'L2', 'IP', etc.)
+
+        Returns:
+            Similarity score (higher = more similar)
+        """
+        # 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,
+                'IP': DistanceStrategy.MAX_INNER_PRODUCT,
+                'DOT': DistanceStrategy.DOT_PRODUCT,
+                'DOT_PRODUCT': DistanceStrategy.DOT_PRODUCT,
+                'MAX_INNER_PRODUCT': DistanceStrategy.MAX_INNER_PRODUCT
+            }
+            strategy = metric_mapping.get(metric.upper(), DistanceStrategy.COSINE)
+        else:
+            strategy = metric
+
+        if strategy == DistanceStrategy.COSINE:
+            # Cosine similarity (returns similarity, not distance)
+            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 in [DistanceStrategy.MAX_INNER_PRODUCT, DistanceStrategy.DOT_PRODUCT]:
+            # Dot product (inner 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,
+        schema: Optional[str] = None,
+        metadata_column: Optional[str] = None,
+        **kwargs
+    ) -> int:
+        """
+        Delete documents from the vector store based on metadata field values.
+
+        Args:
+            documents: List of documents whose metadata values will be used for deletion.
+                    If provided, the pk field values will be extracted from these documents.
+            pk: The metadata field name to use for deletion (default: 'source_type')
+            values: Specific values to delete. Can be a single string or list of strings.
+                If provided, this takes precedence over extracting from documents.
+            table: Override table name
+            schema: Override schema name
+            metadata_column: Override metadata column name
+
+        Returns:
+            int: Number of documents deleted
+
+        Examples:
+            # Delete all documents with source_type 'papers'
+            deleted_count = await store.delete_documents(values='papers')
+
+            # Delete documents with multiple source types
+            deleted_count = await store.delete_documents(values=['papers', 'reports'])
+
+            # Delete based on documents' metadata
+            docs_to_delete = [Document(page_content="test", metadata={"source_type": "papers"})]
+            deleted_count = await store.delete_documents(documents=docs_to_delete)
+
+            # Delete by different metadata field
+            deleted_count = await store.delete_documents(pk='category', values='obsolete')
+        """
+        if not self._connected:
+            await self.connection()
+
+        # Use defaults from instance if not provided
+        table = table or self.table_name
+        schema = schema or self.schema
+        metadata_column = metadata_column or self._document_column or 'cmetadata'
+
+        # Extract values to delete
+        delete_values = []
+
+        if values is not None:
+            # Use provided values
+            if isinstance(values, str):
+                delete_values = [values]
+            else:
+                delete_values = list(values)
+        elif documents:
+            # Extract values from documents metadata
+            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
+
+        # Construct full table name
+        full_table_name = f"{schema}.{table}" if schema != 'public' else table
+
+        deleted_count = 0
+
+        try:
+            async with self.session() as session:
+                for value in delete_values:
+                    # Use JSONB operator to find matching metadata
+                    delete_query = text(f"""
+                        DELETE FROM {full_table_name}
+                        WHERE {metadata_column}->>:pk = :value
+                    """)
+
+                    result = await session.execute(
+                        delete_query,
+                        {"pk": pk, "value": str(value)}
+                    )
+
+                    rows_deleted = result.rowcount
+                    deleted_count += rows_deleted
+
+                    self.logger.info(
+                        f"Deleted {rows_deleted} documents with {pk}='{value}' from {full_table_name}"
+                    )
+
+            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,
+        schema: Optional[str] = None,
+        metadata_column: Optional[str] = None,
+        **kwargs
+    ) -> int:
+        """
+        Delete documents based on multiple metadata field conditions.
+
+        Args:
+            filter_dict: Dictionary of field_name: value(s) pairs for deletion criteria
+            table: Override table name
+            schema: Override schema name
+            metadata_column: Override metadata column name
+
+        Returns:
+            int: Number of documents deleted
+
+        Example:
+            # Delete documents with source_type='papers' AND category='research'
+            deleted = await store.delete_documents_by_filter({
+                'source_type': 'papers',
+                'category': 'research'
+            })
+
+            # Delete documents with source_type in ['papers', 'reports']
+            deleted = await store.delete_documents_by_filter({
+                'source_type': ['papers', 'reports']
+            })
+        """
+        if not self._connected:
+            await self.connection()
+
+        if not filter_dict:
+            raise ValueError("filter_dict cannot be empty")
+
+        # Use defaults from instance if not provided
+        table = table or self.table_name
+        schema = schema or self.schema
+        metadata_column = metadata_column or self._document_column or 'cmetadata'
+
+        # Construct full table name
+        full_table_name = f"{schema}.{table}" if schema != 'public' else table
+
+        # Build WHERE conditions
+        where_conditions = []
+        params = {}
+
+        for field, values in filter_dict.items():
+            if isinstance(values, (list, tuple)):
+                # Handle multiple values with IN operator
+                placeholders = []
+                for i, value in enumerate(values):
+                    param_name = f"{field}_{i}"
+                    placeholders.append(f":{param_name}")
+                    params[param_name] = str(value)
+
+                condition = f"{metadata_column}->>'{field}' IN ({', '.join(placeholders)})"
+                where_conditions.append(condition)
+            else:
+                # Handle single value
+                param_name = f"{field}_single"
+                where_conditions.append(f"{metadata_column}->>'{field}' = :{param_name}")
+                params[param_name] = str(values)
+
+        # Combine conditions with AND
+        where_clause = " AND ".join(where_conditions)
+
+        delete_query = text(f"""
+            DELETE FROM {full_table_name}
+            WHERE {where_clause}
+        """)
+
+        try:
+            async with self.session() as session:
+                result = await session.execute(delete_query, params)
+                deleted_count = result.rowcount
+
+                self.logger.info(
+                    f"Deleted {deleted_count} documents from {full_table_name} "
+                    f"with filter: {filter_dict}"
+                )
+
+                return deleted_count
+
+        except Exception as e:
+            self.logger.error(f"Error deleting documents by filter: {e}")
+            raise RuntimeError(f"Failed to delete documents by filter: {e}") from e
+
+    async def delete_all_documents(
+        self,
+        table: Optional[str] = None,
+        schema: Optional[str] = None,
+        confirm: bool = False,
+        **kwargs
+    ) -> int:
+        """
+        Delete ALL documents from the vector store table.
+
+        WARNING: This will delete all data in the table!
+
+        Args:
+            table: Override table name
+            schema: Override schema name
+            confirm: Must be set to True to proceed with deletion
+
+        Returns:
+            int: Number of documents deleted
+        """
+        if not confirm:
+            raise ValueError(
+                "This operation will delete ALL documents. "
+                "Set confirm=True to proceed."
+            )
+
+        if not self._connected:
+            await self.connection()
+
+        # Use defaults from instance if not provided
+        table = table or self.table_name
+        schema = schema or self.schema
+
+        # Construct full table name
+        full_table_name = f"{schema}.{table}" if schema != 'public' else table
+
+        try:
+            async with self.session() as session:
+                # First count existing documents
+                count_query = text(f"SELECT COUNT(*) FROM {full_table_name}")
+                count_result = await session.execute(count_query)
+                total_docs = count_result.scalar()
+
+                if total_docs == 0:
+                    self.logger.info(f"No documents to delete from {full_table_name}")
+                    return 0
+
+                # Delete all documents
+                delete_query = text(f"DELETE FROM {full_table_name}")
+                result = await session.execute(delete_query)
+                deleted_count = result.rowcount
+
+                self.logger.warning(
+                    f"DELETED ALL {deleted_count} documents from {full_table_name}"
+                )
+
+                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 delete_documents_by_ids(
+        self,
+        document_ids: List[str],
+        table: Optional[str] = None,
+        schema: Optional[str] = None,
+        id_column: Optional[str] = None,
+        **kwargs
+    ) -> int:
+        """
+        Delete documents by their IDs.
+
+        Args:
+            document_ids: List of document IDs to delete
+            table: Override table name
+            schema: Override schema name
+            id_column: Override ID column name
+
+        Returns:
+            int: Number of documents deleted
+
+        Example:
+            deleted_count = await store.delete_documents_by_ids([
+                "doc_1", "doc_2", "doc_3"
+            ])
+        """
+        if not self._connected:
+            await self.connection()
+
+        if not document_ids:
+            self.logger.warning("No document IDs provided for deletion")
+            return 0
+
+        # Use defaults from instance if not provided
+        table = table or self.table_name
+        schema = schema or self.schema
+        id_column = id_column or self._id_column
+
+        # Construct full table name
+        full_table_name = f"{schema}.{table}" if schema != 'public' else table
+
+        # Build parameterized query for multiple IDs
+        placeholders = []
+        params = {}
+        for i, doc_id in enumerate(document_ids):
+            param_name = f"id_{i}"
+            placeholders.append(f":{param_name}")
+            params[param_name] = str(doc_id)
+
+        delete_query = text(f"""
+            DELETE FROM {full_table_name}
+            WHERE {id_column} IN ({', '.join(placeholders)})
+        """)
+
+        try:
+            async with self.session() as session:
+                result = await session.execute(delete_query, params)
+                deleted_count = result.rowcount
+
+                self.logger.info(
+                    f"Deleted {deleted_count} documents by IDs from {full_table_name}"
+                )
+
+                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
+
+    # Additional utility method for safer deletions
+    async def count_documents_by_filter(
+        self,
+        filter_dict: Dict[str, Union[str, List[str]]],
+        table: Optional[str] = None,
+        schema: Optional[str] = None,
+        metadata_column: Optional[str] = None,
+        **kwargs
+    ) -> int:
+        """
+        Count documents that would be affected by a filter (useful before deletion).
+
+        Args:
+            filter_dict: Dictionary of field_name: value(s) pairs for criteria
+            table: Override table name
+            schema: Override schema name
+            metadata_column: Override metadata column name
+
+        Returns:
+            int: Number of documents matching the filter
+        """
+        if not self._connected:
+            await self.connection()
+
+        if not filter_dict:
+            return 0
+
+        # Use defaults from instance if not provided
+        table = table or self.table_name
+        schema = schema or self.schema
+        metadata_column = metadata_column or self._document_column or 'cmetadata'
+
+        # Construct full table name
+        full_table_name = f"{schema}.{table}" if schema != 'public' else table
+
+        # Build WHERE conditions (same logic as delete_documents_by_filter)
+        where_conditions = []
+        params = {}
+
+        for field, values in filter_dict.items():
+            if isinstance(values, (list, tuple)):
+                placeholders = []
+                for i, value in enumerate(values):
+                    param_name = f"{field}_{i}"
+                    placeholders.append(f":{param_name}")
+                    params[param_name] = str(value)
+
+                condition = f"{metadata_column}->>'{field}' IN ({', '.join(placeholders)})"
+                where_conditions.append(condition)
+            else:
+                param_name = f"{field}_single"
+                where_conditions.append(f"{metadata_column}->>'{field}' = :{param_name}")
+                params[param_name] = str(values)
+
+        where_clause = " AND ".join(where_conditions)
+        count_query = text(f"""
+            SELECT COUNT(*) FROM {full_table_name}
+            WHERE {where_clause}
+        """)
+
+        try:
+            async with self.session() as session:
+                result = await session.execute(count_query, params)
+                count = result.scalar()
+
+                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 from_documents(
+        self,
+        documents: List[Document],
+        table: str = None,
+        schema: str = None,
+        embedding_column: str = 'embedding',
+        content_column: str = 'document',
+        metadata_column: str = 'cmetadata',
+        chunk_size: int = 8192,
+        chunk_overlap: int = 200,
+        store_full_document: bool = True,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Add documents using late chunking strategy.
+
+        Args:
+            documents: List of Document objects to process
+            table: Table name
+            schema: Schema name
+            embedding_column: Name of embedding column
+            content_column: Name of content column
+            metadata_column: Name of metadata column
+            chunk_size: Maximum size of each chunk
+            chunk_overlap: Overlap between chunks
+            store_full_document: Whether to store full document alongside chunks
+
+        Returns:
+            Dictionary with processing statistics
+        """
+        if not self._connected:
+            await self.connection()
+
+        table = table or self.table_name
+        schema = schema or self.schema
+
+
+        # Initialize late chunking processor
+        chunking_processor = LateChunkingProcessor(
+            vector_store=self,
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap
+        )
+
+        # Ensure embedding store is initialized
+        if self.embedding_store is None:
+            self.embedding_store = self._define_collection_store(
+                table=table,
+                schema=schema,
+                dimension=self.dimension,
+                id_column=self._id_column,
+                embedding_column=embedding_column,
+                document_column=content_column,
+                metadata_column=metadata_column,
+                text_column=self._text_column,
+            )
+
+        all_inserts = []
+        stats = {
+            'documents_processed': 0,
+            'chunks_created': 0,
+            'full_documents_stored': 0
+        }
+        for doc_idx, document in enumerate(documents):
+            document_id = f"doc_{doc_idx:06d}_{uuid.uuid4().hex[:8]}"
+
+            # Process document with late chunking
+            full_embedding, chunk_infos = await chunking_processor.process_document_late_chunking(
+                document_text=document.page_content,
+                document_id=document_id,
+                metadata=document.metadata
+            )
+            # Store full document if requested
+            if store_full_document:
+                full_doc_metadata = {
+                    **(document.metadata or {}),
+                    'document_id': document_id,
+                    'is_full_document': True,
+                    'total_chunks': len(chunk_infos),
+                    'document_type': 'parent',
+                    'chunking_strategy': 'late_chunking'
+                }
+
+                all_inserts.append({
+                    self._id_column: document_id,
+                    embedding_column: full_embedding.tolist(),
+                    content_column: document.page_content,
+                    metadata_column: full_doc_metadata
+                })
+                stats['full_documents_stored'] += 1
+
+            # Store all chunks
+            for chunk_info in chunk_infos:
+                embed = chunk_info.chunk_embedding if isinstance(chunk_info.chunk_embedding, list) else chunk_info.chunk_embedding.tolist()
+                all_inserts.append({
+                    self._id_column: chunk_info.chunk_id,
+                    embedding_column: embed,
+                    content_column: chunk_info.chunk_text,
+                    metadata_column: chunk_info.metadata
+                })
+                stats['chunks_created'] += 1
+
+            stats['documents_processed'] += 1
+
+        # Bulk insert all data
+        if all_inserts:
+            insert_stmt = insert(self.embedding_store)
+
+            try:
+                async with self.session() as session:
+                    await session.execute(insert_stmt, all_inserts)
+
+                self.logger.info(
+                    f"✅ Late chunking complete: {stats['documents_processed']} documents → "
+                    f"{stats['chunks_created']} chunks + {stats['full_documents_stored']} full docs"
+                )
+
+            except Exception as e:
+                self.logger.error(f"Error in late chunking insert: {e}")
+                raise
+
+        return stats
+
+    async def document_search(
+        self,
+        query: str,
+        table: str = None,
+        schema: str = None,
+        limit: int = 10,
+        search_chunks: bool = True,
+        search_full_docs: bool = False,
+        rerank_with_context: bool = True,
+        context_window: int = 2,
+        **kwargs
+    ) -> List[SearchResult]:
+        """
+        Search with late chunking context awareness.
+
+        Args:
+            query: Search query
+            table: Table name
+            schema: Schema name
+            limit: Number of results
+            search_chunks: Whether to search chunk-level embeddings
+            search_full_docs: Whether to search full document embeddings
+            rerank_with_context: Whether to rerank using surrounding chunks
+            context_window: Number of adjacent chunks to include for context
+
+        Returns:
+            List of SearchResult objects with enhanced context
+        """
+        results = []
+
+        # Search chunks if requested
+        if search_chunks:
+            chunk_filters = {'is_chunk': True}
+            chunk_results = await self.similarity_search(
+                query=query,
+                table=table,
+                schema=schema,
+                limit=limit * 2,  # Get more candidates for reranking
+                metadata_filters=chunk_filters,
+                **kwargs
+            )
+            results.extend(chunk_results)
+
+        # Search full documents if requested
+        if search_full_docs:
+            doc_filters = {'is_full_document': True}
+            doc_results = await self.similarity_search(
+                query=query,
+                table=table,
+                schema=schema,
+                limit=limit,
+                metadata_filters=doc_filters,
+                **kwargs
+            )
+            results.extend(doc_results)
+
+        # Rerank with context if requested
+        if rerank_with_context and search_chunks:
+            results = await self._rerank_with_chunk_context(
+                results=results,
+                query=query,
+                table=table,
+                schema=schema,
+                context_window=context_window
+            )
+
+        # Sort by score and limit
+        results.sort(key=lambda x: x.score)
+        return results[:limit]
+
+    async def _rerank_with_chunk_context(
+        self,
+        results: List[SearchResult],
+        query: str,
+        table: str,
+        schema: str,
+        context_window: int = 2
+    ) -> List[SearchResult]:
+        """
+        Rerank results by including surrounding chunk context.
+        """
+        enhanced_results = []
+
+        for result in results:
+            if not result.metadata.get('is_chunk'):
+                enhanced_results.append(result)
+                continue
+
+            # Get surrounding chunks for context
+            parent_id = result.metadata.get('parent_document_id')
+            chunk_index = result.metadata.get('chunk_index', 0)
+
+            if parent_id:
+                try:
+                    # Find adjacent chunks
+                    context_chunks = await self._get_adjacent_chunks(
+                        parent_id=parent_id,
+                        center_chunk_index=chunk_index,
+                        window_size=context_window,
+                        table=table,
+                        schema=schema
+                    )
+
+                    # Combine text with context
+                    combined_text = self._combine_chunk_context(result, context_chunks)
+
+                    # Re-score with context - ensure embeddings are consistent
+                    context_embedding = self._embed_.embed_query(combined_text)
+                    query_embedding = self._embed_.embed_query(query)
+
+                    # Ensure both embeddings are numpy arrays
+                    if isinstance(context_embedding, list):
+                        context_embedding = np.array(context_embedding, dtype=np.float32)
+                    if isinstance(query_embedding, list):
+                        query_embedding = np.array(query_embedding, dtype=np.float32)
+
+                    # Calculate new similarity score
+                    context_score = self._compute_similarity(
+                        query_embedding, context_embedding, self.distance_strategy
+                    )
+
+                    # Create enhanced result
+                    enhanced_metadata = dict(result.metadata)
+                    enhanced_metadata['context_score'] = context_score
+                    enhanced_metadata['has_context'] = True
+                    enhanced_metadata['context_chunks'] = len(context_chunks)
+
+                    enhanced_result = SearchResult(
+                        id=result.id,
+                        content=combined_text,
+                        metadata=enhanced_metadata,
+                        score=context_score
+                    )
+
+                    enhanced_results.append(enhanced_result)
+
+                except Exception as e:
+                    self.logger.warning(f"Error reranking chunk {result.id}: {e}")
+                    # Fall back to original result if reranking fails
+                    enhanced_results.append(result)
+            else:
+                enhanced_results.append(result)
+
+        return enhanced_results
+
+    async def _get_adjacent_chunks(
+        self,
+        parent_id: str,
+        center_chunk_index: int,
+        window_size: int,
+        table: str,
+        schema: str
+    ) -> List[SearchResult]:
+        """Get adjacent chunks for context."""
+        # Calculate chunk index range
+        start_idx = max(0, center_chunk_index - window_size)
+        end_idx = center_chunk_index + window_size + 1
+
+        # Search for chunks in the range
+        chunk_filters = {
+            'parent_document_id': parent_id,
+            'is_chunk': True
+        }
+
+        # Get all chunks from parent document
+        all_chunks = await self.similarity_search(
+            query="dummy",
+            table=table,
+            schema=schema,
+            limit=1000,  # High limit to get all chunks
+            metadata_filters=chunk_filters
+        )
+
+        # Filter to adjacent chunks
+        adjacent_chunks = [
+            chunk for chunk in all_chunks
+            if start_idx <= chunk.metadata.get('chunk_index', 0) < end_idx
+        ]
+
+        # Sort by chunk index
+        adjacent_chunks.sort(key=lambda x: x.metadata.get('chunk_index', 0))
+
+        return adjacent_chunks
+
+    def _combine_chunk_context(
+        self,
+        center_result: SearchResult,
+        context_chunks: List[SearchResult]
+    ) -> str:
+        """Combine center chunk with surrounding context."""
+        # Sort context chunks by index
+        context_chunks.sort(key=lambda x: x.metadata.get('chunk_index', 0))
+
+        # Combine text
+        combined_parts = []
+        center_idx = center_result.metadata.get('chunk_index', 0)
+
+        for chunk in context_chunks:
+            chunk_idx = chunk.metadata.get('chunk_index', 0)
+            if chunk_idx == center_idx:
+                # Mark the main chunk
+                combined_parts.append(f"[MAIN] {chunk.content} [/MAIN]")
+            else:
+                combined_parts.append(chunk.content)
+
+        return " ... ".join(combined_parts)
+
+    async def collection_exists(self, table: str, schema: str = 'public') -> bool:
+        """
+        Check if a collection (table) exists in the database.
+
+        Args:
+            table: Name of the table to check
+            schema: Schema name (default: 'public')
+
+        Returns:
+            bool: True if the collection exists, False otherwise
+        """
+        if not self._connected:
+            await self.connection()
+
+        async with self.session() as session:
+            query = text(f"""
+                SELECT EXISTS (
+                    SELECT 1 FROM information_schema.tables
+                    WHERE table_schema = :schema AND table_name = :table
+                )
+            """)
+            result = await session.execute(query, {"schema": schema, "table": table})
+            return result.scalar()
+        return False
+
+    async def delete_collection(
+        self,
+        table: str,
+        schema: str = 'public'
+    ) -> None:
+        """
+        Delete a collection (table) from the database.
+
+        Args:
+            table: Name of the table to delete
+            schema: Schema name (default: 'public')
+
+        Raises:
+            RuntimeError: If the collection does not exist or deletion fails
+        """
+        if not self._connected:
+            await self.connection()
+
+        if not await self.collection_exists(table, schema):
+            raise RuntimeError(
+                f"Collection {schema}.{table} does not exist"
+            )
+
+        async with self.session() as session:
+            query = text(
+                f"DROP TABLE IF EXISTS {schema}.{table} CASCADE"
+            )
+            await session.execute(query)
+            self.logger.info(
+                f"Collection {schema}.{table} deleted successfully"
+            )
+
+    async def create_collection(
+        self,
+        table: str,
+        schema: str = 'public',
+        dimension: int = 768,
+        index_type: str = "COSINE",
+        metric_type: str = 'L2',
+        id_column: Optional[str] = None,
+        **kwargs
+    ) -> None:
+        """
+        Create a new collection (table) in the database.
+
+        Args:
+            table: Name of the table to create
+            schema: Schema name (default: 'public')
+            dimension: Embedding dimension (default: 768)
+            index_type: Type of index to create (default: "COSINE")
+            metric_type: Distance metric type (default: 'L2')
+            id_column: Name of the ID column (default: 'id')
+            embedding_column: Name of the embedding column (default: 'embedding')
+            document_column: Name of the document content column (default: 'document')
+            metadata_column: Name of the metadata column (default: 'cmetadata')
+
+        Raises:
+            RuntimeError: If collection creation fails
+        """
+        if not self._connected:
+            await self.connection()
+
+        # Construct full table name
+        full_table_name = f"{schema}.{table}" if schema != 'public' else table
+        self._metric_type: str = metric_type.upper()
+        self._index_type: str = index_type.upper()
+        try:
+            async with self.session() as session:
+                # Check if collection already exists
+                if await self.collection_exists(table, schema):
+                    self.logger.info(
+                        f"Collection {schema}.{table} already exists"
+                    )
+                else:
+                    id_column = id_column or self._id_column or 'id'
+                    # Create the collection:
+                    self.logger.info(f"Creating collection {schema}.{table}...")
+                    create_query = text(f"""
+                        CREATE TABLE {full_table_name} (
+                            {id_column} TEXT PRIMARY KEY
+                        )
+                    """)
+                    await session.execute(create_query)
+                    self.logger.info(
+                        f"Collection {schema}.{table} created successfully"
+                    )
+                # Execute prepare:
+                await self.prepare_embedding_table(
+                    table=table,
+                    schema=schema,
+                    conn=session,
+                    dimension=dimension,
+                    id_column=id_column,
+                    create_all_indexes=True,
+                    **kwargs
+                )
+        except Exception as e:
+            self.logger.error(f"Error creating collection: {e}")
+            raise RuntimeError(
+                f"Failed to create collection: {e}"
+            ) from e