Repository: stanfordnlp/dspy
Pull request #7630 (2025-02-01 00:48:40)

Support status streaming
One interesting and useful feature is status streaming, e.g., we tell users that "we are running the tool XXX", "we finished running the tool XXX" so that users don't wait endlessly until the final output. With callback system, the implementation is relatively straightforward. 

Below is a testing code:

``` 
import asyncio
from typing import Any

from pydantic import BaseModel

import dspy
from dspy.utils.streaming import StatusMessage, StatusMessageProvider

lm = dspy.LM("openai/gpt-4o-mini")
dspy.settings.configure(lm=lm)


class CalendarEvent(BaseModel):
    name: str
    date: str
    participants: list[str]


def write_invitation_letter(participant_name: str, event_info: CalendarEvent):
    if participant_name not in event_info.participants:
        return None
    return f"It's my honor to invite {participant_name} to event {event_info.name} on {event_info.date}"


class InvitationSignature(dspy.Signature):
    participant_name: str = dspy.InputField(desc="The name of the participant to invite")
    event_info: CalendarEvent = dspy.InputField(desc="The information about the event")
    invitation_letter: str = dspy.OutputField(desc="The invitation letter to be sent to the participant")


class CustomStatusMessageProvider(StatusMessageProvider):
    def lm_start_status_message(self, instance: Any):
        return "I am just calling the LLM!"

    def lm_end_status_message(self, outputs: Any):
        return "I have finished calling the LLM!"

    def module_start_status_message(self, instance: Any):
        if isinstance(instance, dspy.ReAct):
            return "I am calling ReAct!"


react = dspy.ReAct(InvitationSignature, tools=[write_invitation_letter])
stream_react = dspy.streamify(react, status_message_provider=CustomStatusMessageProvider())


async def main():
    output = stream_react(
        participant_name="Alice",
        event_info=CalendarEvent(
            name="Science Fair",
            date="Friday",
            participants=["Alice", "Bob"],
        ),
    )
    return_value = None
    async for value in output:
        if isinstance(value, StatusMessage):
            print(value.message)
        else:
            return_value = value
    return return_value


print(asyncio.run(main()))
```

the console output looks like:

```
(dspy) (base) [streaming-status][~/Documents/mlflow_team/dspy]$ python3 simple_react_tmp.py
/Users/chen.qian/miniconda3/envs/dspy/lib/python3.12/site-packages/pydantic/_internal/_config.py:345: UserWarning: Valid config keys have changed in V2:
* 'fields' has been removed
  warnings.warn(message, UserWarning)
I am calling ReAct!
I am just calling the LLM!
I have finished calling the LLM!
Calling tool write_invitation_letter...
Tool calling finished! Querying the LLM with tool calling results...
I am just calling the LLM!
I have finished calling the LLM!
I am just calling the LLM!
I have finished calling the LLM!
Prediction(
    trajectory={'thought_0': 'I need to write an invitation letter for Alice to the Science Fair event.', 'tool_name_0': 'write_invitation_letter', 'tool_args_0': {'participant_name': 'Alice', 'event_info': {'name': 'Science Fair', 'date': 'Friday', 'participants': ['Alice', 'Bob']}}, 'observation_0': "It's my honor to invite Alice to event Science Fair on Friday", 'thought_1': 'I have successfully written the invitation letter for Alice. Now I can finish the task.', 'tool_name_1': 'finish', 'tool_args_1': {}, 'observation_1': 'Completed.'},
    reasoning='The invitation letter is crafted to formally invite Alice to the Science Fair event, highlighting her participation and the date of the event. It emphasizes the honor of inviting her, which adds a personal touch to the invitation.',
    invitation_letter='Dear Alice,\n\nIt is my honor to invite you to the Science Fair taking place this Friday. Your participation is highly valued, and we look forward to seeing your contributions to this exciting event.\n\nBest regards,\n\n[Your Name]'
)
```


Patch:
diff --git a/dspy/utils/__init__.py b/dspy/utils/__init__.py
index b2f2d053..b6e95337 100644
--- a/dspy/utils/__init__.py
+++ b/dspy/utils/__init__.py
@@ -1,5 +1,6 @@
 from dspy.utils.callback import BaseCallback, with_callbacks
 from dspy.utils.dummies import DummyLM, DummyVectorizer, dummy_rm
