diff --git a/mcp_server/src/services/queue_service.py b/mcp_server/src/services/queue_service.py index c64f7497b9..38472b24e3 100644 --- a/mcp_server/src/services/queue_service.py +++ b/mcp_server/src/services/queue_service.py @@ -6,6 +6,9 @@ from datetime import datetime, timezone from typing import Any +from graphiti_core.helpers import validate_excluded_entity_types +from graphiti_core.utils.ontology_utils.entity_types_utils import validate_entity_types + logger = logging.getLogger(__name__) @@ -149,6 +152,9 @@ async def add_episode( if self._graphiti_client is None: raise RuntimeError('Queue service not initialized. Call initialize() first.') + validate_entity_types(entity_types) + validate_excluded_entity_types(excluded_entity_types, entity_types) + async def process_episode(): """Process the episode using the graphiti client.""" try: diff --git a/mcp_server/tests/test_core_parity.py b/mcp_server/tests/test_core_parity.py index abe410cae0..1e5dae63b5 100644 --- a/mcp_server/tests/test_core_parity.py +++ b/mcp_server/tests/test_core_parity.py @@ -15,8 +15,10 @@ import pytest from graphiti_core import Graphiti from graphiti_core.edges import EntityEdge +from graphiti_core.errors import EntityTypeValidationError from graphiti_core.nodes import EntityNode from graphiti_core.search.search_filters import ComparisonOperator, SearchFilters +from pydantic import BaseModel # Add the src directory to the path (mirrors the other unit tests) sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) @@ -155,6 +157,32 @@ def test_invalid_date_raises_value_error(self): class TestQueueServiceThreading: """The queue service must forward every parity param to Graphiti.add_episode.""" + @pytest.mark.asyncio + async def test_add_episode_rejects_invalid_entity_type_before_queueing(self): + class DataObject(BaseModel): + attributes: str | None = None + + client = AsyncMock(spec=Graphiti) + service = QueueService() + await service.initialize(client) + + with pytest.raises( + EntityTypeValidationError, + match='attributes cannot be used as an attribute for DataObject', + ): + await service.add_episode( + group_id='g1', + name='ep', + content='This system has several attributes including age, status, and priority.', + source_description='desc', + episode_type='text', + entity_types={'DataObject': DataObject}, + uuid=None, + ) + + assert 'g1' not in service._episode_queues + client.add_episode.assert_not_called() + @pytest.mark.asyncio async def test_add_episode_forwards_all_params(self): client = AsyncMock(spec=Graphiti) @@ -171,7 +199,10 @@ async def test_add_episode_forwards_all_params(self): content='body', source_description='desc', episode_type='text', - entity_types={'Preference': ENTITY_TYPES['Preference']}, + entity_types={ + 'Preference': ENTITY_TYPES['Preference'], + 'Object': ENTITY_TYPES['Object'], + }, uuid='ep-uuid', reference_time=ref_time, edge_types=edge_types,