LangBot/src/langbot/pkg/rag/service/runtime.py

from __future__ import annotations

import posixpath
import re
from typing import TYPE_CHECKING, Any
from urllib.parse import unquote

import sqlalchemy

from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.entity.persistence import rag as persistence_rag
from langbot.pkg.workspace.errors import WorkspaceNotFoundError

if TYPE_CHECKING:
    from langbot.pkg.core import app


class RAGRuntimeService:
    """Service to handle RAG-related requests from plugins (Runtime).

    This service acts as the bridge between plugin RPC requests and
    LangBot's infrastructure (embedding models, vector databases, file storage).
    """

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

    async def _validate_execution_context(self, execution_context: ExecutionContext) -> None:
        if not isinstance(execution_context, ExecutionContext):
            raise WorkspaceRequiredError('ExecutionContext is required for RAG runtime access')
        if (
            not execution_context.instance_uuid.strip()
            or not execution_context.workspace_uuid.strip()
            or execution_context.placement_generation <= 0
        ):
            raise WorkspaceRequiredError('A complete active ExecutionContext is required')
        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_knowledge_base_uuid(
        self,
        execution_context: ExecutionContext,
        collection_id: str,
    ) -> str:
        """Resolve a plugin logical handle to a Workspace-owned KB UUID."""

        await self._validate_execution_context(execution_context)
        if not isinstance(collection_id, str) or not collection_id.strip():
            raise WorkspaceNotFoundError('Knowledge base not found')
        result = await self.ap.persistence_mgr.execute_async(
            sqlalchemy.select(persistence_rag.KnowledgeBase.uuid)
            .where(persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid)
            .where(
                sqlalchemy.or_(
                    persistence_rag.KnowledgeBase.uuid == collection_id,
                    persistence_rag.KnowledgeBase.collection_id == collection_id,
                )
            )
            .limit(1)
        )
        kb_uuid = result.scalar_one_or_none()
        if kb_uuid is None:
            raise WorkspaceNotFoundError('Knowledge base not found')
        return kb_uuid

    async def vector_upsert(
        self,
        execution_context: ExecutionContext,
        collection_id: str,
        vectors: list[list[float]],
        ids: list[str],
        metadata: list[dict[str, Any]] | None = None,
        documents: list[str] | None = None,
    ) -> None:
        """Handle VECTOR_UPSERT action."""
        knowledge_base_uuid = await self._resolve_knowledge_base_uuid(execution_context, collection_id)
        if len(vectors) != len(ids):
            raise ValueError('vectors and ids must have the same length')
        if metadata is not None and len(metadata) != len(vectors):
            raise ValueError('metadata must have the same length as vectors')
        if documents is not None and len(documents) != len(vectors):
            raise ValueError('documents must have the same length as vectors')
        metadatas = metadata if metadata else [{} for _ in vectors]
        await self.ap.vector_db_mgr.upsert(
            execution_context=execution_context,
            knowledge_base_uuid=knowledge_base_uuid,
            vectors=vectors,
            ids=ids,
            metadata=metadatas,
            documents=documents,
        )

    async def vector_search(
        self,
        execution_context: ExecutionContext,
        collection_id: str,
        query_vector: list[float],
        top_k: int,
        filters: dict[str, Any] | None = None,
        search_type: str = 'vector',
        query_text: str = '',
        vector_weight: float | None = None,
    ) -> list[dict[str, Any]]:
        """Handle VECTOR_SEARCH action."""
        knowledge_base_uuid = await self._resolve_knowledge_base_uuid(execution_context, collection_id)
        return await self.ap.vector_db_mgr.search(
            execution_context=execution_context,
            knowledge_base_uuid=knowledge_base_uuid,
            query_vector=query_vector,
            limit=top_k,
            filter=filters,
            search_type=search_type,
            query_text=query_text,
            vector_weight=vector_weight,
        )

    async def vector_delete(
        self,
        execution_context: ExecutionContext,
        collection_id: str,
        file_ids: list[str] | None = None,
        filters: dict[str, Any] | None = None,
    ) -> int:
        """Handle VECTOR_DELETE action.

        Deletes vectors associated with the given file IDs from the collection.
        Each file_id corresponds to a document whose vectors will be removed.

        Args:
            collection_id: The collection to delete from.
            file_ids: File IDs whose associated vectors should be deleted.
                Each file_id maps to a set of vectors stored with that file_id
                in their metadata.
            filters: Filter-based deletion (not yet supported, will raise).
        """
        knowledge_base_uuid = await self._resolve_knowledge_base_uuid(execution_context, collection_id)
        count = 0
        if file_ids:
            await self.ap.vector_db_mgr.delete_by_file_id(
                execution_context=execution_context,
                knowledge_base_uuid=knowledge_base_uuid,
                file_ids=file_ids,
            )
            count = len(file_ids)
        elif filters:
            count = await self.ap.vector_db_mgr.delete_by_filter(
                execution_context=execution_context,
                knowledge_base_uuid=knowledge_base_uuid,
                filter=filters,
            )
        return count

    async def vector_list(
        self,
        execution_context: ExecutionContext,
        collection_id: str,
        filters: dict[str, Any] | None = None,
        limit: int = 20,
        offset: int = 0,
    ) -> tuple[list[dict[str, Any]], int]:
        """Handle VECTOR_LIST action.

        Args:
            collection_id: The collection to list from.
            filters: Optional metadata filters.
            limit: Maximum number of items to return.
            offset: Number of items to skip.

        Returns:
            Tuple of (items, total).
        """
        knowledge_base_uuid = await self._resolve_knowledge_base_uuid(execution_context, collection_id)
        return await self.ap.vector_db_mgr.list_by_filter(
            execution_context=execution_context,
            knowledge_base_uuid=knowledge_base_uuid,
            filter=filters,
            limit=limit,
            offset=offset,
        )

    async def get_file_stream(
        self,
        execution_context: ExecutionContext,
        storage_path: str,
    ) -> bytes:
        """Handle GET_KNOWLEDEGE_FILE_STREAM action.

        Uses the storage manager abstraction to load file content,
        regardless of the underlying storage provider.
        """
        # Validate storage_path to prevent path traversal
        decoded_path = unquote(storage_path).replace('\\', '/')
        decoded_segments = decoded_path.split('/')
        normalized = posixpath.normpath(decoded_path)
        if (
            not storage_path
            or '\x00' in decoded_path
            or normalized.startswith('/')
            or '..' in decoded_segments
            or '..' in normalized.split('/')
            or re.match(r'^[A-Za-z]:/', normalized)
        ):
            raise ValueError('Invalid storage path')
        await self._validate_execution_context(execution_context)
        result = await self.ap.persistence_mgr.execute_async(
            sqlalchemy.select(persistence_rag.File.uuid)
            .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
            .where(persistence_rag.File.file_name == normalized)
            .limit(1)
        )
        if result.first() is None:
            raise WorkspaceNotFoundError('Knowledge file not found')
        content_bytes = await self.ap.storage_mgr.load_scoped_object_key(
            execution_context,
            normalized,
            expected_owner_type='upload_document',
        )
        return content_bytes if content_bytes else b''