+from dspy.utils.streaming import StatusMessage, StatusMessageProvider, streamify
 
 import os
 import requests
@@ -24,4 +25,7 @@ __all__ = [
     "DummyLM",
     "DummyVectorizer",
     "dummy_rm",
+    "StatusMessage",
+    "StatusMessageProvider",
+    "streamify",
 ]
diff --git a/dspy/utils/asyncify.py b/dspy/utils/asyncify.py
index 0680bc39..7596a03f 100644
--- a/dspy/utils/asyncify.py
+++ b/dspy/utils/asyncify.py
@@ -1,9 +1,10 @@
-from typing import Any, Awaitable, Callable
+from typing import TYPE_CHECKING, Any, Awaitable, Callable
 
 import asyncer
 from anyio import CapacityLimiter
 
-from dspy.primitives.program import Module
+if TYPE_CHECKING:
+    from dspy.primitives.program import Module
 
 _limiter = None
 
@@ -26,7 +27,7 @@ def get_limiter():
     return _limiter
 
 
-def asyncify(program: Module) -> Callable[[Any, Any], Awaitable[Any]]:
+def asyncify(program: "Module") -> Callable[[Any, Any], Awaitable[Any]]:
     """
     Wraps a DSPy program so that it can be called asynchronously. This is useful for running a
     program in parallel with another task (e.g., another DSPy program).
diff --git a/dspy/utils/streaming.py b/dspy/utils/streaming.py
index f51ac11b..d885591c 100644
--- a/dspy/utils/streaming.py
+++ b/dspy/utils/streaming.py
@@ -1,23 +1,178 @@
+import logging
 from asyncio import iscoroutinefunction
-from typing import Any, AsyncGenerator, Awaitable, Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Callable, Dict, Optional
 
 import litellm
 import ujson
 from anyio import create_memory_object_stream, create_task_group
 from anyio.streams.memory import MemoryObjectSendStream
+from asyncer import syncify
 
+from dspy.dsp.utils.settings import settings
 from dspy.primitives.prediction import Prediction
-from dspy.primitives.program import Module
 from dspy.utils.asyncify import asyncify
+from dspy.utils.callback import BaseCallback
 
+logger = logging.getLogger(__name__)
 
-def streamify(program: Module) -> Callable[[Any, Any], Awaitable[Any]]:
+if TYPE_CHECKING:
+    from dspy.primitives.program import Module
+
+
+@dataclass
+class StatusMessage:
+    message: str
+
+
+class StatusMessageProvider:
+    def tool_start_status_message(self, instance: Any, inputs: Dict[str, Any]):
+        return f"Calling tool {instance.name}..."
+
+    def tool_end_status_message(self, outputs: Any):
+        return "Tool calling finished! Querying the LLM with tool calling results..."
+
+    def module_start_status_message(self, instance: Any, inputs: Dict[str, Any]):
+        pass
+
+    def module_end_status_message(self, outputs: Any):
+        pass
+
+    def lm_start_status_message(self, instance: Any, inputs: Dict[str, Any]):
+        pass
+
+    def lm_end_status_message(self, outputs: Any):
+        pass
+
+
+class StatusStreamingCallback(BaseCallback):
+    def __init__(self, status_message_provider: Optional[StatusMessageProvider] = None):
+        self.status_message_provider = status_message_provider or StatusMessageProvider()
+
+    def on_tool_start(
+        self,
+        call_id: str,
+        instance: Any,
+        inputs: Dict[str, Any],
+    ):
+        stream = settings.send_stream
+        if stream is None or instance.name == "finish":
+            return
+
+        @syncify
+        async def send_status():
+            status_message = self.status_message_provider.tool_start_status_message(instance, inputs)
+            if status_message:
+                await stream.send(StatusMessage(status_message))
+
+        send_status()
+
+    def on_tool_end(
+        self,
+        call_id: str,
+        outputs: Optional[Dict[str, Any]],
+        exception: Optional[Exception] = None,
+    ):
+        stream = settings.send_stream
+        if stream is None or outputs == "Completed.":
+            return
+
+        @syncify
+        async def send_status():
+            status_message = self.status_message_provider.tool_end_status_message(outputs)
+            if status_message:
+                await stream.send(StatusMessage(status_message))
+
+        send_status()
+
+    def on_lm_start(
+        self,
+        call_id: str,
+        instance: Any,
+        inputs: Dict[str, Any],
+    ):
+        stream = settings.send_stream
+        if stream is None:
+            return
+
+        @syncify
+        async def send_status():
+            status_message = self.status_message_provider.lm_start_status_message(instance, inputs)
+            if status_message:
+                await stream.send(StatusMessage(status_message))
+
+        send_status()
+
+    def on_lm_end(
+        self,
+        call_id: str,
+        outputs: Optional[Dict[str, Any]],
+        exception: Optional[Exception] = None,
+    ):
+        stream = settings.send_stream
+        if stream is None:
+            return
+
+        @syncify
+        async def send_status():
+            status_message = self.status_message_provider.lm_end_status_message(outputs)
+            if status_message:
+                await stream.send(StatusMessage(status_message))
+
+        send_status()
+
+    def on_module_start(
+        self,
+        call_id: str,
+        instance: Any,
+        inputs: Dict[str, Any],
+    ):
+        stream = settings.send_stream
+        if stream is None:
+            return
+
+        @syncify
+        async def send_status():
+            status_message = self.status_message_provider.module_start_status_message(instance, inputs)
+            if status_message:
+                await stream.send(StatusMessage(status_message))
+
+        send_status()
+
+    def on_module_end(
+        self,
+        call_id: str,
+        outputs: Optional[Dict[str, Any]],
+        exception: Optional[Exception] = None,
+    ):
+        stream = settings.send_stream
+        if stream is None:
+            return
+
+        @syncify
+        async def send_status():
+            status_message = self.status_message_provider.module_end_status_message(outputs)
+            if status_message:
+                await stream.send(StatusMessage(status_message))
+
+        send_status()
+
+
+def streamify(
+    program: "Module",
+    status_message_provider: Optional[StatusMessageProvider] = None,
+) -> Callable[[Any, Any], Awaitable[Any]]:
     """
     Wrap a DSPy program so that it streams its outputs incrementally, rather than returning them
