r2rpy/py/core/providers/database/documents.py

import asyncio
import copy
import csv
import json
import logging
import math
import tempfile
from typing import IO, Any, Optional
from uuid import UUID

import asyncpg
from fastapi import HTTPException

from core.base import (
    DocumentResponse,
    DocumentType,
    GraphConstructionStatus,
    GraphExtractionStatus,
    Handler,
    IngestionStatus,
    R2RException,
    SearchSettings,
)

from .base import PostgresConnectionManager
from .filters import apply_filters

logger = logging.getLogger()


def transform_filter_fields(filters: dict[str, Any]) -> dict[str, Any]:
    """Recursively transform filter field names by replacing 'document_id' with
    'id'. Handles nested logical operators like $and, $or, etc.

    Args:
        filters (dict[str, Any]): The original filters dictionary

    Returns:
        dict[str, Any]: A new dictionary with transformed field names
    """
    if not filters:
        return {}

    transformed = {}

    for key, value in filters.items():
        # Handle logical operators recursively
        if key in ("$and", "$or", "$not"):
            if isinstance(value, list):
                transformed[key] = [
                    transform_filter_fields(item) for item in value
                ]
            else:
                transformed[key] = transform_filter_fields(value)  # type: ignore
            continue

        # Replace 'document_id' with 'id'
        new_key = "id" if key == "document_id" else key

        # Handle nested dictionary cases (e.g., for operators like $eq, $gt, etc.)
        if isinstance(value, dict):
            transformed[new_key] = transform_filter_fields(value)  # type: ignore
        else:
            transformed[new_key] = value

    logger.debug(f"Transformed filters from {filters} to {transformed}")
    return transformed


