mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-27 09:26:22 +02:00
Update to add streaming tests (#600)
This commit is contained in:
parent
f0c95a4c5e
commit
f79d0603f7
9 changed files with 1062 additions and 57 deletions
|
|
@ -59,9 +59,16 @@ class TestGraphRagStreaming:
|
|||
full_text = "Machine learning is a subset of artificial intelligence that focuses on algorithms that learn from data."
|
||||
|
||||
if streaming and chunk_callback:
|
||||
# Simulate streaming chunks
|
||||
# Simulate streaming chunks with end_of_stream flags
|
||||
chunks = []
|
||||
async for chunk in mock_streaming_llm_response():
|
||||
await chunk_callback(chunk)
|
||||
chunks.append(chunk)
|
||||
|
||||
# Send all chunks with end_of_stream=False except the last
|
||||
for i, chunk in enumerate(chunks):
|
||||
is_final = (i == len(chunks) - 1)
|
||||
await chunk_callback(chunk, is_final)
|
||||
|
||||
return full_text
|
||||
else:
|
||||
# Non-streaming response - same text
|
||||
|
|
@ -102,6 +109,9 @@ class TestGraphRagStreaming:
|
|||
assert_streaming_chunks_valid(collector.chunks, min_chunks=1)
|
||||
assert_callback_invoked(AsyncMock(call_count=len(collector.chunks)), min_calls=1)
|
||||
|
||||
# Verify streaming protocol compliance
|
||||
collector.verify_streaming_protocol()
|
||||
|
||||
# Verify full response matches concatenated chunks
|
||||
full_from_chunks = collector.get_full_text()
|
||||
assert result == full_from_chunks
|
||||
|
|
@ -128,7 +138,7 @@ class TestGraphRagStreaming:
|
|||
# Act - Streaming
|
||||
streaming_chunks = []
|
||||
|
||||
async def collect(chunk):
|
||||
async def collect(chunk, end_of_stream):
|
||||
streaming_chunks.append(chunk)
|
||||
|
||||
streaming_result = await graph_rag_streaming.query(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue