mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 16:36:21 +02:00
ffe310af7c
* Fix Metadata/EntityEmbeddings schema migration tail and add regression tests (#776)
The Metadata dataclass dropped its `metadata: list[Triple]` field
and EntityEmbeddings/ChunkEmbeddings settled on a singular
`vector: list[float]` field, but several call sites kept passing
`Metadata(metadata=...)` and `EntityEmbeddings(vectors=...)`. The
bugs were latent until a websocket client first hit
`/api/v1/flow/default/import/entity-contexts`, at which point the
dispatcher TypeError'd on construction.
Production fixes (5 call sites on the same migration tail):
* trustgraph-flow gateway dispatchers entity_contexts_import.py
and graph_embeddings_import.py — drop the stale
Metadata(metadata=...) kwarg; switch graph_embeddings_import
to the singular `vector` wire key.
* trustgraph-base messaging translators knowledge.py and
document_loading.py — fix decode side to read the singular
`"vector"` key, matching what their own encode sides have
always written.
* trustgraph-flow tables/knowledge.py — fix Cassandra row
deserialiser to construct EntityEmbeddings(vector=...)
instead of vectors=.
* trustgraph-flow gateway core_import/core_export — switch the
kg-core msgpack wire format to the singular `"v"`/`"vector"`
key and drop the dead `m["m"]` envelope field that referenced
the removed Metadata.metadata triples list (it was a
guaranteed KeyError on the export side).
Defense-in-depth regression coverage (32 new tests across 7 files):
* tests/contract/test_schema_field_contracts.py — pin the field
set of Metadata, EntityEmbeddings, ChunkEmbeddings,
EntityContext so any future schema rename fails CI loudly
with a clear diff.
* tests/unit/test_translators/test_knowledge_translator_roundtrip.py
and test_document_embeddings_translator_roundtrip.py -
encode→decode round-trip the affected translators end to end,
locking in the singular `"vector"` wire key.
* tests/unit/test_gateway/test_entity_contexts_import_dispatcher.py
and test_graph_embeddings_import_dispatcher.py — exercise the
websocket dispatchers' receive() path with realistic
payloads, the direct regression test for the original
production crash.
* tests/unit/test_gateway/test_core_import_export_roundtrip.py
— pack/unpack the kg-core msgpack format through the real
dispatcher classes (with KnowledgeRequestor mocked),
including a full export→import round-trip.
* tests/unit/test_tables/test_knowledge_table_store.py —
exercise the Cassandra row → schema conversion via __new__ to
bypass the live cluster connection.
Also fixes an unrelated leaked-coroutine RuntimeWarning in
test_gateway/test_service.py::test_run_method_calls_web_run_app: the
mocked aiohttp.web.run_app now closes the coroutine that Api.run() hands
it, mirroring what the real run_app would do, instead of leaving it for
the GC to complain about.
* Fix RabbitMQ request/response race and chunker Flow API drift
Two unrelated regressions surfaced after the v2.2 queue class
refactor. Bundled here because both are small and both block
production.
1. Request/response race against ephemeral RabbitMQ response
queues
Commit feeb92b3 switched response/notify queues to per-subscriber
auto-delete exclusive queues. That fixed orphaned-queue
accumulation but introduced a setup race: Subscriber.start()
created the run() task and returned immediately, while the
underlying RabbitMQ consumer only declared and bound its queue
lazily on the first receive() call. RequestResponse.request()
therefore published the request before any queue was bound to the
matching routing key, and the broker dropped the reply. Symptoms:
"Failed to fetch config on notify" / "Request timeout exception"
repeating roughly every 10s in api-gateway, document-embeddings
and any other service exercising the config notify path.
Fix:
* Add ensure_connected() to the BackendConsumer protocol;
implement it on RabbitMQBackendConsumer (calls _connect
synchronously, declaring and binding the queue) and as a
no-op on PulsarBackendConsumer (Pulsar's client.subscribe is
already synchronous at construction).
* Convert Subscriber's readiness signal from a non-existent
Event to an asyncio.Future created in start(). run() calls
consumer.ensure_connected() immediately after
create_consumer() and sets _ready.set_result(None) on first
successful bind. start() awaits the future via asyncio.wait
so it returns only once the consumer is fully bound. Any
reply published after start() returns is therefore guaranteed
to land in a bound queue.
* First-attempt connection failures call
_ready.set_exception(e) and exit run() so start() unblocks
with the error rather than hanging forever — the existing
higher-level retry pattern in fetch_and_apply_config takes
over from there. Runtime failures after a successful start
still go through the existing retry-with-backoff path.
* Update the two existing graceful-shutdown tests that
monkey-patch Subscriber.run with a custom coroutine to honor
the new contract by signalling _ready themselves.
* Add tests/unit/test_base/test_subscriber_readiness.py with
five regression tests pinning the readiness contract:
ensure_connected must be called before start() returns;
start() must block while ensure_connected runs
(race-condition guard with a threading.Event gate);
first-attempt create_consumer and ensure_connected failures
must propagate to start() instead of hanging;
ensure_connected must run before any receive() call.
2. Chunker Flow parameter lookup using the wrong attribute
trustgraph-base/trustgraph/base/chunking_service.py was reading
flow.parameters.get("chunk-size") and chunk-overlap, but the Flow
class has no `parameters` attribute — parameter lookup is exposed
through Flow.__call__ (flow("chunk-size") returns the resolved
value or None). The exception was caught and logged as a
WARNING, so chunking continued with the default sizes and any
configured chunk-size / chunk-overlap was silently ignored:
chunker - WARNING - Could not parse chunk-size parameter:
'Flow' object has no attribute 'parameters'
The chunker tests didn't catch this because they constructed
mock_flow = MagicMock() and configured
mock_flow.parameters.get.side_effect = ..., which is the same
phantom attribute MagicMock auto-creates on demand. Tests and
production agreed on the wrong API.
Fix: switch chunking_service.py to flow("chunk-size") /
flow("chunk-overlap"). Update both chunker test files to mock the
__call__ side_effect instead of the phantom parameters.get,
merging parameter values into the existing flow() lookup the
on_message tests already used for producer resolution.
280 lines
No EOL
10 KiB
Python
280 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()
|
|
# Flow exposes parameter lookup via __call__: flow("chunk-size")
|
|
mock_flow = MagicMock()
|
|
mock_flow.side_effect = lambda key: {
|
|
"chunk-size": 400, # 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, 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.side_effect = lambda key: {
|
|
"chunk-size": None, # Use default chunk size
|
|
"chunk-overlap": 25 # Override chunk overlap
|
|
}.get(key)
|
|
|
|
# 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.side_effect = lambda key: {
|
|
"chunk-size": 350, # Override chunk size
|
|
"chunk-overlap": 30 # Override chunk overlap
|
|
}.get(key)
|
|
|
|
# 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()
|
|
# 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": 400,
|
|
"chunk-overlap": 40,
|
|
"output": mock_producer,
|
|
"triples": mock_triples_producer,
|
|
}.get(key)
|
|
|
|
# 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.side_effect = lambda key: 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__]) |