mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +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
@@ -2,19 +2,42 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
import io
|
||||
import zipfile
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.core.taskmgr import TaskCapacityError
|
||||
from langbot.pkg.rag.knowledge.kbmgr import RuntimeKnowledgeBase
|
||||
from langbot.pkg.storage.mgr import StorageMgr
|
||||
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
|
||||
|
||||
|
||||
def _make_zip_bytes(entries: dict[str, bytes]) -> bytes:
|
||||
WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
|
||||
CONTEXT = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
placement_generation=2,
|
||||
)
|
||||
|
||||
|
||||
def _upload_key(logical_key: str, *, context: ExecutionContext = CONTEXT) -> str:
|
||||
return StorageMgr.scoped_object_key(
|
||||
context,
|
||||
owner_type='upload_document',
|
||||
owner='account:test',
|
||||
key=logical_key,
|
||||
)
|
||||
|
||||
|
||||
def _make_zip_bytes(entries: dict[str, bytes], *, compression: int = zipfile.ZIP_STORED) -> bytes:
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, 'w') as zf:
|
||||
with zipfile.ZipFile(buffer, 'w', compression=compression) as zf:
|
||||
for name, content in entries.items():
|
||||
zf.writestr(name, content)
|
||||
zf.mkdir('emptydir')
|
||||
@@ -25,26 +48,39 @@ def _make_app() -> Mock:
|
||||
app = Mock()
|
||||
app.logger = Mock()
|
||||
app.task_mgr = Mock()
|
||||
app.storage_mgr = Mock()
|
||||
app.storage_mgr.storage_provider = Mock()
|
||||
app.storage_mgr.storage_provider.exists = AsyncMock(return_value=True)
|
||||
app.storage_mgr.storage_provider.load = AsyncMock()
|
||||
app.storage_mgr.storage_provider.save = AsyncMock()
|
||||
app.storage_mgr.storage_provider.size = AsyncMock(return_value=123)
|
||||
app.storage_mgr.storage_provider.delete = AsyncMock()
|
||||
storage_mgr = StorageMgr(app)
|
||||
storage_mgr.storage_provider = Mock()
|
||||
storage_mgr.storage_provider.exists = AsyncMock(return_value=True)
|
||||
storage_mgr.storage_provider.load = AsyncMock()
|
||||
storage_mgr.storage_provider.load_bounded = AsyncMock()
|
||||
storage_mgr.storage_provider.save = AsyncMock()
|
||||
storage_mgr.storage_provider.size = AsyncMock(return_value=123)
|
||||
storage_mgr.storage_provider.delete = AsyncMock()
|
||||
app.storage_mgr = storage_mgr
|
||||
app.persistence_mgr = Mock()
|
||||
app.persistence_mgr.execute_async = AsyncMock()
|
||||
app.persistence_mgr.execute_async = AsyncMock(return_value=SimpleNamespace(rowcount=1))
|
||||
app.plugin_connector = Mock()
|
||||
app.plugin_connector.require_workspace_context = AsyncMock(side_effect=lambda context: context)
|
||||
app.workspace_service = SimpleNamespace(
|
||||
get_execution_binding=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
instance_uuid=CONTEXT.instance_uuid,
|
||||
workspace_uuid=CONTEXT.workspace_uuid,
|
||||
placement_generation=CONTEXT.placement_generation,
|
||||
)
|
||||
)
|
||||
)
|
||||
return app
|
||||
|
||||
|
||||
def _make_kb(plugin_id: str | None = 'author/engine') -> RuntimeKnowledgeBase:
|
||||
kb_entity = Mock()
|
||||
kb_entity.uuid = 'test-kb-uuid'
|
||||
kb_entity.workspace_uuid = WORKSPACE_A
|
||||
kb_entity.collection_id = 'test-collection'
|
||||
kb_entity.creation_settings = {}
|
||||
kb_entity.knowledge_engine_plugin_id = plugin_id
|
||||
return RuntimeKnowledgeBase(_make_app(), kb_entity)
|
||||
return RuntimeKnowledgeBase(_make_app(), kb_entity, CONTEXT)
|
||||
|
||||
|
||||
class TestStoreFile:
|
||||
@@ -58,33 +94,89 @@ class TestStoreFile:
|
||||
|
||||
kb.ap.task_mgr.create_user_task = Mock(side_effect=create_user_task)
|
||||
|
||||
task_id = await kb.store_file('documents/test.pdf')
|
||||
object_key = _upload_key('documents/test.pdf')
|
||||
task_id = await kb.store_file(CONTEXT, object_key)
|
||||
|
||||
assert task_id == 'task-1'
|
||||
kb.ap.storage_mgr.storage_provider.exists.assert_awaited_once_with('documents/test.pdf')
|
||||
kb.ap.storage_mgr.storage_provider.exists.assert_awaited_once_with(object_key)
|
||||
kb.ap.persistence_mgr.execute_async.assert_awaited_once()
|
||||
call_kwargs = kb.ap.task_mgr.create_user_task.call_args.kwargs
|
||||
assert call_kwargs['kind'] == 'knowledge-operation'
|
||||
assert call_kwargs['name'] == 'knowledge-store-file-documents/test.pdf'
|
||||
assert call_kwargs['label'] == 'Store file documents/test.pdf'
|
||||
assert call_kwargs['name'] == f'knowledge-store-file-{object_key}'
|
||||
assert call_kwargs['label'] == f'Store file {object_key}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_file_raises_when_source_file_missing(self):
|
||||
kb = _make_kb()
|
||||
kb.ap.storage_mgr.storage_provider.exists = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(Exception, match='File missing.pdf not found'):
|
||||
await kb.store_file('missing.pdf')
|
||||
object_key = _upload_key('missing.pdf')
|
||||
with pytest.raises(WorkspaceNotFoundError, match='Upload not found'):
|
||||
await kb.store_file(CONTEXT, object_key)
|
||||
|
||||
kb.ap.persistence_mgr.execute_async.assert_not_awaited()
|
||||
kb.ap.task_mgr.create_user_task.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_file_rejects_cross_workspace_upload_key(self):
|
||||
kb = _make_kb()
|
||||
other_context = ExecutionContext(
|
||||
instance_uuid=CONTEXT.instance_uuid,
|
||||
workspace_uuid='00000000-0000-0000-0000-00000000000b',
|
||||
placement_generation=CONTEXT.placement_generation,
|
||||
)
|
||||
|
||||
with pytest.raises(WorkspaceNotFoundError, match='Upload not found'):
|
||||
await kb.store_file(CONTEXT, _upload_key('stolen.pdf', context=other_context))
|
||||
|
||||
kb.ap.storage_mgr.storage_provider.exists.assert_not_awaited()
|
||||
kb.ap.persistence_mgr.execute_async.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_file_rejects_stale_generation_and_wrong_owner_type(self):
|
||||
kb = _make_kb()
|
||||
stale_context = ExecutionContext(
|
||||
instance_uuid=CONTEXT.instance_uuid,
|
||||
workspace_uuid=CONTEXT.workspace_uuid,
|
||||
placement_generation=CONTEXT.placement_generation + 1,
|
||||
)
|
||||
plugin_key = StorageMgr.scoped_object_key(
|
||||
CONTEXT,
|
||||
owner_type='plugin_config',
|
||||
owner='plugin:test',
|
||||
key='config.pdf',
|
||||
)
|
||||
|
||||
for object_key in (_upload_key('stale.pdf', context=stale_context), plugin_key, 'raw.pdf'):
|
||||
with pytest.raises(WorkspaceNotFoundError, match='Upload not found'):
|
||||
await kb.store_file(CONTEXT, object_key)
|
||||
|
||||
kb.ap.storage_mgr.storage_provider.exists.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_file_rolls_back_pending_record_when_task_capacity_is_exhausted(self):
|
||||
kb = _make_kb()
|
||||
object_key = _upload_key('queued.pdf')
|
||||
|
||||
def reject(coro, **_kwargs):
|
||||
coro.close()
|
||||
raise TaskCapacityError('capacity')
|
||||
|
||||
kb.ap.task_mgr.create_user_task.side_effect = reject
|
||||
|
||||
with pytest.raises(TaskCapacityError, match='capacity'):
|
||||
await kb.store_file(CONTEXT, object_key)
|
||||
|
||||
statements = [str(call.args[0]) for call in kb.ap.persistence_mgr.execute_async.await_args_list]
|
||||
assert any(statement.startswith('INSERT') for statement in statements)
|
||||
assert any(statement.startswith('DELETE') for statement in statements)
|
||||
|
||||
|
||||
class TestStoreZipFile:
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_zip_file_extracts_supported_files_and_skips_noise(self):
|
||||
kb = _make_kb()
|
||||
kb.ap.storage_mgr.storage_provider.load = AsyncMock(
|
||||
kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
|
||||
return_value=_make_zip_bytes(
|
||||
{
|
||||
'doc1.pdf': b'pdf',
|
||||
@@ -99,63 +191,131 @@ class TestStoreZipFile:
|
||||
)
|
||||
kb.store_file = AsyncMock(side_effect=['task-pdf', 'task-txt', 'task-md', 'task-html'])
|
||||
|
||||
task_id = await kb._store_zip_file('archive.zip', parser_plugin_id='parser/plugin')
|
||||
zip_key = _upload_key('archive.zip')
|
||||
task_id = await kb._store_zip_file(CONTEXT, zip_key, parser_plugin_id='parser/plugin')
|
||||
|
||||
assert task_id == 'task-pdf'
|
||||
assert kb.ap.storage_mgr.storage_provider.save.await_count == 4
|
||||
saved_names = [call.args[0] for call in kb.ap.storage_mgr.storage_provider.save.await_args_list]
|
||||
assert any(name.startswith('doc1_') and name.endswith('.pdf') for name in saved_names)
|
||||
assert any(name.startswith('doc2_') and name.endswith('.txt') for name in saved_names)
|
||||
assert any(name.startswith('subdir_doc3_') and name.endswith('.md') for name in saved_names)
|
||||
assert any(name.startswith('page_') and name.endswith('.html') for name in saved_names)
|
||||
assert not any('image' in name for name in saved_names)
|
||||
assert not any('hidden' in name for name in saved_names)
|
||||
assert not any('__MACOSX' in name for name in saved_names)
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with('archive.zip')
|
||||
assert {name.rsplit('.', 1)[-1] for name in saved_names} == {'pdf', 'txt', 'md', 'html'}
|
||||
for name in saved_names:
|
||||
StorageMgr.require_scoped_object_key(
|
||||
CONTEXT,
|
||||
name,
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
forwarded_keys = [call.args[1] for call in kb.store_file.await_args_list]
|
||||
assert forwarded_keys == saved_names
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(zip_key)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_zip_file_raises_when_no_supported_files(self):
|
||||
kb = _make_kb()
|
||||
kb.ap.storage_mgr.storage_provider.load = AsyncMock(
|
||||
kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
|
||||
return_value=_make_zip_bytes({'image.png': b'png', 'video.mp4': b'video'})
|
||||
)
|
||||
kb.store_file = AsyncMock()
|
||||
|
||||
with pytest.raises(Exception, match='No supported files found'):
|
||||
await kb._store_zip_file('archive.zip')
|
||||
await kb._store_zip_file(CONTEXT, _upload_key('archive.zip'))
|
||||
|
||||
kb.store_file.assert_not_awaited()
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with('archive.zip')
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip'))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_zip_file_rejects_too_many_documents_before_extracting(self):
|
||||
kb = _make_kb()
|
||||
kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
|
||||
return_value=_make_zip_bytes({f'doc-{index}.txt': b'text' for index in range(9)})
|
||||
)
|
||||
kb.store_file = AsyncMock()
|
||||
|
||||
with pytest.raises(ValueError, match='too many supported documents'):
|
||||
await kb._store_zip_file(CONTEXT, _upload_key('archive.zip'))
|
||||
|
||||
kb.store_file.assert_not_awaited()
|
||||
kb.ap.storage_mgr.storage_provider.save.assert_not_awaited()
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip'))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_zip_file_rejects_extreme_compression_ratio_before_extracting(self):
|
||||
kb = _make_kb()
|
||||
kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
|
||||
return_value=_make_zip_bytes(
|
||||
{'bomb.txt': b'A' * (1024 * 1024)},
|
||||
compression=zipfile.ZIP_DEFLATED,
|
||||
)
|
||||
)
|
||||
kb.store_file = AsyncMock()
|
||||
|
||||
with pytest.raises(ValueError, match='compression-ratio limit'):
|
||||
await kb._store_zip_file(CONTEXT, _upload_key('archive.zip'))
|
||||
|
||||
kb.store_file.assert_not_awaited()
|
||||
kb.ap.storage_mgr.storage_provider.save.assert_not_awaited()
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip'))
|
||||
|
||||
|
||||
class TestStoreFileTask:
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_file_task_opens_uow_before_first_database_helper(self):
|
||||
kb = _make_kb()
|
||||
active_workspace = contextvars.ContextVar('rag_task_workspace', default=None)
|
||||
observed = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def tenant_uow(workspace_uuid):
|
||||
token = active_workspace.set(workspace_uuid)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
active_workspace.reset(token)
|
||||
|
||||
kb.ap.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
|
||||
kb.ap.persistence_mgr.tenant_uow = tenant_uow
|
||||
|
||||
async def assert_execution_context(_context):
|
||||
observed.append(active_workspace.get())
|
||||
|
||||
kb._assert_execution_context = AsyncMock(side_effect=assert_execution_context)
|
||||
kb._set_file_status = AsyncMock(side_effect=[True, True])
|
||||
kb._ingest_document = AsyncMock(return_value={'status': 'completed'})
|
||||
object_key = _upload_key('scoped.pdf')
|
||||
file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
|
||||
|
||||
await kb._store_file_task(CONTEXT, file_obj, Mock())
|
||||
|
||||
assert observed[0] == WORKSPACE_A
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_file_task_marks_completed_and_cleans_storage(self):
|
||||
kb = _make_kb()
|
||||
kb._ingest_document = AsyncMock(return_value={'status': 'completed'})
|
||||
file_obj = SimpleNamespace(uuid='file-uuid', file_name='test.pdf', extension='pdf')
|
||||
object_key = _upload_key('test.pdf')
|
||||
file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
|
||||
task_context = Mock()
|
||||
|
||||
await kb._store_file_task(file_obj, task_context)
|
||||
await kb._store_file_task(CONTEXT, file_obj, task_context)
|
||||
|
||||
task_context.set_current_action.assert_called_once_with('Processing file')
|
||||
kb.ap.storage_mgr.storage_provider.size.assert_awaited_once_with('test.pdf')
|
||||
kb.ap.storage_mgr.storage_provider.size.assert_awaited_once_with(object_key)
|
||||
kb._ingest_document.assert_awaited_once()
|
||||
assert kb.ap.persistence_mgr.execute_async.await_count == 2
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with('test.pdf')
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(object_key)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_file_task_marks_failed_and_cleans_storage(self):
|
||||
kb = _make_kb()
|
||||
kb._ingest_document = AsyncMock(return_value={'status': 'failed', 'error_message': 'parser failed'})
|
||||
file_obj = SimpleNamespace(uuid='file-uuid', file_name='bad.pdf', extension='pdf')
|
||||
object_key = _upload_key('bad.pdf')
|
||||
file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
|
||||
task_context = Mock()
|
||||
|
||||
with pytest.raises(Exception, match='parser failed'):
|
||||
await kb._store_file_task(file_obj, task_context)
|
||||
await kb._store_file_task(CONTEXT, file_obj, task_context)
|
||||
|
||||
assert kb.ap.persistence_mgr.execute_async.await_count == 2
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with('bad.pdf')
|
||||
kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(object_key)
|
||||
|
||||
|
||||
class TestDeleteDocument:
|
||||
@@ -163,7 +323,7 @@ class TestDeleteDocument:
|
||||
async def test_delete_document_returns_false_when_no_plugin_id(self):
|
||||
kb = _make_kb(plugin_id=None)
|
||||
|
||||
result = await kb._delete_document('doc-id')
|
||||
result = await kb._delete_document(CONTEXT, 'doc-id')
|
||||
|
||||
assert result is False
|
||||
|
||||
@@ -172,7 +332,7 @@ class TestDeleteDocument:
|
||||
kb = _make_kb()
|
||||
kb.ap.plugin_connector.call_rag_delete_document = AsyncMock(return_value=True)
|
||||
|
||||
result = await kb._delete_document('doc-id')
|
||||
result = await kb._delete_document(CONTEXT, 'doc-id')
|
||||
|
||||
assert result is True
|
||||
kb.ap.plugin_connector.call_rag_delete_document.assert_awaited_once_with(
|
||||
@@ -184,7 +344,7 @@ class TestDeleteDocument:
|
||||
kb = _make_kb()
|
||||
kb.ap.plugin_connector.call_rag_delete_document = AsyncMock(side_effect=Exception('plugin error'))
|
||||
|
||||
result = await kb._delete_document('doc-id')
|
||||
result = await kb._delete_document(CONTEXT, 'doc-id')
|
||||
|
||||
assert result is False
|
||||
kb.ap.logger.error.assert_called_once()
|
||||
|
||||
+601
-492
File diff suppressed because it is too large
Load Diff
@@ -1,474 +1,467 @@
|
||||
"""Tests for RAGRuntimeService.
|
||||
|
||||
Tests the service that handles RAG-related requests from plugins,
|
||||
using mocked vector_db_mgr and storage_mgr.
|
||||
"""
|
||||
"""Tenant-aware tests for the plugin-facing RAG runtime service."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.utils.import_isolation import isolated_sys_modules
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
|
||||
|
||||
|
||||
class TestRAGRuntimeServiceVectorUpsert:
|
||||
"""Tests for vector_upsert method."""
|
||||
WORKSPACE_UUID = '00000000-0000-0000-0000-00000000000a'
|
||||
CONTEXT = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid=WORKSPACE_UUID,
|
||||
placement_generation=4,
|
||||
)
|
||||
|
||||
def _create_mock_app(self):
|
||||
"""Create mock app with vector_db_mgr and storage_mgr."""
|
||||
mock_app = MagicMock()
|
||||
mock_app.vector_db_mgr = MagicMock()
|
||||
mock_app.vector_db_mgr.upsert = AsyncMock()
|
||||
mock_app.storage_mgr = MagicMock()
|
||||
mock_app.storage_mgr.storage_provider = MagicMock()
|
||||
mock_app.storage_mgr.storage_provider.load = AsyncMock(return_value=b'content')
|
||||
return mock_app
|
||||
|
||||
def _make_rag_import_mocks(self):
|
||||
"""Create mocks needed for importing RAG service."""
|
||||
return {
|
||||
'langbot.pkg.core.app': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.rag': MagicMock(),
|
||||
}
|
||||
class _ScalarResult:
|
||||
def __init__(self, value):
|
||||
self.value = value
|
||||
|
||||
def scalar_one_or_none(self):
|
||||
return self.value
|
||||
|
||||
def first(self):
|
||||
return None if self.value is None else (self.value,)
|
||||
|
||||
|
||||
def _app(*, kb_uuid='kb-a', file_exists=True):
|
||||
persistence_results = [_ScalarResult(kb_uuid)]
|
||||
if file_exists is not None:
|
||||
persistence_results.append(_ScalarResult('file-a' if file_exists else None))
|
||||
return SimpleNamespace(
|
||||
workspace_service=SimpleNamespace(
|
||||
get_execution_binding=AsyncMock(return_value=SimpleNamespace(instance_uuid='instance-a'))
|
||||
),
|
||||
persistence_mgr=SimpleNamespace(execute_async=AsyncMock(side_effect=persistence_results)),
|
||||
vector_db_mgr=SimpleNamespace(
|
||||
upsert=AsyncMock(),
|
||||
search=AsyncMock(return_value=[{'id': 'chunk-a'}]),
|
||||
delete_by_file_id=AsyncMock(),
|
||||
delete_by_filter=AsyncMock(return_value=3),
|
||||
list_by_filter=AsyncMock(return_value=([{'id': 'chunk-a'}], 1)),
|
||||
),
|
||||
storage_mgr=SimpleNamespace(
|
||||
load_scoped_object_key=AsyncMock(return_value=b'content'),
|
||||
storage_provider=SimpleNamespace(load=AsyncMock(return_value=b'content')),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_upsert_resolves_canonical_kb_and_forwards_trusted_context():
|
||||
app = _app()
|
||||
service = RAGRuntimeService(app)
|
||||
|
||||
await service.vector_upsert(
|
||||
CONTEXT,
|
||||
'logical-collection',
|
||||
[[0.1, 0.2]],
|
||||
['chunk-a'],
|
||||
metadata=[{'file_id': 'file-a'}],
|
||||
documents=['hello'],
|
||||
)
|
||||
|
||||
app.vector_db_mgr.upsert.assert_awaited_once_with(
|
||||
execution_context=CONTEXT,
|
||||
knowledge_base_uuid='kb-a',
|
||||
vectors=[[0.1, 0.2]],
|
||||
ids=['chunk-a'],
|
||||
metadata=[{'file_id': 'file-a'}],
|
||||
documents=['hello'],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_upsert_rejects_mismatched_lengths():
|
||||
service = RAGRuntimeService(_app())
|
||||
with pytest.raises(ValueError, match='vectors and ids'):
|
||||
await service.vector_upsert(CONTEXT, 'kb-a', [[0.1]], ['a', 'b'])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_or_cross_workspace_collection_is_not_forwarded():
|
||||
app = _app(kb_uuid=None)
|
||||
service = RAGRuntimeService(app)
|
||||
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await service.vector_search(CONTEXT, 'kb-from-other-workspace', [0.1], 5)
|
||||
app.vector_db_mgr.search.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_search_forwards_all_search_options():
|
||||
app = _app()
|
||||
service = RAGRuntimeService(app)
|
||||
|
||||
result = await service.vector_search(
|
||||
CONTEXT,
|
||||
'kb-a',
|
||||
[0.1, 0.2],
|
||||
7,
|
||||
filters={'file_id': 'file-a'},
|
||||
search_type='hybrid',
|
||||
query_text='hello',
|
||||
vector_weight=0.7,
|
||||
)
|
||||
|
||||
assert result == [{'id': 'chunk-a'}]
|
||||
app.vector_db_mgr.search.assert_awaited_once_with(
|
||||
execution_context=CONTEXT,
|
||||
knowledge_base_uuid='kb-a',
|
||||
query_vector=[0.1, 0.2],
|
||||
limit=7,
|
||||
filter={'file_id': 'file-a'},
|
||||
search_type='hybrid',
|
||||
query_text='hello',
|
||||
vector_weight=0.7,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_delete_and_list_stay_on_canonical_kb():
|
||||
delete_app = _app()
|
||||
delete_service = RAGRuntimeService(delete_app)
|
||||
assert await delete_service.vector_delete(CONTEXT, 'logical', file_ids=['file-a']) == 1
|
||||
delete_app.vector_db_mgr.delete_by_file_id.assert_awaited_once_with(
|
||||
execution_context=CONTEXT,
|
||||
knowledge_base_uuid='kb-a',
|
||||
file_ids=['file-a'],
|
||||
)
|
||||
|
||||
filter_app = _app()
|
||||
filter_service = RAGRuntimeService(filter_app)
|
||||
assert await filter_service.vector_delete(CONTEXT, 'logical', filters={'page': 1}) == 3
|
||||
filter_app.vector_db_mgr.delete_by_filter.assert_awaited_once_with(
|
||||
execution_context=CONTEXT,
|
||||
knowledge_base_uuid='kb-a',
|
||||
filter={'page': 1},
|
||||
)
|
||||
|
||||
list_app = _app()
|
||||
list_service = RAGRuntimeService(list_app)
|
||||
assert await list_service.vector_list(CONTEXT, 'logical', {'page': 1}, 10, 2) == (
|
||||
[{'id': 'chunk-a'}],
|
||||
1,
|
||||
)
|
||||
list_app.vector_db_mgr.list_by_filter.assert_awaited_once_with(
|
||||
execution_context=CONTEXT,
|
||||
knowledge_base_uuid='kb-a',
|
||||
filter={'page': 1},
|
||||
limit=10,
|
||||
offset=2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_stream_requires_workspace_owned_file():
|
||||
app = _app(file_exists=True)
|
||||
app.persistence_mgr.execute_async = AsyncMock(return_value=_ScalarResult('file-a'))
|
||||
service = RAGRuntimeService(app)
|
||||
assert await service.get_file_stream(CONTEXT, 'nested/file.pdf') == b'content'
|
||||
app.storage_mgr.load_scoped_object_key.assert_awaited_once_with(
|
||||
CONTEXT,
|
||||
'nested/file.pdf',
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
|
||||
missing_app = _app(file_exists=False)
|
||||
missing_app.persistence_mgr.execute_async = AsyncMock(return_value=_ScalarResult(None))
|
||||
missing_service = RAGRuntimeService(missing_app)
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await missing_service.get_file_stream(CONTEXT, 'other.pdf')
|
||||
missing_app.storage_mgr.load_scoped_object_key.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'unsafe_path',
|
||||
[
|
||||
'',
|
||||
'../secret.txt',
|
||||
'/absolute/path.txt',
|
||||
'..\\secret.txt',
|
||||
'nested\\..\\secret.txt',
|
||||
'%2e%2e/secret.txt',
|
||||
'nested/%2e%2e/secret.txt',
|
||||
'C:\\secret.txt',
|
||||
'safe/\x00file.txt',
|
||||
],
|
||||
)
|
||||
async def test_file_stream_rejects_unsafe_paths_before_storage(unsafe_path):
|
||||
app = _app(file_exists=True)
|
||||
service = RAGRuntimeService(app)
|
||||
with pytest.raises(ValueError, match='Invalid storage path'):
|
||||
await service.get_file_stream(CONTEXT, unsafe_path)
|
||||
app.storage_mgr.load_scoped_object_key.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_rejects_wrong_instance_binding():
|
||||
app = _app()
|
||||
app.workspace_service.get_execution_binding.return_value = SimpleNamespace(instance_uuid='instance-b')
|
||||
service = RAGRuntimeService(app)
|
||||
|
||||
with pytest.raises(Exception, match='another LangBot instance'):
|
||||
await service.vector_search(CONTEXT, 'kb-a', [0.1], 1)
|
||||
app.persistence_mgr.execute_async.assert_not_awaited()
|
||||
|
||||
|
||||
# The following classes preserve the pre-tenancy regression scenarios. They
|
||||
# intentionally exercise the same inputs through the new trusted
|
||||
# ExecutionContext and assert the canonical Workspace-owned KB forwarded to the
|
||||
# vector layer.
|
||||
class TestRAGRuntimeServiceVectorUpsertRegression:
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_upsert_basic(self):
|
||||
"""Basic vector upsert delegates to vector_db_mgr."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
service = RAGRuntimeService(app)
|
||||
vectors = [[0.1, 0.2], [0.3, 0.4]]
|
||||
ids = ['id1', 'id2']
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
await service.vector_upsert(CONTEXT, 'test_collection', vectors, ids)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
vectors = [[0.1, 0.2], [0.3, 0.4]]
|
||||
ids = ['id1', 'id2']
|
||||
|
||||
await service.vector_upsert(
|
||||
collection_id='test_collection',
|
||||
vectors=vectors,
|
||||
ids=ids,
|
||||
)
|
||||
|
||||
mock_app.vector_db_mgr.upsert.assert_called_once()
|
||||
call_args = mock_app.vector_db_mgr.upsert.call_args
|
||||
assert call_args.kwargs['collection_name'] == 'test_collection'
|
||||
assert call_args.kwargs['vectors'] == vectors
|
||||
assert call_args.kwargs['ids'] == ids
|
||||
# Default metadata is empty dicts
|
||||
assert call_args.kwargs['metadata'] == [{} for _ in vectors]
|
||||
app.vector_db_mgr.upsert.assert_awaited_once_with(
|
||||
execution_context=CONTEXT,
|
||||
knowledge_base_uuid='kb-a',
|
||||
vectors=vectors,
|
||||
ids=ids,
|
||||
metadata=[{}, {}],
|
||||
documents=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_upsert_with_metadata(self):
|
||||
"""Vector upsert with provided metadata."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
metadata = [{'file_id': 'abc', 'page': 1}]
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
await RAGRuntimeService(app).vector_upsert(
|
||||
CONTEXT,
|
||||
'test',
|
||||
[[0.1, 0.2]],
|
||||
['id1'],
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
vectors = [[0.1, 0.2]]
|
||||
ids = ['id1']
|
||||
metadata = [{'file_id': 'abc', 'page': 1}]
|
||||
|
||||
await service.vector_upsert(
|
||||
collection_id='test',
|
||||
vectors=vectors,
|
||||
ids=ids,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
call_args = mock_app.vector_db_mgr.upsert.call_args
|
||||
assert call_args.kwargs['metadata'] == metadata
|
||||
assert app.vector_db_mgr.upsert.await_args.kwargs['metadata'] == metadata
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_upsert_with_documents(self):
|
||||
"""Vector upsert with documents for full-text search."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
documents = ['This is a test document']
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
vectors = [[0.1, 0.2]]
|
||||
ids = ['id1']
|
||||
documents = ['This is a test document']
|
||||
|
||||
await service.vector_upsert(
|
||||
collection_id='test',
|
||||
vectors=vectors,
|
||||
ids=ids,
|
||||
documents=documents,
|
||||
)
|
||||
|
||||
call_args = mock_app.vector_db_mgr.upsert.call_args
|
||||
assert call_args.kwargs['documents'] == documents
|
||||
|
||||
|
||||
class TestRAGRuntimeServiceVectorSearch:
|
||||
"""Tests for vector_search method."""
|
||||
|
||||
def _create_mock_app(self):
|
||||
"""Create mock app."""
|
||||
mock_app = MagicMock()
|
||||
mock_app.vector_db_mgr = MagicMock()
|
||||
mock_app.vector_db_mgr.search = AsyncMock(
|
||||
return_value=[
|
||||
{'id': 'id1', 'distance': 0.1, 'metadata': {'file_id': 'abc'}},
|
||||
{'id': 'id2', 'distance': 0.2, 'metadata': {'file_id': 'def'}},
|
||||
]
|
||||
await RAGRuntimeService(app).vector_upsert(
|
||||
CONTEXT,
|
||||
'test',
|
||||
[[0.1, 0.2]],
|
||||
['id1'],
|
||||
documents=documents,
|
||||
)
|
||||
return mock_app
|
||||
|
||||
def _make_rag_import_mocks(self):
|
||||
return {
|
||||
'langbot.pkg.core.app': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.rag': MagicMock(),
|
||||
}
|
||||
assert app.vector_db_mgr.upsert.await_args.kwargs['documents'] == documents
|
||||
|
||||
|
||||
class TestRAGRuntimeServiceVectorSearchRegression:
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_search_basic(self):
|
||||
"""Basic vector search delegates to vector_db_mgr."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
query_vector = [0.1, 0.2, 0.3]
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
result = await RAGRuntimeService(app).vector_search(
|
||||
CONTEXT,
|
||||
'test',
|
||||
query_vector,
|
||||
5,
|
||||
)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
query_vector = [0.1, 0.2, 0.3]
|
||||
|
||||
result = await service.vector_search(
|
||||
collection_id='test',
|
||||
query_vector=query_vector,
|
||||
top_k=5,
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
mock_app.vector_db_mgr.search.assert_called_once()
|
||||
call_args = mock_app.vector_db_mgr.search.call_args
|
||||
assert call_args.kwargs['collection_name'] == 'test'
|
||||
assert call_args.kwargs['query_vector'] == query_vector
|
||||
assert call_args.kwargs['limit'] == 5
|
||||
assert result == [{'id': 'chunk-a'}]
|
||||
app.vector_db_mgr.search.assert_awaited_once_with(
|
||||
execution_context=CONTEXT,
|
||||
knowledge_base_uuid='kb-a',
|
||||
query_vector=query_vector,
|
||||
limit=5,
|
||||
filter=None,
|
||||
search_type='vector',
|
||||
query_text='',
|
||||
vector_weight=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_search_with_filters(self):
|
||||
"""Vector search with metadata filters."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
filters = {'file_id': 'abc'}
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
await RAGRuntimeService(app).vector_search(
|
||||
CONTEXT,
|
||||
'test',
|
||||
[0.1, 0.2],
|
||||
10,
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
filters = {'file_id': 'abc'}
|
||||
|
||||
await service.vector_search(
|
||||
collection_id='test',
|
||||
query_vector=[0.1, 0.2],
|
||||
top_k=10,
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
call_args = mock_app.vector_db_mgr.search.call_args
|
||||
assert call_args.kwargs['filter'] == filters
|
||||
assert app.vector_db_mgr.search.await_args.kwargs['filter'] == filters
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_search_hybrid_mode(self):
|
||||
"""Vector search with hybrid search type."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
await RAGRuntimeService(app).vector_search(
|
||||
CONTEXT,
|
||||
'test',
|
||||
[0.1, 0.2],
|
||||
10,
|
||||
search_type='hybrid',
|
||||
query_text='search query',
|
||||
vector_weight=0.7,
|
||||
)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
await service.vector_search(
|
||||
collection_id='test',
|
||||
query_vector=[0.1, 0.2],
|
||||
top_k=10,
|
||||
search_type='hybrid',
|
||||
query_text='search query',
|
||||
vector_weight=0.7,
|
||||
)
|
||||
|
||||
call_args = mock_app.vector_db_mgr.search.call_args
|
||||
assert call_args.kwargs['search_type'] == 'hybrid'
|
||||
assert call_args.kwargs['query_text'] == 'search query'
|
||||
assert call_args.kwargs['vector_weight'] == 0.7
|
||||
kwargs = app.vector_db_mgr.search.await_args.kwargs
|
||||
assert kwargs['search_type'] == 'hybrid'
|
||||
assert kwargs['query_text'] == 'search query'
|
||||
assert kwargs['vector_weight'] == 0.7
|
||||
|
||||
|
||||
class TestRAGRuntimeServiceVectorDelete:
|
||||
"""Tests for vector_delete method."""
|
||||
|
||||
def _create_mock_app(self):
|
||||
mock_app = MagicMock()
|
||||
mock_app.vector_db_mgr = MagicMock()
|
||||
mock_app.vector_db_mgr.delete_by_file_id = AsyncMock()
|
||||
mock_app.vector_db_mgr.delete_by_filter = AsyncMock(return_value=5)
|
||||
return mock_app
|
||||
|
||||
def _make_rag_import_mocks(self):
|
||||
return {
|
||||
'langbot.pkg.core.app': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.rag': MagicMock(),
|
||||
}
|
||||
|
||||
class TestRAGRuntimeServiceVectorDeleteRegression:
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_delete_by_file_ids(self):
|
||||
"""Delete by file_ids delegates to delete_by_file_id."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
result = await RAGRuntimeService(app).vector_delete(
|
||||
CONTEXT,
|
||||
'test',
|
||||
file_ids=['file1', 'file2', 'file3'],
|
||||
)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
result = await service.vector_delete(
|
||||
collection_id='test',
|
||||
file_ids=['file1', 'file2', 'file3'],
|
||||
)
|
||||
|
||||
assert result == 3 # Returns count of file_ids
|
||||
mock_app.vector_db_mgr.delete_by_file_id.assert_called_once()
|
||||
call_args = mock_app.vector_db_mgr.delete_by_file_id.call_args
|
||||
assert call_args.kwargs['collection_name'] == 'test'
|
||||
assert call_args.kwargs['file_ids'] == ['file1', 'file2', 'file3']
|
||||
assert result == 3
|
||||
app.vector_db_mgr.delete_by_file_id.assert_awaited_once_with(
|
||||
execution_context=CONTEXT,
|
||||
knowledge_base_uuid='kb-a',
|
||||
file_ids=['file1', 'file2', 'file3'],
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_delete_by_filters(self):
|
||||
"""Delete by filters delegates to delete_by_filter."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
filters = {'status': 'deleted'}
|
||||
app.vector_db_mgr.delete_by_filter.return_value = 5
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
result = await RAGRuntimeService(app).vector_delete(
|
||||
CONTEXT,
|
||||
'test',
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
filters = {'status': 'deleted'}
|
||||
|
||||
result = await service.vector_delete(
|
||||
collection_id='test',
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
assert result == 5 # Returns count from delete_by_filter
|
||||
mock_app.vector_db_mgr.delete_by_filter.assert_called_once()
|
||||
call_args = mock_app.vector_db_mgr.delete_by_filter.call_args
|
||||
assert call_args.kwargs['collection_name'] == 'test'
|
||||
assert call_args.kwargs['filter'] == filters
|
||||
assert result == 5
|
||||
assert app.vector_db_mgr.delete_by_filter.await_args.kwargs['filter'] == filters
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_delete_no_params(self):
|
||||
"""Delete with no params returns 0."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
result = await RAGRuntimeService(app).vector_delete(CONTEXT, 'test')
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
result = await service.vector_delete(collection_id='test')
|
||||
|
||||
assert result == 0
|
||||
mock_app.vector_db_mgr.delete_by_file_id.assert_not_called()
|
||||
mock_app.vector_db_mgr.delete_by_filter.assert_not_called()
|
||||
assert result == 0
|
||||
app.vector_db_mgr.delete_by_file_id.assert_not_awaited()
|
||||
app.vector_db_mgr.delete_by_filter.assert_not_awaited()
|
||||
|
||||
|
||||
class TestRAGRuntimeServiceVectorList:
|
||||
"""Tests for vector_list method."""
|
||||
|
||||
def _create_mock_app(self):
|
||||
mock_app = MagicMock()
|
||||
mock_app.vector_db_mgr = MagicMock()
|
||||
mock_app.vector_db_mgr.list_by_filter = AsyncMock(
|
||||
return_value=([{'id': 'id1', 'metadata': {'file_id': 'abc'}}], 10)
|
||||
)
|
||||
return mock_app
|
||||
|
||||
def _make_rag_import_mocks(self):
|
||||
return {
|
||||
'langbot.pkg.core.app': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.rag': MagicMock(),
|
||||
}
|
||||
|
||||
class TestRAGRuntimeServiceVectorListRegression:
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_list_basic(self):
|
||||
"""Basic vector list delegates to vector_db_mgr."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
items, total = await RAGRuntimeService(app).vector_list(CONTEXT, 'test')
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
items, total = await service.vector_list(
|
||||
collection_id='test',
|
||||
)
|
||||
|
||||
assert len(items) == 1
|
||||
assert total == 10
|
||||
mock_app.vector_db_mgr.list_by_filter.assert_called_once()
|
||||
call_args = mock_app.vector_db_mgr.list_by_filter.call_args
|
||||
assert call_args.kwargs['collection_name'] == 'test'
|
||||
assert call_args.kwargs['limit'] == 20 # Default
|
||||
assert call_args.kwargs['offset'] == 0 # Default
|
||||
assert items == [{'id': 'chunk-a'}]
|
||||
assert total == 1
|
||||
app.vector_db_mgr.list_by_filter.assert_awaited_once_with(
|
||||
execution_context=CONTEXT,
|
||||
knowledge_base_uuid='kb-a',
|
||||
filter=None,
|
||||
limit=20,
|
||||
offset=0,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_list_with_pagination(self):
|
||||
"""Vector list with custom pagination."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
await RAGRuntimeService(app).vector_list(CONTEXT, 'test', limit=50, offset=100)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
await service.vector_list(
|
||||
collection_id='test',
|
||||
limit=50,
|
||||
offset=100,
|
||||
)
|
||||
|
||||
call_args = mock_app.vector_db_mgr.list_by_filter.call_args
|
||||
assert call_args.kwargs['limit'] == 50
|
||||
assert call_args.kwargs['offset'] == 100
|
||||
kwargs = app.vector_db_mgr.list_by_filter.await_args.kwargs
|
||||
assert kwargs['limit'] == 50
|
||||
assert kwargs['offset'] == 100
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_list_with_filters(self):
|
||||
"""Vector list with metadata filters."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app()
|
||||
filters = {'file_id': 'abc'}
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
await RAGRuntimeService(app).vector_list(CONTEXT, 'test', filters=filters)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
filters = {'file_id': 'abc'}
|
||||
|
||||
await service.vector_list(
|
||||
collection_id='test',
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
call_args = mock_app.vector_db_mgr.list_by_filter.call_args
|
||||
assert call_args.kwargs['filter'] == filters
|
||||
assert app.vector_db_mgr.list_by_filter.await_args.kwargs['filter'] == filters
|
||||
|
||||
|
||||
class TestRAGRuntimeServiceGetFileStream:
|
||||
"""Tests for get_file_stream method."""
|
||||
|
||||
def _create_mock_app(self):
|
||||
mock_app = MagicMock()
|
||||
mock_app.vector_db_mgr = MagicMock()
|
||||
mock_app.storage_mgr = MagicMock()
|
||||
mock_app.storage_mgr.storage_provider = MagicMock()
|
||||
mock_app.storage_mgr.storage_provider.load = AsyncMock(return_value=b'file content')
|
||||
return mock_app
|
||||
|
||||
def _make_rag_import_mocks(self):
|
||||
return {
|
||||
'langbot.pkg.core.app': MagicMock(),
|
||||
'langbot_plugin.api.entities.builtin.rag': MagicMock(),
|
||||
}
|
||||
|
||||
class TestRAGRuntimeServiceGetFileStreamRegression:
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_file_stream_basic(self):
|
||||
"""Get file stream loads from storage."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app(file_exists=True)
|
||||
app.persistence_mgr.execute_async = AsyncMock(return_value=_ScalarResult('file-a'))
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
result = await RAGRuntimeService(app).get_file_stream(
|
||||
CONTEXT,
|
||||
'knowledge/files/doc.pdf',
|
||||
)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
result = await service.get_file_stream('knowledge/files/doc.pdf')
|
||||
|
||||
assert result == b'file content'
|
||||
mock_app.storage_mgr.storage_provider.load.assert_called_once_with('knowledge/files/doc.pdf')
|
||||
assert result == b'content'
|
||||
app.storage_mgr.load_scoped_object_key.assert_awaited_once_with(
|
||||
CONTEXT,
|
||||
'knowledge/files/doc.pdf',
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_file_stream_empty_result(self):
|
||||
"""Empty file returns empty bytes."""
|
||||
mock_app = self._create_mock_app()
|
||||
mock_app.storage_mgr.storage_provider.load = AsyncMock(return_value=None)
|
||||
app = _app(file_exists=True)
|
||||
app.persistence_mgr.execute_async = AsyncMock(return_value=_ScalarResult('file-a'))
|
||||
app.storage_mgr.load_scoped_object_key.return_value = None
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
result = await RAGRuntimeService(app).get_file_stream(CONTEXT, 'nonexistent.pdf')
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
result = await service.get_file_stream('nonexistent.pdf')
|
||||
|
||||
assert result == b''
|
||||
assert result == b''
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_file_stream_normalizes_safe_path(self):
|
||||
"""Safe relative paths are normalized before loading."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app(file_exists=True)
|
||||
app.persistence_mgr.execute_async = AsyncMock(return_value=_ScalarResult('file-a'))
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
result = await RAGRuntimeService(app).get_file_stream(
|
||||
CONTEXT,
|
||||
'knowledge/./files/doc.pdf',
|
||||
)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
result = await service.get_file_stream('knowledge/./files/doc.pdf')
|
||||
|
||||
assert result == b'file content'
|
||||
mock_app.storage_mgr.storage_provider.load.assert_called_once_with('knowledge/files/doc.pdf')
|
||||
assert result == b'content'
|
||||
app.storage_mgr.load_scoped_object_key.assert_awaited_once_with(
|
||||
CONTEXT,
|
||||
'knowledge/files/doc.pdf',
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_file_stream_path_traversal_blocked(self):
|
||||
"""Path traversal attacks are blocked."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app(file_exists=True)
|
||||
service = RAGRuntimeService(app)
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
# Absolute path should raise ValueError
|
||||
with pytest.raises(ValueError, match='Invalid storage path'):
|
||||
await service.get_file_stream('/etc/passwd')
|
||||
|
||||
# Path traversal should raise ValueError
|
||||
with pytest.raises(ValueError, match='Invalid storage path'):
|
||||
await service.get_file_stream('knowledge/../../../etc/passwd')
|
||||
with pytest.raises(ValueError, match='Invalid storage path'):
|
||||
await service.get_file_stream(CONTEXT, '/etc/passwd')
|
||||
with pytest.raises(ValueError, match='Invalid storage path'):
|
||||
await service.get_file_stream(CONTEXT, 'knowledge/../../../etc/passwd')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
@@ -486,36 +479,22 @@ class TestRAGRuntimeServiceGetFileStream:
|
||||
],
|
||||
)
|
||||
async def test_get_file_stream_rejects_unsafe_paths(self, storage_path: str):
|
||||
"""Unsafe runtime file paths are rejected before storage load."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app(file_exists=True)
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
with pytest.raises(ValueError, match='Invalid storage path'):
|
||||
await RAGRuntimeService(app).get_file_stream(CONTEXT, storage_path)
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
with pytest.raises(ValueError, match='Invalid storage path'):
|
||||
await service.get_file_stream(storage_path)
|
||||
|
||||
mock_app.storage_mgr.storage_provider.load.assert_not_called()
|
||||
app.storage_mgr.load_scoped_object_key.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_file_stream_normalizes_path(self):
|
||||
"""Valid paths with .. in filename (not traversal) should work."""
|
||||
mock_app = self._create_mock_app()
|
||||
app = _app(file_exists=True)
|
||||
app.persistence_mgr.execute_async = AsyncMock(return_value=_ScalarResult('file-a'))
|
||||
|
||||
mocks = self._make_rag_import_mocks()
|
||||
await RAGRuntimeService(app).get_file_stream(CONTEXT, 'knowledge/files/test.pdf')
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
|
||||
service = RAGRuntimeService(mock_app)
|
||||
|
||||
# Path that contains '..' as part of filename (not traversal)
|
||||
# This should NOT raise - posixpath.normpath handles this
|
||||
# But the current implementation checks '..' in split('/')
|
||||
# Let's test a simple valid path
|
||||
await service.get_file_stream('knowledge/files/test.pdf')
|
||||
mock_app.storage_mgr.storage_provider.load.assert_called()
|
||||
app.storage_mgr.load_scoped_object_key.assert_awaited_once_with(
|
||||
CONTEXT,
|
||||
'knowledge/files/test.pdf',
|
||||
expected_owner_type='upload_document',
|
||||
)
|
||||
|
||||
@@ -0,0 +1,390 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
||||
|
||||
from langbot.pkg.api.http.authz import WorkspaceRequiredError
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.api.http.service.knowledge import KnowledgeService
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.rag import File, KnowledgeBase
|
||||
from langbot.pkg.entity.persistence.workspace import Workspace, WorkspaceExecutionState
|
||||
from langbot.pkg.rag.knowledge.kbmgr import RAGManager
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
from langbot.pkg.vector.mgr import VectorDBManager
|
||||
from langbot.pkg.vector.vdbs.pgvector_db import PgVectorDatabase
|
||||
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
|
||||
from langbot.pkg.workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy
|
||||
from langbot.pkg.workspace.service import WorkspaceService
|
||||
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
INSTANCE_UUID = 'instance-rag-isolation'
|
||||
WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
|
||||
WORKSPACE_B = '00000000-0000-0000-0000-00000000000b'
|
||||
|
||||
|
||||
class _PersistenceManager:
|
||||
def __init__(self, engine):
|
||||
self.engine = engine
|
||||
|
||||
def get_db_engine(self):
|
||||
return self.engine
|
||||
|
||||
async def execute_async(self, *args, **kwargs):
|
||||
async with self.engine.connect() as connection:
|
||||
result = await connection.execute(*args, **kwargs)
|
||||
await connection.commit()
|
||||
return result
|
||||
|
||||
@asynccontextmanager
|
||||
async def tenant_uow(self, _workspace_uuid):
|
||||
# This lightweight fixture does not emulate PostgreSQL RLS; production
|
||||
# persistence tests cover the transaction-bound unit of work itself.
|
||||
async with AsyncSession(self.engine, expire_on_commit=False) as session, session.begin():
|
||||
yield SimpleNamespace(session=session)
|
||||
|
||||
@staticmethod
|
||||
def serialize_model(model, row, masked_columns=()):
|
||||
return {
|
||||
column.name: (
|
||||
getattr(row, column.name).isoformat()
|
||||
if isinstance(getattr(row, column.name), datetime.datetime)
|
||||
else getattr(row, column.name)
|
||||
)
|
||||
for column in model.__table__.columns
|
||||
if column.name not in masked_columns
|
||||
}
|
||||
|
||||
|
||||
class _RecordingVectorDatabase:
|
||||
def __init__(self):
|
||||
self.collections: list[str] = []
|
||||
self.metadatas: list[list[dict]] = []
|
||||
self.calls: list[tuple[str, str]] = []
|
||||
|
||||
async def add_embeddings(self, *, collection, ids, embeddings_list, metadatas, documents):
|
||||
self.collections.append(collection)
|
||||
self.metadatas.append(metadatas)
|
||||
self.calls.append(('upsert', collection))
|
||||
|
||||
async def search(self, *, collection, **_kwargs):
|
||||
self.calls.append(('search', collection))
|
||||
return {'ids': [[]], 'distances': [[]], 'metadatas': [[]]}
|
||||
|
||||
async def delete_by_file_id(self, collection, _file_id):
|
||||
self.calls.append(('delete_by_file_id', collection))
|
||||
|
||||
async def delete_collection(self, collection):
|
||||
self.calls.append(('delete_collection', collection))
|
||||
|
||||
async def delete_by_filter(self, collection, _filter):
|
||||
self.calls.append(('delete_by_filter', collection))
|
||||
return 1
|
||||
|
||||
async def list_by_filter(self, collection, _filter, _limit, _offset):
|
||||
self.calls.append(('list_by_filter', collection))
|
||||
return [], 0
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def tenant_rag(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "rag-tenant.db"}')
|
||||
async with engine.begin() as connection:
|
||||
await connection.run_sync(Base.metadata.create_all)
|
||||
await connection.execute(
|
||||
sqlalchemy.insert(Workspace),
|
||||
[
|
||||
{
|
||||
'uuid': WORKSPACE_A,
|
||||
'instance_uuid': INSTANCE_UUID,
|
||||
'name': 'Workspace A',
|
||||
'slug': 'workspace-a',
|
||||
'source': 'local',
|
||||
},
|
||||
{
|
||||
'uuid': WORKSPACE_B,
|
||||
'instance_uuid': INSTANCE_UUID,
|
||||
'name': 'Workspace B',
|
||||
'slug': 'workspace-b',
|
||||
'source': 'cloud_projection',
|
||||
},
|
||||
],
|
||||
)
|
||||
await connection.execute(
|
||||
sqlalchemy.insert(WorkspaceExecutionState),
|
||||
[
|
||||
{
|
||||
'workspace_uuid': WORKSPACE_A,
|
||||
'instance_uuid': INSTANCE_UUID,
|
||||
'active_generation': 3,
|
||||
'state': 'active',
|
||||
'source': 'local',
|
||||
'write_fenced': False,
|
||||
},
|
||||
{
|
||||
'workspace_uuid': WORKSPACE_B,
|
||||
'instance_uuid': INSTANCE_UUID,
|
||||
'active_generation': 3,
|
||||
'state': 'active',
|
||||
'source': 'cloud',
|
||||
'write_fenced': False,
|
||||
},
|
||||
],
|
||||
)
|
||||
await connection.execute(
|
||||
sqlalchemy.insert(KnowledgeBase),
|
||||
[
|
||||
{
|
||||
'uuid': 'kb-a',
|
||||
'workspace_uuid': WORKSPACE_A,
|
||||
'name': 'Same Knowledge Base',
|
||||
'description': 'A',
|
||||
'knowledge_engine_plugin_id': 'author/engine',
|
||||
'collection_id': 'kb-a',
|
||||
'creation_settings': {},
|
||||
'retrieval_settings': {},
|
||||
},
|
||||
{
|
||||
'uuid': 'kb-b',
|
||||
'workspace_uuid': WORKSPACE_B,
|
||||
'name': 'Same Knowledge Base',
|
||||
'description': 'B',
|
||||
'knowledge_engine_plugin_id': 'author/engine',
|
||||
'collection_id': 'kb-b',
|
||||
'creation_settings': {},
|
||||
'retrieval_settings': {},
|
||||
},
|
||||
],
|
||||
)
|
||||
await connection.execute(
|
||||
sqlalchemy.insert(File),
|
||||
[
|
||||
{
|
||||
'uuid': 'file-a',
|
||||
'workspace_uuid': WORKSPACE_A,
|
||||
'kb_id': 'kb-a',
|
||||
'file_name': 'a.pdf',
|
||||
'extension': 'pdf',
|
||||
},
|
||||
{
|
||||
'uuid': 'file-b',
|
||||
'workspace_uuid': WORKSPACE_B,
|
||||
'kb_id': 'kb-b',
|
||||
'file_name': 'b.pdf',
|
||||
'extension': 'pdf',
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
app = SimpleNamespace()
|
||||
app.persistence_mgr = _PersistenceManager(engine)
|
||||
app.logger = Mock()
|
||||
app.workspace_policy = SingleWorkspacePolicy()
|
||||
app.plugin_connector = SimpleNamespace(
|
||||
is_enable_plugin=False,
|
||||
rag_on_kb_create=AsyncMock(),
|
||||
rag_on_kb_delete=AsyncMock(),
|
||||
)
|
||||
app.workspace_service = WorkspaceService(app, instance_uuid=INSTANCE_UUID)
|
||||
app.rag_mgr = RAGManager(app)
|
||||
await app.rag_mgr.initialize()
|
||||
app.knowledge_service = KnowledgeService(app)
|
||||
|
||||
yield app, engine
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
def _context(workspace_uuid: str) -> ExecutionContext:
|
||||
return ExecutionContext(
|
||||
instance_uuid=INSTANCE_UUID,
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=3,
|
||||
)
|
||||
|
||||
|
||||
async def test_context_is_mandatory_and_same_names_are_isolated(tenant_rag):
|
||||
app, _engine = tenant_rag
|
||||
|
||||
with pytest.raises(WorkspaceRequiredError):
|
||||
await app.knowledge_service.get_knowledge_bases(None)
|
||||
|
||||
bases_a = await app.knowledge_service.get_knowledge_bases(_context(WORKSPACE_A))
|
||||
bases_b = await app.knowledge_service.get_knowledge_bases(_context(WORKSPACE_B))
|
||||
assert [(item['uuid'], item['name']) for item in bases_a] == [('kb-a', 'Same Knowledge Base')]
|
||||
assert [(item['uuid'], item['name']) for item in bases_b] == [('kb-b', 'Same Knowledge Base')]
|
||||
|
||||
|
||||
async def test_cross_workspace_uuid_and_file_guessing_return_not_found(tenant_rag):
|
||||
app, engine = tenant_rag
|
||||
context_a = _context(WORKSPACE_A)
|
||||
|
||||
assert await app.knowledge_service.get_knowledge_base(context_a, 'kb-b') is None
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await app.knowledge_service.update_knowledge_base(context_a, 'kb-b', {'name': 'stolen'})
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await app.knowledge_service.delete_knowledge_base(context_a, 'kb-b')
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await app.knowledge_service.get_files_by_knowledge_base(context_a, 'kb-b')
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await app.knowledge_service.delete_file(context_a, 'kb-b', 'file-b')
|
||||
|
||||
async with engine.connect() as connection:
|
||||
assert (
|
||||
await connection.scalar(sqlalchemy.select(KnowledgeBase.name).where(KnowledgeBase.uuid == 'kb-b'))
|
||||
== 'Same Knowledge Base'
|
||||
)
|
||||
assert await connection.scalar(sqlalchemy.select(File.uuid).where(File.uuid == 'file-b')) == 'file-b'
|
||||
|
||||
|
||||
async def test_runtime_rejects_cross_workspace_collection_reference(tenant_rag):
|
||||
app, _engine = tenant_rag
|
||||
app.vector_db_mgr = SimpleNamespace(upsert=AsyncMock())
|
||||
service = RAGRuntimeService(app)
|
||||
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await service.vector_upsert(
|
||||
_context(WORKSPACE_A),
|
||||
'kb-b',
|
||||
vectors=[[0.1, 0.2]],
|
||||
ids=['chunk-1'],
|
||||
)
|
||||
app.vector_db_mgr.upsert.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_physical_vector_handles_do_not_collide_across_workspaces(tenant_rag):
|
||||
app, _engine = tenant_rag
|
||||
database = _RecordingVectorDatabase()
|
||||
manager = VectorDBManager(app)
|
||||
manager.vector_db = database
|
||||
|
||||
context_a = _context(WORKSPACE_A)
|
||||
context_b = _context(WORKSPACE_B)
|
||||
assert manager.physical_collection_name(context_a, 'same-kb-id') != manager.physical_collection_name(
|
||||
context_b,
|
||||
'same-kb-id',
|
||||
)
|
||||
|
||||
await manager.upsert(context_a, 'kb-a', [[0.1]], ['a'], metadata=[{'source': 'client'}])
|
||||
await manager.upsert(context_b, 'kb-b', [[0.2]], ['b'], metadata=[{'source': 'client'}])
|
||||
assert len(set(database.collections)) == 2
|
||||
assert database.metadatas[0][0]['_langbot_workspace_uuid'] == WORKSPACE_A
|
||||
assert database.metadatas[1][0]['_langbot_workspace_uuid'] == WORKSPACE_B
|
||||
|
||||
|
||||
async def test_migrated_local_kb_keeps_legacy_collection_for_every_vector_operation(tenant_rag):
|
||||
app, engine = tenant_rag
|
||||
legacy_collection = 'legacy-collection-kb-a'
|
||||
async with engine.begin() as connection:
|
||||
await connection.execute(
|
||||
sqlalchemy.update(KnowledgeBase)
|
||||
.where(KnowledgeBase.uuid == 'kb-a')
|
||||
.values(
|
||||
collection_id=legacy_collection,
|
||||
legacy_vector_collection=True,
|
||||
)
|
||||
)
|
||||
|
||||
database = _RecordingVectorDatabase()
|
||||
manager = VectorDBManager(app)
|
||||
manager.vector_db = database
|
||||
context = _context(WORKSPACE_A)
|
||||
|
||||
await manager.upsert(context, 'kb-a', [[0.1]], ['chunk-a'])
|
||||
await manager.search(context, 'kb-a', [0.1], 3)
|
||||
await manager.delete_by_file_id(context, 'kb-a', ['file-a'])
|
||||
assert await manager.delete_by_filter(context, 'kb-a', {'file_id': 'file-a'}) == 1
|
||||
assert await manager.list_by_filter(context, 'kb-a', {'file_id': 'file-a'}) == ([], 0)
|
||||
await manager.delete_collection(context, 'kb-a')
|
||||
|
||||
assert database.calls == [
|
||||
('upsert', legacy_collection),
|
||||
('search', legacy_collection),
|
||||
('delete_by_file_id', legacy_collection),
|
||||
('delete_by_filter', legacy_collection),
|
||||
('list_by_filter', legacy_collection),
|
||||
('delete_collection', legacy_collection),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('deny_by', ['projected_workspace', 'multi_workspace_policy'])
|
||||
async def test_legacy_marker_is_ignored_outside_single_local_workspace(tenant_rag, deny_by):
|
||||
app, engine = tenant_rag
|
||||
workspace_uuid = WORKSPACE_B if deny_by == 'projected_workspace' else WORKSPACE_A
|
||||
kb_uuid = 'kb-b' if deny_by == 'projected_workspace' else 'kb-a'
|
||||
legacy_collection = f'legacy-{deny_by}'
|
||||
async with engine.begin() as connection:
|
||||
await connection.execute(
|
||||
sqlalchemy.update(KnowledgeBase)
|
||||
.where(KnowledgeBase.uuid == kb_uuid)
|
||||
.values(
|
||||
collection_id=legacy_collection,
|
||||
legacy_vector_collection=True,
|
||||
)
|
||||
)
|
||||
if deny_by == 'multi_workspace_policy':
|
||||
app.workspace_policy = CloudWorkspacePolicy()
|
||||
|
||||
database = _RecordingVectorDatabase()
|
||||
manager = VectorDBManager(app)
|
||||
manager.vector_db = database
|
||||
context = _context(workspace_uuid)
|
||||
await manager.search(context, kb_uuid, [0.1], 3)
|
||||
|
||||
assert database.calls == [('search', manager.physical_collection_name(context, kb_uuid))]
|
||||
assert database.calls[0][1] != legacy_collection
|
||||
app.logger.warning.assert_called_once()
|
||||
|
||||
|
||||
async def test_stale_generation_is_rejected_before_vector_access(tenant_rag):
|
||||
app, _engine = tenant_rag
|
||||
database = _RecordingVectorDatabase()
|
||||
manager = VectorDBManager(app)
|
||||
manager.vector_db = database
|
||||
stale = ExecutionContext(
|
||||
instance_uuid=INSTANCE_UUID,
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
placement_generation=2,
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match='generation'):
|
||||
await manager.upsert(stale, 'kb-a', [[0.1]], ['a'])
|
||||
assert database.collections == []
|
||||
|
||||
|
||||
async def test_pgvector_first_write_binds_dimension_and_later_mismatch_fails(tenant_rag):
|
||||
app, engine = tenant_rag
|
||||
manager = VectorDBManager(app)
|
||||
pgvector = object.__new__(PgVectorDatabase)
|
||||
pgvector.allowed_dimensions = frozenset({1, 2})
|
||||
pgvector.add_embeddings = AsyncMock()
|
||||
pgvector.search = AsyncMock(return_value={'ids': [[]], 'distances': [[]], 'metadatas': [[]]})
|
||||
manager.vector_db = pgvector
|
||||
context = _context(WORKSPACE_A)
|
||||
|
||||
await manager.upsert(context, 'kb-a', [[0.1]], ['chunk-a'])
|
||||
scope = pgvector.add_embeddings.await_args.kwargs['scope']
|
||||
assert scope.workspace_uuid == WORKSPACE_A
|
||||
assert scope.knowledge_base_uuid == 'kb-a'
|
||||
assert scope.embedding_dimension == 1
|
||||
|
||||
async with engine.connect() as connection:
|
||||
selected_dimension = await connection.scalar(
|
||||
sqlalchemy.select(KnowledgeBase.embedding_dimension).where(
|
||||
KnowledgeBase.workspace_uuid == WORKSPACE_A,
|
||||
KnowledgeBase.uuid == 'kb-a',
|
||||
)
|
||||
)
|
||||
assert selected_dimension == 1
|
||||
|
||||
with pytest.raises(ValueError, match='dimension is 1, not 2'):
|
||||
await manager.upsert(context, 'kb-a', [[0.1, 0.2]], ['chunk-b'])
|
||||
with pytest.raises(ValueError, match='not enabled'):
|
||||
await manager.search(context, 'kb-a', [0.1, 0.2, 0.3], 3)
|
||||
Reference in New Issue
Block a user