class PostgresDocumentsHandler(Handler):
    TABLE_NAME = "documents"

    def __init__(
        self,
        project_name: str,
        connection_manager: PostgresConnectionManager,
        dimension: int | float,
    ):
        self.dimension = dimension
        super().__init__(project_name, connection_manager)

    async def create_tables(self):
        logger.info(
            f"Creating table, if it does not exist: {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}"
        )

        vector_dim = (
            "" if math.isnan(self.dimension) else f"({self.dimension})"
        )
        vector_type = f"vector{vector_dim}"

        try:
            query = f"""
            CREATE TABLE IF NOT EXISTS {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)} (
                id UUID PRIMARY KEY,
                collection_ids UUID[],
                owner_id UUID,
                type TEXT,
                metadata JSONB,
                title TEXT,
                summary TEXT NULL,
                summary_embedding {vector_type} NULL,
                version TEXT,
                size_in_bytes INT,
                ingestion_status TEXT DEFAULT 'pending',
                extraction_status TEXT DEFAULT 'pending',
                created_at TIMESTAMPTZ DEFAULT NOW(),
                updated_at TIMESTAMPTZ DEFAULT NOW(),
                ingestion_attempt_number INT DEFAULT 0,
                raw_tsvector tsvector GENERATED ALWAYS AS (
                    setweight(to_tsvector('english', COALESCE(title, '')), 'A') ||
                    setweight(to_tsvector('english', COALESCE(summary, '')), 'B') ||
                    setweight(to_tsvector('english', COALESCE((metadata->>'description')::text, '')), 'C')
                ) STORED,
                total_tokens INT DEFAULT 0
            );
            CREATE INDEX IF NOT EXISTS idx_collection_ids_{self.project_name}
            ON {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)} USING GIN (collection_ids);

            -- Full text search index
            CREATE INDEX IF NOT EXISTS idx_doc_search_{self.project_name}
            ON {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}
            USING GIN (raw_tsvector);
            """
            await self.connection_manager.execute_query(query)

            # ---------------------------------------------------------------
            # Now check if total_tokens column exists in the 'documents' table
            # ---------------------------------------------------------------
            # 1) See what columns exist
            # column_check_query = f"""
            # SELECT column_name
            # FROM information_schema.columns
            # WHERE table_name = '{self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}'
            # AND table_schema = CURRENT_SCHEMA()
            # """
            # existing_columns = await self.connection_manager.fetch_query(column_check_query)
            # 2) Parse the table name for schema checks
            table_full_name = self._get_table_name(
                PostgresDocumentsHandler.TABLE_NAME
            )
            parsed_schema = "public"
            parsed_table_name = table_full_name
            if "." in table_full_name:
                parts = table_full_name.split(".", maxsplit=1)
                parsed_schema = parts[0].replace('"', "").strip()
                parsed_table_name = parts[1].replace('"', "").strip()
            else:
                parsed_table_name = parsed_table_name.replace('"', "").strip()

            # 3) Check columns
            column_check_query = f"""
            SELECT column_name
            FROM information_schema.columns
            WHERE table_name = '{parsed_table_name}'
            AND table_schema = '{parsed_schema}'
            """
            existing_columns = await self.connection_manager.fetch_query(
                column_check_query
            )

            existing_column_names = {
                row["column_name"] for row in existing_columns
            }

            if "total_tokens" not in existing_column_names:
                # 2) If missing, see if the table already has data
                # doc_count_query = f"SELECT COUNT(*) FROM {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}"
                # doc_count = await self.connection_manager.fetchval(doc_count_query)
                doc_count_query = f"SELECT COUNT(*) AS doc_count FROM {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}"
                row = await self.connection_manager.fetchrow_query(
                    doc_count_query
                )
                if row is None:
                    doc_count = 0
                else:
                    doc_count = row[
                        "doc_count"
                    ]  # or row[0] if you prefer positional indexing

                if doc_count > 0:
                    # We already have documents, but no total_tokens column
                    # => ask user to run r2r db migrate
                    logger.warning(
                        "Adding the missing 'total_tokens' column to the 'documents' table, this will impact existing files."
                    )

                create_tokens_col = f"""
                ALTER TABLE {table_full_name}
                ADD COLUMN total_tokens INT DEFAULT 0
                """
                await self.connection_manager.execute_query(create_tokens_col)

        except Exception as e:
            logger.warning(f"Error {e} when creating document table.")
            raise e

    async def upsert_documents_overview(
        self, documents_overview: DocumentResponse | list[DocumentResponse]
    ) -> None:
        if isinstance(documents_overview, DocumentResponse):
            documents_overview = [documents_overview]

        # TODO: make this an arg
        max_retries = 20
        for document in documents_overview:
            retries = 0
            while retries < max_retries:
                try:
                    async with (
                        self.connection_manager.pool.get_connection() as conn  # type: ignore
                    ):
                        async with conn.transaction():
                            # Lock the row for update
                            check_query = f"""
                            SELECT ingestion_attempt_number, ingestion_status FROM {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}
                            WHERE id = $1 FOR UPDATE
                            """
                            existing_doc = await conn.fetchrow(
                                check_query, document.id
                            )

                            db_entry = document.convert_to_db_entry()

                            if existing_doc:
                                db_version = existing_doc[
                                    "ingestion_attempt_number"
                                ]
                                db_status = existing_doc["ingestion_status"]
                                new_version = db_entry[
                                    "ingestion_attempt_number"
                                ]

                                # Only increment version if status is changing to 'success' or if it's a new version
                                if (
                                    db_status != "success"
                                    and db_entry["ingestion_status"]
                                    == "success"
                                ) or (new_version > db_version):
                                    new_attempt_number = db_version + 1
                                else:
                                    new_attempt_number = db_version

                                db_entry["ingestion_attempt_number"] = (
                                    new_attempt_number
                                )

                                update_query = f"""
                                UPDATE {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}
                                SET collection_ids = $1,
                                    owner_id = $2,
                                    type = $3,
                                    metadata = $4,
                                    title = $5,
                                    version = $6,
                                    size_in_bytes = $7,
                                    ingestion_status = $8,
                                    extraction_status = $9,
                                    updated_at = $10,
                                    ingestion_attempt_number = $11,
                                    summary = $12,
                                    summary_embedding = $13,
                                    total_tokens = $14
                                WHERE id = $15
                                """

                                await conn.execute(
                                    update_query,
                                    db_entry["collection_ids"],
                                    db_entry["owner_id"],
                                    db_entry["document_type"],
                                    db_entry["metadata"],
                                    db_entry["title"],
                                    db_entry["version"],
                                    db_entry["size_in_bytes"],
                                    db_entry["ingestion_status"],
                                    db_entry["extraction_status"],
                                    db_entry["updated_at"],
                                    db_entry["ingestion_attempt_number"],
                                    db_entry["summary"],
                                    db_entry["summary_embedding"],
                                    db_entry[
                                        "total_tokens"
                                    ],  # pass the new field here
                                    document.id,
                                )
                            else:
                                insert_query = f"""
                                INSERT INTO {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}
                                (id, collection_ids, owner_id, type, metadata, title, version,
                                size_in_bytes, ingestion_status, extraction_status, created_at,
                                updated_at, ingestion_attempt_number, summary, summary_embedding, total_tokens)
                                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
                                """
                                await conn.execute(
                                    insert_query,
                                    db_entry["id"],
                                    db_entry["collection_ids"],
                                    db_entry["owner_id"],
                                    db_entry["document_type"],
                                    db_entry["metadata"],
                                    db_entry["title"],
                                    db_entry["version"],
                                    db_entry["size_in_bytes"],
                                    db_entry["ingestion_status"],
                                    db_entry["extraction_status"],
                                    db_entry["created_at"],
                                    db_entry["updated_at"],
                                    db_entry["ingestion_attempt_number"],
                                    db_entry["summary"],
                                    db_entry["summary_embedding"],
                                    db_entry["total_tokens"],
                                )

                    break  # Success, exit the retry loop
                except (
                    asyncpg.exceptions.UniqueViolationError,
                    asyncpg.exceptions.DeadlockDetectedError,
                ) as e:
                    retries += 1
                    if retries == max_retries:
                        logger.error(
                            f"Failed to update document {document.id} after {max_retries} attempts. Error: {str(e)}"
                        )
                        raise
                    else:
                        wait_time = 0.1 * (2**retries)  # Exponential backoff
                        await asyncio.sleep(wait_time)

    async def delete(
        self, document_id: UUID, version: Optional[str] = None
    ) -> None:
        query = f"""
        DELETE FROM {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}
        WHERE id = $1
        """

        params = [str(document_id)]

        if version:
            query += " AND version = $2"
            params.append(version)

        await self.connection_manager.execute_query(query=query, params=params)

    async def _get_status_from_table(
        self,
        ids: list[UUID],
        table_name: str,
        status_type: str,
        column_name: str,
    ):
        """Get the workflow status for a given document or list of documents.

        Args:
            ids (list[UUID]): The document IDs.
            table_name (str): The table name.
            status_type (str): The type of status to retrieve.

        Returns:
            The workflow status for the given document or list of documents.
        """
        query = f"""
            SELECT {status_type} FROM {self._get_table_name(table_name)}
            WHERE {column_name} = ANY($1)
        """
        return [
            row[status_type]
            for row in await self.connection_manager.fetch_query(query, [ids])
        ]

    async def _get_ids_from_table(
        self,
        status: list[str],
        table_name: str,
        status_type: str,
        collection_id: Optional[UUID] = None,
    ):
        """Get the IDs from a given table.

        Args:
            status (str | list[str]): The status or list of statuses to retrieve.
            table_name (str): The table name.
            status_type (str): The type of status to retrieve.
        """
        query = f"""
            SELECT id FROM {self._get_table_name(table_name)}
            WHERE {status_type} = ANY($1) and $2 = ANY(collection_ids)
        """
        records = await self.connection_manager.fetch_query(
            query, [status, collection_id]
        )
        return [record["id"] for record in records]

    async def _set_status_in_table(
        self,
        ids: list[UUID],
        status: str,
        table_name: str,
        status_type: str,
        column_name: str,
    ):
        """Set the workflow status for a given document or list of documents.

        Args:
            ids (list[UUID]): The document IDs.
            status (str): The status to set.
            table_name (str): The table name.
            status_type (str): The type of status to set.
            column_name (str): The column name in the table to update.
        """
        query = f"""
            UPDATE {self._get_table_name(table_name)}
            SET {status_type} = $1
            WHERE {column_name} = Any($2)
        """
        await self.connection_manager.execute_query(query, [status, ids])

    def _get_status_model(self, status_type: str):
        """Get the status model for a given status type.

        Args:
            status_type (str): The type of status to retrieve.

        Returns:
            The status model for the given status type.
        """
        if status_type == "ingestion":
            return IngestionStatus
        elif status_type == "extraction_status":
            return GraphExtractionStatus
        elif status_type in {"graph_cluster_status", "graph_sync_status"}:
            return GraphConstructionStatus
        else:
            raise R2RException(
                status_code=400, message=f"Invalid status type: {status_type}"
            )

    async def get_workflow_status(
        self, id: UUID | list[UUID], status_type: str
    ):
        """Get the workflow status for a given document or list of documents.

        Args:
            id (UUID | list[UUID]): The document ID or list of document IDs.
            status_type (str): The type of status to retrieve.

        Returns:
            The workflow status for the given document or list of documents.
        """

        ids = [id] if isinstance(id, UUID) else id
        out_model = self._get_status_model(status_type)
        result = await self._get_status_from_table(
            ids,
            out_model.table_name(),
            status_type,
            out_model.id_column(),
        )

        result = [out_model[status.upper()] for status in result]
        return result[0] if isinstance(id, UUID) else result

    async def set_workflow_status(
        self, id: UUID | list[UUID], status_type: str, status: str
    ):
        """Set the workflow status for a given document or list of documents.

        Args:
            id (UUID | list[UUID]): The document ID or list of document IDs.
            status_type (str): The type of status to set.
            status (str): The status to set.
        """
        ids = [id] if isinstance(id, UUID) else id
        out_model = self._get_status_model(status_type)

        return await self._set_status_in_table(
            ids,
            status,
            out_model.table_name(),
            status_type,
            out_model.id_column(),
        )

    async def get_document_ids_by_status(
        self,
        status_type: str,
        status: str | list[str],
        collection_id: Optional[UUID] = None,
    ):
        """Get the IDs for a given status.

        Args:
            ids_key (str): The key to retrieve the IDs.
            status_type (str): The type of status to retrieve.
            status (str | list[str]): The status or list of statuses to retrieve.
        """

        if isinstance(status, str):
            status = [status]

        out_model = self._get_status_model(status_type)
        return await self._get_ids_from_table(
            status, out_model.table_name(), status_type, collection_id
        )

    async def get_documents_overview(
        self,
        offset: int,
        limit: int,
        filter_user_ids: Optional[list[UUID]] = None,
        filter_document_ids: Optional[list[UUID]] = None,
        filter_collection_ids: Optional[list[UUID]] = None,
        include_summary_embedding: Optional[bool] = True,
        filters: Optional[dict[str, Any]] = None,
        sort_order: str = "DESC",
        owner_only: bool = False,
    ) -> dict[str, Any]:
        """Fetch overviews of documents with optional offset/limit pagination.

        You can use either:
          - Traditional filters: `filter_user_ids`, `filter_document_ids`, `filter_collection_ids`
          - A `filters` dict (e.g., like we do in semantic search), which will be passed to `apply_filters`.

        If both the `filters` dict and any of the traditional filter arguments are provided,
        this method will raise an error.
        """

        filters = copy.deepcopy(filters)
        filters = transform_filter_fields(filters)  # type: ignore

        # Safety check: We do not allow mixing the old filter arguments with the new `filters` dict.
        # This keeps the query logic unambiguous.
        if filters and any(
            [
                filter_user_ids,
                filter_document_ids,
                filter_collection_ids,
            ]
        ):
            raise HTTPException(
                status_code=400,
                detail=(
                    "Cannot use both the 'filters' dictionary "
                    "and the 'filter_*_ids' parameters simultaneously."
                ),
            )

        conditions = []
        params: list[Any] = []
        param_index = 1

        # -------------------------------------------
        # 1) If using the new `filters` dict approach
        # -------------------------------------------
        if filters:
            # Apply the filters to generate a WHERE clause
            filter_condition, filter_params = apply_filters(
                filters, params, mode="condition_only"
            )
            if filter_condition:
                conditions.append(filter_condition)
            # Make sure we keep adding to the same params list
            params.extend(filter_params)
            param_index += len(filter_params)

        # -------------------------------------------
        # 2) If using the old filter_*_ids approach
        # -------------------------------------------
        else:
            # Handle document IDs with AND
            if filter_document_ids:
                conditions.append(f"id = ANY(${param_index})")
                params.append(filter_document_ids)
                param_index += 1

            # For owner/collection filters, we used OR logic previously
            # so we combine them into a single sub-condition in parentheses
            owner_conditions = []
            collection_conditions = []

            if filter_user_ids:
                owner_conditions.append(f"owner_id = ANY(${param_index})")
                params.append(filter_user_ids)
                param_index += 1

            if filter_collection_ids:
                collection_conditions.append(
                    f"collection_ids && ${param_index}"
                )
                params.append(filter_collection_ids)
                param_index += 1

            if owner_only:
                if owner_conditions:
                    conditions.append(f"({' OR '.join(owner_conditions)})")

                if collection_conditions:
                    conditions.append(
                        f"({' OR '.join(collection_conditions)})"
                    )
            elif (
                combined_conditions := owner_conditions + collection_conditions
            ):
                conditions.append(f"({' OR '.join(combined_conditions)})")

        # -------------------------
        # Build the full query
        # -------------------------
        base_query = (
            f"FROM {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}"
        )
        if conditions:
            # Combine everything with AND
            base_query += " WHERE " + " AND ".join(conditions)

        # Construct SELECT fields (including total_entries via window function)
        select_fields = """
            SELECT
                id,
                collection_ids,
                owner_id,
                type,
                metadata,
                title,
                version,
                size_in_bytes,
                ingestion_status,
                extraction_status,
                created_at,
                updated_at,
                summary,
                summary_embedding,
                total_tokens,
                COUNT(*) OVER() AS total_entries
        """

        query = f"""
            {select_fields}
            {base_query}
            ORDER BY created_at {sort_order}
            OFFSET ${param_index}
        """
        params.append(offset)
        param_index += 1

        if limit != -1:
            query += f" LIMIT ${param_index}"
            params.append(limit)
            param_index += 1

        try:
            results = await self.connection_manager.fetch_query(query, params)
            total_entries = results[0]["total_entries"] if results else 0

            documents = []
            for row in results:
                # Safely handle the embedding
                embedding = None
                if (
                    "summary_embedding" in row
                    and row["summary_embedding"] is not None
                ):
                    try:
                        # The embedding is stored as a string like "[0.1, 0.2, ...]"
                        embedding_str = row["summary_embedding"]
                        if embedding_str.startswith(
                            "["
                        ) and embedding_str.endswith("]"):
                            embedding = [
                                float(x)
                                for x in embedding_str[1:-1].split(",")
                                if x
                            ]
                    except Exception as e:
                        logger.warning(
                            f"Failed to parse embedding for document {row['id']}: {e}"
                        )

                documents.append(
                    DocumentResponse(
                        id=row["id"],
                        collection_ids=row["collection_ids"],
                        owner_id=row["owner_id"],
                        document_type=DocumentType(row["type"]),
                        metadata=json.loads(row["metadata"]),
                        title=row["title"],
                        version=row["version"],
                        size_in_bytes=row["size_in_bytes"],
                        ingestion_status=IngestionStatus(
                            row["ingestion_status"]
                        ),
                        extraction_status=GraphExtractionStatus(
                            row["extraction_status"]
                        ),
                        created_at=row["created_at"],
                        updated_at=row["updated_at"],
                        summary=row["summary"] if "summary" in row else None,
                        summary_embedding=(
                            embedding if include_summary_embedding else None
                        ),
                        total_tokens=row["total_tokens"],
                    )
                )
            return {"results": documents, "total_entries": total_entries}
        except Exception as e:
            logger.error(f"Error in get_documents_overview: {str(e)}")
            raise HTTPException(
                status_code=500,
                detail="Database query failed",
            ) from e

    async def update_document_metadata(
        self,
        document_id: UUID,
        metadata: list[dict],
        overwrite: bool = False,
    ) -> DocumentResponse:
        """
        Update the metadata of a document, either by appending to existing metadata or overwriting it.
        Accepts a list of metadata dictionaries.
        """

        doc_result = await self.get_documents_overview(
            offset=0,
            limit=1,
            filter_document_ids=[document_id],
        )

        if not doc_result["results"]:
            raise HTTPException(
                status_code=404,
                detail=f"Document with ID {document_id} not found",
            )

        existing_doc = doc_result["results"][0]

        if overwrite:
            combined_metadata: dict[str, Any] = {}
            for meta_item in metadata:
                combined_metadata |= meta_item
            existing_doc.metadata = combined_metadata
        else:
            for meta_item in metadata:
                existing_doc.metadata.update(meta_item)

        await self.upsert_documents_overview(existing_doc)

        return existing_doc

    async def semantic_document_search(
        self, query_embedding: list[float], search_settings: SearchSettings
    ) -> list[DocumentResponse]:
        """Search documents using semantic similarity with their summary
        embeddings."""

        where_clauses = ["summary_embedding IS NOT NULL"]
        params: list[str | int | bytes] = [str(query_embedding)]

        vector_dim = (
            "" if math.isnan(self.dimension) else f"({self.dimension})"
        )
        filters = copy.deepcopy(search_settings.filters)
        if filters:
            filter_condition, params = apply_filters(
                transform_filter_fields(filters), params, mode="condition_only"
            )
            if filter_condition:
                where_clauses.append(filter_condition)

        where_clause = " AND ".join(where_clauses)

        query = f"""
        WITH document_scores AS (
            SELECT
                id,
                collection_ids,
                owner_id,
                type,
                metadata,
                title,
                version,
                size_in_bytes,
                ingestion_status,
                extraction_status,
                created_at,
                updated_at,
                summary,
                summary_embedding,
                total_tokens,
                (summary_embedding <=> $1::vector({vector_dim})) as semantic_distance
            FROM {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}
            WHERE {where_clause}
            ORDER BY semantic_distance ASC
            LIMIT ${len(params) + 1}
            OFFSET ${len(params) + 2}
        )
        SELECT *,
            1.0 - semantic_distance as semantic_score
        FROM document_scores
        """

        params.extend([search_settings.limit, search_settings.offset])

        results = await self.connection_manager.fetch_query(query, params)

        return [
            DocumentResponse(
                id=row["id"],
                collection_ids=row["collection_ids"],
                owner_id=row["owner_id"],
                document_type=DocumentType(row["type"]),
                metadata={
                    **(
                        json.loads(row["metadata"])
                        if search_settings.include_metadatas
                        else {}
                    ),
                    "search_score": float(row["semantic_score"]),
                    "search_type": "semantic",
                },
                title=row["title"],
                version=row["version"],
                size_in_bytes=row["size_in_bytes"],
                ingestion_status=IngestionStatus(row["ingestion_status"]),
                extraction_status=GraphExtractionStatus(
                    row["extraction_status"]
                ),
                created_at=row["created_at"],
                updated_at=row["updated_at"],
                summary=row["summary"],
                summary_embedding=[
                    float(x)
                    for x in row["summary_embedding"][1:-1].split(",")
                    if x
                ],
                total_tokens=row["total_tokens"],
            )
            for row in results
        ]

    async def full_text_document_search(
        self, query_text: str, search_settings: SearchSettings
    ) -> list[DocumentResponse]:
        """Enhanced full-text search using generated tsvector."""

        where_clauses = ["raw_tsvector @@ websearch_to_tsquery('english', $1)"]
        params: list[str | int | bytes] = [query_text]

        filters = copy.deepcopy(search_settings.filters)
        if filters:
            filter_condition, params = apply_filters(
                transform_filter_fields(filters), params, mode="condition_only"
            )
            if filter_condition:
                where_clauses.append(filter_condition)

        where_clause = " AND ".join(where_clauses)

        query = f"""
        WITH document_scores AS (
            SELECT
                id,
                collection_ids,
                owner_id,
                type,
                metadata,
                title,
                version,
                size_in_bytes,
                ingestion_status,
                extraction_status,
                created_at,
                updated_at,
                summary,
                summary_embedding,
                total_tokens,
                ts_rank_cd(raw_tsvector, websearch_to_tsquery('english', $1), 32) as text_score
            FROM {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)}
            WHERE {where_clause}
            ORDER BY text_score DESC
            LIMIT ${len(params) + 1}
            OFFSET ${len(params) + 2}
        )
        SELECT * FROM document_scores
        """

        params.extend([search_settings.limit, search_settings.offset])

        results = await self.connection_manager.fetch_query(query, params)

        return [
            DocumentResponse(
                id=row["id"],
                collection_ids=row["collection_ids"],
                owner_id=row["owner_id"],
                document_type=DocumentType(row["type"]),
                metadata={
                    **(
                        json.loads(row["metadata"])
                        if search_settings.include_metadatas
                        else {}
                    ),
                    "search_score": float(row["text_score"]),
                    "search_type": "full_text",
                },
                title=row["title"],
                version=row["version"],
                size_in_bytes=row["size_in_bytes"],
                ingestion_status=IngestionStatus(row["ingestion_status"]),
                extraction_status=GraphExtractionStatus(
                    row["extraction_status"]
                ),
                created_at=row["created_at"],
                updated_at=row["updated_at"],
                summary=row["summary"],
                summary_embedding=(
                    [
                        float(x)
                        for x in row["summary_embedding"][1:-1].split(",")
                        if x
                    ]
                    if row["summary_embedding"]
                    else None
                ),
                total_tokens=row["total_tokens"],
            )
            for row in results
        ]

    async def hybrid_document_search(
        self,
        query_text: str,
        query_embedding: list[float],
        search_settings: SearchSettings,
    ) -> list[DocumentResponse]:
        """Search documents using both semantic and full-text search with RRF
        fusion."""

        # Get more results than needed for better fusion
        extended_settings = copy.deepcopy(search_settings)
        extended_settings.limit = search_settings.limit * 3

        # Get results from both search methods
        semantic_results = await self.semantic_document_search(
            query_embedding, extended_settings
        )
        full_text_results = await self.full_text_document_search(
            query_text, extended_settings
        )

        # Combine results using RRF
        doc_scores: dict[str, dict] = {}

        # Process semantic results
        for rank, result in enumerate(semantic_results, 1):
            doc_id = str(result.id)
            doc_scores[doc_id] = {
                "semantic_rank": rank,
                "full_text_rank": len(full_text_results)
                + 1,  # Default rank if not found
                "data": result,
            }

        # Process full-text results
        for rank, result in enumerate(full_text_results, 1):
            doc_id = str(result.id)
            if doc_id in doc_scores:
                doc_scores[doc_id]["full_text_rank"] = rank
            else:
                doc_scores[doc_id] = {
                    "semantic_rank": len(semantic_results)
                    + 1,  # Default rank if not found
                    "full_text_rank": rank,
                    "data": result,
                }

        # Calculate RRF scores using hybrid search settings
        rrf_k = search_settings.hybrid_settings.rrf_k
        semantic_weight = search_settings.hybrid_settings.semantic_weight
        full_text_weight = search_settings.hybrid_settings.full_text_weight

        for scores in doc_scores.values():
            semantic_score = 1 / (rrf_k + scores["semantic_rank"])
            full_text_score = 1 / (rrf_k + scores["full_text_rank"])

            # Weighted combination
            combined_score = (
                semantic_score * semantic_weight
                + full_text_score * full_text_weight
            ) / (semantic_weight + full_text_weight)

            scores["final_score"] = combined_score

        # Sort by final score and apply offset/limit
        sorted_results = sorted(
            doc_scores.values(), key=lambda x: x["final_score"], reverse=True
        )[
            search_settings.offset : search_settings.offset
            + search_settings.limit
        ]

        return [
            DocumentResponse(
                **{
                    **result["data"].__dict__,
                    "metadata": {
                        **(
                            result["data"].metadata
                            if search_settings.include_metadatas
                            else {}
                        ),
                        "search_score": result["final_score"],
                        "semantic_rank": result["semantic_rank"],
                        "full_text_rank": result["full_text_rank"],
                        "search_type": "hybrid",
                    },
                }
            )
            for result in sorted_results
        ]

    async def search_documents(
        self,
        query_text: str,
        query_embedding: Optional[list[float]] = None,
        settings: Optional[SearchSettings] = None,
    ) -> list[DocumentResponse]:
        """Main search method that delegates to the appropriate search method
        based on settings."""
        if settings is None:
            settings = SearchSettings()

        if (
            settings.use_semantic_search and settings.use_fulltext_search
        ) or settings.use_hybrid_search:
            if query_embedding is None:
                raise ValueError(
                    "query_embedding is required for hybrid search"
                )
            return await self.hybrid_document_search(
                query_text, query_embedding, settings
            )
        elif settings.use_semantic_search:
            if query_embedding is None:
                raise ValueError(
                    "query_embedding is required for vector search"
                )
            return await self.semantic_document_search(
                query_embedding, settings
            )
        else:
            return await self.full_text_document_search(query_text, settings)

    async def export_to_csv(
        self,
        columns: Optional[list[str]] = None,
        filters: Optional[dict] = None,
        include_header: bool = True,
    ) -> tuple[str, IO]:
        """Creates a CSV file from the PostgreSQL data and returns the path to
        the temp file."""
        valid_columns = {
            "id",
            "collection_ids",
            "owner_id",
            "type",
            "metadata",
            "title",
            "summary",
            "version",
            "size_in_bytes",
            "ingestion_status",
            "extraction_status",
            "created_at",
            "updated_at",
            "total_tokens",
        }
        filters = copy.deepcopy(filters)
        filters = transform_filter_fields(filters)  # type: ignore

        if not columns:
            columns = list(valid_columns)
        elif invalid_cols := set(columns) - valid_columns:
            raise ValueError(f"Invalid columns: {invalid_cols}")

        select_stmt = f"""
            SELECT
                id::text,
                collection_ids::text,
                owner_id::text,
                type::text,
                metadata::text AS metadata,
                title,
                summary,
                version,
                size_in_bytes,
                ingestion_status,
                extraction_status,
                to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') AS created_at,
                to_char(updated_at, 'YYYY-MM-DD HH24:MI:SS') AS updated_at,
                total_tokens
            FROM {self._get_table_name(self.TABLE_NAME)}
        """

        conditions = []
        params: list[Any] = []
        param_index = 1

        if filters:
            for field, value in filters.items():
                if field not in valid_columns:
                    continue

                if isinstance(value, dict):
                    for op, val in value.items():
                        if op == "$eq":
                            conditions.append(f"{field} = ${param_index}")
                            params.append(val)
                            param_index += 1
                        elif op == "$gt":
                            conditions.append(f"{field} > ${param_index}")
                            params.append(val)
                            param_index += 1
                        elif op == "$lt":
                            conditions.append(f"{field} < ${param_index}")
                            params.append(val)
                            param_index += 1
                else:
                    # Direct equality
                    conditions.append(f"{field} = ${param_index}")
                    params.append(value)
                    param_index += 1

        if conditions:
            select_stmt = f"{select_stmt} WHERE {' AND '.join(conditions)}"

        select_stmt = f"{select_stmt} ORDER BY created_at DESC"

        temp_file = None
        try:
            temp_file = tempfile.NamedTemporaryFile(
                mode="w", delete=True, suffix=".csv"
            )
            writer = csv.writer(temp_file, quoting=csv.QUOTE_ALL)

            async with self.connection_manager.pool.get_connection() as conn:  # type: ignore
                async with conn.transaction():
                    cursor = await conn.cursor(select_stmt, *params)

                    if include_header:
                        writer.writerow(columns)

                    chunk_size = 1000
                    while True:
                        rows = await cursor.fetch(chunk_size)
                        if not rows:
                            break
                        for row in rows:
                            row_dict = {
                                "id": row[0],
                                "collection_ids": row[1],
                                "owner_id": row[2],
                                "type": row[3],
                                "metadata": row[4],
                                "title": row[5],
                                "summary": row[6],
                                "version": row[7],
                                "size_in_bytes": row[8],
                                "ingestion_status": row[9],
                                "extraction_status": row[10],
                                "created_at": row[11],
                                "updated_at": row[12],
                                "total_tokens": row[13],
                            }
                            writer.writerow([row_dict[col] for col in columns])

            temp_file.flush()
            return temp_file.name, temp_file

        except Exception as e:
            if temp_file:
                temp_file.close()
            raise HTTPException(
                status_code=500,
                detail=f"Failed to export data: {str(e)}",
            ) from e