feat: Add async streaming support in HuggingFaceLocalChatGenerator

[mdrazak2001]

Related Issues

• fixes In HuggingFaceLocalChatGenerator add an async version of HFTokenStreamHandler and update type signature for async streaming callback #9391

Proposed Changes:

• added async version of hugging face local streaming handler.

How did you test it?

• added a unit tes

Notes for the reviewer

Checklist

• I have read the contributors guidelines and the code of conduct
• I have updated the related issue with new insights and changes
• I added unit tests and updated the docstrings
• I've used one of the conventional commit types for my PR title: fix:, feat:, build:, chore:, ci:, docs:, style:, refactor:, perf:, test: and added ! in case the PR includes breaking changes.
• I documented my code
• I ran pre-commit hooks and fixed any issue

[coveralls]

Pull Request Test Coverage Report for Build 15588007009

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

• For more information on this, see Tracking coverage changes with pull request builds.
• To avoid this issue with future PRs, see these Recommended CI Configurations.
• For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

• 0 of 0 changed or added relevant lines in 0 files are covered.
• 44 unchanged lines in 2 files lost coverage.
• Overall coverage increased (+0.06%) to 90.277%

---

Files with Coverage Reduction	New Missed Lines	%
utils/hf.py	10	86.77%
components/generators/chat/hugging_face_local.py	34	83.63%

Totals
Change from base Build 15580713019:	0.06%
Covered Lines:	11551
Relevant Lines:	12795

---

💛 - Coveralls

[sjrl]

To make mypy happy could we add an assert here asserting that streaming_callback is of type AsyncStreamingCallbackT?

or update AsyncHFTokenStreamingHandler such that the type hint for stream_handler is StreamingCallbackT

[mpangrazzi]

I've left a comment below!

[sjrl]

Could we also add an integration test for this? So an async version of test_live_run with streaming?

[mpangrazzi]

I would also simplify this, removing asyncio.Event usage and reuse an already available mock.
Something like:

@pytest.mark.asyncio
async def test_run_async_with_streaming_callback(self, model_info_mock, mock_pipeline_with_tokenizer):
    streaming_chunks = []

    async def streaming_callback(chunk: StreamingChunk) -> None:
        streaming_chunks.append(chunk)

    # Create a mock that simulates streaming behavior
    def mock_pipeline_call(*args, **kwargs):
        streamer = kwargs.get("streamer")
        if streamer:
            # Simulate streaming chunks
            streamer.on_finalized_text("Berlin", stream_end=False)
            streamer.on_finalized_text(" is cool", stream_end=True)
        return [{"generated_text": "Berlin is cool"}]

    # Setup the mock pipeline with streaming simulation
    mock_pipeline_with_tokenizer.side_effect = mock_pipeline_call

    generator = HuggingFaceLocalChatGenerator(model="test-model", streaming_callback=streaming_callback)
    generator.pipeline = mock_pipeline_with_tokenizer

    messages = [ChatMessage.from_user("Test message")]
    response = await generator.run_async(messages)

    # Verify streaming chunks were collected
    assert len(streaming_chunks) == 2
    assert streaming_chunks[0].content == "Berlin"
    assert streaming_chunks[1].content == " is cool\n"

    # Verify the final response
    assert isinstance(response, dict)
    assert "replies" in response
    assert len(response["replies"]) == 1
    assert isinstance(response["replies"][0], ChatMessage)
    assert response["replies"][0].text == "Berlin is cool"

WDYT?

[mdrazak2001]

@sjrl: added a live integration test
@mpangrazzi:
refactored the flow a bit:

1. AsyncHFTokenStreamingHandler

• created a asyncio queue that will store chunks as and when the synchronous on_finalized_text is triggered by the hf generator and without blocking.
• added a process_queue coroutine function that'll ingest chunks from the asyncio queue, if its empty itll wait, if it fetches the chunk itll invoke the stream_handler, wait till operation is completed and give main loop control.

2. HuggingFaceLocalChatGenerator._run_streaming_async

• added a asyncio task to process the queue in background

