From c81d952f9f73aad9127c6f12f93d84d48687e741 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 26 Nov 2025 19:33:40 +0000 Subject: [PATCH] Prompt tests + LLM tests --- .../test_prompt_streaming_integration.py | 389 ++++++++++++++++++ ...t_text_completion_streaming_integration.py | 365 ++++++++++++++++ tests/test_parser_debug.py | 51 +++ 3 files changed, 805 insertions(+) create mode 100644 tests/integration/test_prompt_streaming_integration.py create mode 100644 tests/integration/test_text_completion_streaming_integration.py create mode 100644 tests/test_parser_debug.py diff --git a/tests/integration/test_prompt_streaming_integration.py b/tests/integration/test_prompt_streaming_integration.py new file mode 100644 index 00000000..65b3d68a --- /dev/null +++ b/tests/integration/test_prompt_streaming_integration.py @@ -0,0 +1,389 @@ +""" +Integration tests for Prompt Service Streaming Functionality + +These tests verify the streaming behavior of the Prompt service, +testing how it coordinates between templates and text completion streaming. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock +from trustgraph.prompt.template.service import Processor +from trustgraph.schema import PromptRequest, PromptResponse, TextCompletionResponse +from tests.utils.streaming_assertions import ( + assert_streaming_chunks_valid, + assert_callback_invoked, +) + + +@pytest.mark.integration +class TestPromptStreaming: + """Integration tests for Prompt service streaming""" + + @pytest.fixture + def mock_flow_context_streaming(self): + """Mock flow context with streaming text completion support""" + context = MagicMock() + + # Mock text completion client with streaming + text_completion_client = AsyncMock() + + async def streaming_request(request, recipient=None, timeout=600): + """Simulate streaming text completion""" + if request.streaming and recipient: + # Simulate streaming chunks + chunks = [ + "Machine", " learning", " is", " a", " field", + " of", " artificial", " intelligence", "." + ] + + for i, chunk_text in enumerate(chunks): + is_final = (i == len(chunks) - 1) + response = TextCompletionResponse( + response=chunk_text, + error=None, + end_of_stream=is_final + ) + final = await recipient(response) + if final: + break + + # Final empty chunk + await recipient(TextCompletionResponse( + response="", + error=None, + end_of_stream=True + )) + + text_completion_client.request = streaming_request + + # Mock response producer + response_producer = AsyncMock() + + def context_router(service_name): + if service_name == "text-completion-request": + return text_completion_client + elif service_name == "response": + return response_producer + else: + return AsyncMock() + + context.side_effect = context_router + return context + + @pytest.fixture + def mock_prompt_manager(self): + """Mock PromptManager with simple template""" + manager = MagicMock() + + async def invoke_template(kind, input_vars, llm_function): + """Simulate template invocation""" + # Call the LLM function with simple prompts + system = "You are a helpful assistant." + prompt = f"Question: {input_vars.get('question', 'test')}" + result = await llm_function(system, prompt) + return result + + manager.invoke = invoke_template + return manager + + @pytest.fixture + def prompt_processor_streaming(self, mock_prompt_manager): + """Create Prompt processor with streaming support""" + processor = MagicMock() + processor.manager = mock_prompt_manager + processor.config_key = "prompt" + + # Bind the actual on_request method + processor.on_request = Processor.on_request.__get__(processor, Processor) + + return processor + + @pytest.mark.asyncio + async def test_prompt_streaming_basic(self, prompt_processor_streaming, mock_flow_context_streaming): + """Test basic prompt streaming functionality""" + # Arrange + request = PromptRequest( + id="kg_prompt", + terms={"question": '"What is machine learning?"'}, + streaming=True + ) + + message = MagicMock() + message.value.return_value = request + message.properties.return_value = {"id": "test-123"} + + consumer = MagicMock() + + # Act + await prompt_processor_streaming.on_request( + message, consumer, mock_flow_context_streaming + ) + + # Assert + # Verify response producer was called multiple times (for streaming chunks) + response_producer = mock_flow_context_streaming("response") + assert response_producer.send.call_count > 0 + + # Verify streaming chunks were sent + calls = response_producer.send.call_args_list + assert len(calls) > 1 # Should have multiple chunks + + # Check that responses have end_of_stream flag + for call in calls: + response = call.args[0] + assert isinstance(response, PromptResponse) + assert hasattr(response, 'end_of_stream') + + # Last response should have end_of_stream=True + last_call = calls[-1] + last_response = last_call.args[0] + assert last_response.end_of_stream is True + + @pytest.mark.asyncio + async def test_prompt_streaming_non_streaming_mode(self, prompt_processor_streaming, + mock_flow_context_streaming): + """Test prompt service in non-streaming mode""" + # Arrange + request = PromptRequest( + id="kg_prompt", + terms={"question": '"What is AI?"'}, + streaming=False # Non-streaming + ) + + message = MagicMock() + message.value.return_value = request + message.properties.return_value = {"id": "test-456"} + + consumer = MagicMock() + + # Mock non-streaming text completion + text_completion_client = mock_flow_context_streaming("text-completion-request") + + async def non_streaming_text_completion(system, prompt, streaming=False): + return "AI is the simulation of human intelligence in machines." + + text_completion_client.text_completion = non_streaming_text_completion + + # Act + await prompt_processor_streaming.on_request( + message, consumer, mock_flow_context_streaming + ) + + # Assert + # Verify response producer was called once (non-streaming) + response_producer = mock_flow_context_streaming("response") + # Note: In non-streaming mode, the service sends a single response + assert response_producer.send.call_count >= 1 + + @pytest.mark.asyncio + async def test_prompt_streaming_chunk_forwarding(self, prompt_processor_streaming, + mock_flow_context_streaming): + """Test that prompt service forwards chunks immediately""" + # Arrange + request = PromptRequest( + id="test_prompt", + terms={"question": '"Test query"'}, + streaming=True + ) + + message = MagicMock() + message.value.return_value = request + message.properties.return_value = {"id": "test-789"} + + consumer = MagicMock() + + # Act + await prompt_processor_streaming.on_request( + message, consumer, mock_flow_context_streaming + ) + + # Assert + # Verify chunks were forwarded with proper structure + response_producer = mock_flow_context_streaming("response") + calls = response_producer.send.call_args_list + + for call in calls: + response = call.args[0] + # Each response should have text and end_of_stream fields + assert hasattr(response, 'text') + assert hasattr(response, 'end_of_stream') + + @pytest.mark.asyncio + async def test_prompt_streaming_error_handling(self, prompt_processor_streaming): + """Test error handling during streaming""" + # Arrange + context = MagicMock() + + # Mock text completion client that raises an error + text_completion_client = AsyncMock() + + async def failing_request(request, recipient=None, timeout=600): + if recipient: + # Send error response + error_response = TextCompletionResponse( + response="", + error=MagicMock(message="Text completion error"), + end_of_stream=True + ) + await recipient(error_response) + + text_completion_client.request = failing_request + + def context_router(service_name): + if service_name == "text-completion-request": + return text_completion_client + elif service_name == "response": + return AsyncMock() + else: + return AsyncMock() + + context.side_effect = context_router + + request = PromptRequest( + id="test_prompt", + terms={"question": '"Test"'}, + streaming=True + ) + + message = MagicMock() + message.value.return_value = request + message.properties.return_value = {"id": "test-error"} + + consumer = MagicMock() + + # Act & Assert + with pytest.raises(RuntimeError) as exc_info: + await prompt_processor_streaming.on_request(message, consumer, context) + + assert "Text completion error" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_prompt_streaming_preserves_message_id(self, prompt_processor_streaming, + mock_flow_context_streaming): + """Test that message IDs are preserved through streaming""" + # Arrange + message_id = "unique-test-id-12345" + + request = PromptRequest( + id="test_prompt", + terms={"question": '"Test"'}, + streaming=True + ) + + message = MagicMock() + message.value.return_value = request + message.properties.return_value = {"id": message_id} + + consumer = MagicMock() + + # Act + await prompt_processor_streaming.on_request( + message, consumer, mock_flow_context_streaming + ) + + # Assert + # Verify all responses were sent with the correct message ID + response_producer = mock_flow_context_streaming("response") + calls = response_producer.send.call_args_list + + for call in calls: + properties = call.kwargs.get('properties') + assert properties is not None + assert properties['id'] == message_id + + @pytest.mark.asyncio + async def test_prompt_streaming_empty_response_handling(self, prompt_processor_streaming): + """Test handling of empty responses during streaming""" + # Arrange + context = MagicMock() + + # Mock text completion that sends empty chunks + text_completion_client = AsyncMock() + + async def empty_streaming_request(request, recipient=None, timeout=600): + if request.streaming and recipient: + # Send empty chunk followed by final marker + await recipient(TextCompletionResponse( + response="", + error=None, + end_of_stream=False + )) + await recipient(TextCompletionResponse( + response="", + error=None, + end_of_stream=True + )) + + text_completion_client.request = empty_streaming_request + response_producer = AsyncMock() + + def context_router(service_name): + if service_name == "text-completion-request": + return text_completion_client + elif service_name == "response": + return response_producer + else: + return AsyncMock() + + context.side_effect = context_router + + request = PromptRequest( + id="test_prompt", + terms={"question": '"Test"'}, + streaming=True + ) + + message = MagicMock() + message.value.return_value = request + message.properties.return_value = {"id": "test-empty"} + + consumer = MagicMock() + + # Act + await prompt_processor_streaming.on_request(message, consumer, context) + + # Assert + # Should still send responses even if empty (including final marker) + assert response_producer.send.call_count > 0 + + # Last response should have end_of_stream=True + last_call = response_producer.send.call_args_list[-1] + last_response = last_call.args[0] + assert last_response.end_of_stream is True + + @pytest.mark.asyncio + async def test_prompt_streaming_concatenation_matches_complete(self, prompt_processor_streaming, + mock_flow_context_streaming): + """Test that streaming chunks concatenate to form complete response""" + # Arrange + request = PromptRequest( + id="test_prompt", + terms={"question": '"What is ML?"'}, + streaming=True + ) + + message = MagicMock() + message.value.return_value = request + message.properties.return_value = {"id": "test-concat"} + + consumer = MagicMock() + + # Act + await prompt_processor_streaming.on_request( + message, consumer, mock_flow_context_streaming + ) + + # Assert + # Collect all response texts + response_producer = mock_flow_context_streaming("response") + calls = response_producer.send.call_args_list + + chunk_texts = [] + for call in calls: + response = call.args[0] + if response.text and not response.end_of_stream: + chunk_texts.append(response.text) + + # Verify chunks concatenate to expected result + full_text = "".join(chunk_texts) + assert full_text == "Machine learning is a field of artificial intelligence." diff --git a/tests/integration/test_text_completion_streaming_integration.py b/tests/integration/test_text_completion_streaming_integration.py new file mode 100644 index 00000000..7c9ff18a --- /dev/null +++ b/tests/integration/test_text_completion_streaming_integration.py @@ -0,0 +1,365 @@ +""" +Integration tests for Text Completion Streaming Functionality + +These tests verify the streaming behavior of the Text Completion service, +testing token-by-token response delivery through the complete pipeline. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock +from openai.types.chat import ChatCompletionChunk +from openai.types.chat.chat_completion_chunk import Choice as StreamChoice, ChoiceDelta + +from trustgraph.model.text_completion.openai.llm import Processor +from trustgraph.base import LlmChunk +from tests.utils.streaming_assertions import ( + assert_streaming_chunks_valid, + assert_callback_invoked, +) + + +@pytest.mark.integration +class TestTextCompletionStreaming: + """Integration tests for Text Completion streaming""" + + @pytest.fixture + def mock_streaming_openai_client(self, mock_streaming_llm_response): + """Mock OpenAI client with streaming support""" + client = MagicMock() + + def create_streaming_completion(**kwargs): + """Generator that yields streaming chunks""" + if kwargs.get('stream', False): + # Simulate OpenAI streaming response + chunks_text = [ + "Machine", " learning", " is", " a", " subset", + " of", " AI", " that", " enables", " computers", + " to", " learn", " from", " data", "." + ] + + for text in chunks_text: + delta = ChoiceDelta(content=text, role=None) + choice = StreamChoice(index=0, delta=delta, finish_reason=None) + chunk = ChatCompletionChunk( + id="chatcmpl-streaming", + choices=[choice], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion.chunk" + ) + yield chunk + else: + # Non-streaming response (shouldn't happen in these tests) + raise ValueError("Expected streaming mode") + + client.chat.completions.create.return_value = create_streaming_completion() + return client + + @pytest.fixture + def text_completion_processor_streaming(self, mock_streaming_openai_client): + """Create text completion processor with streaming support""" + processor = MagicMock() + processor.default_model = "gpt-3.5-turbo" + processor.temperature = 0.7 + processor.max_output = 1024 + processor.openai = mock_streaming_openai_client + + # Bind the actual streaming method + processor.generate_content_stream = Processor.generate_content_stream.__get__( + processor, Processor + ) + + return processor + + @pytest.mark.asyncio + async def test_text_completion_streaming_basic(self, text_completion_processor_streaming, + streaming_chunk_collector): + """Test basic text completion streaming functionality""" + # Arrange + system_prompt = "You are a helpful assistant." + user_prompt = "What is machine learning?" + collector = streaming_chunk_collector() + + # Act - Collect all chunks + chunks = [] + async for chunk in text_completion_processor_streaming.generate_content_stream( + system_prompt, user_prompt + ): + chunks.append(chunk) + if chunk.text: + await collector.collect(chunk.text) + + # Assert + assert len(chunks) > 1 # Should have multiple chunks + + # Verify all chunks are LlmChunk objects + for chunk in chunks: + assert isinstance(chunk, LlmChunk) + assert chunk.model == "gpt-3.5-turbo" + + # Verify last chunk has is_final=True + assert chunks[-1].is_final is True + + # Verify we got meaningful content + full_text = collector.get_full_text() + assert "machine" in full_text.lower() or "learning" in full_text.lower() + + @pytest.mark.asyncio + async def test_text_completion_streaming_chunk_structure(self, text_completion_processor_streaming): + """Test that streaming chunks have correct structure""" + # Arrange + system_prompt = "You are a helpful assistant." + user_prompt = "Explain AI." + + # Act + chunks = [] + async for chunk in text_completion_processor_streaming.generate_content_stream( + system_prompt, user_prompt + ): + chunks.append(chunk) + + # Assert - Verify chunk structure + for i, chunk in enumerate(chunks[:-1]): # All except last + assert isinstance(chunk, LlmChunk) + assert chunk.text is not None + assert chunk.model == "gpt-3.5-turbo" + assert chunk.is_final is False + + # Last chunk should be final marker + final_chunk = chunks[-1] + assert final_chunk.is_final is True + assert final_chunk.model == "gpt-3.5-turbo" + + @pytest.mark.asyncio + async def test_text_completion_streaming_concatenation(self, text_completion_processor_streaming): + """Test that chunks concatenate to form complete response""" + # Arrange + system_prompt = "You are a helpful assistant." + user_prompt = "What is AI?" + + # Act - Collect all chunk texts + chunk_texts = [] + async for chunk in text_completion_processor_streaming.generate_content_stream( + system_prompt, user_prompt + ): + if chunk.text and not chunk.is_final: + chunk_texts.append(chunk.text) + + # Assert + full_text = "".join(chunk_texts) + assert len(full_text) > 0 + assert len(chunk_texts) > 1 # Should have multiple chunks + + # Verify completeness - should be a coherent sentence + assert full_text == "Machine learning is a subset of AI that enables computers to learn from data." + + @pytest.mark.asyncio + async def test_text_completion_streaming_final_marker(self, text_completion_processor_streaming): + """Test that final chunk properly marks end of stream""" + # Arrange + system_prompt = "You are a helpful assistant." + user_prompt = "Test query" + + # Act + chunks = [] + async for chunk in text_completion_processor_streaming.generate_content_stream( + system_prompt, user_prompt + ): + chunks.append(chunk) + + # Assert + # Should have at least content chunks + final marker + assert len(chunks) >= 2 + + # Only the last chunk should have is_final=True + for chunk in chunks[:-1]: + assert chunk.is_final is False + + assert chunks[-1].is_final is True + + @pytest.mark.asyncio + async def test_text_completion_streaming_model_parameter(self, mock_streaming_openai_client): + """Test that model parameter is preserved in streaming""" + # Arrange + processor = MagicMock() + processor.default_model = "gpt-4" + processor.temperature = 0.5 + processor.max_output = 2048 + processor.openai = mock_streaming_openai_client + processor.generate_content_stream = Processor.generate_content_stream.__get__( + processor, Processor + ) + + # Act + chunks = [] + async for chunk in processor.generate_content_stream("System", "Prompt"): + chunks.append(chunk) + + # Assert + # Verify OpenAI was called with correct model + call_args = mock_streaming_openai_client.chat.completions.create.call_args + assert call_args.kwargs['model'] == "gpt-4" + assert call_args.kwargs['temperature'] == 0.5 + assert call_args.kwargs['max_tokens'] == 2048 + assert call_args.kwargs['stream'] is True + + # Verify chunks have correct model + for chunk in chunks: + assert chunk.model == "gpt-4" + + @pytest.mark.asyncio + async def test_text_completion_streaming_temperature_parameter(self, mock_streaming_openai_client): + """Test that temperature parameter is applied in streaming""" + # Arrange + temperatures = [0.0, 0.5, 1.0, 1.5] + + for temp in temperatures: + processor = MagicMock() + processor.default_model = "gpt-3.5-turbo" + processor.temperature = temp + processor.max_output = 1024 + processor.openai = mock_streaming_openai_client + processor.generate_content_stream = Processor.generate_content_stream.__get__( + processor, Processor + ) + + # Act + chunks = [] + async for chunk in processor.generate_content_stream("System", "Prompt"): + chunks.append(chunk) + if chunk.is_final: + break + + # Assert + call_args = mock_streaming_openai_client.chat.completions.create.call_args + assert call_args.kwargs['temperature'] == temp + + # Reset mock for next iteration + mock_streaming_openai_client.reset_mock() + + @pytest.mark.asyncio + async def test_text_completion_streaming_error_propagation(self): + """Test that errors during streaming are properly propagated""" + # Arrange + mock_client = MagicMock() + + def failing_stream(**kwargs): + yield from [] + raise Exception("Streaming error") + + mock_client.chat.completions.create.return_value = failing_stream() + + processor = MagicMock() + processor.default_model = "gpt-3.5-turbo" + processor.temperature = 0.7 + processor.max_output = 1024 + processor.openai = mock_client + processor.generate_content_stream = Processor.generate_content_stream.__get__( + processor, Processor + ) + + # Act & Assert + with pytest.raises(Exception) as exc_info: + async for chunk in processor.generate_content_stream("System", "Prompt"): + pass + + assert "Streaming error" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_text_completion_streaming_empty_chunks_filtered(self, mock_streaming_openai_client): + """Test that empty chunks are handled correctly""" + # Arrange - Mock that returns some empty chunks + def create_streaming_with_empties(**kwargs): + chunks_text = ["Hello", "", " world", "", "!"] + + for text in chunks_text: + delta = ChoiceDelta(content=text if text else None, role=None) + choice = StreamChoice(index=0, delta=delta, finish_reason=None) + chunk = ChatCompletionChunk( + id="chatcmpl-streaming", + choices=[choice], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion.chunk" + ) + yield chunk + + mock_streaming_openai_client.chat.completions.create.return_value = create_streaming_with_empties() + + processor = MagicMock() + processor.default_model = "gpt-3.5-turbo" + processor.temperature = 0.7 + processor.max_output = 1024 + processor.openai = mock_streaming_openai_client + processor.generate_content_stream = Processor.generate_content_stream.__get__( + processor, Processor + ) + + # Act + chunks = [] + async for chunk in processor.generate_content_stream("System", "Prompt"): + chunks.append(chunk) + + # Assert - Only non-empty chunks should be yielded (plus final marker) + text_chunks = [c for c in chunks if not c.is_final] + assert len(text_chunks) == 3 # "Hello", " world", "!" + assert "".join(c.text for c in text_chunks) == "Hello world!" + + @pytest.mark.asyncio + async def test_text_completion_streaming_prompt_construction(self, mock_streaming_openai_client): + """Test that system and user prompts are correctly combined for streaming""" + # Arrange + processor = MagicMock() + processor.default_model = "gpt-3.5-turbo" + processor.temperature = 0.7 + processor.max_output = 1024 + processor.openai = mock_streaming_openai_client + processor.generate_content_stream = Processor.generate_content_stream.__get__( + processor, Processor + ) + + system_prompt = "You are an expert." + user_prompt = "Explain quantum physics." + + # Act + chunks = [] + async for chunk in processor.generate_content_stream(system_prompt, user_prompt): + chunks.append(chunk) + if chunk.is_final: + break + + # Assert - Verify prompts were combined correctly + call_args = mock_streaming_openai_client.chat.completions.create.call_args + messages = call_args.kwargs['messages'] + assert len(messages) == 1 + + message_content = messages[0]['content'][0]['text'] + assert system_prompt in message_content + assert user_prompt in message_content + assert message_content.startswith(system_prompt) + + @pytest.mark.asyncio + async def test_text_completion_streaming_chunk_count(self, text_completion_processor_streaming): + """Test that streaming produces expected number of chunks""" + # Arrange + system_prompt = "You are a helpful assistant." + user_prompt = "Test" + + # Act + chunks = [] + async for chunk in text_completion_processor_streaming.generate_content_stream( + system_prompt, user_prompt + ): + chunks.append(chunk) + + # Assert + # Should have 15 content chunks + 1 final marker = 16 total + assert len(chunks) == 16 + + # 15 content chunks + content_chunks = [c for c in chunks if not c.is_final] + assert len(content_chunks) == 15 + + # 1 final marker + final_chunks = [c for c in chunks if c.is_final] + assert len(final_chunks) == 1 diff --git a/tests/test_parser_debug.py b/tests/test_parser_debug.py new file mode 100644 index 00000000..2ad7867f --- /dev/null +++ b/tests/test_parser_debug.py @@ -0,0 +1,51 @@ +""" +Debug script to understand StreamingReActParser behavior +""" +import sys +sys.path.insert(0, '/home/mark/work/trustgraph.ai/trustgraph/trustgraph-flow') + +from trustgraph.agent.react.streaming_parser import StreamingReActParser + +# Test the parser with different chunking strategies +full_text = """Thought: I need to search for information about machine learning. +Action: knowledge_query +Args: { + "question": "What is machine learning?" +}""" + +print("=" * 60) +print("TEST 1: Send as single chunk") +print("=" * 60) + +parser1 = StreamingReActParser() +parser1.feed(full_text) +parser1.finalize() +result1 = parser1.get_result() +print(f"Result: {result1}") +print(f"Action name: '{result1.name}'") +print(f"Arguments: {result1.arguments}") + +print("\n" + "=" * 60) +print("TEST 2: Send line-by-line chunks") +print("=" * 60) + +parser2 = StreamingReActParser() +chunks = [ + "Thought: I need to search for information about machine learning.\n", + "Action: knowledge_query\n", + "Args: {\n", + ' "question": "What is machine learning?"\n', + "}" +] + +for i, chunk in enumerate(chunks): + print(f"\nChunk {i}: {repr(chunk)}") + print(f" Before: state={parser2.state}, action_buffer='{parser2.action_buffer}'") + parser2.feed(chunk) + print(f" After: state={parser2.state}, action_buffer='{parser2.action_buffer}'") + +parser2.finalize() +result2 = parser2.get_result() +print(f"\nFinal Result: {result2}") +print(f"Action name: '{result2.name}'") +print(f"Arguments: {result2.arguments}")