Repository: stanfordnlp/dspy

Issue:
[Bug] Usage tracking does not work with Predict
### What happened?

Using the code sample provided in PR #8014 , I am not able to get usage tracking when replacing ChainOfThought with Predict. It always returns None. See below for simple example.

### Steps to reproduce

```
import dspy

dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=False), track_usage=True)

predict = dspy.Predict("question -> answer")
output = predict(question="What is the capital of France?")
print(output.get_lm_usage())
```

### DSPy version

2.6.22

Discussion:

Solution patch:
diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py
index 023ade472a..f77bc40ecb 100644
--- a/dspy/clients/lm.py
+++ b/dspy/clients/lm.py
@@ -231,35 +231,25 @@ def dump_state(self):
         return {key: getattr(self, key) for key in state_keys} | self.kwargs
 
 
-def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}):
-    retry_kwargs = dict(
-        retry_policy=_get_litellm_retry_policy(num_retries),
-        retry_strategy="exponential_backoff_retry",
-        # In LiteLLM version 1.55.3 (the first version that supports retry_policy as an argument
-        # to completion()), the default value of max_retries is non-zero for certain providers, and
-        # max_retries is stacked on top of the retry_policy. To avoid this, we set max_retries=0
-        max_retries=0,
-    )
-
+def _get_stream_completion_fn(
+    request: Dict[str, Any],
+    retry_kwargs: Dict[str, Any],
+    cache_kwargs: Dict[str, Any],
+    sync=True,
+):
     stream = dspy.settings.send_stream
     caller_predict = dspy.settings.caller_predict
+
     if stream is None:
-        # If `streamify` is not used, or if the exact predict doesn't need to be streamed,
-        # we can just return the completion without streaming.
-        return litellm.completion(
-            cache=cache,
-            **retry_kwargs,
-            **request,
-        )
+        return None
 
     # The stream is already opened, and will be closed by the caller.
     stream = cast(MemoryObjectSendStream, stream)
     caller_predict_id = id(caller_predict) if caller_predict else None
 
-    @syncify
-    async def stream_completion():
+    async def stream_completion(request: Dict[str, Any], retry_kwargs: Dict[str, Any], cache_kwargs: Dict[str, Any]):
         response = await litellm.acompletion(
-            cache=cache,
+            cache=cache_kwargs,
             stream=True,
             **retry_kwargs,
             **request,
@@ -273,6 +263,37 @@ async def stream_completion():
             await stream.send(chunk)
         return litellm.stream_chunk_builder(chunks)
 
+    def sync_stream_completion():
+        syncified_stream_completion = syncify(stream_completion)
+        return syncified_stream_completion(request, retry_kwargs, cache_kwargs)
+
+    async def async_stream_completion():
+        return await stream_completion(request, retry_kwargs, cache_kwargs)
+
+    if sync:
+        return sync_stream_completion
+    else:
+        return async_stream_completion
+
+
+def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}):
+    retry_kwargs = dict(
+        retry_policy=_get_litellm_retry_policy(num_retries),
+        retry_strategy="exponential_backoff_retry",
+        # In LiteLLM version 1.55.3 (the first version that supports retry_policy as an argument
+        # to completion()), the default value of max_retries is non-zero for certain providers, and
+        # max_retries is stacked on top of the retry_policy. To avoid this, we set max_retries=0
+        max_retries=0,
+    )
+
+    stream_completion = _get_stream_completion_fn(request, retry_kwargs, cache, sync=True)
+    if stream_completion is None:
+        return litellm.completion(
+            cache=cache,
+            **retry_kwargs,
+            **request,
+        )
+
     return stream_completion()
 
 