[mpangrazzi]

I've added some comments for possible improvements. Let me know if they are clear enough!

[mpangrazzi]

What about adding a more explicit check here? (can apply also on line 427):

if stop_words:
    for stop_word in stop_words:
        replies = [reply.replace(stop_word, "").rstrip() for reply in replies]

[mpangrazzi]

Can we use select_streaming_callback utility here? (we used it in other generators)

You can get it from:

from haystack.dataclasses.streaming_chunk import select_streaming_callback

so we can avoid assert usage!

[mpangrazzi]

This cleanup logic can be a bit more robust: we can add a short timeout so we can ensure queue is drained:

finally:
    try:
        await asyncio.wait_for(queue_processor, timeout=0.1)
    except asyncio.TimeoutError:
        queue_processor.cancel()
        with suppress(asyncio.CancelledError):
            await queue_processor

WDYT?

[mpangrazzi]

I think that this test is a bit over-engineered. What about something simpler like the following? No need to use e.g. asyncio.Event to check when the streaming is done.

@pytest.mark.integration
@pytest.mark.slow
@pytest.mark.flaky(reruns=3, reruns_delay=10)
@pytest.mark.asyncio
async def test_live_run_async_with_streaming(self, monkeypatch):
    monkeypatch.delenv("HF_API_TOKEN", raising=False)

    streaming_chunks = []

    async def streaming_callback(chunk: StreamingChunk) -> None:
        streaming_chunks.append(chunk)


    llm = HuggingFaceLocalChatGenerator(
        model="Qwen/Qwen2.5-0.5B-Instruct",
        generation_kwargs={"max_new_tokens": 50},
        streaming_callback=streaming_callback,
    )
    llm.warm_up()

    response = await llm.run_async(
        messages=[ChatMessage.from_user("Please create a summary about the following topic: Capital of France")]
    )

    # Verify that the response is not None
    assert len(streaming_chunks) > 0
    assert "replies" in response
    assert isinstance(response["replies"][0], ChatMessage)
    assert response["replies"][0].text is not None

    # Verify that the response contains the word "Paris"
    assert "Paris" in response["replies"][0].text

    # Verify streaming chunks contain actual content
    total_streamed_content = "".join(chunk.content for chunk in streaming_chunks)
    assert len(total_streamed_content.strip()) > 0
    assert "Paris" in total_streamed_content

WDYT?

[mpangrazzi]

I would also simplify this, removing asyncio.Event usage and reuse an already available mock.
Something like:

@pytest.mark.asyncio
async def test_run_async_with_streaming_callback(self, model_info_mock, mock_pipeline_with_tokenizer):
    streaming_chunks = []

    async def streaming_callback(chunk: StreamingChunk) -> None:
        streaming_chunks.append(chunk)

    # Create a mock that simulates streaming behavior
    def mock_pipeline_call(*args, **kwargs):
        streamer = kwargs.get("streamer")
        if streamer:
            # Simulate streaming chunks
            streamer.on_finalized_text("Berlin", stream_end=False)
            streamer.on_finalized_text(" is cool", stream_end=True)
        return [{"generated_text": "Berlin is cool"}]

    # Setup the mock pipeline with streaming simulation
    mock_pipeline_with_tokenizer.side_effect = mock_pipeline_call

    generator = HuggingFaceLocalChatGenerator(model="test-model", streaming_callback=streaming_callback)
    generator.pipeline = mock_pipeline_with_tokenizer

    messages = [ChatMessage.from_user("Test message")]
    response = await generator.run_async(messages)

    # Verify streaming chunks were collected
    assert len(streaming_chunks) == 2
    assert streaming_chunks[0].content == "Berlin"
    assert streaming_chunks[1].content == " is cool\n"

    # Verify the final response
    assert isinstance(response, dict)
    assert "replies" in response
    assert len(response["replies"]) == 1
    assert isinstance(response["replies"][0], ChatMessage)
    assert response["replies"][0].text == "Berlin is cool"

WDYT?

[mpangrazzi]

LGTM - we can address other minor nits later!