mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
24f0190ce7
Adds a RabbitMQ backend as an alternative to Pulsar, selectable via PUBSUB_BACKEND=rabbitmq. Both backends implement the same PubSubBackend protocol — no application code changes needed to switch. RabbitMQ topology: - Single topic exchange per topicspace (e.g. 'tg') - Routing key derived from queue class and topic name - Shared consumers: named queue bound to exchange (competing, round-robin) - Exclusive consumers: anonymous auto-delete queue (broadcast, each gets every message). Used by Subscriber and config push consumer. - Thread-local producer connections (pika is not thread-safe) - Push-based consumption via basic_consume with process_data_events for heartbeat processing Consumer model changes: - Consumer class creates one backend consumer per concurrent task (required for pika thread safety, harmless for Pulsar) - Consumer class accepts consumer_type parameter - Subscriber passes consumer_type='exclusive' for broadcast semantics - Config push consumer uses consumer_type='exclusive' so every processor instance receives config updates - handle_one_from_queue receives consumer as parameter for correct per-connection ack/nack LibrarianClient: - New shared client class replacing duplicated librarian request-response code across 6+ services (chunking, decoders, RAG, etc.) - Uses stream-document instead of get-document-content for fetching document content in 1MB chunks (avoids broker message size limits) - Standalone object (self.librarian = LibrarianClient(...)) not a mixin - get-document-content marked deprecated in schema and OpenAPI spec Serialisation: - Extracted dataclass_to_dict/dict_to_dataclass to shared serialization.py (used by both Pulsar and RabbitMQ backends) Librarian queues: - Changed from flow class (persistent) back to request/response class now that stream-document eliminates large single messages - API upload chunk size reduced from 5MB to 3MB to stay under broker limits after base64 encoding Factory and CLI: - get_pubsub() handles 'rabbitmq' backend with RabbitMQ connection params - add_pubsub_args() includes RabbitMQ options (host, port, credentials) - add_pubsub_args(standalone=True) defaults to localhost for CLI tools - init_trustgraph skips Pulsar admin setup for non-Pulsar backends - tg-dump-queues and tg-monitor-prompts use backend abstraction - BaseClient and ConfigClient accept generic pubsub config
279 lines
No EOL
10 KiB
Python
279 lines
No EOL
10 KiB
Python
"""
|
|
Unit tests for trustgraph.chunking.token
|
|
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.token.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 TestTokenChunkerSimple(IsolatedAsyncioTestCase):
|
|
"""Test Token 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': 300,
|
|
'chunk_overlap': 20,
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
# Act
|
|
processor = Processor(**config)
|
|
|
|
# Assert
|
|
assert processor.default_chunk_size == 300
|
|
assert processor.default_chunk_overlap == 20
|
|
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': 250, # Default chunk size
|
|
'chunk_overlap': 15,
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
processor = Processor(**config)
|
|
|
|
# Mock message and flow
|
|
mock_message = MagicMock()
|
|
mock_consumer = MagicMock()
|
|
mock_flow = MagicMock()
|
|
mock_flow.parameters.get.side_effect = lambda param: {
|
|
"chunk-size": 400, # Override chunk size
|
|
"chunk-overlap": None # Use default chunk overlap
|
|
}.get(param)
|
|
|
|
# Act
|
|
chunk_size, chunk_overlap = await processor.chunk_document(
|
|
mock_message, mock_consumer, mock_flow, 250, 15
|
|
)
|
|
|
|
# Assert
|
|
assert chunk_size == 400 # Should use overridden value
|
|
assert chunk_overlap == 15 # 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': 250,
|
|
'chunk_overlap': 15, # 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.parameters.get.side_effect = lambda param: {
|
|
"chunk-size": None, # Use default chunk size
|
|
"chunk-overlap": 25 # Override chunk overlap
|
|
}.get(param)
|
|
|
|
# Act
|
|
chunk_size, chunk_overlap = await processor.chunk_document(
|
|
mock_message, mock_consumer, mock_flow, 250, 15
|
|
)
|
|
|
|
# Assert
|
|
assert chunk_size == 250 # Should use default value
|
|
assert chunk_overlap == 25 # 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': 250,
|
|
'chunk_overlap': 15,
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
processor = Processor(**config)
|
|
|
|
# Mock message and flow
|
|
mock_message = MagicMock()
|
|
mock_consumer = MagicMock()
|
|
mock_flow = MagicMock()
|
|
mock_flow.parameters.get.side_effect = lambda param: {
|
|
"chunk-size": 350, # Override chunk size
|
|
"chunk-overlap": 30 # Override chunk overlap
|
|
}.get(param)
|
|
|
|
# Act
|
|
chunk_size, chunk_overlap = await processor.chunk_document(
|
|
mock_message, mock_consumer, mock_flow, 250, 15
|
|
)
|
|
|
|
# Assert
|
|
assert chunk_size == 350 # Should use overridden value
|
|
assert chunk_overlap == 30 # Should use overridden value
|
|
|
|
@patch('trustgraph.base.librarian_client.Consumer')
|
|
@patch('trustgraph.base.librarian_client.Producer')
|
|
@patch('trustgraph.chunking.token.chunker.TokenTextSplitter')
|
|
@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 token chunk content"
|
|
mock_splitter.create_documents.return_value = [mock_document]
|
|
mock_splitter_class.return_value = mock_splitter
|
|
|
|
config = {
|
|
'id': 'test-chunker',
|
|
'chunk_size': 250,
|
|
'chunk_overlap': 15,
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
processor = Processor(**config)
|
|
|
|
# Mock save_child_document to avoid librarian producer interactions
|
|
processor.librarian.save_child_document = AsyncMock(return_value="chunk-id")
|
|
|
|
# Mock message with TextDocument
|
|
mock_message = MagicMock()
|
|
mock_text_doc = MagicMock()
|
|
mock_text_doc.metadata = Metadata(
|
|
id="test-doc-456",
|
|
user="test-user",
|
|
collection="test-collection"
|
|
)
|
|
mock_text_doc.text = b"This is test document content for token chunking"
|
|
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()
|
|
mock_flow = MagicMock()
|
|
mock_flow.parameters.get.side_effect = lambda param: {
|
|
"chunk-size": 400,
|
|
"chunk-overlap": 40,
|
|
}.get(param)
|
|
mock_flow.side_effect = lambda name: {
|
|
"output": mock_producer,
|
|
"triples": mock_triples_producer,
|
|
}.get(name)
|
|
|
|
# Act
|
|
await processor.on_message(mock_message, mock_consumer, mock_flow)
|
|
|
|
# Assert
|
|
# Verify TokenTextSplitter was called with overridden parameters (last call)
|
|
expected_call = [
|
|
('encoding_name', 'cl100k_base'),
|
|
('chunk_size', 400),
|
|
('chunk_overlap', 40)
|
|
]
|
|
actual_last_call = mock_splitter_class.call_args_list[-1]
|
|
assert actual_last_call.kwargs['encoding_name'] == "cl100k_base"
|
|
assert actual_last_call.kwargs['chunk_size'] == 400
|
|
assert actual_last_call.kwargs['chunk_overlap'] == 40
|
|
|
|
# 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': 250,
|
|
'chunk_overlap': 15,
|
|
'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.parameters.get.return_value = None # No overrides
|
|
|
|
# Act
|
|
chunk_size, chunk_overlap = await processor.chunk_document(
|
|
mock_message, mock_consumer, mock_flow, 250, 15
|
|
)
|
|
|
|
# Assert
|
|
assert chunk_size == 250 # Should use default value
|
|
assert chunk_overlap == 15 # Should use default value
|
|
|
|
@patch('trustgraph.base.librarian_client.Consumer')
|
|
@patch('trustgraph.base.librarian_client.Producer')
|
|
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
|
def test_token_chunker_uses_different_defaults(self, mock_producer, mock_consumer):
|
|
"""Test that token chunker has different defaults than recursive chunker"""
|
|
# Arrange & Act
|
|
config = {
|
|
'id': 'test-chunker',
|
|
'concurrency': 1,
|
|
'taskgroup': AsyncMock()
|
|
}
|
|
|
|
processor = Processor(**config)
|
|
|
|
# Assert - Token chunker should have different defaults
|
|
assert processor.default_chunk_size == 250 # Token chunker default
|
|
assert processor.default_chunk_overlap == 15 # Token chunker default
|
|
|
|
|
|
if __name__ == '__main__':
|
|
pytest.main([__file__]) |