-    all at once.
+    all at once. It also provides status messages to the user to indicate the progress of the program, and users
+    can implement their own status message provider to customize the status messages and what module to generate
+    status messages for.
 
     Args:
         program: The DSPy program to wrap with streaming functionality.
+        status_message_provider: A custom status message generator to use instead of the default one. Users can
+            implement their own status message generator to customize the status messages and what module to generate
+            status messages for.
 
     Returns:
         A function that takes the same arguments as the original program, but returns an async
@@ -26,27 +181,68 @@ def streamify(program: Module) -> Callable[[Any, Any], Awaitable[Any]]:
     Example:
 
     ```python
-    class TestSignature(dspy.Signature):
-        input_text: str = dspy.InputField()
-        output_text: str = dspy.OutputField()
+    import asyncio
+    import dspy
 
     # Create the program and wrap it with streaming functionality
-    program = dspy.streamify(dspy.Predict(TestSignature))
+    program = dspy.streamify(dspy.Predict("q->a"))
 
     # Use the program with streaming output
     async def use_streaming():
-        output_stream = program(input_text="Test")
-        async for value in output_stream:
-            print(value)  # Print each streamed value incrementally
+        output = program(q="Why did a chicken cross the kitchen?")
+        return_value = None
+        async for value in output:
+            if isinstance(value, dspy.Prediction):
+                return_value = value
+            else:
+                print(value)
+        return return_value
+
+    output = asyncio.run(use_streaming())
+    print(output)
     ```
-    """
+
+    Example with custom status message provider:
+    ```python
+    import asyncio
     import dspy
 
+    class MyStatusMessageProvider(StatusMessageProvider):
+        def module_start_status_message(self, instance, inputs):
+            return f"Predicting..."
+
+        def tool_end_status_message(self, outputs):
+            return f"Tool calling finished with output: {outputs}!"
+
+    # Create the program and wrap it with streaming functionality
+    program = dspy.streamify(dspy.Predict("q->a"), status_message_provider=MyStatusMessageProvider())
+
+    # Use the program with streaming output
+    async def use_streaming():
+        output = program(q="Why did a chicken cross the kitchen?")
+        return_value = None
+        async for value in output:
+            if isinstance(value, dspy.Prediction):
+                return_value = value
+            else:
+                print(value)
+        return return_value
+
+    output = asyncio.run(use_streaming())
+    print(output)
+    ```
+    """
     if not iscoroutinefunction(program):
         program = asyncify(program)
 
