mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-21 20:06:06 +00:00
feat(tenancy): implement workspace isolation
This commit is contained in:
Junyan Qin
@@ -9,7 +9,27 @@ from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.rag.knowledge.kbmgr import RuntimeKnowledgeBase
|
||||
from langbot.pkg.storage.mgr import StorageMgr
|
||||
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
|
||||
|
||||
|
||||
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]) -> bytes:
|
||||
@@ -25,26 +45,38 @@ 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.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.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,27 +90,65 @@ 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()
|
||||
|
||||
|
||||
class TestStoreZipFile:
|
||||
@pytest.mark.asyncio
|
||||
@@ -99,19 +169,22 @@ 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):
|
||||
@@ -122,10 +195,10 @@ class TestStoreZipFile:
|
||||
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'))
|
||||
|
||||
|
||||
class TestStoreFileTask:
|
||||
@@ -133,29 +206,31 @@ class TestStoreFileTask:
|
||||
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 +238,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 +247,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 +259,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()
|
||||
|
||||
+548
-495
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,350 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import 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.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
|
||||
|
||||
@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 == []
|
||||
Reference in New Issue
Block a user