If you are from Haystack, this blog is not for you!
Need for this blog
Love the haystack team and I hope they roll out a RC soon on for the or .
If not, follow the code snippet and I'll try to explain why we have it implemented this way.
async def run_pipeline(self, pipeline: AsyncPipeline, input_data: Dict[str, Any]) -> AsyncGenerator[str, None]:
request_collector = ChunkCollector() # code snippet is below
loop = asyncio.get_running_loop()
# Create sync wrapper for async callback
async def async_callback(chunk):
await collect_chunk(request_collector.queue, chunk) # code snippet is below
def sync_callback(chunk):
# Use run_coroutine_threadsafe instead of create_task
future = asyncio.run_coroutine_threadsafe(async_callback(chunk), loop)
try:
# Wait for the coroutine to complete
future.result()
except Exception as e:
print(f"Error in sync_callback: {str(e)}")
# Set callbacks using sync wrapper
input_data["generator"]["streaming_callback"] = sync_callback
async def pipeline_runner():
try:
async for _ in pipeline.run(input_data):
pass
finally:
await request_collector.queue.put(None)
# Create task for pipeline
pipeline_task = asyncio.create_task(pipeline_runner())
try:
# Start yielding chunks
async for chunk in request_collector.generator():
yield chunk
finally:
# Ensure pipeline task is cleaned up
if not pipeline_task.done():
pipeline_task.cancel()
try:
await pipeline_task
except asyncio.CancelledError:
pass
Some Q/A's for you
* Why dont we do a direct async callback without wrapping?
async def callback(chunk):
await collect_chunk(request_collector.queue, chunk)
input_data["generator"]["streaming_callback"] = callback
The generator is calling the callback synchronously, but we're passing an async function. So, we need a sync wrapper around our async callback, instead of trying to await an async generator.
* Why dont we just create a task?
def sync_callback(chunk):
asyncio.create_task(async_callback(chunk))
The callback is being called from a different thread where there's no event loop. So, we need a thread-safe way to schedule the callback.
Chunking in SSE Format
We need to define the request_collector which handles the queue, stores the chunks and also yeilds the chunks from the queue(in SSE format).
from typing import AsyncGenerator
import uuid
import json
from asyncio import Queue
from haystack.dataclasses import StreamingChunk
class ChunkCollector:
"""Collects and queues streaming chunks."""
def __init__(self):
self.queue = Queue()
async def generator(self) -> AsyncGenerator[str, None]:
"""Yields chunks from the queue."""
# Send initial metadata event
yield 'event: metadata\n' + f'data: {{"run_id": "{uuid.uuid4()}"}}\n\n'
while True:
chunk = await self.queue.get()
if chunk is None:
# Send end event
yield 'event: end\n\n'
break
# Send data event
yield f'event: data\ndata: {json.dumps(chunk)}\n\n'
async def collect_chunk(queue: Queue, chunk: StreamingChunk):
"""
Collect chunks and store them in the queue.
:param queue: Queue to store the chunks
:param chunk: StreamingChunk to be collected
"""
if chunk and chunk.content:
await queue.put(chunk.content)
Frontend
You can directcly use EventSource or fetch. For this tutorial, let's use here as sockets will be just 👌🏽
If you found this blog helpful, just send a good vibe my way—whether it’s my (research taking off || side project getting some honey || landing some gigs) ✌🏽!
SOCIAL SHARE CARD GENERATOR