Repository: stanfordnlp/dspy (python)

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

Pull request description:
Fix predict usage tracking
fix #8145 

It actually indicates a bigger problem that our `dspy.Predict`'s `__call__` is not using the `Module.__call__`. We override the `__call__` in `dspy.Predict` to disallow positional args, but it should still utilize the parent's `__call__`.

We also improve the error message on users passing positional args to `dspy.Predict`. 

Before:

```
TypeError: Predict.__call__() takes 1 positional argument but 2 were given
```

After:

```
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?')
```


Patch:
diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py
index 023ade47..f77bc40e 100644
--- a/dspy/clients/lm.py
+++ b/dspy/clients/lm.py
@@ -231,35 +231,25 @@ class LM(BaseLM):
         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 @@ def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cac
             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 bde63261..b4a94a36 100644
--- a/dspy/predict/predict.py
+++ b/dspy/predict/predict.py
@@ -11,7 +11,6 @@ from dspy.predict.parameter import Parameter
 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 @@ class Predict(Module, Parameter):
 
         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 @@ class Predict(Module, Parameter):
             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 @@ class Predict(Module, Parameter):
         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 a6a5ae88..3edea2d5 100644
--- a/dspy/primitives/tool.py
+++ b/dspy/primitives/tool.py
@@ -146,6 +146,7 @@ class Tool:
             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 bda50341..9a3cf74c 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 @@ class StatusStreamingCallback(BaseCallback):
         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 @@ class StatusStreamingCallback(BaseCallback):
         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 @@ class StatusStreamingCallback(BaseCallback):
         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 @@ class StatusStreamingCallback(BaseCallback):
         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 @@ class StatusStreamingCallback(BaseCallback):
         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 @@ class StatusStreamingCallback(BaseCallback):
         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 1a91af6d..c99f2a15 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 @@ def streamify(
     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 7dfc4fa4..1efdf621 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 3461de47..c6052606 100644
--- a/tests/predict/test_predict.py
+++ b/tests/predict/test_predict.py
@@ -11,6 +11,7 @@ import dspy
 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 @@ def test_call_predict_with_chat_history(adapter_type):
     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 c09f4806..3cf2cfad 100644
--- a/tests/streaming/test_streaming.py
+++ b/tests/streaming/test_streaming.py
@@ -160,3 +160,67 @@ async def test_stream_listener():
 
     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"