mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
d35473f7f7
Introduces `workspace` as the isolation boundary for config, flows,
library, and knowledge data. Removes `user` as a schema-level field
throughout the code, API specs, and tests; workspace provides the
same separation more cleanly at the trusted flow.workspace layer
rather than through client-supplied message fields.
Design
------
- IAM tech spec (docs/tech-specs/iam.md) documents current state,
proposed auth/access model, and migration direction.
- Data ownership model (docs/tech-specs/data-ownership-model.md)
captures the workspace/collection/flow hierarchy.
Schema + messaging
------------------
- Drop `user` field from AgentRequest/Step, GraphRagQuery,
DocumentRagQuery, Triples/Graph/Document/Row EmbeddingsRequest,
Sparql/Rows/Structured QueryRequest, ToolServiceRequest.
- Keep collection/workspace routing via flow.workspace at the
service layer.
- Translators updated to not serialise/deserialise user.
API specs
---------
- OpenAPI schemas and path examples cleaned of user fields.
- Websocket async-api messages updated.
- Removed the unused parameters/User.yaml.
Services + base
---------------
- Librarian, collection manager, knowledge, config: all operations
scoped by workspace. Config client API takes workspace as first
positional arg.
- `flow.workspace` set at flow start time by the infrastructure;
no longer pass-through from clients.
- Tool service drops user-personalisation passthrough.
CLI + SDK
---------
- tg-init-workspace and workspace-aware import/export.
- All tg-* commands drop user args; accept --workspace.
- Python API/SDK (flow, socket_client, async_*, explainability,
library) drop user kwargs from every method signature.
MCP server
----------
- All tool endpoints drop user parameters; socket_manager no longer
keyed per user.
Flow service
------------
- Closure-based topic cleanup on flow stop: only delete topics
whose blueprint template was parameterised AND no remaining
live flow (across all workspaces) still resolves to that topic.
Three scopes fall out naturally from template analysis:
* {id} -> per-flow, deleted on stop
* {blueprint} -> per-blueprint, kept while any flow of the
same blueprint exists
* {workspace} -> per-workspace, kept while any flow in the
workspace exists
* literal -> global, never deleted (e.g. tg.request.librarian)
Fixes a bug where stopping a flow silently destroyed the global
librarian exchange, wedging all library operations until manual
restart.
RabbitMQ backend
----------------
- heartbeat=60, blocked_connection_timeout=300. Catches silently
dead connections (broker restart, orphaned channels, network
partitions) within ~2 heartbeat windows, so the consumer
reconnects and re-binds its queue rather than sitting forever
on a zombie connection.
Tests
-----
- Full test refresh: unit, integration, contract, provenance.
- Dropped user-field assertions and constructor kwargs across
~100 test files.
- Renamed user-collection isolation tests to workspace-collection.
257 lines
No EOL
9.5 KiB
Python
257 lines
No EOL
9.5 KiB
Python
"""
|
|
Unit tests for trustgraph.chunking.recursive
|
|
Testing parameter override functionality for chunk-size and chunk-overlap
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
from unittest import IsolatedAsyncioTestCase
|
|
|
|
# Import the service under test
|
|
from trustgraph.chunking.recursive.chunker import Processor
|
|
from trustgraph.schema import TextDocument, Chunk, Metadata
|
|
|
|
|
|
class MockAsyncProcessor:
|
|
def __init__(self, **params):
|
|
self.config_handlers = []
|
|
self.id = params.get('id', 'test-service')
|
|
self.specifications = []
|
|
self.pubsub = MagicMock()
|
|
self.taskgroup = params.get('taskgroup', MagicMock())
|
|
|
|
|
|
class TestRecursiveChunkerSimple(IsolatedAsyncioTestCase):
|
|
"""Test Recursive chunker functionality"""
|
|
|
|
@patch('trustgraph.base.librarian_client.Consumer')
|
|
@patch('trustgraph.base.librarian_client.Producer')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
|
def test_processor_initialization_basic(self, mock_producer, mock_consumer):
|
|
"""Test basic processor initialization"""
|
|
# Arrange
|
|
config = {
|
|
'id': 'test-chunker',
|
|
'chunk_size': 1500,
|
|
'chunk_overlap': 150,
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
# Act
|
|
processor = Processor(**config)
|
|
|
|
# Assert
|
|
assert processor.default_chunk_size == 1500
|
|
assert processor.default_chunk_overlap == 150
|
|
assert hasattr(processor, 'text_splitter')
|
|
|
|
# Verify parameter specs are registered
|
|
param_specs = [spec for spec in processor.specifications
|
|
if hasattr(spec, 'name') and spec.name in ['chunk-size', 'chunk-overlap']]
|
|
assert len(param_specs) == 2
|
|
|
|
@patch('trustgraph.base.librarian_client.Consumer')
|
|
@patch('trustgraph.base.librarian_client.Producer')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
|
async def test_chunk_document_with_chunk_size_override(self, mock_producer, mock_consumer):
|
|
"""Test chunk_document with chunk-size parameter override"""
|
|
# Arrange
|
|
config = {
|
|
'id': 'test-chunker',
|
|
'chunk_size': 1000, # Default chunk size
|
|
'chunk_overlap': 100,
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
processor = Processor(**config)
|
|
|
|
# Mock message and flow
|
|
mock_message = MagicMock()
|
|
mock_consumer = MagicMock()
|
|
# Flow exposes parameter lookup via __call__: flow("chunk-size")
|
|
mock_flow = MagicMock()
|
|
mock_flow.side_effect = lambda key: {
|
|
"chunk-size": 2000, # Override chunk size
|
|
"chunk-overlap": None # Use default chunk overlap
|
|
}.get(key)
|
|
|
|
# Act
|
|
chunk_size, chunk_overlap = await processor.chunk_document(
|
|
mock_message, mock_consumer, mock_flow, 1000, 100
|
|
)
|
|
|
|
# Assert
|
|
assert chunk_size == 2000 # Should use overridden value
|
|
assert chunk_overlap == 100 # Should use default value
|
|
|
|
@patch('trustgraph.base.librarian_client.Consumer')
|
|
@patch('trustgraph.base.librarian_client.Producer')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
|
async def test_chunk_document_with_chunk_overlap_override(self, mock_producer, mock_consumer):
|
|
"""Test chunk_document with chunk-overlap parameter override"""
|
|
# Arrange
|
|
config = {
|
|
'id': 'test-chunker',
|
|
'chunk_size': 1000,
|
|
'chunk_overlap': 100, # Default chunk overlap
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
processor = Processor(**config)
|
|
|
|
# Mock message and flow
|
|
mock_message = MagicMock()
|
|
mock_consumer = MagicMock()
|
|
mock_flow = MagicMock()
|
|
mock_flow.side_effect = lambda key: {
|
|
"chunk-size": None, # Use default chunk size
|
|
"chunk-overlap": 200 # Override chunk overlap
|
|
}.get(key)
|
|
|
|
# Act
|
|
chunk_size, chunk_overlap = await processor.chunk_document(
|
|
mock_message, mock_consumer, mock_flow, 1000, 100
|
|
)
|
|
|
|
# Assert
|
|
assert chunk_size == 1000 # Should use default value
|
|
assert chunk_overlap == 200 # Should use overridden value
|
|
|
|
@patch('trustgraph.base.librarian_client.Consumer')
|
|
@patch('trustgraph.base.librarian_client.Producer')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
|
async def test_chunk_document_with_both_parameters_override(self, mock_producer, mock_consumer):
|
|
"""Test chunk_document with both chunk-size and chunk-overlap overrides"""
|
|
# Arrange
|
|
config = {
|
|
'id': 'test-chunker',
|
|
'chunk_size': 1000,
|
|
'chunk_overlap': 100,
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
processor = Processor(**config)
|
|
|
|
# Mock message and flow
|
|
mock_message = MagicMock()
|
|
mock_consumer = MagicMock()
|
|
mock_flow = MagicMock()
|
|
mock_flow.side_effect = lambda key: {
|
|
"chunk-size": 1500, # Override chunk size
|
|
"chunk-overlap": 150 # Override chunk overlap
|
|
}.get(key)
|
|
|
|
# Act
|
|
chunk_size, chunk_overlap = await processor.chunk_document(
|
|
mock_message, mock_consumer, mock_flow, 1000, 100
|
|
)
|
|
|
|
# Assert
|
|
assert chunk_size == 1500 # Should use overridden value
|
|
assert chunk_overlap == 150 # Should use overridden value
|
|
|
|
@patch('trustgraph.base.librarian_client.Consumer')
|
|
@patch('trustgraph.base.librarian_client.Producer')
|
|
@patch('trustgraph.chunking.recursive.chunker.RecursiveCharacterTextSplitter')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
|
async def test_on_message_uses_flow_parameters(self, mock_splitter_class, mock_producer, mock_consumer):
|
|
"""Test that on_message method uses parameters from flow"""
|
|
# Arrange
|
|
mock_splitter = MagicMock()
|
|
mock_document = MagicMock()
|
|
mock_document.page_content = "Test chunk content"
|
|
mock_splitter.create_documents.return_value = [mock_document]
|
|
mock_splitter_class.return_value = mock_splitter
|
|
|
|
config = {
|
|
'id': 'test-chunker',
|
|
'chunk_size': 1000,
|
|
'chunk_overlap': 100,
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
processor = Processor(**config)
|
|
|
|
# Mock save_child_document to avoid waiting for librarian response
|
|
processor.librarian.save_child_document = AsyncMock(return_value="mock-doc-id")
|
|
|
|
# Mock message with TextDocument
|
|
mock_message = MagicMock()
|
|
mock_text_doc = MagicMock()
|
|
mock_text_doc.metadata = Metadata(
|
|
id="test-doc-123",
|
|
collection="test-collection"
|
|
)
|
|
mock_text_doc.text = b"This is test document content"
|
|
mock_text_doc.document_id = "" # No librarian fetch needed
|
|
mock_message.value.return_value = mock_text_doc
|
|
|
|
# Mock consumer and flow with parameter overrides
|
|
mock_consumer = MagicMock()
|
|
mock_producer = AsyncMock()
|
|
mock_triples_producer = AsyncMock()
|
|
# Flow.__call__ resolves parameters and producers/consumers from the
|
|
# same dict — merge both kinds here.
|
|
mock_flow = MagicMock()
|
|
mock_flow.side_effect = lambda key: {
|
|
"chunk-size": 1500,
|
|
"chunk-overlap": 150,
|
|
"output": mock_producer,
|
|
"triples": mock_triples_producer,
|
|
}.get(key)
|
|
|
|
# Act
|
|
await processor.on_message(mock_message, mock_consumer, mock_flow)
|
|
|
|
# Assert
|
|
# Verify RecursiveCharacterTextSplitter was called with overridden parameters (last call)
|
|
actual_last_call = mock_splitter_class.call_args_list[-1]
|
|
assert actual_last_call.kwargs['chunk_size'] == 1500
|
|
assert actual_last_call.kwargs['chunk_overlap'] == 150
|
|
assert actual_last_call.kwargs['length_function'] == len
|
|
assert actual_last_call.kwargs['is_separator_regex'] == False
|
|
|
|
# Verify chunk was sent to output
|
|
mock_producer.send.assert_called_once()
|
|
sent_chunk = mock_producer.send.call_args[0][0]
|
|
assert isinstance(sent_chunk, Chunk)
|
|
|
|
@patch('trustgraph.base.librarian_client.Consumer')
|
|
@patch('trustgraph.base.librarian_client.Producer')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
|
async def test_chunk_document_with_no_overrides(self, mock_producer, mock_consumer):
|
|
"""Test chunk_document when no parameters are overridden (flow returns None)"""
|
|
# Arrange
|
|
config = {
|
|
'id': 'test-chunker',
|
|
'chunk_size': 1000,
|
|
'chunk_overlap': 100,
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
processor = Processor(**config)
|
|
|
|
# Mock message and flow that returns None for all parameters
|
|
mock_message = MagicMock()
|
|
mock_consumer = MagicMock()
|
|
mock_flow = MagicMock()
|
|
mock_flow.side_effect = lambda key: None # No overrides
|
|
|
|
# Act
|
|
chunk_size, chunk_overlap = await processor.chunk_document(
|
|
mock_message, mock_consumer, mock_flow, 1000, 100
|
|
)
|
|
|
|
# Assert
|
|
assert chunk_size == 1000 # Should use default value
|
|
assert chunk_overlap == 100 # Should use default value
|
|
|
|
|
|
if __name__ == '__main__':
|
|
pytest.main([__file__]) |