@@ -314,11 +335,15 @@ async def alitellm_completion(request: Dict[str, Any], num_retries: int, cache={
         max_retries=0,
     )
 
-    return await litellm.acompletion(
-        cache=cache,
-        **retry_kwargs,
-        **request,
-    )
+    stream_completion = _get_stream_completion_fn(request, retry_kwargs, cache, sync=False)
+    if stream_completion is None:
+        return await litellm.acompletion(
+            cache=cache,
+            **retry_kwargs,
+            **request,
+        )
+
+    return await stream_completion()
 
 
 async def alitellm_text_completion(
diff --git a/dspy/predict/predict.py b/dspy/predict/predict.py
index bde6326165..b4a94a3672 100644
--- a/dspy/predict/predict.py
+++ b/dspy/predict/predict.py
@@ -11,7 +11,6 @@
 from dspy.primitives.prediction import Prediction
 from dspy.primitives.program import Module
 from dspy.signatures.signature import ensure_signature
-from dspy.utils.callback import with_callbacks
 
 logger = logging.getLogger(__name__)
 
@@ -71,13 +70,24 @@ def load_state(self, state):
 
         return self
 
-    @with_callbacks
-    def __call__(self, **kwargs):
-        return self.forward(**kwargs)
+    def __call__(self, *args, **kwargs):
+        if args:
+            raise ValueError(
+                "Positional arguments are not allowed when calling `dspy.Predict`, must use keyword arguments "
+                "that match your signature input fields. For example: "
+                "dspy.Predict('question -> answer')(question='What is the capital of France?')"
+            )
+
+        return super().__call__(**kwargs)
 
-    @with_callbacks
-    async def acall(self, **kwargs):
-        return await self.aforward(**kwargs)
+    async def acall(self, *args, **kwargs):
+        if args:
+            raise ValueError(
+                "Positional arguments are not allowed when calling `dspy.Predict`, must use keyword arguments "
+                "that match your signature input fields. For example: "
+                "dspy.Predict('question -> answer').acall(question='What is the capital of France?')"
+            )
+        return await super().acall(**kwargs)
 
     def _forward_preprocess(self, **kwargs):
         # Extract the three privileged keyword arguments.
@@ -114,17 +124,20 @@ def _forward_postprocess(self, completions, signature, **kwargs):
             trace.append((self, {**kwargs}, pred))
         return pred
 
+    def _should_stream(self):
+        stream_listeners = settings.stream_listeners or []
+        should_stream = settings.send_stream is not None
+        if should_stream and len(stream_listeners) > 0:
+            should_stream = any(stream_listener.predict == self for stream_listener in stream_listeners)
+
+        return should_stream
+
     def forward(self, **kwargs):
         lm, config, signature, demos, kwargs = self._forward_preprocess(**kwargs)
 
         adapter = settings.adapter or ChatAdapter()
 
-        stream_listeners = settings.stream_listeners or []
-        stream = settings.send_stream is not None
-        if stream and len(stream_listeners) > 0:
-            stream = any(stream_listener.predict == self for stream_listener in stream_listeners)
-
-        if stream:
+        if self._should_stream():
             with settings.context(caller_predict=self):
                 completions = adapter(lm, lm_kwargs=config, signature=signature, demos=demos, inputs=kwargs)
         else:
@@ -137,7 +150,13 @@ async def aforward(self, **kwargs):
         lm, config, signature, demos, kwargs = self._forward_preprocess(**kwargs)
 
         adapter = settings.adapter or ChatAdapter()
-        completions = await adapter.acall(lm, lm_kwargs=config, signature=signature, demos=demos, inputs=kwargs)
+        if self._should_stream():
+            with settings.context(caller_predict=self):
+                completions = await adapter.acall(lm, lm_kwargs=config, signature=signature, demos=demos, inputs=kwargs)
+        else:
+            with settings.context(send_stream=None):
+                completions = await adapter.acall(lm, lm_kwargs=config, signature=signature, demos=demos, inputs=kwargs)
+
         return self._forward_postprocess(completions, signature, **kwargs)
 
     def update_config(self, **kwargs):
diff --git a/dspy/primitives/tool.py b/dspy/primitives/tool.py
index a6a5ae887c..3edea2d52c 100644
--- a/dspy/primitives/tool.py
+++ b/dspy/primitives/tool.py
@@ -146,6 +146,7 @@ def __call__(self, **kwargs):
             raise ValueError("You are calling `__call__` on an async tool, please use `acall` instead.")
         return result
 
+    @with_callbacks
     async def acall(self, **kwargs):
         parsed_kwargs = self._validate_and_parse_args(**kwargs)
         result = self.func(**parsed_kwargs)
diff --git a/dspy/streaming/messages.py b/dspy/streaming/messages.py
index bda50341bc..9a3cf74c6c 100644
--- a/dspy/streaming/messages.py
+++ b/dspy/streaming/messages.py
@@ -1,8 +1,7 @@
+import asyncio
 from dataclasses import dataclass
 from typing import Any, Dict, Optional
 
-from asyncer import syncify
-
 from dspy.dsp.utils.settings import settings
 from dspy.utils.callback import BaseCallback
 
@@ -21,6 +20,39 @@ class StatusMessage:
     message: str
 
 
+def sync_send_to_stream(stream, message):
+    """Send message to stream in a sync context, regardless of whether the caller is async or not."""
+    # Try to get current event loop, create one if none exists
+    try:
+        loop = asyncio.get_event_loop()
+    except RuntimeError:
+        # "There is no current event loop in thread" error
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+    # If we're in an async context
+    if loop.is_running():
+        # In an async context, we need to use an approach that doesn't block
+        # Create a new thread and run a new event loop there
+        import concurrent.futures
+
+        def run_async_in_new_loop():
+            new_loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(new_loop)
+            try:
+                return new_loop.run_until_complete(stream.send(message))
+            finally:
+                new_loop.close()
+
+        # Run the function in a separate thread
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+            future = executor.submit(run_async_in_new_loop)
+            return future.result()  # This shouldn't hang now
+    else:
+        # We're in a sync context, use run_until_complete
+        return loop.run_until_complete(stream.send(message))
+
+
 class StatusMessageProvider:
     """Provides customizable status message streaming for DSPy programs.
 
@@ -80,13 +112,9 @@ def on_tool_start(
         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()
+        status_message = self.status_message_provider.tool_start_status_message(instance, inputs)
+        if status_message:
+            sync_send_to_stream(stream, StatusMessage(status_message))
 
     def on_tool_end(
         self,
@@ -98,13 +126,9 @@ def on_tool_end(
         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()
+        status_message = self.status_message_provider.tool_end_status_message(outputs)
+        if status_message:
+            sync_send_to_stream(stream, StatusMessage(status_message))
 
     def on_lm_start(
         self,
@@ -116,13 +140,9 @@ def on_lm_start(
         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()
+        status_message = self.status_message_provider.lm_start_status_message(instance, inputs)
+        if status_message:
+            sync_send_to_stream(stream, StatusMessage(status_message))
 
     def on_lm_end(
         self,
@@ -134,13 +154,9 @@ def on_lm_end(
         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()
+        status_message = self.status_message_provider.lm_end_status_message(outputs)
+        if status_message:
+            sync_send_to_stream(stream, StatusMessage(status_message))
 
     def on_module_start(
         self,
@@ -152,13 +168,9 @@ def on_module_start(
         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()
+        status_message = self.status_message_provider.module_start_status_message(instance, inputs)
+        if status_message:
+            sync_send_to_stream(stream, StatusMessage(status_message))
 
     def on_module_end(
         self,
@@ -170,10 +182,6 @@ def on_module_end(
         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()
+        status_message = self.status_message_provider.module_end_status_message(outputs)
+        if status_message:
+            sync_send_to_stream(stream, StatusMessage(status_message))
diff --git a/dspy/streaming/streamify.py b/dspy/streaming/streamify.py
index 1a91af6d29..c99f2a154a 100644
--- a/dspy/streaming/streamify.py
+++ b/dspy/streaming/streamify.py
@@ -25,6 +25,7 @@ def streamify(
     status_message_provider: Optional[StatusMessageProvider] = None,
     stream_listeners: Optional[List[StreamListener]] = None,
     include_final_prediction_in_output_stream: bool = True,
+    is_async_program: bool = False,
 ) -> Callable[[Any, Any], Awaitable[Any]]:
     """
     Wrap a DSPy program so that it streams its outputs incrementally, rather than returning them
@@ -43,6 +44,8 @@ def streamify(
             useful when `stream_listeners` is provided. If `False`, the final prediction will not be included in the
             output stream. When the program hit cache, or no listeners captured anything, the final prediction will
             still be included in the output stream even if this is `False`.
+        is_async_program: Whether the program is async. If `False`, the program will be wrapped with `asyncify`,
+            otherwise the program will be called with `acall`.
 
     Returns:
         A function that takes the same arguments as the original program, but returns an async
@@ -146,7 +149,9 @@ async def use_streaming():
     else:
         predict_id_to_listener = {}
 
-    if not iscoroutinefunction(program):
+    if is_async_program:
+        program = program.acall
+    elif not iscoroutinefunction(program):
         program = asyncify(program)
 
     callbacks = settings.callbacks
diff --git a/dspy/streaming/streaming_listener.py b/dspy/streaming/streaming_listener.py
index 7dfc4fa40f..1efdf62133 100644
--- a/dspy/streaming/streaming_listener.py
+++ b/dspy/streaming/streaming_listener.py
@@ -148,6 +148,7 @@ def find_predictor_for_stream_listeners(program: "Module", stream_listeners: Lis
     predict_id_to_listener = defaultdict(list)
     for listener in stream_listeners:
         if listener.predict:
+            predict_id_to_listener[id(listener.predict)].append(listener)
             continue
         if listener.signature_field_name not in field_name_to_named_predictor:
             raise ValueError(


Test patch:
diff --git a/tests/predict/test_predict.py b/tests/predict/test_predict.py
index 3461de479d..c605260683 100644
--- a/tests/predict/test_predict.py
+++ b/tests/predict/test_predict.py
@@ -11,6 +11,7 @@
 from dspy import Predict, Signature
 from dspy.utils.dummies import DummyLM
 from unittest.mock import patch, MagicMock, Mock
+from litellm import ModelResponse
 
 
 def test_initialization_with_string_signature():
@@ -491,6 +492,30 @@ class MySignature(dspy.Signature):
     assert "are you sure that's correct" in messages[3]["content"]
 
 
+def test_lm_usage():
+    program = Predict("question -> answer")
+    dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=False), track_usage=True)
+    with patch(
+        "dspy.clients.lm.litellm_completion",
+        return_value=ModelResponse(
+            choices=[{"message": {"content": "[[ ## answer ## ]]\nParis"}}],
+            usage={"total_tokens": 10},
+        ),
+    ):
+        result = program(question="What is the capital of France?")
+        assert result.answer == "Paris"
+        assert result.get_lm_usage()["openai/gpt-4o-mini"]["total_tokens"] == 10
+
+
+def test_positional_arguments():
+    program = Predict("question -> answer")
+    with pytest.raises(ValueError) as e:
+        program("What is the capital of France?")
+    assert "Positional arguments are not allowed when calling `dspy.Predict`, must use keyword arguments" in str(
+        e.value
+    )
+
+
 @pytest.mark.parametrize("adapter_type", ["chat", "json"])
 def test_field_constraints(adapter_type):
     class SpyLM(dspy.LM):
diff --git a/tests/streaming/test_streaming.py b/tests/streaming/test_streaming.py
index c09f480632..3cf2cfaddc 100644
--- a/tests/streaming/test_streaming.py
+++ b/tests/streaming/test_streaming.py
@@ -160,3 +160,67 @@ def __call__(self, x: str, **kwargs):
 
     assert all_chunks[-1].predict_name == "predict2"
     assert all_chunks[-1].signature_field_name == "judgement"
+
+
+@pytest.mark.anyio
+async def test_default_status_streaming_in_async_program():
+    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")
+
+        async def acall(self, x: str):
+            question = await self.generate_question.acall(x=x)
+            return await self.predict.acall(question=question)
+
+    lm = dspy.utils.DummyLM([{"answer": "red"}, {"answer": "blue"}])
+    with dspy.context(lm=lm):
+        program = dspy.streamify(MyProgram(), is_async_program=True)
+        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.skipif(not os.getenv("OPENAI_API_KEY"), reason="OpenAI API key not found in environment variables")
+@pytest.mark.anyio
+async def test_stream_listener_in_async_program():
+    class MyProgram(dspy.Module):
+        def __init__(self):
+            self.predict1 = dspy.Predict("question->answer")
+            self.predict2 = dspy.Predict("question, answer->judgement")
+
+        async def acall(self, x: str, **kwargs):
+            answer = await self.predict1.acall(question=x, **kwargs)
+            judgement = await self.predict2.acall(question=x, answer=answer, **kwargs)
+            return judgement
+
+    # Turn off the cache to ensure the stream is produced.
+    dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=False))
+    my_program = MyProgram()
+    program = dspy.streamify(
+        my_program,
+        is_async_program=True,
+        stream_listeners=[
+            dspy.streaming.StreamListener(signature_field_name="answer"),
+            dspy.streaming.StreamListener(signature_field_name="judgement"),
+        ],
+        include_final_prediction_in_output_stream=False,
+    )
+    output = program(x="why did a chicken cross the kitchen?")
+    all_chunks = []
+    async for value in output:
+        if isinstance(value, dspy.streaming.StreamResponse):
+            all_chunks.append(value)
+
+    assert all_chunks[0].predict_name == "predict1"
+    assert all_chunks[0].signature_field_name == "answer"
+
+    assert all_chunks[-1].predict_name == "predict2"
+    assert all_chunks[-1].signature_field_name == "judgement"