From 65e297266ff653b4f6f9773f5448b7f1d69b19b7 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 26 Nov 2025 17:51:07 +0000 Subject: [PATCH] Streaming tests --- tests/integration/conftest.py | 200 ++++++++++ .../test_agent_streaming_integration.py | 354 ++++++++++++++++++ ...test_document_rag_streaming_integration.py | 273 ++++++++++++++ .../integration/test_graph_rag_integration.py | 269 +++++++++++++ .../test_graph_rag_streaming_integration.py | 248 ++++++++++++ tests/utils/__init__.py | 29 ++ tests/utils/streaming_assertions.py | 218 +++++++++++ 7 files changed, 1591 insertions(+) create mode 100644 tests/integration/test_agent_streaming_integration.py create mode 100644 tests/integration/test_document_rag_streaming_integration.py create mode 100644 tests/integration/test_graph_rag_integration.py create mode 100644 tests/integration/test_graph_rag_streaming_integration.py create mode 100644 tests/utils/__init__.py create mode 100644 tests/utils/streaming_assertions.py diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 0f47077c..af5dda5b 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -382,6 +382,206 @@ def sample_kg_triples(): ] +# Streaming test fixtures + +@pytest.fixture +def mock_streaming_llm_response(): + """Mock streaming LLM response with realistic chunks""" + async def _generate_chunks(): + """Generate realistic streaming chunks""" + chunks = [ + "Machine", + " learning", + " is", + " a", + " subset", + " of", + " artificial", + " intelligence", + " that", + " focuses", + " on", + " algorithms", + " that", + " learn", + " from", + " data", + "." + ] + for chunk in chunks: + yield chunk + return _generate_chunks + + +@pytest.fixture +def sample_streaming_agent_response(): + """Sample streaming agent response chunks""" + return [ + { + "chunk_type": "thought", + "content": "I need to search", + "end_of_message": False, + "end_of_dialog": False + }, + { + "chunk_type": "thought", + "content": " for information", + "end_of_message": False, + "end_of_dialog": False + }, + { + "chunk_type": "thought", + "content": " about machine learning.", + "end_of_message": True, + "end_of_dialog": False + }, + { + "chunk_type": "action", + "content": "knowledge_query", + "end_of_message": True, + "end_of_dialog": False + }, + { + "chunk_type": "observation", + "content": "Machine learning is", + "end_of_message": False, + "end_of_dialog": False + }, + { + "chunk_type": "observation", + "content": " a subset of AI.", + "end_of_message": True, + "end_of_dialog": False + }, + { + "chunk_type": "final-answer", + "content": "Machine learning", + "end_of_message": False, + "end_of_dialog": False + }, + { + "chunk_type": "final-answer", + "content": " is a subset", + "end_of_message": False, + "end_of_dialog": False + }, + { + "chunk_type": "final-answer", + "content": " of artificial intelligence.", + "end_of_message": True, + "end_of_dialog": True + } + ] + + +@pytest.fixture +def streaming_chunk_collector(): + """Helper to collect streaming chunks for assertions""" + class ChunkCollector: + def __init__(self): + self.chunks = [] + self.complete = False + + async def collect(self, chunk): + """Async callback to collect chunks""" + self.chunks.append(chunk) + + def get_full_text(self): + """Concatenate all chunk content""" + return "".join(self.chunks) + + def get_chunk_types(self): + """Get list of chunk types if chunks are dicts""" + if self.chunks and isinstance(self.chunks[0], dict): + return [c.get("chunk_type") for c in self.chunks] + return [] + + return ChunkCollector + + +@pytest.fixture +def mock_streaming_prompt_response(): + """Mock streaming prompt service response""" + async def _generate_prompt_chunks(): + """Generate streaming chunks for prompt responses""" + chunks = [ + "Based on the", + " provided context,", + " here is", + " the answer:", + " Machine learning", + " enables computers", + " to learn", + " from data." + ] + for chunk in chunks: + yield chunk + return _generate_prompt_chunks + + +@pytest.fixture +def sample_rag_streaming_chunks(): + """Sample RAG streaming response chunks""" + return [ + { + "chunk": "Based on", + "end_of_stream": False + }, + { + "chunk": " the knowledge", + "end_of_stream": False + }, + { + "chunk": " graph,", + "end_of_stream": False + }, + { + "chunk": " machine learning", + "end_of_stream": False + }, + { + "chunk": " is a subset", + "end_of_stream": False + }, + { + "chunk": " of AI.", + "end_of_stream": False + }, + { + "chunk": None, + "end_of_stream": True, + "response": "Based on the knowledge graph, machine learning is a subset of AI." + } + ] + + +@pytest.fixture +def streaming_error_scenarios(): + """Common error scenarios for streaming tests""" + return { + "connection_drop": { + "exception": ConnectionError, + "message": "Connection lost during streaming", + "chunks_before_error": 5 + }, + "timeout": { + "exception": TimeoutError, + "message": "Streaming timeout exceeded", + "chunks_before_error": 10 + }, + "rate_limit": { + "exception": Exception, + "message": "Rate limit exceeded", + "chunks_before_error": 3 + }, + "invalid_chunk": { + "exception": ValueError, + "message": "Invalid chunk format", + "chunks_before_error": 7 + } + } + + # Test markers for integration tests pytestmark = pytest.mark.integration diff --git a/tests/integration/test_agent_streaming_integration.py b/tests/integration/test_agent_streaming_integration.py new file mode 100644 index 00000000..b119eeb6 --- /dev/null +++ b/tests/integration/test_agent_streaming_integration.py @@ -0,0 +1,354 @@ +""" +Integration tests for Agent Manager Streaming Functionality + +These tests verify the streaming behavior of the Agent service, testing +chunk-by-chunk delivery of thoughts, actions, observations, and final answers. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock +from trustgraph.agent.react.agent_manager import AgentManager +from trustgraph.agent.react.tools import KnowledgeQueryImpl +from trustgraph.agent.react.types import Tool, Argument +from tests.utils.streaming_assertions import ( + assert_agent_streaming_chunks, + assert_streaming_chunks_valid, + assert_callback_invoked, + assert_chunk_types_valid, +) + + +@pytest.mark.integration +class TestAgentStreaming: + """Integration tests for Agent streaming functionality""" + + @pytest.fixture + def mock_prompt_client_streaming(self): + """Mock prompt client with streaming support""" + client = AsyncMock() + + async def agent_react_streaming(variables, timeout=600, streaming=False, chunk_callback=None): + if streaming and chunk_callback: + # Simulate streaming response with chunks + chunks = [ + "Thought: I need to", + " search for", + " information about", + " machine learning.", + "\n", + "Action: knowledge_query\n", + "Args: {\n", + ' "question": "What is machine learning?"\n', + "}" + ] + + full_text = "" + for chunk in chunks: + full_text += chunk + await chunk_callback(chunk) + + return full_text + else: + # Non-streaming response + return """Thought: I need to search for information about machine learning. +Action: knowledge_query +Args: { + "question": "What is machine learning?" +}""" + + client.agent_react.side_effect = agent_react_streaming + return client + + @pytest.fixture + def mock_flow_context(self, mock_prompt_client_streaming): + """Mock flow context with streaming prompt client""" + context = MagicMock() + + # Mock graph RAG client + graph_rag_client = AsyncMock() + graph_rag_client.rag.return_value = "Machine learning is a subset of AI." + + def context_router(service_name): + if service_name == "prompt-request": + return mock_prompt_client_streaming + elif service_name == "graph-rag-request": + return graph_rag_client + else: + return AsyncMock() + + context.side_effect = context_router + return context + + @pytest.fixture + def sample_tools(self): + """Sample tool configuration""" + return { + "knowledge_query": Tool( + name="knowledge_query", + description="Query the knowledge graph", + arguments=[ + Argument( + name="question", + type="string", + description="The question to ask" + ) + ], + implementation=KnowledgeQueryImpl, + config={} + ) + } + + @pytest.fixture + def agent_manager(self, mock_flow_context, sample_tools): + """Create AgentManager instance with streaming support""" + return AgentManager( + prompt_client=mock_flow_context("prompt-request"), + tools=sample_tools, + verbose=True + ) + + @pytest.mark.asyncio + async def test_agent_streaming_thought_chunks(self, agent_manager, mock_flow_context): + """Test that thought chunks are streamed correctly""" + # Arrange + thought_chunks = [] + + async def think_callback(chunk): + thought_chunks.append(chunk) + + # Act + await agent_manager.react( + question="What is machine learning?", + history=[], + think_callback=think_callback, + observe_callback=AsyncMock(), + context=mock_flow_context, + streaming=True + ) + + # Assert + assert len(thought_chunks) > 0 + assert_streaming_chunks_valid(thought_chunks, min_chunks=1) + + # Verify thought content makes sense + full_thought = "".join(thought_chunks) + assert "search" in full_thought.lower() or "information" in full_thought.lower() + + @pytest.mark.asyncio + async def test_agent_streaming_observation_chunks(self, agent_manager, mock_flow_context): + """Test that observation chunks are streamed correctly""" + # Arrange + observation_chunks = [] + + async def observe_callback(chunk): + observation_chunks.append(chunk) + + # Act + await agent_manager.react( + question="What is machine learning?", + history=[], + think_callback=AsyncMock(), + observe_callback=observe_callback, + context=mock_flow_context, + streaming=True + ) + + # Assert + # Note: Observations come from tool execution, which may or may not be streamed + # depending on the tool implementation + # For now, verify callback was set up + assert observe_callback is not None + + @pytest.mark.asyncio + async def test_agent_streaming_vs_non_streaming(self, agent_manager, mock_flow_context): + """Test that streaming and non-streaming produce equivalent results""" + # Arrange + question = "What is machine learning?" + history = [] + + # Act - Non-streaming + non_streaming_result = await agent_manager.react( + question=question, + history=history, + think_callback=None, + observe_callback=None, + context=mock_flow_context, + streaming=False + ) + + # Act - Streaming + thought_chunks = [] + observation_chunks = [] + + async def think_cb(chunk): + thought_chunks.append(chunk) + + async def observe_cb(chunk): + observation_chunks.append(chunk) + + streaming_result = await agent_manager.react( + question=question, + history=history, + think_callback=think_cb, + observe_callback=observe_cb, + context=mock_flow_context, + streaming=True + ) + + # Assert - Results should be equivalent (or both valid) + assert non_streaming_result is not None + assert streaming_result is not None + + @pytest.mark.asyncio + async def test_agent_streaming_callback_invocation(self, agent_manager, mock_flow_context): + """Test that callbacks are invoked with correct parameters""" + # Arrange + think_callback = AsyncMock() + observe_callback = AsyncMock() + + # Act + await agent_manager.react( + question="What is machine learning?", + history=[], + think_callback=think_callback, + observe_callback=observe_callback, + context=mock_flow_context, + streaming=True + ) + + # Assert - Think callback should be invoked + assert think_callback.call_count > 0 + + # Verify all callback invocations had string arguments + for call in think_callback.call_args_list: + assert len(call.args) > 0 + assert isinstance(call.args[0], str) + + @pytest.mark.asyncio + async def test_agent_streaming_without_callbacks(self, agent_manager, mock_flow_context): + """Test streaming parameter without callbacks (should work gracefully)""" + # Arrange & Act + result = await agent_manager.react( + question="What is machine learning?", + history=[], + think_callback=None, + observe_callback=None, + context=mock_flow_context, + streaming=True # Streaming enabled but no callbacks + ) + + # Assert - Should complete without error + assert result is not None + + @pytest.mark.asyncio + async def test_agent_streaming_with_conversation_history(self, agent_manager, mock_flow_context, + sample_agent_responses): + """Test streaming with existing conversation history""" + # Arrange + history = sample_agent_responses[:1] # Use first response as history + think_callback = AsyncMock() + + # Act + result = await agent_manager.react( + question="Tell me more about neural networks", + history=history, + think_callback=think_callback, + observe_callback=AsyncMock(), + context=mock_flow_context, + streaming=True + ) + + # Assert + assert result is not None + assert think_callback.call_count > 0 + + @pytest.mark.asyncio + async def test_agent_streaming_error_propagation(self, agent_manager, mock_flow_context): + """Test that errors during streaming are properly propagated""" + # Arrange + mock_prompt_client = mock_flow_context("prompt-request") + mock_prompt_client.agent_react.side_effect = Exception("Prompt service error") + + think_callback = AsyncMock() + observe_callback = AsyncMock() + + # Act & Assert + with pytest.raises(Exception) as exc_info: + await agent_manager.react( + question="test question", + history=[], + think_callback=think_callback, + observe_callback=observe_callback, + context=mock_flow_context, + streaming=True + ) + + assert "Prompt service error" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_agent_streaming_multi_step_reasoning(self, agent_manager, mock_flow_context, + mock_prompt_client_streaming): + """Test streaming through multi-step reasoning process""" + # Arrange - Mock a multi-step response + step_responses = [ + """Thought: I need to search for basic information. +Action: knowledge_query +Args: {"question": "What is AI?"}""", + """Thought: Now I can answer the question. +Final Answer: AI is the simulation of human intelligence in machines.""" + ] + + call_count = 0 + + async def multi_step_agent_react(variables, timeout=600, streaming=False, chunk_callback=None): + nonlocal call_count + response = step_responses[min(call_count, len(step_responses) - 1)] + call_count += 1 + + if streaming and chunk_callback: + for chunk in response.split(): + await chunk_callback(chunk + " ") + return response + return response + + mock_prompt_client_streaming.agent_react.side_effect = multi_step_agent_react + + think_callback = AsyncMock() + observe_callback = AsyncMock() + + # Act + result = await agent_manager.react( + question="What is artificial intelligence?", + history=[], + think_callback=think_callback, + observe_callback=observe_callback, + context=mock_flow_context, + streaming=True + ) + + # Assert + assert result is not None + assert think_callback.call_count > 0 + + @pytest.mark.asyncio + async def test_agent_streaming_preserves_tool_config(self, agent_manager, mock_flow_context): + """Test that streaming preserves tool configuration and context""" + # Arrange + think_callback = AsyncMock() + observe_callback = AsyncMock() + + # Act + await agent_manager.react( + question="What is machine learning?", + history=[], + think_callback=think_callback, + observe_callback=observe_callback, + context=mock_flow_context, + streaming=True + ) + + # Assert - Verify prompt client was called with streaming + mock_prompt_client = mock_flow_context("prompt-request") + call_args = mock_prompt_client.agent_react.call_args + assert call_args.kwargs['streaming'] is True + assert call_args.kwargs['chunk_callback'] is not None diff --git a/tests/integration/test_document_rag_streaming_integration.py b/tests/integration/test_document_rag_streaming_integration.py new file mode 100644 index 00000000..45601b39 --- /dev/null +++ b/tests/integration/test_document_rag_streaming_integration.py @@ -0,0 +1,273 @@ +""" +Integration tests for DocumentRAG streaming functionality + +These tests verify the streaming behavior of DocumentRAG, testing token-by-token +response delivery through the complete pipeline. +""" + +import pytest +from unittest.mock import AsyncMock +from trustgraph.retrieval.document_rag.document_rag import DocumentRag +from tests.utils.streaming_assertions import ( + assert_streaming_chunks_valid, + assert_callback_invoked, +) + + +@pytest.mark.integration +class TestDocumentRagStreaming: + """Integration tests for DocumentRAG streaming""" + + @pytest.fixture + def mock_embeddings_client(self): + """Mock embeddings client""" + client = AsyncMock() + client.embed.return_value = [[0.1, 0.2, 0.3, 0.4, 0.5]] + return client + + @pytest.fixture + def mock_doc_embeddings_client(self): + """Mock document embeddings client""" + client = AsyncMock() + client.query.return_value = [ + "Machine learning is a subset of AI.", + "Deep learning uses neural networks.", + "Supervised learning needs labeled data." + ] + return client + + @pytest.fixture + def mock_streaming_prompt_client(self, mock_streaming_llm_response): + """Mock prompt client with streaming support""" + client = AsyncMock() + + async def document_prompt_side_effect(query, documents, timeout=600, streaming=False, chunk_callback=None): + if streaming and chunk_callback: + # Simulate streaming chunks + full_text = "" + async for chunk in mock_streaming_llm_response(): + full_text += chunk + await chunk_callback(chunk) + return full_text + else: + # Non-streaming response + return "Machine learning is a subset of artificial intelligence." + + client.document_prompt.side_effect = document_prompt_side_effect + return client + + @pytest.fixture + def document_rag_streaming(self, mock_embeddings_client, mock_doc_embeddings_client, + mock_streaming_prompt_client): + """Create DocumentRag instance with streaming support""" + return DocumentRag( + embeddings_client=mock_embeddings_client, + doc_embeddings_client=mock_doc_embeddings_client, + prompt_client=mock_streaming_prompt_client, + verbose=True + ) + + @pytest.mark.asyncio + async def test_document_rag_streaming_basic(self, document_rag_streaming, streaming_chunk_collector): + """Test basic DocumentRAG streaming functionality""" + # Arrange + query = "What is machine learning?" + collector = streaming_chunk_collector() + + # Act + result = await document_rag_streaming.query( + query=query, + user="test_user", + collection="test_collection", + doc_limit=10, + streaming=True, + chunk_callback=collector.collect + ) + + # Assert + assert_streaming_chunks_valid(collector.chunks, min_chunks=1) + assert_callback_invoked(AsyncMock(call_count=len(collector.chunks)), min_calls=1) + + # Verify full response matches concatenated chunks + full_from_chunks = collector.get_full_text() + assert result == full_from_chunks + + # Verify content is reasonable + assert len(result) > 0 + + @pytest.mark.asyncio + async def test_document_rag_streaming_vs_non_streaming(self, document_rag_streaming): + """Test that streaming and non-streaming produce equivalent results""" + # Arrange + query = "What is machine learning?" + user = "test_user" + collection = "test_collection" + doc_limit = 10 + + # Act - Non-streaming + non_streaming_result = await document_rag_streaming.query( + query=query, + user=user, + collection=collection, + doc_limit=doc_limit, + streaming=False + ) + + # Act - Streaming + streaming_chunks = [] + + async def collect(chunk): + streaming_chunks.append(chunk) + + streaming_result = await document_rag_streaming.query( + query=query, + user=user, + collection=collection, + doc_limit=doc_limit, + streaming=True, + chunk_callback=collect + ) + + # Assert - Results should be equivalent + assert streaming_result == non_streaming_result + assert len(streaming_chunks) > 0 + assert "".join(streaming_chunks) == streaming_result + + @pytest.mark.asyncio + async def test_document_rag_streaming_callback_invocation(self, document_rag_streaming): + """Test that chunk callback is invoked correctly""" + # Arrange + callback = AsyncMock() + + # Act + result = await document_rag_streaming.query( + query="test query", + user="test_user", + collection="test_collection", + doc_limit=5, + streaming=True, + chunk_callback=callback + ) + + # Assert + assert callback.call_count > 0 + assert result is not None + + # Verify all callback invocations had string arguments + for call in callback.call_args_list: + assert isinstance(call.args[0], str) + + @pytest.mark.asyncio + async def test_document_rag_streaming_without_callback(self, document_rag_streaming): + """Test streaming parameter without callback (should fall back to non-streaming)""" + # Arrange & Act + result = await document_rag_streaming.query( + query="test query", + user="test_user", + collection="test_collection", + doc_limit=5, + streaming=True, + chunk_callback=None # No callback provided + ) + + # Assert - Should complete without error + assert result is not None + assert isinstance(result, str) + + @pytest.mark.asyncio + async def test_document_rag_streaming_with_no_documents(self, document_rag_streaming, + mock_doc_embeddings_client): + """Test streaming with no documents found""" + # Arrange + mock_doc_embeddings_client.query.return_value = [] # No documents + callback = AsyncMock() + + # Act + result = await document_rag_streaming.query( + query="unknown topic", + user="test_user", + collection="test_collection", + doc_limit=10, + streaming=True, + chunk_callback=callback + ) + + # Assert - Should still produce streamed response + assert result is not None + assert callback.call_count > 0 + + @pytest.mark.asyncio + async def test_document_rag_streaming_error_propagation(self, document_rag_streaming, + mock_embeddings_client): + """Test that errors during streaming are properly propagated""" + # Arrange + mock_embeddings_client.embed.side_effect = Exception("Embeddings error") + callback = AsyncMock() + + # Act & Assert + with pytest.raises(Exception) as exc_info: + await document_rag_streaming.query( + query="test query", + user="test_user", + collection="test_collection", + doc_limit=5, + streaming=True, + chunk_callback=callback + ) + + assert "Embeddings error" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_document_rag_streaming_with_different_doc_limits(self, document_rag_streaming, + mock_doc_embeddings_client): + """Test streaming with various document limits""" + # Arrange + callback = AsyncMock() + doc_limits = [1, 5, 10, 20] + + for limit in doc_limits: + # Reset mocks + mock_doc_embeddings_client.reset_mock() + callback.reset_mock() + + # Act + result = await document_rag_streaming.query( + query="test query", + user="test_user", + collection="test_collection", + doc_limit=limit, + streaming=True, + chunk_callback=callback + ) + + # Assert + assert result is not None + assert callback.call_count > 0 + + # Verify doc_limit was passed correctly + call_args = mock_doc_embeddings_client.query.call_args + assert call_args.kwargs['limit'] == limit + + @pytest.mark.asyncio + async def test_document_rag_streaming_preserves_user_collection(self, document_rag_streaming, + mock_doc_embeddings_client): + """Test that streaming preserves user/collection isolation""" + # Arrange + callback = AsyncMock() + user = "test_user_123" + collection = "test_collection_456" + + # Act + await document_rag_streaming.query( + query="test query", + user=user, + collection=collection, + doc_limit=10, + streaming=True, + chunk_callback=callback + ) + + # Assert - Verify user/collection were passed to document embeddings client + call_args = mock_doc_embeddings_client.query.call_args + assert call_args.kwargs['user'] == user + assert call_args.kwargs['collection'] == collection diff --git a/tests/integration/test_graph_rag_integration.py b/tests/integration/test_graph_rag_integration.py new file mode 100644 index 00000000..a0608819 --- /dev/null +++ b/tests/integration/test_graph_rag_integration.py @@ -0,0 +1,269 @@ +""" +Integration tests for GraphRAG retrieval system + +These tests verify the end-to-end functionality of the GraphRAG system, +testing the coordination between embeddings, graph retrieval, triple querying, and prompt services. +Following the TEST_STRATEGY.md approach for integration testing. + +NOTE: This is the first integration test file for GraphRAG (previously had only unit tests). +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock +from trustgraph.retrieval.graph_rag.graph_rag import GraphRag + + +@pytest.mark.integration +class TestGraphRagIntegration: + """Integration tests for GraphRAG system coordination""" + + @pytest.fixture + def mock_embeddings_client(self): + """Mock embeddings client that returns realistic vector embeddings""" + client = AsyncMock() + client.embed.return_value = [ + [0.1, 0.2, 0.3, 0.4, 0.5], # Realistic 5-dimensional embedding + ] + return client + + @pytest.fixture + def mock_graph_embeddings_client(self): + """Mock graph embeddings client that returns realistic entities""" + client = AsyncMock() + client.query.return_value = [ + "http://trustgraph.ai/e/machine-learning", + "http://trustgraph.ai/e/artificial-intelligence", + "http://trustgraph.ai/e/neural-networks" + ] + return client + + @pytest.fixture + def mock_triples_client(self): + """Mock triples client that returns realistic knowledge graph triples""" + client = AsyncMock() + + # Mock different queries return different triples + async def query_side_effect(s=None, p=None, o=None, limit=None, user=None, collection=None): + # Mock label queries + if p == "http://www.w3.org/2000/01/rdf-schema#label": + if s == "http://trustgraph.ai/e/machine-learning": + return [MagicMock(s=s, p=p, o="Machine Learning")] + elif s == "http://trustgraph.ai/e/artificial-intelligence": + return [MagicMock(s=s, p=p, o="Artificial Intelligence")] + elif s == "http://trustgraph.ai/e/neural-networks": + return [MagicMock(s=s, p=p, o="Neural Networks")] + return [] + + # Mock relationship queries + if s == "http://trustgraph.ai/e/machine-learning": + return [ + MagicMock( + s="http://trustgraph.ai/e/machine-learning", + p="http://trustgraph.ai/is_subset_of", + o="http://trustgraph.ai/e/artificial-intelligence" + ), + MagicMock( + s="http://trustgraph.ai/e/machine-learning", + p="http://www.w3.org/2000/01/rdf-schema#label", + o="Machine Learning" + ) + ] + + return [] + + client.query.side_effect = query_side_effect + return client + + @pytest.fixture + def mock_prompt_client(self): + """Mock prompt client that generates realistic responses""" + client = AsyncMock() + client.kg_prompt.return_value = ( + "Machine learning is a subset of artificial intelligence that enables computers " + "to learn from data without being explicitly programmed. It uses algorithms " + "and statistical models to find patterns in data." + ) + return client + + @pytest.fixture + def graph_rag(self, mock_embeddings_client, mock_graph_embeddings_client, + mock_triples_client, mock_prompt_client): + """Create GraphRag instance with mocked dependencies""" + return GraphRag( + embeddings_client=mock_embeddings_client, + graph_embeddings_client=mock_graph_embeddings_client, + triples_client=mock_triples_client, + prompt_client=mock_prompt_client, + verbose=True + ) + + @pytest.mark.asyncio + async def test_graph_rag_end_to_end_flow(self, graph_rag, mock_embeddings_client, + mock_graph_embeddings_client, mock_triples_client, + mock_prompt_client): + """Test complete GraphRAG pipeline from query to response""" + # Arrange + query = "What is machine learning?" + user = "test_user" + collection = "ml_knowledge" + entity_limit = 50 + triple_limit = 30 + + # Act + result = await graph_rag.query( + query=query, + user=user, + collection=collection, + entity_limit=entity_limit, + triple_limit=triple_limit + ) + + # Assert - Verify service coordination + + # 1. Should compute embeddings for query + mock_embeddings_client.embed.assert_called_once_with(query) + + # 2. Should query graph embeddings to find relevant entities + mock_graph_embeddings_client.query.assert_called_once() + call_args = mock_graph_embeddings_client.query.call_args + assert call_args.kwargs['vectors'] == [[0.1, 0.2, 0.3, 0.4, 0.5]] + assert call_args.kwargs['limit'] == entity_limit + assert call_args.kwargs['user'] == user + assert call_args.kwargs['collection'] == collection + + # 3. Should query triples to build knowledge subgraph + assert mock_triples_client.query.call_count > 0 + + # 4. Should call prompt with knowledge graph + mock_prompt_client.kg_prompt.assert_called_once() + call_args = mock_prompt_client.kg_prompt.call_args + assert call_args.args[0] == query # First arg is query + assert isinstance(call_args.args[1], list) # Second arg is kg (list of triples) + + # Verify final response + assert result is not None + assert isinstance(result, str) + assert "machine learning" in result.lower() + + @pytest.mark.asyncio + async def test_graph_rag_with_different_limits(self, graph_rag, mock_embeddings_client, + mock_graph_embeddings_client): + """Test GraphRAG with various entity and triple limits""" + # Arrange + query = "Explain neural networks" + test_configs = [ + {"entity_limit": 10, "triple_limit": 10}, + {"entity_limit": 50, "triple_limit": 30}, + {"entity_limit": 100, "triple_limit": 100}, + ] + + for config in test_configs: + # Reset mocks + mock_embeddings_client.reset_mock() + mock_graph_embeddings_client.reset_mock() + + # Act + await graph_rag.query( + query=query, + user="test_user", + collection="test_collection", + entity_limit=config["entity_limit"], + triple_limit=config["triple_limit"] + ) + + # Assert + call_args = mock_graph_embeddings_client.query.call_args + assert call_args.kwargs['limit'] == config["entity_limit"] + + @pytest.mark.asyncio + async def test_graph_rag_error_propagation(self, graph_rag, mock_embeddings_client): + """Test that errors from underlying services are properly propagated""" + # Arrange + mock_embeddings_client.embed.side_effect = Exception("Embeddings service error") + + # Act & Assert + with pytest.raises(Exception) as exc_info: + await graph_rag.query( + query="test query", + user="test_user", + collection="test_collection" + ) + + assert "Embeddings service error" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_graph_rag_with_empty_knowledge_graph(self, graph_rag, mock_graph_embeddings_client, + mock_triples_client, mock_prompt_client): + """Test GraphRAG handles empty knowledge graph gracefully""" + # Arrange + mock_graph_embeddings_client.query.return_value = [] # No entities found + mock_triples_client.query.return_value = [] # No triples found + + # Act + result = await graph_rag.query( + query="unknown topic", + user="test_user", + collection="test_collection" + ) + + # Assert + # Should still call prompt client with empty knowledge graph + mock_prompt_client.kg_prompt.assert_called_once() + call_args = mock_prompt_client.kg_prompt.call_args + assert isinstance(call_args.args[1], list) # kg should be a list + assert result is not None + + @pytest.mark.asyncio + async def test_graph_rag_label_caching(self, graph_rag, mock_triples_client): + """Test that label lookups are cached to reduce redundant queries""" + # Arrange + query = "What is machine learning?" + + # First query + await graph_rag.query( + query=query, + user="test_user", + collection="test_collection" + ) + + first_call_count = mock_triples_client.query.call_count + mock_triples_client.reset_mock() + + # Second identical query + await graph_rag.query( + query=query, + user="test_user", + collection="test_collection" + ) + + second_call_count = mock_triples_client.query.call_count + + # Assert - Second query should make fewer triple queries due to caching + # Note: This is a weak assertion because caching behavior depends on + # implementation details, but it verifies the concept + assert second_call_count >= 0 # Should complete without errors + + @pytest.mark.asyncio + async def test_graph_rag_multi_user_isolation(self, graph_rag, mock_graph_embeddings_client): + """Test that different users/collections are properly isolated""" + # Arrange + query = "test query" + user1, collection1 = "user1", "collection1" + user2, collection2 = "user2", "collection2" + + # Act + await graph_rag.query(query=query, user=user1, collection=collection1) + await graph_rag.query(query=query, user=user2, collection=collection2) + + # Assert - Both users should have separate queries + assert mock_graph_embeddings_client.query.call_count == 2 + + # Verify first call + first_call = mock_graph_embeddings_client.query.call_args_list[0] + assert first_call.kwargs['user'] == user1 + assert first_call.kwargs['collection'] == collection1 + + # Verify second call + second_call = mock_graph_embeddings_client.query.call_args_list[1] + assert second_call.kwargs['user'] == user2 + assert second_call.kwargs['collection'] == collection2 diff --git a/tests/integration/test_graph_rag_streaming_integration.py b/tests/integration/test_graph_rag_streaming_integration.py new file mode 100644 index 00000000..d12ec232 --- /dev/null +++ b/tests/integration/test_graph_rag_streaming_integration.py @@ -0,0 +1,248 @@ +""" +Integration tests for GraphRAG streaming functionality + +These tests verify the streaming behavior of GraphRAG, testing token-by-token +response delivery through the complete pipeline. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock +from trustgraph.retrieval.graph_rag.graph_rag import GraphRag +from tests.utils.streaming_assertions import ( + assert_streaming_chunks_valid, + assert_rag_streaming_chunks, + assert_streaming_content_matches, + assert_callback_invoked, +) + + +@pytest.mark.integration +class TestGraphRagStreaming: + """Integration tests for GraphRAG streaming""" + + @pytest.fixture + def mock_embeddings_client(self): + """Mock embeddings client""" + client = AsyncMock() + client.embed.return_value = [[0.1, 0.2, 0.3, 0.4, 0.5]] + return client + + @pytest.fixture + def mock_graph_embeddings_client(self): + """Mock graph embeddings client""" + client = AsyncMock() + client.query.return_value = [ + "http://trustgraph.ai/e/machine-learning", + ] + return client + + @pytest.fixture + def mock_triples_client(self): + """Mock triples client with minimal responses""" + client = AsyncMock() + + async def query_side_effect(s=None, p=None, o=None, limit=None, user=None, collection=None): + if p == "http://www.w3.org/2000/01/rdf-schema#label": + return [MagicMock(s=s, p=p, o="Machine Learning")] + return [] + + client.query.side_effect = query_side_effect + return client + + @pytest.fixture + def mock_streaming_prompt_client(self, mock_streaming_llm_response): + """Mock prompt client with streaming support""" + client = AsyncMock() + + async def kg_prompt_side_effect(query, kg, timeout=600, streaming=False, chunk_callback=None): + if streaming and chunk_callback: + # Simulate streaming chunks + full_text = "" + async for chunk in mock_streaming_llm_response(): + full_text += chunk + await chunk_callback(chunk) + return full_text + else: + # Non-streaming response + return "Machine learning is a subset of artificial intelligence." + + client.kg_prompt.side_effect = kg_prompt_side_effect + return client + + @pytest.fixture + def graph_rag_streaming(self, mock_embeddings_client, mock_graph_embeddings_client, + mock_triples_client, mock_streaming_prompt_client): + """Create GraphRag instance with streaming support""" + return GraphRag( + embeddings_client=mock_embeddings_client, + graph_embeddings_client=mock_graph_embeddings_client, + triples_client=mock_triples_client, + prompt_client=mock_streaming_prompt_client, + verbose=True + ) + + @pytest.mark.asyncio + async def test_graph_rag_streaming_basic(self, graph_rag_streaming, streaming_chunk_collector): + """Test basic GraphRAG streaming functionality""" + # Arrange + query = "What is machine learning?" + collector = streaming_chunk_collector() + + # Act + result = await graph_rag_streaming.query( + query=query, + user="test_user", + collection="test_collection", + streaming=True, + chunk_callback=collector.collect + ) + + # Assert + assert_streaming_chunks_valid(collector.chunks, min_chunks=1) + assert_callback_invoked(AsyncMock(call_count=len(collector.chunks)), min_calls=1) + + # Verify full response matches concatenated chunks + full_from_chunks = collector.get_full_text() + assert result == full_from_chunks + + # Verify content is reasonable + assert "machine" in result.lower() or "learning" in result.lower() + + @pytest.mark.asyncio + async def test_graph_rag_streaming_vs_non_streaming(self, graph_rag_streaming): + """Test that streaming and non-streaming produce equivalent results""" + # Arrange + query = "What is machine learning?" + user = "test_user" + collection = "test_collection" + + # Act - Non-streaming + non_streaming_result = await graph_rag_streaming.query( + query=query, + user=user, + collection=collection, + streaming=False + ) + + # Act - Streaming + streaming_chunks = [] + + async def collect(chunk): + streaming_chunks.append(chunk) + + streaming_result = await graph_rag_streaming.query( + query=query, + user=user, + collection=collection, + streaming=True, + chunk_callback=collect + ) + + # Assert - Results should be equivalent + assert streaming_result == non_streaming_result + assert len(streaming_chunks) > 0 + assert "".join(streaming_chunks) == streaming_result + + @pytest.mark.asyncio + async def test_graph_rag_streaming_callback_invocation(self, graph_rag_streaming): + """Test that chunk callback is invoked correctly""" + # Arrange + callback = AsyncMock() + + # Act + result = await graph_rag_streaming.query( + query="test query", + user="test_user", + collection="test_collection", + streaming=True, + chunk_callback=callback + ) + + # Assert + assert callback.call_count > 0 + assert result is not None + + # Verify all callback invocations had string arguments + for call in callback.call_args_list: + assert isinstance(call.args[0], str) + + @pytest.mark.asyncio + async def test_graph_rag_streaming_without_callback(self, graph_rag_streaming): + """Test streaming parameter without callback (should fall back to non-streaming)""" + # Arrange & Act + result = await graph_rag_streaming.query( + query="test query", + user="test_user", + collection="test_collection", + streaming=True, + chunk_callback=None # No callback provided + ) + + # Assert - Should complete without error + assert result is not None + assert isinstance(result, str) + + @pytest.mark.asyncio + async def test_graph_rag_streaming_with_empty_kg(self, graph_rag_streaming, + mock_graph_embeddings_client): + """Test streaming with empty knowledge graph""" + # Arrange + mock_graph_embeddings_client.query.return_value = [] # No entities + callback = AsyncMock() + + # Act + result = await graph_rag_streaming.query( + query="unknown topic", + user="test_user", + collection="test_collection", + streaming=True, + chunk_callback=callback + ) + + # Assert - Should still produce streamed response + assert result is not None + assert callback.call_count > 0 + + @pytest.mark.asyncio + async def test_graph_rag_streaming_error_propagation(self, graph_rag_streaming, + mock_embeddings_client): + """Test that errors during streaming are properly propagated""" + # Arrange + mock_embeddings_client.embed.side_effect = Exception("Embeddings error") + callback = AsyncMock() + + # Act & Assert + with pytest.raises(Exception) as exc_info: + await graph_rag_streaming.query( + query="test query", + user="test_user", + collection="test_collection", + streaming=True, + chunk_callback=callback + ) + + assert "Embeddings error" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_graph_rag_streaming_preserves_parameters(self, graph_rag_streaming, + mock_graph_embeddings_client): + """Test that streaming preserves all query parameters""" + # Arrange + callback = AsyncMock() + entity_limit = 25 + triple_limit = 15 + + # Act + await graph_rag_streaming.query( + query="test query", + user="test_user", + collection="test_collection", + entity_limit=entity_limit, + triple_limit=triple_limit, + streaming=True, + chunk_callback=callback + ) + + # Assert - Verify parameters were passed to underlying services + call_args = mock_graph_embeddings_client.query.call_args + assert call_args.kwargs['limit'] == entity_limit diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py new file mode 100644 index 00000000..985bcbf1 --- /dev/null +++ b/tests/utils/__init__.py @@ -0,0 +1,29 @@ +"""Test utilities for TrustGraph tests""" + +from .streaming_assertions import ( + assert_streaming_chunks_valid, + assert_streaming_sequence, + assert_agent_streaming_chunks, + assert_rag_streaming_chunks, + assert_streaming_completion, + assert_streaming_content_matches, + assert_no_empty_chunks, + assert_streaming_error_handled, + assert_chunk_types_valid, + assert_streaming_latency_acceptable, + assert_callback_invoked, +) + +__all__ = [ + "assert_streaming_chunks_valid", + "assert_streaming_sequence", + "assert_agent_streaming_chunks", + "assert_rag_streaming_chunks", + "assert_streaming_completion", + "assert_streaming_content_matches", + "assert_no_empty_chunks", + "assert_streaming_error_handled", + "assert_chunk_types_valid", + "assert_streaming_latency_acceptable", + "assert_callback_invoked", +] diff --git a/tests/utils/streaming_assertions.py b/tests/utils/streaming_assertions.py new file mode 100644 index 00000000..cc9164ed --- /dev/null +++ b/tests/utils/streaming_assertions.py @@ -0,0 +1,218 @@ +""" +Streaming test assertion helpers + +Provides reusable assertion functions for validating streaming behavior +across different TrustGraph services. +""" + +from typing import List, Dict, Any, Optional + + +def assert_streaming_chunks_valid(chunks: List[Any], min_chunks: int = 1): + """ + Assert that streaming chunks are valid and non-empty. + + Args: + chunks: List of streaming chunks + min_chunks: Minimum number of expected chunks + """ + assert len(chunks) >= min_chunks, f"Expected at least {min_chunks} chunks, got {len(chunks)}" + assert all(chunk is not None for chunk in chunks), "All chunks should be non-None" + + +def assert_streaming_sequence(chunks: List[Dict[str, Any]], expected_sequence: List[str], key: str = "chunk_type"): + """ + Assert that streaming chunks follow an expected sequence. + + Args: + chunks: List of chunk dictionaries + expected_sequence: Expected sequence of chunk types/values + key: Dictionary key to check (default: "chunk_type") + """ + actual_sequence = [chunk.get(key) for chunk in chunks if key in chunk] + assert actual_sequence == expected_sequence, \ + f"Expected sequence {expected_sequence}, got {actual_sequence}" + + +def assert_agent_streaming_chunks(chunks: List[Dict[str, Any]]): + """ + Assert that agent streaming chunks have valid structure. + + Validates: + - All chunks have chunk_type field + - All chunks have content field + - All chunks have end_of_message field + - All chunks have end_of_dialog field + - Last chunk has end_of_dialog=True + + Args: + chunks: List of agent streaming chunk dictionaries + """ + assert len(chunks) > 0, "Expected at least one chunk" + + for i, chunk in enumerate(chunks): + assert "chunk_type" in chunk, f"Chunk {i} missing chunk_type" + assert "content" in chunk, f"Chunk {i} missing content" + assert "end_of_message" in chunk, f"Chunk {i} missing end_of_message" + assert "end_of_dialog" in chunk, f"Chunk {i} missing end_of_dialog" + + # Validate chunk_type values + valid_types = ["thought", "action", "observation", "final-answer"] + assert chunk["chunk_type"] in valid_types, \ + f"Invalid chunk_type '{chunk['chunk_type']}' at index {i}" + + # Last chunk should signal end of dialog + assert chunks[-1]["end_of_dialog"] is True, \ + "Last chunk should have end_of_dialog=True" + + +def assert_rag_streaming_chunks(chunks: List[Dict[str, Any]]): + """ + Assert that RAG streaming chunks have valid structure. + + Validates: + - All chunks except last have chunk field + - All chunks have end_of_stream field + - Last chunk has end_of_stream=True + - Last chunk may have response field with complete text + + Args: + chunks: List of RAG streaming chunk dictionaries + """ + assert len(chunks) > 0, "Expected at least one chunk" + + for i, chunk in enumerate(chunks): + assert "end_of_stream" in chunk, f"Chunk {i} missing end_of_stream" + + if i < len(chunks) - 1: + # Non-final chunks should have chunk content and end_of_stream=False + assert "chunk" in chunk, f"Chunk {i} missing chunk field" + assert chunk["end_of_stream"] is False, \ + f"Non-final chunk {i} should have end_of_stream=False" + else: + # Final chunk should have end_of_stream=True + assert chunk["end_of_stream"] is True, \ + "Last chunk should have end_of_stream=True" + + +def assert_streaming_completion(chunks: List[Dict[str, Any]], expected_complete_flag: str = "end_of_stream"): + """ + Assert that streaming completed properly. + + Args: + chunks: List of streaming chunk dictionaries + expected_complete_flag: Name of the completion flag field + """ + assert len(chunks) > 0, "Expected at least one chunk" + + # Check that all but last chunk have completion flag = False + for i, chunk in enumerate(chunks[:-1]): + assert chunk.get(expected_complete_flag) is False, \ + f"Non-final chunk {i} should have {expected_complete_flag}=False" + + # Check that last chunk has completion flag = True + assert chunks[-1].get(expected_complete_flag) is True, \ + f"Final chunk should have {expected_complete_flag}=True" + + +def assert_streaming_content_matches(chunks: List, expected_content: str, content_key: str = "chunk"): + """ + Assert that concatenated streaming chunks match expected content. + + Args: + chunks: List of streaming chunks (strings or dicts) + expected_content: Expected complete content after concatenation + content_key: Dictionary key for content (used if chunks are dicts) + """ + if isinstance(chunks[0], dict): + # Extract content from chunk dictionaries + content_chunks = [ + chunk.get(content_key, "") + for chunk in chunks + if chunk.get(content_key) is not None + ] + actual_content = "".join(content_chunks) + else: + # Chunks are already strings + actual_content = "".join(chunks) + + assert actual_content == expected_content, \ + f"Expected content '{expected_content}', got '{actual_content}'" + + +def assert_no_empty_chunks(chunks: List[Dict[str, Any]], content_key: str = "content"): + """ + Assert that no chunks have empty content (except final chunk if it's completion marker). + + Args: + chunks: List of streaming chunk dictionaries + content_key: Dictionary key for content + """ + for i, chunk in enumerate(chunks[:-1]): + content = chunk.get(content_key) + assert content is not None and len(content) > 0, \ + f"Chunk {i} has empty content" + + +def assert_streaming_error_handled(chunks: List[Dict[str, Any]], error_flag: str = "error"): + """ + Assert that streaming error was properly signaled. + + Args: + chunks: List of streaming chunk dictionaries + error_flag: Name of the error flag field + """ + # Check that at least one chunk has error flag + has_error = any(chunk.get(error_flag) is not None for chunk in chunks) + assert has_error, "Expected error flag in at least one chunk" + + # If last chunk has error, should also have completion flag + if chunks[-1].get(error_flag): + # Check for completion flags (either end_of_stream or end_of_dialog) + completion_flags = ["end_of_stream", "end_of_dialog"] + has_completion = any(chunks[-1].get(flag) is True for flag in completion_flags) + assert has_completion, \ + "Error chunk should have completion flag set to True" + + +def assert_chunk_types_valid(chunks: List[Dict[str, Any]], valid_types: List[str], type_key: str = "chunk_type"): + """ + Assert that all chunk types are from a valid set. + + Args: + chunks: List of streaming chunk dictionaries + valid_types: List of valid chunk type values + type_key: Dictionary key for chunk type + """ + for i, chunk in enumerate(chunks): + chunk_type = chunk.get(type_key) + assert chunk_type in valid_types, \ + f"Chunk {i} has invalid type '{chunk_type}', expected one of {valid_types}" + + +def assert_streaming_latency_acceptable(chunk_timestamps: List[float], max_gap_seconds: float = 5.0): + """ + Assert that streaming latency between chunks is acceptable. + + Args: + chunk_timestamps: List of timestamps when chunks were received + max_gap_seconds: Maximum acceptable gap between chunks in seconds + """ + assert len(chunk_timestamps) > 1, "Need at least 2 timestamps to check latency" + + for i in range(1, len(chunk_timestamps)): + gap = chunk_timestamps[i] - chunk_timestamps[i-1] + assert gap <= max_gap_seconds, \ + f"Gap between chunks {i-1} and {i} is {gap:.2f}s, exceeds max {max_gap_seconds}s" + + +def assert_callback_invoked(mock_callback, min_calls: int = 1): + """ + Assert that a streaming callback was invoked minimum number of times. + + Args: + mock_callback: AsyncMock callback object + min_calls: Minimum number of expected calls + """ + assert mock_callback.call_count >= min_calls, \ + f"Expected callback to be called at least {min_calls} times, was called {mock_callback.call_count} times"