LangBot/src/langbot/pkg/vector/mgr.py

from __future__ import annotations

import uuid

import sqlalchemy

from ..api.http.authz import WorkspaceRequiredError
from ..api.http.context import ExecutionContext
from ..core import app
from ..entity.persistence import rag as persistence_rag
from ..entity.persistence import workspace as persistence_workspace
from ..workspace.errors import WorkspaceNotFoundError
from .vdb import VectorDatabase, SearchType


class VectorDBManager:
    ap: app.Application
    vector_db: VectorDatabase = None

    def __init__(self, ap: app.Application):
        self.ap = ap

    async def initialize(self):
        kb_config = self.ap.instance_config.data.get('vdb')
        if kb_config:
            vdb_type = kb_config.get('use')

            if vdb_type == 'chroma':
                from .vdbs.chroma import ChromaVectorDatabase

                self.vector_db = ChromaVectorDatabase(self.ap)
                self.ap.logger.info('Initialized Chroma vector database backend.')

            elif vdb_type == 'qdrant':
                from .vdbs.qdrant import QdrantVectorDatabase

                self.vector_db = QdrantVectorDatabase(self.ap)
                self.ap.logger.info('Initialized Qdrant vector database backend.')
            elif vdb_type == 'seekdb':
                from .vdbs.seekdb import SeekDBVectorDatabase

                self.vector_db = SeekDBVectorDatabase(self.ap)
                self.ap.logger.info('Initialized SeekDB vector database backend.')

            elif vdb_type == 'valkey_search':
                from .vdbs.valkey_search import ValkeySearchVectorDatabase

                self.vector_db = ValkeySearchVectorDatabase(self.ap)
                self.ap.logger.info('Initialized Valkey Search vector database backend.')

            elif vdb_type == 'milvus':
                from .vdbs.milvus import MilvusVectorDatabase

                # Get Milvus configuration
                milvus_config = kb_config.get('milvus', {})
                uri = milvus_config.get('uri', './data/milvus.db')
                token = milvus_config.get('token')
                db_name = milvus_config.get('db_name', 'default')
                self.vector_db = MilvusVectorDatabase(self.ap, uri=uri, token=token, db_name=db_name)
                self.ap.logger.info('Initialized Milvus vector database backend.')

            elif vdb_type == 'pgvector':
                from .vdbs.pgvector_db import PgVectorDatabase

                # Get pgvector configuration
                pgvector_config = kb_config.get('pgvector', {})
                use_business_database = pgvector_config.get('use_business_database', False)
                allowed_dimensions = pgvector_config.get(
                    'allowed_dimensions',
                    [384, 512, 768, 1024, 1536],
                )
                common_options = {
                    'use_business_database': use_business_database,
                    'allowed_dimensions': allowed_dimensions,
                }
                if use_business_database:
                    self.vector_db = PgVectorDatabase(self.ap, **common_options)
                    self.ap.logger.info('Initialized pgvector on the shared business PostgreSQL database.')
                    return
                connection_string = pgvector_config.get('connection_string')
                if connection_string:
                    self.vector_db = PgVectorDatabase(
                        self.ap,
                        connection_string=connection_string,
                        **common_options,
                    )
                else:
                    # Use individual parameters
                    host = pgvector_config.get('host', 'localhost')
                    port = pgvector_config.get('port', 5432)
                    database = pgvector_config.get('database', 'langbot')
                    user = pgvector_config.get('user', 'postgres')
                    password = pgvector_config.get('password', 'postgres')
                    self.vector_db = PgVectorDatabase(
                        self.ap,
                        host=host,
                        port=port,
                        database=database,
                        user=user,
                        password=password,
                        **common_options,
                    )
                self.ap.logger.info('Initialized pgvector database backend.')

            else:
                from .vdbs.chroma import ChromaVectorDatabase

                self.vector_db = ChromaVectorDatabase(self.ap)
                self.ap.logger.warning('No valid vector database backend configured, defaulting to Chroma.')
        else:
            from .vdbs.chroma import ChromaVectorDatabase

            self.vector_db = ChromaVectorDatabase(self.ap)
            self.ap.logger.warning('No vector database backend configured, defaulting to Chroma.')

    async def shutdown(self) -> None:
        """Release the active vector backend deterministically."""

        vector_db = self.vector_db
        self.vector_db = None
        if vector_db is not None:
            await vector_db.close()

    def get_supported_search_types(self) -> list[str]:
        """Return the search types supported by the current VDB backend."""
        if self.vector_db is None:
            return [SearchType.VECTOR.value]
        return [st.value for st in self.vector_db.supported_search_types()]

    @staticmethod
    def physical_collection_name(
        execution_context: ExecutionContext,
        knowledge_base_uuid: str,
    ) -> str:
        """Derive an opaque physical collection from trusted tenant identity.

        Vector backends have different collection-name constraints, so the
        instance, Workspace and knowledge-base identifiers are encoded through
        UUIDv5 instead of being concatenated into a client-visible handle.
        Placement generation is deliberately not part of the name: generation
        fencing rejects stale work while preserving data across placements.
        """

        if not isinstance(execution_context, ExecutionContext):
            raise WorkspaceRequiredError('ExecutionContext is required for vector access')
        instance_uuid = execution_context.instance_uuid.strip()
        workspace_uuid = execution_context.workspace_uuid.strip()
        kb_uuid = knowledge_base_uuid.strip() if isinstance(knowledge_base_uuid, str) else ''
        if not instance_uuid or not workspace_uuid or not kb_uuid:
            raise WorkspaceRequiredError('Instance, Workspace and knowledge-base context are required')
        if execution_context.placement_generation <= 0:
            raise WorkspaceRequiredError('A positive placement generation is required')

        collection_uuid = uuid.uuid5(
            uuid.NAMESPACE_URL,
            f'langbot:knowledge-vector:{instance_uuid}:{workspace_uuid}:{kb_uuid}',
        )
        return f'lb_{collection_uuid.hex}'

    async def _validate_execution_context(self, execution_context: ExecutionContext) -> None:
        """Validate the active placement before touching a vector backend."""

        # Also performs structural validation before accessing app services.
        self.physical_collection_name(execution_context, 'context-validation')
        workspace_service = getattr(self.ap, 'workspace_service', None)
        if workspace_service is None:
            raise WorkspaceRequiredError('Workspace execution service is unavailable')
        binding = await workspace_service.get_execution_binding(
            execution_context.workspace_uuid,
            expected_generation=execution_context.placement_generation,
        )
        if binding.instance_uuid != execution_context.instance_uuid:
            raise WorkspaceRequiredError('ExecutionContext belongs to another LangBot instance')

    async def _resolve_physical_collection_name(
        self,
        execution_context: ExecutionContext,
        knowledge_base_uuid: str,
    ) -> str:
        """Resolve a scoped collection or an explicitly migrated OSS handle.

        Legacy handles are server-owned migration state, not caller input.
        They are honored only for the one local Workspace under the OSS
        single-Workspace policy.  A projected/cloud Workspace always gets the
        opaque tenant-derived collection, even if its database row was
        incorrectly marked as legacy.
        """

        await self._validate_execution_context(execution_context)
        async with self.ap.persistence_mgr.tenant_uow(execution_context.workspace_uuid):
            result = await self.ap.persistence_mgr.execute_async(
                sqlalchemy.select(
                    persistence_rag.KnowledgeBase.collection_id,
                    persistence_rag.KnowledgeBase.legacy_vector_collection,
                    persistence_workspace.Workspace.source,
                )
                .join(
                    persistence_workspace.Workspace,
                    persistence_workspace.Workspace.uuid == persistence_rag.KnowledgeBase.workspace_uuid,
                )
                .where(
                    persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
                    persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
                    persistence_workspace.Workspace.instance_uuid == execution_context.instance_uuid,
                )
                .limit(1)
            )
        row = result.first()
        if row is None:
            raise WorkspaceNotFoundError('Knowledge base not found')

        collection_id, legacy_vector_collection, workspace_source = row
        if legacy_vector_collection:
            policy = getattr(self.ap, 'workspace_policy', None)
            is_single_workspace = policy is not None and not getattr(
                policy,
                'multi_workspace_enabled',
                True,
            )
            is_local_workspace = workspace_source == persistence_workspace.WorkspaceSource.LOCAL.value
            if is_single_workspace and is_local_workspace and isinstance(collection_id, str) and collection_id.strip():
                return collection_id
            self.ap.logger.warning(
                'Ignored a legacy vector collection marker outside the local single-Workspace compatibility boundary.'
            )

        return self.physical_collection_name(execution_context, knowledge_base_uuid)

    def _pgvector_database(self):
        from .vdbs.pgvector_db import PgVectorDatabase

        return self.vector_db if isinstance(self.vector_db, PgVectorDatabase) else None

    async def _resolve_pgvector_scope(
        self,
        execution_context: ExecutionContext,
        knowledge_base_uuid: str,
        *,
        expected_dimension: int | None,
        initialize_dimension: bool,
    ):
        """Bind and verify the server-owned knowledge-base vector dimension."""

        from .vdbs.pgvector_db import PgVectorScope

        pgvector = self._pgvector_database()
        if pgvector is None:  # pragma: no cover - private call invariant
            raise RuntimeError('pgvector scope requested for another vector backend')
        if expected_dimension is not None and expected_dimension not in pgvector.allowed_dimensions:
            raise ValueError(f'Embedding dimension {expected_dimension} is not enabled for this deployment')

        async with self.ap.persistence_mgr.tenant_uow(execution_context.workspace_uuid):
            query = sqlalchemy.select(persistence_rag.KnowledgeBase.embedding_dimension).where(
                persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
                persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
            )
            current_dimension = (await self.ap.persistence_mgr.execute_async(query)).scalar_one_or_none()
            if current_dimension is None and expected_dimension is not None and initialize_dimension:
                await self.ap.persistence_mgr.execute_async(
                    sqlalchemy.update(persistence_rag.KnowledgeBase)
                    .where(
                        persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
                        persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
                        persistence_rag.KnowledgeBase.embedding_dimension.is_(None),
                    )
                    .values(embedding_dimension=expected_dimension)
                )
                current_dimension = (await self.ap.persistence_mgr.execute_async(query)).scalar_one_or_none()

            if expected_dimension is not None and current_dimension != expected_dimension:
                if current_dimension is None:
                    raise ValueError('Knowledge base has no selected pgvector embedding dimension')
                raise ValueError(f'Knowledge base embedding dimension is {current_dimension}, not {expected_dimension}')

        return PgVectorScope(
            workspace_uuid=execution_context.workspace_uuid,
            knowledge_base_uuid=knowledge_base_uuid,
            embedding_dimension=current_dimension,
        )

    async def upsert(
        self,
        execution_context: ExecutionContext,
        knowledge_base_uuid: str,
        vectors: list[list[float]],
        ids: list[str],
        metadata: list[dict] | None = None,
        documents: list[str] | None = None,
    ):
        """Upsert vectors into a server-derived tenant collection."""

        collection_name = await self._resolve_physical_collection_name(
            execution_context,
            knowledge_base_uuid,
        )
        source_metadata = metadata or [{} for _ in vectors]
        scoped_metadata = [
            {
                **item,
                '_langbot_instance_uuid': execution_context.instance_uuid,
                '_langbot_workspace_uuid': execution_context.workspace_uuid,
                '_langbot_knowledge_base_uuid': knowledge_base_uuid,
            }
            for item in source_metadata
        ]
        pgvector = self._pgvector_database()
        if pgvector is not None:
            if not vectors:
                return
            scope = await self._resolve_pgvector_scope(
                execution_context,
                knowledge_base_uuid,
                expected_dimension=len(vectors[0]),
                initialize_dimension=True,
            )
            await pgvector.add_embeddings(
                collection=collection_name,
                ids=ids,
                embeddings_list=vectors,
                metadatas=scoped_metadata,
                documents=documents,
                scope=scope,
            )
            return
        await self.vector_db.add_embeddings(
            collection=collection_name,
            ids=ids,
            embeddings_list=vectors,
            metadatas=scoped_metadata,
            documents=documents,
        )

    async def search(
        self,
        execution_context: ExecutionContext,
        knowledge_base_uuid: str,
        query_vector: list[float],
        limit: int,
        filter: dict | None = None,
        search_type: str = 'vector',
        query_text: str = '',
        vector_weight: float | None = None,
    ) -> list[dict]:
        """Proxy: Search vectors.

        Returns a list of dicts with keys: 'id', 'distance', 'metadata'.
        The underlying VectorDatabase.search returns Chroma-style format:
        { 'ids': [['id1']], 'distances': [[0.1]], 'metadatas': [[{}]] }
        """
        collection_name = await self._resolve_physical_collection_name(
            execution_context,
            knowledge_base_uuid,
        )
        pgvector = self._pgvector_database()
        if pgvector is not None:
            scope = await self._resolve_pgvector_scope(
                execution_context,
                knowledge_base_uuid,
                expected_dimension=len(query_vector),
                initialize_dimension=False,
            )
            results = await pgvector.search(
                collection=collection_name,
                query_embedding=query_vector,
                k=limit,
                search_type=search_type,
                query_text=query_text,
                filter=filter,
                vector_weight=vector_weight,
                scope=scope,
            )
        else:
            results = await self.vector_db.search(
                collection=collection_name,
                query_embedding=query_vector,
                k=limit,
                search_type=search_type,
                query_text=query_text,
                filter=filter,
                vector_weight=vector_weight,
            )

        if not results or 'ids' not in results or not results['ids']:
            return []

        # Flatten nested lists (Chroma returns batch-style: list of lists)
        raw_ids = results['ids']
        raw_dists = results.get('distances', [])
        raw_metas = results.get('metadatas', [])

        r_ids = raw_ids[0] if raw_ids and isinstance(raw_ids[0], list) else raw_ids
        r_dists = raw_dists[0] if raw_dists and isinstance(raw_dists[0], list) else raw_dists
        r_metas = raw_metas[0] if raw_metas and isinstance(raw_metas[0], list) else raw_metas

        parsed_results = []
        for i, id_val in enumerate(r_ids):
            parsed_results.append(
                {
                    'id': id_val,
                    'distance': r_dists[i] if r_dists and i < len(r_dists) else 0.0,
                    'metadata': r_metas[i] if r_metas and i < len(r_metas) else {},
                }
            )

        return parsed_results

    async def delete_by_file_id(
        self,
        execution_context: ExecutionContext,
        knowledge_base_uuid: str,
        file_ids: list[str],
    ):
        """Proxy: Delete vectors by file_id (metadata-level identifier).

        This delegates to VectorDatabase.delete_by_file_id which removes
        all vectors associated with the given file IDs.
        """
        collection_name = await self._resolve_physical_collection_name(
            execution_context,
            knowledge_base_uuid,
        )
        pgvector = self._pgvector_database()
        scope = None
        if pgvector is not None:
            scope = await self._resolve_pgvector_scope(
                execution_context,
                knowledge_base_uuid,
                expected_dimension=None,
                initialize_dimension=False,
            )
        for file_id in file_ids:
            if pgvector is not None:
                await pgvector.delete_by_file_id(collection_name, file_id, scope=scope)
            else:
                await self.vector_db.delete_by_file_id(collection_name, file_id)

    async def delete_collection(
        self,
        execution_context: ExecutionContext,
        knowledge_base_uuid: str,
    ):
        """Delete one server-derived tenant collection."""

        collection_name = await self._resolve_physical_collection_name(
            execution_context,
            knowledge_base_uuid,
        )
        pgvector = self._pgvector_database()
        if pgvector is not None:
            scope = await self._resolve_pgvector_scope(
                execution_context,
                knowledge_base_uuid,
                expected_dimension=None,
                initialize_dimension=False,
            )
            await pgvector.delete_collection(collection_name, scope=scope)
        else:
            await self.vector_db.delete_collection(collection_name)

    async def delete_by_filter(
        self,
        execution_context: ExecutionContext,
        knowledge_base_uuid: str,
        filter: dict,
    ) -> int:
        """Proxy: Delete vectors by metadata filter.

        Returns:
            Number of deleted vectors (best-effort; some backends return 0).
        """
        collection_name = await self._resolve_physical_collection_name(
            execution_context,
            knowledge_base_uuid,
        )
        pgvector = self._pgvector_database()
        if pgvector is not None:
            scope = await self._resolve_pgvector_scope(
                execution_context,
                knowledge_base_uuid,
                expected_dimension=None,
                initialize_dimension=False,
            )
            return await pgvector.delete_by_filter(collection_name, filter, scope=scope)
        return await self.vector_db.delete_by_filter(collection_name, filter)

    async def list_by_filter(
        self,
        execution_context: ExecutionContext,
        knowledge_base_uuid: str,
        filter: dict | None = None,
        limit: int = 20,
        offset: int = 0,
    ) -> tuple[list[dict], int]:
        """Proxy: List vectors by metadata filter with pagination.

        Returns:
            Tuple of (items, total).
        """
        collection_name = await self._resolve_physical_collection_name(
            execution_context,
            knowledge_base_uuid,
        )
        pgvector = self._pgvector_database()
        if pgvector is not None:
            scope = await self._resolve_pgvector_scope(
                execution_context,
                knowledge_base_uuid,
                expected_dimension=None,
                initialize_dimension=False,
            )
            return await pgvector.list_by_filter(collection_name, filter, limit, offset, scope=scope)
        return await self.vector_db.list_by_filter(collection_name, filter, limit, offset)