mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
Update to add streaming tests (#600)
This commit is contained in:
cybermaggedon
GitHub
parent
f0c95a4c5e
commit
f79d0603f7
9 changed files with 1062 additions and 57 deletions
260
tests/unit/test_base/test_prompt_client_streaming.py
Normal file
260
tests/unit/test_base/test_prompt_client_streaming.py
Normal file
|
|
@ -0,0 +1,260 @@
|
|||
"""
|
||||
Unit tests for PromptClient streaming callback behavior.
|
||||
|
||||
These tests verify that the prompt client correctly passes the end_of_stream
|
||||
flag to chunk callbacks, ensuring proper streaming protocol compliance.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
from trustgraph.base.prompt_client import PromptClient
|
||||
from trustgraph.schema import PromptResponse
|
||||
|
||||
|
||||
class TestPromptClientStreamingCallback:
|
||||
"""Test PromptClient streaming callback behavior"""
|
||||
|
||||
@pytest.fixture
|
||||
def prompt_client(self):
|
||||
"""Create a PromptClient with mocked dependencies"""
|
||||
# Mock all the required initialization parameters
|
||||
with patch.object(PromptClient, '__init__', lambda self: None):
|
||||
client = PromptClient()
|
||||
return client
|
||||
|
||||
@pytest.fixture
|
||||
def mock_request_response(self):
|
||||
"""Create a mock request/response handler"""
|
||||
async def mock_request(request, recipient=None, timeout=600):
|
||||
if recipient:
|
||||
# Simulate streaming responses
|
||||
responses = [
|
||||
PromptResponse(text="Hello", object=None, error=None, end_of_stream=False),
|
||||
PromptResponse(text=" world", object=None, error=None, end_of_stream=False),
|
||||
PromptResponse(text="!", object=None, error=None, end_of_stream=False),
|
||||
PromptResponse(text="", object=None, error=None, end_of_stream=True),
|
||||
]
|
||||
for resp in responses:
|
||||
should_stop = await recipient(resp)
|
||||
if should_stop:
|
||||
break
|
||||
else:
|
||||
# Non-streaming response
|
||||
return PromptResponse(text="Hello world!", object=None, error=None)
|
||||
|
||||
return mock_request
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_callback_receives_chunk_and_end_of_stream(self, prompt_client, mock_request_response):
|
||||
"""Test that callback receives both chunk text and end_of_stream flag"""
|
||||
# Arrange
|
||||
prompt_client.request = mock_request_response
|
||||
|
||||
callback = AsyncMock()
|
||||
|
||||
# Act
|
||||
await prompt_client.prompt(
|
||||
id="test-prompt",
|
||||
variables={"query": "test"},
|
||||
streaming=True,
|
||||
chunk_callback=callback
|
||||
)
|
||||
|
||||
# Assert - callback should be called with (chunk, end_of_stream) signature
|
||||
assert callback.call_count == 4
|
||||
|
||||
# Verify first chunk: text + end_of_stream=False
|
||||
assert callback.call_args_list[0] == call("Hello", False)
|
||||
|
||||
# Verify second chunk
|
||||
assert callback.call_args_list[1] == call(" world", False)
|
||||
|
||||
# Verify third chunk
|
||||
assert callback.call_args_list[2] == call("!", False)
|
||||
|
||||
# Verify final chunk: empty text + end_of_stream=True
|
||||
assert callback.call_args_list[3] == call("", True)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_callback_receives_empty_final_chunk(self, prompt_client, mock_request_response):
|
||||
"""Test that empty final chunks are passed to callback"""
|
||||
# Arrange
|
||||
prompt_client.request = mock_request_response
|
||||
|
||||
chunks_received = []
|
||||
|
||||
async def collect_chunks(chunk, end_of_stream):
|
||||
chunks_received.append((chunk, end_of_stream))
|
||||
|
||||
# Act
|
||||
await prompt_client.prompt(
|
||||
id="test-prompt",
|
||||
variables={"query": "test"},
|
||||
streaming=True,
|
||||
chunk_callback=collect_chunks
|
||||
)
|
||||
|
||||
# Assert - should receive the empty final chunk
|
||||
final_chunk = chunks_received[-1]
|
||||
assert final_chunk == ("", True), "Final chunk should be empty string with end_of_stream=True"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_callback_signature_with_non_empty_final_chunk(self, prompt_client):
|
||||
"""Test callback signature when LLM sends non-empty final chunk"""
|
||||
# Arrange
|
||||
async def mock_request_non_empty_final(request, recipient=None, timeout=600):
|
||||
if recipient:
|
||||
# Some LLMs send content in the final chunk
|
||||
responses = [
|
||||
PromptResponse(text="Hello", object=None, error=None, end_of_stream=False),
|
||||
PromptResponse(text=" world!", object=None, error=None, end_of_stream=True),
|
||||
]
|
||||
for resp in responses:
|
||||
should_stop = await recipient(resp)
|
||||
if should_stop:
|
||||
break
|
||||
|
||||
prompt_client.request = mock_request_non_empty_final
|
||||
|
||||
callback = AsyncMock()
|
||||
|
||||
# Act
|
||||
await prompt_client.prompt(
|
||||
id="test-prompt",
|
||||
variables={"query": "test"},
|
||||
streaming=True,
|
||||
chunk_callback=callback
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert callback.call_count == 2
|
||||
assert callback.call_args_list[0] == call("Hello", False)
|
||||
assert callback.call_args_list[1] == call(" world!", True)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_callback_not_called_without_text(self, prompt_client):
|
||||
"""Test that callback is not called for responses without text"""
|
||||
# Arrange
|
||||
async def mock_request_no_text(request, recipient=None, timeout=600):
|
||||
if recipient:
|
||||
# Response with only end_of_stream, no text
|
||||
responses = [
|
||||
PromptResponse(text="Content", object=None, error=None, end_of_stream=False),
|
||||
PromptResponse(text=None, object=None, error=None, end_of_stream=True),
|
||||
]
|
||||
for resp in responses:
|
||||
should_stop = await recipient(resp)
|
||||
if should_stop:
|
||||
break
|
||||
|
||||
prompt_client.request = mock_request_no_text
|
||||
|
||||
callback = AsyncMock()
|
||||
|
||||
# Act
|
||||
await prompt_client.prompt(
|
||||
id="test-prompt",
|
||||
variables={"query": "test"},
|
||||
streaming=True,
|
||||
chunk_callback=callback
|
||||
)
|
||||
|
||||
# Assert - callback should only be called once (for "Content")
|
||||
assert callback.call_count == 1
|
||||
assert callback.call_args_list[0] == call("Content", False)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_synchronous_callback_also_receives_end_of_stream(self, prompt_client):
|
||||
"""Test that synchronous callbacks also receive end_of_stream parameter"""
|
||||
# Arrange
|
||||
async def mock_request(request, recipient=None, timeout=600):
|
||||
if recipient:
|
||||
responses = [
|
||||
PromptResponse(text="test", object=None, error=None, end_of_stream=False),
|
||||
PromptResponse(text="", object=None, error=None, end_of_stream=True),
|
||||
]
|
||||
for resp in responses:
|
||||
should_stop = await recipient(resp)
|
||||
if should_stop:
|
||||
break
|
||||
|
||||
prompt_client.request = mock_request
|
||||
|
||||
callback = MagicMock() # Synchronous mock
|
||||
|
||||
# Act
|
||||
await prompt_client.prompt(
|
||||
id="test-prompt",
|
||||
variables={"query": "test"},
|
||||
streaming=True,
|
||||
chunk_callback=callback
|
||||
)
|
||||
|
||||
# Assert - synchronous callback should also get both parameters
|
||||
assert callback.call_count == 2
|
||||
assert callback.call_args_list[0] == call("test", False)
|
||||
assert callback.call_args_list[1] == call("", True)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kg_prompt_passes_parameters_to_callback(self, prompt_client):
|
||||
"""Test that kg_prompt correctly passes streaming parameters"""
|
||||
# Arrange
|
||||
async def mock_request(request, recipient=None, timeout=600):
|
||||
if recipient:
|
||||
responses = [
|
||||
PromptResponse(text="Answer", object=None, error=None, end_of_stream=False),
|
||||
PromptResponse(text="", object=None, error=None, end_of_stream=True),
|
||||
]
|
||||
for resp in responses:
|
||||
should_stop = await recipient(resp)
|
||||
if should_stop:
|
||||
break
|
||||
|
||||
prompt_client.request = mock_request
|
||||
|
||||
callback = AsyncMock()
|
||||
|
||||
# Act
|
||||
await prompt_client.kg_prompt(
|
||||
query="What is machine learning?",
|
||||
kg=[("subject", "predicate", "object")],
|
||||
streaming=True,
|
||||
chunk_callback=callback
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert callback.call_count == 2
|
||||
assert callback.call_args_list[0] == call("Answer", False)
|
||||
assert callback.call_args_list[1] == call("", True)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_prompt_passes_parameters_to_callback(self, prompt_client):
|
||||
"""Test that document_prompt correctly passes streaming parameters"""
|
||||
# Arrange
|
||||
async def mock_request(request, recipient=None, timeout=600):
|
||||
if recipient:
|
||||
responses = [
|
||||
PromptResponse(text="Summary", object=None, error=None, end_of_stream=False),
|
||||
PromptResponse(text="", object=None, error=None, end_of_stream=True),
|
||||
]
|
||||
for resp in responses:
|
||||
should_stop = await recipient(resp)
|
||||
if should_stop:
|
||||
break
|
||||
|
||||
prompt_client.request = mock_request
|
||||
|
||||
callback = AsyncMock()
|
||||
|
||||
# Act
|
||||
await prompt_client.document_prompt(
|
||||
query="Summarize this",
|
||||
documents=["doc1", "doc2"],
|
||||
streaming=True,
|
||||
chunk_callback=callback
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert callback.call_count == 2
|
||||
assert callback.call_args_list[0] == call("Summary", False)
|
||||
assert callback.call_args_list[1] == call("", True)
|
||||
Loading…
Add table
Add a link
Reference in a new issue