mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
404 lines
14 KiB
Python
404 lines
14 KiB
Python
"""
|
|
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
|
|
from trustgraph.schema import Error
|
|
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 with proper Error schema
|
|
error_response = TextCompletionResponse(
|
|
response="",
|
|
error=Error(message="Text completion error", type="processing_error"),
|
|
end_of_stream=True
|
|
)
|
|
await recipient(error_response)
|
|
|
|
text_completion_client.request = failing_request
|
|
|
|
# Mock response producer to capture error response
|
|
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-error"}
|
|
|
|
consumer = MagicMock()
|
|
|
|
# Act - The service catches errors and sends error responses, doesn't raise
|
|
await prompt_processor_streaming.on_request(message, consumer, context)
|
|
|
|
# Assert - Verify error response was sent
|
|
assert response_producer.send.call_count > 0
|
|
|
|
# Check that at least one response contains an error
|
|
error_sent = False
|
|
for call in response_producer.send.call_args_list:
|
|
response = call.args[0]
|
|
if hasattr(response, 'error') and response.error:
|
|
error_sent = True
|
|
assert "Text completion error" in response.error.message
|
|
break
|
|
|
|
assert error_sent, "Expected error response to be sent"
|
|
|
|
@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"
|