mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-09-01 15:17:15 +00:00
feat(tenancy): add Workspace multi-tenant foundation (#2353)
* Document multi-tenant workspace architecture * Add OSS and commercial workspace boundaries * docs: redesign multi-tenant workspace architecture * feat(tenancy): implement workspace isolation * docs(tenancy): record verification evidence * docs(tenancy): revise single-instance SaaS topology * docs(tenancy): refine architecture options * docs: finalize cloud v2 multi-tenant decisions * feat(tenancy): establish cloud isolation foundations * feat(tenancy): harden shared cloud runtime boundaries * docs(tenancy): record final isolation verification * fix(tenancy): close isolation and permission gaps * docs(tenancy): record final isolation verification * feat(tenancy): connect cloud workspace control plane * fix(build): install git for pinned SDK * docs(cloud): update control plane verification * chore: update multi-tenant SDK pin * fix(cloud): skip legacy model sync during startup * test(cloud): preserve minimal model manager fixtures * fix(cloud): preserve authenticated account context * fix(cloud): reuse authenticated account for user info * feat(cloud): complete Workspace settings navigation * test(web): cover Workspace dropdown menu * feat(web): place workspace controls in sidebar * refactor(web): streamline workspace controls * style(web): format workspace layout test * fix(cloud): surface runtime and workspace plan status * fix(plugin): keep runtime identity stable across restarts * fix(ui): widen and center workspace switcher * fix(ui): hide roles from workspace switcher * fix(ui): align workspace switcher with sidebar entries * feat(workspace): add in-product collaboration and direct Cloud launch * style: format collaboration changes * fix(workspace): bind collaboration APIs to tenant UoW * fix(cloud): preserve Core-owned collaboration state * test(cloud): require Space identity for invite registration * feat(cloud): complete secure invitation experience * style(web): format invitation flows * fix(cloud): recover box runtime without unscoped skill reload * feat(oss): enforce invitation account and owner billing flows * style: format OSS account service * test(oss): cover invitation logout handoff * fix(oss): resolve workspace owner in scoped session * feat(cloud): harden multi-tenant runtime resources * fix(cloud): bound runtime restart storms * fix(cloud): eliminate periodic runtime CPU spikes * fix(cloud): enforce instance capacity ceilings * fix(cloud): scope public login capability discovery * fix(cloud): bound tenant maintenance and monitoring work * fix(runtime): bound tenant resource amplification * fix(deps): pin green multi-tenant plugin SDK * fix(cloud): handle unavailable skill capability * fix(security): require authentication for image file endpoint (H-2) - Changed /api/v1/files/image from AuthType.NONE to USER_TOKEN_OR_API_KEY - Added Permission.RESOURCE_VIEW requirement - Prevents unauthenticated cross-tenant file access via leaked keys - Fixes HIGH severity finding from multi-tenant security review docs: add comprehensive database migration guide - Complete migration steps for OSS → multi-tenant - Backup, execution, verification procedures - Rollback scenarios and recovery plans - Performance tuning recommendations * test: add comprehensive cross-tenant isolation tests Added 7 critical test scenarios for multi-tenant boundaries: - Cross-tenant bot access prevention - Viewer role read-only enforcement - Removed member immediate access revocation - Model provider credential isolation - WebSocket message isolation - Invitation token workspace scoping - Multi-workspace context validation These tests address P0-2 coverage gaps for: - workspaces.py (membership & invitation flows) - user.py (authentication & authorization) - websocket_chat.py (real-time isolation) - plugins.py (resource access control) docs: finalize database migration guide * fix(security): resolve M-1, M-2, M-3 security findings M-1: WebSocket authorization TOCTOU race (FIXED) - Changed _revalidate_websocket_authorization to return RequestContext - Ensures validated context is used immediately without race window - Prevents removed members from sending messages during revalidation gap M-2: Model Manager cache workspace isolation (VERIFIED) - Confirmed _CacheKey already uses 4-tuple: (instance, workspace, generation, resource) - Cache is properly scoped per workspace, no cross-tenant leakage possible - No code change needed, documented as working correctly M-3: Invitation lock workspace scoping (FIXED) - Changed lock key from token_digest to workspace_uuid:token_digest - Prevents DoS where attacker locks token in Workspace A to block Workspace B - Locks now isolated per workspace All MEDIUM severity findings from security review now resolved. * fix(cloud): unblock tenant CI and enforce knowledge quotas * fix(tenancy): scope rerank model sync --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
RockChinQ
@@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
import abc
|
||||
|
||||
from langbot.pkg.core import app
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot_plugin.api.entities.builtin.rag import context as rag_context
|
||||
|
||||
|
||||
@@ -22,10 +23,16 @@ class KnowledgeBaseInterface(metaclass=abc.ABCMeta):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def retrieve(self, query: str, settings: dict | None = None) -> list[rag_context.RetrievalResultEntry]:
|
||||
async def retrieve(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
query: str,
|
||||
settings: dict | None = None,
|
||||
) -> list[rag_context.RetrievalResultEntry]:
|
||||
"""Retrieve relevant documents from the knowledge base
|
||||
|
||||
Args:
|
||||
execution_context: Trusted active Workspace placement.
|
||||
query: The query string
|
||||
settings: Optional per-request retrieval settings overrides
|
||||
|
||||
@@ -50,6 +57,6 @@ class KnowledgeBaseInterface(metaclass=abc.ABCMeta):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def dispose(self):
|
||||
async def dispose(self, execution_context: ExecutionContext):
|
||||
"""Clean up resources"""
|
||||
pass
|
||||
|
||||
@@ -1,48 +1,130 @@
|
||||
from __future__ import annotations
|
||||
import asyncio
|
||||
import io
|
||||
import mimetypes
|
||||
import os.path
|
||||
import traceback
|
||||
import uuid
|
||||
import zipfile
|
||||
import io
|
||||
from typing import Any
|
||||
from langbot.pkg.core import app
|
||||
|
||||
import sqlalchemy
|
||||
|
||||
|
||||
from langbot.pkg.entity.persistence import rag as persistence_rag
|
||||
from langbot.pkg.core import taskmgr
|
||||
from langbot_plugin.api.entities.builtin.rag import context as rag_context
|
||||
|
||||
from langbot.pkg.api.http.authz import WorkspaceRequiredError
|
||||
from langbot.pkg.api.http.context import ExecutionContext, RequestContext
|
||||
from langbot.pkg.api.http.service.tenant import TenantContext, require_workspace_uuid
|
||||
from langbot.pkg.core import app, taskmgr
|
||||
from langbot.pkg.core.task_boundary import run_in_workspace_uow
|
||||
from langbot.pkg.entity.persistence import rag as persistence_rag
|
||||
from langbot.pkg.workspace.errors import WorkspaceInvariantError, WorkspaceNotFoundError
|
||||
|
||||
from .base import KnowledgeBaseInterface
|
||||
|
||||
|
||||
_MAX_ZIP_ARCHIVE_ENTRIES = 1024
|
||||
_MAX_ZIP_DOCUMENTS = 8
|
||||
_MAX_ZIP_FILE_BYTES = 10 * 1024 * 1024
|
||||
_MAX_ZIP_UNCOMPRESSED_BYTES = 40 * 1024 * 1024
|
||||
_MAX_ZIP_COMPRESSION_RATIO = 100
|
||||
|
||||
|
||||
class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
ap: app.Application
|
||||
|
||||
knowledge_base_entity: persistence_rag.KnowledgeBase
|
||||
|
||||
def __init__(self, ap: app.Application, knowledge_base_entity: persistence_rag.KnowledgeBase):
|
||||
def __init__(
|
||||
self,
|
||||
ap: app.Application,
|
||||
knowledge_base_entity: persistence_rag.KnowledgeBase,
|
||||
execution_context: ExecutionContext,
|
||||
):
|
||||
super().__init__(ap)
|
||||
self.knowledge_base_entity = knowledge_base_entity
|
||||
self.execution_context = execution_context
|
||||
|
||||
async def initialize(self):
|
||||
pass
|
||||
|
||||
async def _assert_execution_context(self, execution_context: ExecutionContext) -> None:
|
||||
"""Reject stale or cross-Workspace runtime access."""
|
||||
|
||||
if not isinstance(execution_context, ExecutionContext):
|
||||
raise WorkspaceRequiredError('ExecutionContext is required for knowledge runtime access')
|
||||
if (
|
||||
execution_context.instance_uuid != self.execution_context.instance_uuid
|
||||
or execution_context.workspace_uuid != self.execution_context.workspace_uuid
|
||||
or execution_context.placement_generation != self.execution_context.placement_generation
|
||||
):
|
||||
raise WorkspaceNotFoundError('Knowledge base not found')
|
||||
if self.knowledge_base_entity.workspace_uuid != execution_context.workspace_uuid:
|
||||
raise WorkspaceNotFoundError('Knowledge base not found')
|
||||
binding = await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
if binding.instance_uuid != execution_context.instance_uuid:
|
||||
raise WorkspaceNotFoundError('Knowledge base not found')
|
||||
|
||||
async def _require_plugin_runtime_context(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
) -> ExecutionContext:
|
||||
"""Fence every singleton Plugin Runtime call to this runtime KB."""
|
||||
|
||||
await self._assert_execution_context(execution_context)
|
||||
return await self.ap.plugin_connector.require_workspace_context(execution_context)
|
||||
|
||||
def _require_upload_object_key(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
object_key: str,
|
||||
) -> None:
|
||||
"""Reject raw, cross-Workspace, stale, or non-upload object keys."""
|
||||
|
||||
try:
|
||||
self.ap.storage_mgr.require_scoped_object_key(
|
||||
execution_context,
|
||||
object_key,
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
except (WorkspaceRequiredError, ValueError) as exc:
|
||||
raise WorkspaceNotFoundError('Upload not found') from exc
|
||||
|
||||
async def _store_file_task(
|
||||
self, file: persistence_rag.File, task_context: taskmgr.TaskContext, parser_plugin_id: str | None = None
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
file: persistence_rag.File,
|
||||
task_context: taskmgr.TaskContext,
|
||||
parser_plugin_id: str | None = None,
|
||||
):
|
||||
await run_in_workspace_uow(
|
||||
self.ap,
|
||||
execution_context.workspace_uuid,
|
||||
lambda: self._assert_execution_context(execution_context),
|
||||
)
|
||||
self._require_upload_object_key(execution_context, file.file_name)
|
||||
try:
|
||||
# set file status to processing
|
||||
await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.update(persistence_rag.File)
|
||||
.where(persistence_rag.File.uuid == file.uuid)
|
||||
.values(status='processing')
|
||||
)
|
||||
status_visible = False
|
||||
for retry_delay in (0.0, 0.01, 0.05, 0.1):
|
||||
if retry_delay:
|
||||
await asyncio.sleep(retry_delay)
|
||||
if await self._set_file_status(execution_context, file.uuid, 'processing'):
|
||||
status_visible = True
|
||||
break
|
||||
if not status_visible:
|
||||
raise WorkspaceNotFoundError('Knowledge file was not committed before its background task started')
|
||||
|
||||
task_context.set_current_action('Processing file')
|
||||
|
||||
# Get file size from storage
|
||||
file_size = await self.ap.storage_mgr.storage_provider.size(file.file_name)
|
||||
file_size = await self.ap.storage_mgr.size_scoped_object_key(
|
||||
execution_context,
|
||||
file.file_name,
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
|
||||
# Detect MIME type from extension
|
||||
mime_type, _ = mimetypes.guess_type(file.file_name)
|
||||
@@ -53,16 +135,22 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
parsed_content = None
|
||||
if parser_plugin_id:
|
||||
task_context.set_current_action('Parsing file')
|
||||
file_bytes = await self.ap.storage_mgr.storage_provider.load(file.file_name)
|
||||
file_bytes = await self.ap.storage_mgr.load_scoped_object_key(
|
||||
execution_context,
|
||||
file.file_name,
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
parse_context = {
|
||||
'mime_type': mime_type,
|
||||
'filename': file.file_name,
|
||||
'metadata': {},
|
||||
}
|
||||
await self._require_plugin_runtime_context(execution_context)
|
||||
parsed_content = await self.ap.plugin_connector.call_parser(parser_plugin_id, parse_context, file_bytes)
|
||||
|
||||
# Call plugin to ingest document
|
||||
result = await self._ingest_document(
|
||||
execution_context,
|
||||
{
|
||||
'document_id': file.uuid,
|
||||
'filename': file.file_name,
|
||||
@@ -80,44 +168,94 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
raise Exception(error_msg)
|
||||
|
||||
# set file status to completed
|
||||
await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.update(persistence_rag.File)
|
||||
.where(persistence_rag.File.uuid == file.uuid)
|
||||
.values(status='completed')
|
||||
)
|
||||
if not await self._set_file_status(execution_context, file.uuid, 'completed'):
|
||||
raise WorkspaceNotFoundError('Knowledge file not found')
|
||||
|
||||
except Exception as e:
|
||||
self.ap.logger.error(f'Error storing file {file.uuid}: {e}')
|
||||
traceback.print_exc()
|
||||
# set file status to failed
|
||||
await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.update(persistence_rag.File)
|
||||
.where(persistence_rag.File.uuid == file.uuid)
|
||||
.values(status='failed')
|
||||
)
|
||||
# A stale placement is fenced from all writes, including failure
|
||||
# status updates from an old background task.
|
||||
try:
|
||||
if not await self._set_file_status(execution_context, file.uuid, 'failed'):
|
||||
raise WorkspaceNotFoundError('Knowledge file not found')
|
||||
except Exception:
|
||||
self.ap.logger.warning(f'Skipping stale RAG task status update for file {file.uuid}')
|
||||
|
||||
raise
|
||||
finally:
|
||||
# delete file from storage
|
||||
await self.ap.storage_mgr.storage_provider.delete(file.file_name)
|
||||
# An old background task must not touch an upload after its
|
||||
# placement generation has been fenced off.
|
||||
try:
|
||||
await self._assert_execution_context(execution_context)
|
||||
await self.ap.storage_mgr.delete_scoped_object_key(
|
||||
execution_context,
|
||||
file.file_name,
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
except (WorkspaceRequiredError, WorkspaceNotFoundError):
|
||||
self.ap.logger.warning(f'Skipping stale RAG upload cleanup for file {file.uuid}')
|
||||
|
||||
async def store_file(self, file_id: str, parser_plugin_id: str | None = None) -> str:
|
||||
async def _set_file_status(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
file_uuid: str,
|
||||
status: str,
|
||||
) -> bool:
|
||||
"""Commit one detached-task status transition in its own tenant UoW."""
|
||||
|
||||
async def update() -> bool:
|
||||
await self._assert_execution_context(execution_context)
|
||||
result = await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.update(persistence_rag.File)
|
||||
.where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
|
||||
.where(persistence_rag.File.uuid == file_uuid)
|
||||
.values(status=status)
|
||||
)
|
||||
return getattr(result, 'rowcount', 0) > 0
|
||||
|
||||
persistence_mgr = self.ap.persistence_mgr
|
||||
managed_mode = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) in {
|
||||
'cloud_runtime',
|
||||
'oss_compat',
|
||||
}
|
||||
tenant_uow = getattr(persistence_mgr, 'tenant_uow', None)
|
||||
if managed_mode:
|
||||
if not callable(tenant_uow):
|
||||
raise RuntimeError('Knowledge tasks require an explicit tenant UoW')
|
||||
async with tenant_uow(execution_context.workspace_uuid):
|
||||
return await update()
|
||||
return await update()
|
||||
|
||||
async def store_file(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
file_id: str,
|
||||
parser_plugin_id: str | None = None,
|
||||
) -> str:
|
||||
await self._assert_execution_context(execution_context)
|
||||
self._require_upload_object_key(execution_context, file_id)
|
||||
# pre checking
|
||||
if not await self.ap.storage_mgr.storage_provider.exists(file_id):
|
||||
raise Exception(f'File {file_id} not found')
|
||||
if not await self.ap.storage_mgr.exists_scoped_object_key(
|
||||
execution_context,
|
||||
file_id,
|
||||
expected_owner_type='upload_document',
|
||||
):
|
||||
raise WorkspaceNotFoundError('Upload not found')
|
||||
|
||||
file_name = file_id
|
||||
_, ext = os.path.splitext(file_name)
|
||||
extension = ext.lstrip('.').lower() if ext else ''
|
||||
|
||||
if extension == 'zip':
|
||||
return await self._store_zip_file(file_id, parser_plugin_id=parser_plugin_id)
|
||||
return await self._store_zip_file(execution_context, file_id, parser_plugin_id=parser_plugin_id)
|
||||
|
||||
file_uuid = str(uuid.uuid4())
|
||||
kb_id = self.knowledge_base_entity.uuid
|
||||
|
||||
file_obj_data = {
|
||||
'uuid': file_uuid,
|
||||
'workspace_uuid': execution_context.workspace_uuid,
|
||||
'kb_id': kb_id,
|
||||
'file_name': file_name,
|
||||
'extension': extension,
|
||||
@@ -130,20 +268,47 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
|
||||
# run background task asynchronously
|
||||
ctx = taskmgr.TaskContext.new()
|
||||
wrapper = self.ap.task_mgr.create_user_task(
|
||||
self._store_file_task(file_obj, task_context=ctx, parser_plugin_id=parser_plugin_id),
|
||||
kind='knowledge-operation',
|
||||
name=f'knowledge-store-file-{file_id}',
|
||||
label=f'Store file {file_id}',
|
||||
context=ctx,
|
||||
)
|
||||
try:
|
||||
wrapper = self.ap.task_mgr.create_user_task(
|
||||
self._store_file_task(
|
||||
execution_context,
|
||||
file_obj,
|
||||
task_context=ctx,
|
||||
parser_plugin_id=parser_plugin_id,
|
||||
),
|
||||
kind='knowledge-operation',
|
||||
name=f'knowledge-store-file-{file_id}',
|
||||
label=f'Store file {file_id}',
|
||||
context=ctx,
|
||||
instance_uuid=execution_context.instance_uuid,
|
||||
workspace_uuid=execution_context.workspace_uuid,
|
||||
placement_generation=execution_context.placement_generation,
|
||||
)
|
||||
except taskmgr.TaskCapacityError:
|
||||
await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.delete(persistence_rag.File)
|
||||
.where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
|
||||
.where(persistence_rag.File.uuid == file_uuid)
|
||||
)
|
||||
raise
|
||||
return wrapper.id
|
||||
|
||||
async def _store_zip_file(self, zip_file_id: str, parser_plugin_id: str | None = None) -> str:
|
||||
async def _store_zip_file(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
zip_file_id: str,
|
||||
parser_plugin_id: str | None = None,
|
||||
) -> str:
|
||||
"""Handle ZIP file by extracting each document and storing them separately."""
|
||||
await self._assert_execution_context(execution_context)
|
||||
self._require_upload_object_key(execution_context, zip_file_id)
|
||||
self.ap.logger.info(f'Processing ZIP file: {zip_file_id}')
|
||||
|
||||
zip_bytes = await self.ap.storage_mgr.storage_provider.load(zip_file_id)
|
||||
zip_bytes = await self.ap.storage_mgr.load_scoped_object_key(
|
||||
execution_context,
|
||||
zip_file_id,
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
|
||||
supported_extensions = {'txt', 'pdf', 'docx', 'md', 'html'}
|
||||
stored_file_tasks = []
|
||||
@@ -151,9 +316,21 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
try:
|
||||
# use utf-8 encoding
|
||||
with zipfile.ZipFile(io.BytesIO(zip_bytes), 'r', metadata_encoding='utf-8') as zip_ref:
|
||||
if len(zip_ref.filelist) > _MAX_ZIP_ARCHIVE_ENTRIES:
|
||||
raise ValueError('ZIP archive contains too many entries')
|
||||
|
||||
supported_files: list[zipfile.ZipInfo] = []
|
||||
total_uncompressed_bytes = 0
|
||||
for file_info in zip_ref.filelist:
|
||||
# skip directories and hidden files
|
||||
if file_info.is_dir() or file_info.filename.startswith('.'):
|
||||
normalized_name = file_info.filename.replace('\\', '/').strip('/')
|
||||
path_parts = normalized_name.split('/')
|
||||
if (
|
||||
file_info.is_dir()
|
||||
or not normalized_name
|
||||
or any(part.startswith('.') for part in path_parts)
|
||||
or '__MACOSX' in path_parts
|
||||
):
|
||||
continue
|
||||
|
||||
_, file_ext = os.path.splitext(file_info.filename)
|
||||
@@ -161,29 +338,60 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
if file_extension not in supported_extensions:
|
||||
self.ap.logger.debug(f'Skipping unsupported file in ZIP: {file_info.filename}')
|
||||
continue
|
||||
if file_info.flag_bits & 0x1:
|
||||
raise ValueError('Encrypted ZIP entries are not supported')
|
||||
if file_info.file_size > _MAX_ZIP_FILE_BYTES:
|
||||
raise ValueError(f'ZIP document exceeds the file size limit: {file_info.filename}')
|
||||
if (
|
||||
file_info.file_size
|
||||
and file_info.file_size > max(file_info.compress_size, 1) * _MAX_ZIP_COMPRESSION_RATIO
|
||||
):
|
||||
raise ValueError(f'ZIP document exceeds the compression-ratio limit: {file_info.filename}')
|
||||
total_uncompressed_bytes += file_info.file_size
|
||||
if total_uncompressed_bytes > _MAX_ZIP_UNCOMPRESSED_BYTES:
|
||||
raise ValueError('ZIP documents exceed the uncompressed size limit')
|
||||
supported_files.append(file_info)
|
||||
if len(supported_files) > _MAX_ZIP_DOCUMENTS:
|
||||
raise ValueError('ZIP archive contains too many supported documents')
|
||||
|
||||
for file_info in supported_files:
|
||||
try:
|
||||
file_content = zip_ref.read(file_info.filename)
|
||||
file_content = await asyncio.to_thread(zip_ref.read, file_info)
|
||||
|
||||
base_name = file_info.filename.replace('/', '_').replace('\\', '_')
|
||||
file_stem, file_ext = os.path.splitext(base_name)
|
||||
extension = file_ext.lstrip('.')
|
||||
|
||||
if file_stem.startswith('__MACOSX'):
|
||||
continue
|
||||
|
||||
extracted_file_id = file_stem + '_' + str(uuid.uuid4())[:8] + '.' + extension
|
||||
# save file to storage
|
||||
extracted_object_key = await self.ap.storage_mgr.save_scoped(
|
||||
execution_context,
|
||||
owner_type='upload_document',
|
||||
owner=f'knowledge-base:{self.knowledge_base_entity.uuid}',
|
||||
key=extracted_file_id,
|
||||
value=file_content,
|
||||
)
|
||||
|
||||
await self.ap.storage_mgr.storage_provider.save(extracted_file_id, file_content)
|
||||
|
||||
task_id = await self.store_file(extracted_file_id, parser_plugin_id=parser_plugin_id)
|
||||
try:
|
||||
task_id = await self.store_file(
|
||||
execution_context,
|
||||
extracted_object_key,
|
||||
parser_plugin_id=parser_plugin_id,
|
||||
)
|
||||
except Exception:
|
||||
await self.ap.storage_mgr.delete_scoped_object_key(
|
||||
execution_context,
|
||||
extracted_object_key,
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
raise
|
||||
stored_file_tasks.append(task_id)
|
||||
|
||||
self.ap.logger.info(
|
||||
f'Extracted and stored file from ZIP: {file_info.filename} -> {extracted_file_id}'
|
||||
f'Extracted and stored file from ZIP: {file_info.filename} -> {extracted_object_key}'
|
||||
)
|
||||
|
||||
except taskmgr.TaskCapacityError:
|
||||
raise
|
||||
except Exception as e:
|
||||
self.ap.logger.warning(f'Failed to extract file {file_info.filename} from ZIP: {e}')
|
||||
continue
|
||||
@@ -197,20 +405,33 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
return stored_file_tasks[0] if stored_file_tasks else ''
|
||||
finally:
|
||||
try:
|
||||
await self.ap.storage_mgr.storage_provider.delete(zip_file_id)
|
||||
await self._assert_execution_context(execution_context)
|
||||
await self.ap.storage_mgr.delete_scoped_object_key(
|
||||
execution_context,
|
||||
zip_file_id,
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except (WorkspaceRequiredError, WorkspaceNotFoundError):
|
||||
self.ap.logger.warning(f'Skipping stale RAG ZIP cleanup for upload {zip_file_id}')
|
||||
except Exception as e:
|
||||
self.ap.logger.warning(f'Failed to cleanup ZIP file {zip_file_id}: {e}')
|
||||
|
||||
async def retrieve(self, query: str, settings: dict | None = None) -> list[rag_context.RetrievalResultEntry]:
|
||||
async def retrieve(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
query: str,
|
||||
settings: dict | None = None,
|
||||
) -> list[rag_context.RetrievalResultEntry]:
|
||||
await self._assert_execution_context(execution_context)
|
||||
# Merge stored retrieval_settings with per-request overrides
|
||||
stored = self.knowledge_base_entity.retrieval_settings or {}
|
||||
merged = {**stored, **(settings or {})}
|
||||
if 'top_k' not in merged:
|
||||
merged['top_k'] = 5 # fallback default
|
||||
|
||||
response = await self._retrieve(query, merged)
|
||||
response = await self._retrieve(execution_context, query, merged)
|
||||
|
||||
results_data = response.get('results', [])
|
||||
entries = []
|
||||
@@ -221,12 +442,25 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
entries.append(r)
|
||||
return entries
|
||||
|
||||
async def delete_file(self, file_id: str):
|
||||
await self._delete_document(file_id)
|
||||
async def delete_file(self, execution_context: ExecutionContext, file_id: str):
|
||||
await self._assert_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.kb_id == self.knowledge_base_entity.uuid)
|
||||
.where(persistence_rag.File.uuid == file_id)
|
||||
.limit(1)
|
||||
)
|
||||
if result.first() is None:
|
||||
raise WorkspaceNotFoundError('Knowledge file not found')
|
||||
await self._delete_document(execution_context, file_id)
|
||||
|
||||
# Also cleanup DB record
|
||||
await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.delete(persistence_rag.File).where(persistence_rag.File.uuid == file_id)
|
||||
sqlalchemy.delete(persistence_rag.File)
|
||||
.where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
|
||||
.where(persistence_rag.File.kb_id == self.knowledge_base_entity.uuid)
|
||||
.where(persistence_rag.File.uuid == file_id)
|
||||
)
|
||||
|
||||
def get_uuid(self) -> str:
|
||||
@@ -241,14 +475,16 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
"""Get the Knowledge Engine plugin ID"""
|
||||
return self.knowledge_base_entity.knowledge_engine_plugin_id or ''
|
||||
|
||||
async def dispose(self):
|
||||
async def dispose(self, execution_context: ExecutionContext):
|
||||
"""Dispose the knowledge base, notifying the plugin to cleanup."""
|
||||
await self._on_kb_delete()
|
||||
await self._assert_execution_context(execution_context)
|
||||
await self._on_kb_delete(execution_context)
|
||||
|
||||
# ========== Plugin Communication Methods ==========
|
||||
|
||||
async def _on_kb_create(self) -> None:
|
||||
async def _on_kb_create(self, execution_context: ExecutionContext) -> None:
|
||||
"""Notify plugin about KB creation."""
|
||||
await self._assert_execution_context(execution_context)
|
||||
plugin_id = self.knowledge_base_entity.knowledge_engine_plugin_id
|
||||
if not plugin_id:
|
||||
return
|
||||
@@ -258,17 +494,20 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
self.ap.logger.info(
|
||||
f'Calling RAG plugin {plugin_id}: on_knowledge_base_create(kb_id={self.knowledge_base_entity.uuid})'
|
||||
)
|
||||
await self._require_plugin_runtime_context(execution_context)
|
||||
await self.ap.plugin_connector.rag_on_kb_create(plugin_id, self.knowledge_base_entity.uuid, config)
|
||||
except Exception as e:
|
||||
self.ap.logger.error(f'Failed to notify plugin {plugin_id} on KB create: {e}')
|
||||
raise
|
||||
|
||||
async def _on_kb_delete(self) -> None:
|
||||
async def _on_kb_delete(self, execution_context: ExecutionContext) -> None:
|
||||
"""Notify plugin about KB deletion."""
|
||||
await self._assert_execution_context(execution_context)
|
||||
plugin_id = self.knowledge_base_entity.knowledge_engine_plugin_id
|
||||
if not plugin_id:
|
||||
return
|
||||
|
||||
await self._require_plugin_runtime_context(execution_context)
|
||||
try:
|
||||
self.ap.logger.info(
|
||||
f'Calling RAG plugin {plugin_id}: on_knowledge_base_delete(kb_id={self.knowledge_base_entity.uuid})'
|
||||
@@ -279,11 +518,13 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
|
||||
async def _ingest_document(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
file_metadata: dict[str, Any],
|
||||
storage_path: str,
|
||||
parsed_content: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Call plugin to ingest document."""
|
||||
await self._assert_execution_context(execution_context)
|
||||
kb = self.knowledge_base_entity
|
||||
plugin_id = kb.knowledge_engine_plugin_id
|
||||
if not plugin_id:
|
||||
@@ -306,6 +547,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
'parsed_content': parsed_content,
|
||||
}
|
||||
|
||||
await self._require_plugin_runtime_context(execution_context)
|
||||
try:
|
||||
result = await self.ap.plugin_connector.call_rag_ingest(plugin_id, context_data)
|
||||
return result
|
||||
@@ -315,6 +557,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
|
||||
async def _retrieve(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
query: str,
|
||||
settings: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
@@ -324,6 +567,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
ValueError: If no RAG plugin is configured for this KB.
|
||||
Exception: If the plugin retrieval call fails.
|
||||
"""
|
||||
await self._assert_execution_context(execution_context)
|
||||
kb = self.knowledge_base_entity
|
||||
plugin_id = kb.knowledge_engine_plugin_id
|
||||
if not plugin_id:
|
||||
@@ -333,25 +577,28 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
# for plugins that need it. Do NOT move them into filters, as filters
|
||||
# are passed directly to vector_search by some plugins (e.g. LangRAG)
|
||||
# and would cause empty results when the metadata field doesn't exist.
|
||||
filters = settings.pop('filters', {})
|
||||
plugin_settings = dict(settings)
|
||||
filters = plugin_settings.pop('filters', {})
|
||||
|
||||
retrieval_context = {
|
||||
'query': query,
|
||||
'knowledge_base_id': kb.uuid,
|
||||
'collection_id': kb.collection_id or kb.uuid,
|
||||
'retrieval_settings': settings,
|
||||
'retrieval_settings': plugin_settings,
|
||||
'creation_settings': kb.creation_settings or {},
|
||||
'filters': filters,
|
||||
}
|
||||
|
||||
await self._require_plugin_runtime_context(execution_context)
|
||||
result = await self.ap.plugin_connector.call_rag_retrieve(
|
||||
plugin_id,
|
||||
retrieval_context,
|
||||
)
|
||||
return result
|
||||
|
||||
async def _delete_document(self, document_id: str) -> bool:
|
||||
async def _delete_document(self, execution_context: ExecutionContext, document_id: str) -> bool:
|
||||
"""Call plugin to delete document."""
|
||||
await self._assert_execution_context(execution_context)
|
||||
kb = self.knowledge_base_entity
|
||||
plugin_id = kb.knowledge_engine_plugin_id
|
||||
if not plugin_id:
|
||||
@@ -359,6 +606,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
|
||||
self.ap.logger.info(f'Calling RAG plugin {plugin_id}: delete_document(doc_id={document_id})')
|
||||
|
||||
await self._require_plugin_runtime_context(execution_context)
|
||||
try:
|
||||
return await self.ap.plugin_connector.call_rag_delete_document(plugin_id, document_id, kb.uuid)
|
||||
except Exception as e:
|
||||
@@ -369,29 +617,115 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
|
||||
class RAGManager:
|
||||
ap: app.Application
|
||||
|
||||
knowledge_bases: dict[str, KnowledgeBaseInterface]
|
||||
knowledge_bases: dict[tuple[str, str], RuntimeKnowledgeBase]
|
||||
|
||||
def __init__(self, ap: app.Application):
|
||||
self.ap = ap
|
||||
self.knowledge_bases = {}
|
||||
self._scope_generations: dict[tuple[str, str], int] = {}
|
||||
self._knowledge_keys_by_scope: dict[
|
||||
tuple[str, str],
|
||||
set[tuple[str, str]],
|
||||
] = {}
|
||||
|
||||
def _cache_runtime(
|
||||
self,
|
||||
runtime: RuntimeKnowledgeBase,
|
||||
) -> None:
|
||||
context = runtime.execution_context
|
||||
self._observe_execution_context(context)
|
||||
key = (
|
||||
context.workspace_uuid,
|
||||
runtime.get_uuid(),
|
||||
)
|
||||
self.knowledge_bases[key] = runtime
|
||||
scope = (context.instance_uuid, context.workspace_uuid)
|
||||
self._knowledge_keys_by_scope.setdefault(scope, set()).add(key)
|
||||
|
||||
def _pop_runtime(
|
||||
self,
|
||||
context: ExecutionContext,
|
||||
kb_uuid: str,
|
||||
) -> RuntimeKnowledgeBase | None:
|
||||
key = (context.workspace_uuid, kb_uuid)
|
||||
runtime = self.knowledge_bases.pop(key, None)
|
||||
scope = (context.instance_uuid, context.workspace_uuid)
|
||||
keys = self._knowledge_keys_by_scope.get(scope)
|
||||
if keys is not None:
|
||||
keys.discard(key)
|
||||
if not keys:
|
||||
self._knowledge_keys_by_scope.pop(scope, None)
|
||||
self._scope_generations.pop(scope, None)
|
||||
return runtime
|
||||
|
||||
def _observe_execution_context(self, context: ExecutionContext) -> None:
|
||||
scope = (context.instance_uuid, context.workspace_uuid)
|
||||
previous_generation = self._scope_generations.get(scope)
|
||||
if previous_generation is not None and context.placement_generation < previous_generation:
|
||||
raise WorkspaceInvariantError('RAG runtime placement generation rolled back')
|
||||
if previous_generation == context.placement_generation:
|
||||
return
|
||||
if previous_generation is not None:
|
||||
for key in self._knowledge_keys_by_scope.pop(scope, ()):
|
||||
self.knowledge_bases.pop(key, None)
|
||||
self._scope_generations[scope] = context.placement_generation
|
||||
|
||||
async def initialize(self):
|
||||
await self.load_knowledge_bases_from_db()
|
||||
|
||||
async def get_all_knowledge_base_details(self) -> list[dict]:
|
||||
async def _to_execution_context(
|
||||
self,
|
||||
context: RequestContext | ExecutionContext,
|
||||
*,
|
||||
_binding_validated: bool = False,
|
||||
) -> ExecutionContext:
|
||||
if isinstance(context, RequestContext):
|
||||
execution_context = ExecutionContext.from_request(context)
|
||||
elif isinstance(context, ExecutionContext):
|
||||
execution_context = context
|
||||
else:
|
||||
raise WorkspaceRequiredError('RequestContext or ExecutionContext is required')
|
||||
|
||||
if not _binding_validated:
|
||||
binding = await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
if binding.instance_uuid != execution_context.instance_uuid:
|
||||
raise WorkspaceNotFoundError('Workspace not found')
|
||||
scope = (
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
)
|
||||
if scope in self._scope_generations:
|
||||
self._observe_execution_context(execution_context)
|
||||
return execution_context
|
||||
|
||||
async def _get_engine_map(self, context: TenantContext) -> dict[str, dict]:
|
||||
engine_map: dict[str, dict] = {}
|
||||
connector = getattr(self.ap, 'plugin_connector', None)
|
||||
if connector is not None and connector.is_enable_plugin:
|
||||
await connector.require_workspace_context(context)
|
||||
try:
|
||||
engines = await connector.list_knowledge_engines()
|
||||
engine_map = {engine['plugin_id']: engine for engine in engines}
|
||||
except Exception as e:
|
||||
self.ap.logger.warning(f'Failed to list Knowledge Engines: {e}')
|
||||
return engine_map
|
||||
|
||||
async def get_all_knowledge_base_details(self, context: TenantContext) -> list[dict]:
|
||||
"""Get all knowledge bases with enriched Knowledge Engine details."""
|
||||
workspace_uuid = require_workspace_uuid(context)
|
||||
# 1. Get raw KBs from DB
|
||||
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_rag.KnowledgeBase))
|
||||
result = await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.select(persistence_rag.KnowledgeBase).where(
|
||||
persistence_rag.KnowledgeBase.workspace_uuid == workspace_uuid
|
||||
)
|
||||
)
|
||||
knowledge_bases = result.all()
|
||||
|
||||
# 2. Get all available Knowledge Engines for enrichment
|
||||
engine_map = {}
|
||||
if self.ap.plugin_connector.is_enable_plugin:
|
||||
try:
|
||||
engines = await self.ap.plugin_connector.list_knowledge_engines()
|
||||
engine_map = {e['plugin_id']: e for e in engines}
|
||||
except Exception as e:
|
||||
self.ap.logger.warning(f'Failed to list Knowledge Engines: {e}')
|
||||
engine_map = await self._get_engine_map(context)
|
||||
|
||||
# 3. Serialize and enrich
|
||||
kb_list = []
|
||||
@@ -402,10 +736,13 @@ class RAGManager:
|
||||
|
||||
return kb_list
|
||||
|
||||
async def get_knowledge_base_details(self, kb_uuid: str) -> dict | None:
|
||||
async def get_knowledge_base_details(self, context: TenantContext, kb_uuid: str) -> dict | None:
|
||||
"""Get specific knowledge base with enriched Knowledge Engine details."""
|
||||
workspace_uuid = require_workspace_uuid(context)
|
||||
result = await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.select(persistence_rag.KnowledgeBase).where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
|
||||
sqlalchemy.select(persistence_rag.KnowledgeBase)
|
||||
.where(persistence_rag.KnowledgeBase.workspace_uuid == workspace_uuid)
|
||||
.where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
|
||||
)
|
||||
kb = result.first()
|
||||
if not kb:
|
||||
@@ -414,13 +751,7 @@ class RAGManager:
|
||||
kb_dict = self.ap.persistence_mgr.serialize_model(persistence_rag.KnowledgeBase, kb)
|
||||
|
||||
# Fetch engines
|
||||
engine_map = {}
|
||||
if self.ap.plugin_connector.is_enable_plugin:
|
||||
try:
|
||||
engines = await self.ap.plugin_connector.list_knowledge_engines()
|
||||
engine_map = {e['plugin_id']: e for e in engines}
|
||||
except Exception as e:
|
||||
self.ap.logger.warning(f'Failed to list Knowledge Engines: {e}')
|
||||
engine_map = await self._get_engine_map(context)
|
||||
|
||||
self._enrich_kb_dict(kb_dict, engine_map)
|
||||
return kb_dict
|
||||
@@ -465,6 +796,7 @@ class RAGManager:
|
||||
|
||||
async def create_knowledge_base(
|
||||
self,
|
||||
context: RequestContext | ExecutionContext,
|
||||
name: str,
|
||||
knowledge_engine_plugin_id: str,
|
||||
creation_settings: dict,
|
||||
@@ -472,8 +804,10 @@ class RAGManager:
|
||||
description: str = '',
|
||||
) -> persistence_rag.KnowledgeBase:
|
||||
"""Create a new knowledge base using a RAG plugin."""
|
||||
execution_context = await self._to_execution_context(context)
|
||||
# Validate that the Knowledge Engine plugin exists
|
||||
if self.ap.plugin_connector.is_enable_plugin:
|
||||
await self.ap.plugin_connector.require_workspace_context(execution_context)
|
||||
try:
|
||||
engines = await self.ap.plugin_connector.list_knowledge_engines()
|
||||
engine_ids = [e.get('plugin_id') for e in engines]
|
||||
@@ -490,6 +824,7 @@ class RAGManager:
|
||||
|
||||
kb_data = {
|
||||
'uuid': kb_uuid,
|
||||
'workspace_uuid': execution_context.workspace_uuid,
|
||||
'name': name,
|
||||
'description': description,
|
||||
'knowledge_engine_plugin_id': knowledge_engine_plugin_id,
|
||||
@@ -505,15 +840,17 @@ class RAGManager:
|
||||
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_rag.KnowledgeBase).values(kb_data))
|
||||
|
||||
# Load into Runtime
|
||||
runtime_kb = await self.load_knowledge_base(kb)
|
||||
runtime_kb = await self.load_knowledge_base(execution_context, kb)
|
||||
|
||||
# Notify Plugin — rollback DB record and runtime entry on failure
|
||||
try:
|
||||
await runtime_kb._on_kb_create()
|
||||
await runtime_kb._on_kb_create(execution_context)
|
||||
except Exception:
|
||||
self.knowledge_bases.pop(kb_uuid, None)
|
||||
self._pop_runtime(execution_context, kb_uuid)
|
||||
await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.delete(persistence_rag.KnowledgeBase).where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
|
||||
sqlalchemy.delete(persistence_rag.KnowledgeBase)
|
||||
.where(persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid)
|
||||
.where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
|
||||
)
|
||||
raise
|
||||
|
||||
@@ -524,14 +861,52 @@ class RAGManager:
|
||||
self.ap.logger.info('Loading knowledge bases from db...')
|
||||
|
||||
self.knowledge_bases = {}
|
||||
self._scope_generations = {}
|
||||
self._knowledge_keys_by_scope = {}
|
||||
|
||||
# Load knowledge bases
|
||||
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
|
||||
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
||||
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
|
||||
if cloud_runtime:
|
||||
if not callable(list_bindings) or not callable(tenant_uow):
|
||||
raise RuntimeError('Cloud knowledge loading requires explicit instance discovery and tenant UoWs')
|
||||
for binding in await list_bindings():
|
||||
async with tenant_uow(binding.workspace_uuid):
|
||||
result = await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.select(persistence_rag.KnowledgeBase)
|
||||
.where(persistence_rag.KnowledgeBase.workspace_uuid == binding.workspace_uuid)
|
||||
.order_by(persistence_rag.KnowledgeBase.uuid)
|
||||
)
|
||||
for knowledge_base in result.all():
|
||||
try:
|
||||
await self.load_knowledge_base(
|
||||
ExecutionContext(
|
||||
instance_uuid=binding.instance_uuid,
|
||||
workspace_uuid=binding.workspace_uuid,
|
||||
placement_generation=binding.placement_generation,
|
||||
),
|
||||
knowledge_base,
|
||||
_binding_validated=True,
|
||||
)
|
||||
except Exception as e:
|
||||
self.ap.logger.error(
|
||||
f'Error loading knowledge base {knowledge_base.uuid}: {e}\n{traceback.format_exc()}'
|
||||
)
|
||||
return
|
||||
|
||||
# Compatibility path for isolated manager tests and older embedders.
|
||||
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_rag.KnowledgeBase))
|
||||
knowledge_bases = result.all()
|
||||
|
||||
for knowledge_base in knowledge_bases:
|
||||
try:
|
||||
await self.load_knowledge_base(knowledge_base)
|
||||
binding = await self.ap.workspace_service.get_execution_binding(knowledge_base.workspace_uuid)
|
||||
execution_context = ExecutionContext(
|
||||
instance_uuid=binding.instance_uuid,
|
||||
workspace_uuid=binding.workspace_uuid,
|
||||
placement_generation=binding.placement_generation,
|
||||
)
|
||||
await self.load_knowledge_base(execution_context, knowledge_base)
|
||||
except Exception as e:
|
||||
self.ap.logger.error(
|
||||
f'Error loading knowledge base {knowledge_base.uuid}: {e}\n{traceback.format_exc()}'
|
||||
@@ -539,7 +914,10 @@ class RAGManager:
|
||||
|
||||
async def load_knowledge_base(
|
||||
self,
|
||||
context: RequestContext | ExecutionContext,
|
||||
knowledge_base_entity: persistence_rag.KnowledgeBase | sqlalchemy.Row | dict,
|
||||
*,
|
||||
_binding_validated: bool = False,
|
||||
) -> RuntimeKnowledgeBase:
|
||||
if isinstance(knowledge_base_entity, sqlalchemy.Row):
|
||||
# Safe access to _mapping for SQLAlchemy 1.4+
|
||||
@@ -551,23 +929,48 @@ class RAGManager:
|
||||
}
|
||||
knowledge_base_entity = persistence_rag.KnowledgeBase(**filtered_dict)
|
||||
|
||||
runtime_knowledge_base = RuntimeKnowledgeBase(ap=self.ap, knowledge_base_entity=knowledge_base_entity)
|
||||
execution_context = await self._to_execution_context(
|
||||
context,
|
||||
_binding_validated=_binding_validated,
|
||||
)
|
||||
if knowledge_base_entity.workspace_uuid != execution_context.workspace_uuid:
|
||||
raise WorkspaceNotFoundError('Knowledge base not found')
|
||||
runtime_knowledge_base = RuntimeKnowledgeBase(
|
||||
ap=self.ap,
|
||||
knowledge_base_entity=knowledge_base_entity,
|
||||
execution_context=execution_context,
|
||||
)
|
||||
|
||||
await runtime_knowledge_base.initialize()
|
||||
|
||||
self.knowledge_bases[runtime_knowledge_base.get_uuid()] = runtime_knowledge_base
|
||||
self._cache_runtime(runtime_knowledge_base)
|
||||
|
||||
return runtime_knowledge_base
|
||||
|
||||
async def get_knowledge_base_by_uuid(self, kb_uuid: str) -> KnowledgeBaseInterface | None:
|
||||
return self.knowledge_bases.get(kb_uuid)
|
||||
async def get_knowledge_base_by_uuid(
|
||||
self,
|
||||
context: RequestContext | ExecutionContext,
|
||||
kb_uuid: str,
|
||||
) -> RuntimeKnowledgeBase | None:
|
||||
execution_context = await self._to_execution_context(context)
|
||||
return self.knowledge_bases.get((execution_context.workspace_uuid, kb_uuid))
|
||||
|
||||
async def remove_knowledge_base_from_runtime(self, kb_uuid: str):
|
||||
self.knowledge_bases.pop(kb_uuid, None)
|
||||
async def remove_knowledge_base_from_runtime(
|
||||
self,
|
||||
context: RequestContext | ExecutionContext,
|
||||
kb_uuid: str,
|
||||
) -> None:
|
||||
execution_context = await self._to_execution_context(context)
|
||||
self._pop_runtime(execution_context, kb_uuid)
|
||||
|
||||
async def delete_knowledge_base(self, kb_uuid: str):
|
||||
kb = self.knowledge_bases.pop(kb_uuid, None)
|
||||
async def delete_knowledge_base(
|
||||
self,
|
||||
context: RequestContext | ExecutionContext,
|
||||
kb_uuid: str,
|
||||
) -> None:
|
||||
execution_context = await self._to_execution_context(context)
|
||||
kb = self._pop_runtime(execution_context, kb_uuid)
|
||||
if kb is not None:
|
||||
await kb.dispose()
|
||||
await kb.dispose(execution_context)
|
||||
else:
|
||||
self.ap.logger.warning(f'Knowledge base {kb_uuid} not found in runtime, skipping plugin notification')
|
||||
|
||||
@@ -5,6 +5,13 @@ 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
|
||||
|
||||
@@ -19,8 +26,54 @@ class RAGRuntimeService:
|
||||
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],
|
||||
@@ -28,9 +81,17 @@ class RAGRuntimeService:
|
||||
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(
|
||||
collection_name=collection_id,
|
||||
execution_context=execution_context,
|
||||
knowledge_base_uuid=knowledge_base_uuid,
|
||||
vectors=vectors,
|
||||
ids=ids,
|
||||
metadata=metadatas,
|
||||
@@ -39,6 +100,7 @@ class RAGRuntimeService:
|
||||
|
||||
async def vector_search(
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
collection_id: str,
|
||||
query_vector: list[float],
|
||||
top_k: int,
|
||||
@@ -48,8 +110,10 @@ class RAGRuntimeService:
|
||||
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(
|
||||
collection_name=collection_id,
|
||||
execution_context=execution_context,
|
||||
knowledge_base_uuid=knowledge_base_uuid,
|
||||
query_vector=query_vector,
|
||||
limit=top_k,
|
||||
filter=filters,
|
||||
@@ -59,7 +123,11 @@ class RAGRuntimeService:
|
||||
)
|
||||
|
||||
async def vector_delete(
|
||||
self, collection_id: str, file_ids: list[str] | None = None, filters: dict[str, Any] | None = None
|
||||
self,
|
||||
execution_context: ExecutionContext,
|
||||
collection_id: str,
|
||||
file_ids: list[str] | None = None,
|
||||
filters: dict[str, Any] | None = None,
|
||||
) -> int:
|
||||
"""Handle VECTOR_DELETE action.
|
||||
|
||||
@@ -73,16 +141,26 @@ class RAGRuntimeService:
|
||||
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(collection_name=collection_id, file_ids=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(collection_name=collection_id, filter=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,
|
||||
@@ -99,14 +177,20 @@ class RAGRuntimeService:
|
||||
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(
|
||||
collection_name=collection_id,
|
||||
execution_context=execution_context,
|
||||
knowledge_base_uuid=knowledge_base_uuid,
|
||||
filter=filters,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
async def get_file_stream(self, storage_path: str) -> bytes:
|
||||
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,
|
||||
@@ -125,5 +209,18 @@ class RAGRuntimeService:
|
||||
or re.match(r'^[A-Za-z]:/', normalized)
|
||||
):
|
||||
raise ValueError('Invalid storage path')
|
||||
content_bytes = await self.ap.storage_mgr.storage_provider.load(normalized)
|
||||
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''
|
||||
|
||||
Reference in New Issue
Block a user