+    callbacks = settings.callbacks
+    status_streaming_callback = StatusStreamingCallback(status_message_provider)
+    if not any(isinstance(c, StatusStreamingCallback) for c in callbacks):
+        callbacks.append(status_streaming_callback)
+    settings.configure(callbacks=callbacks)
+
     async def generator(args, kwargs, stream: MemoryObjectSendStream):
-        with dspy.settings.context(send_stream=stream):
+        with settings.context(send_stream=stream):
             prediction = await program(*args, **kwargs)
 
         await stream.send(prediction)


Test patch:
diff --git a/tests/utils/test_streaming.py b/tests/utils/test_streaming.py
index 2b9c9fd5..3e890780 100644
--- a/tests/utils/test_streaming.py
+++ b/tests/utils/test_streaming.py
@@ -1,7 +1,7 @@
 import pytest
 
 import dspy
-from dspy.utils.streaming import streaming_response
+from dspy.utils.streaming import StatusMessage, StatusMessageProvider, streaming_response
 from ..test_utils.server import litellm_test_server
 
 
@@ -58,3 +58,67 @@ async def test_streaming_response_yields_expected_response_chunks(litellm_test_s
         assert all(chunk.startswith("data: ") for chunk in output_chunks)
         assert 'data: {"prediction":{"output_text":"Hello!"}}\n\n' in output_chunks
         assert output_chunks[-1] == "data: [DONE]\n\n"
+
+
+@pytest.mark.anyio
+async def test_default_status_streaming():
+    class MyProgram(dspy.Module):
+        def __init__(self):
+            self.generate_question = dspy.Tool(lambda x: f"What color is the {x}?", name="generate_question")
+            self.predict = dspy.Predict("question->answer")
+
+        def __call__(self, x: str):
+            question = self.generate_question(x=x)
+            return self.predict(question=question)
+
+    lm = dspy.utils.DummyLM([{"answer": "red"}, {"answer": "blue"}])
+    with dspy.context(lm=lm):
+        program = dspy.streamify(MyProgram())
+        output = program("sky")
+
+        status_messages = []
+        async for value in output:
+            if isinstance(value, StatusMessage):
+                status_messages.append(value)
+
+    assert len(status_messages) == 2
+    assert status_messages[0].message == "Calling tool generate_question..."
+    assert status_messages[1].message == "Tool calling finished! Querying the LLM with tool calling results..."
+
+
+@pytest.mark.anyio
+async def test_custom_status_streaming():
+    class MyProgram(dspy.Module):
+        def __init__(self):
+            self.generate_question = dspy.Tool(lambda x: f"What color is the {x}?", name="generate_question")
+            self.predict = dspy.Predict("question->answer")
+
+        def __call__(self, x: str):
+            question = self.generate_question(x=x)
+            return self.predict(question=question)
+
+    class MyStatusMessageProvider(StatusMessageProvider):
+        def tool_start_status_message(self, instance, inputs):
+            return f"Tool starting!"
+
+        def tool_end_status_message(self, outputs):
+            return "Tool finished!"
+
+        def module_start_status_message(self, instance, inputs):
+            if isinstance(instance, dspy.Predict):
+                return "Predict starting!"
+
+    lm = dspy.utils.DummyLM([{"answer": "red"}, {"answer": "blue"}])
+    with dspy.context(lm=lm):
+        program = dspy.streamify(MyProgram(), status_message_provider=MyStatusMessageProvider())
+        output = program("sky")
+
+        status_messages = []
+        async for value in output:
+            if isinstance(value, StatusMessage):
+                status_messages.append(value)
+
+        assert len(status_messages) == 3
+        assert status_messages[0].message == "Tool starting!"
+        assert status_messages[1].message == "Tool finished!"
+        assert status_messages[2].message == "Predict starting!"