remove extra predict field

Author: chenmoneygithub

dspy/adapters/base.py

@@ -18,19 +18,21 @@ def __init_subclass__(cls, **kwargs) -> None:
         cls.format = with_callbacks(cls.format)
         cls.parse = with_callbacks(cls.parse)
 
-    def __call__(self, lm, lm_kwargs, signature, demos, inputs, predict=None):
+    def __call__(self, lm, lm_kwargs, signature, demos, inputs):
         inputs_ = self.format(signature, demos, inputs)
         inputs_ = dict(prompt=inputs_) if isinstance(inputs_, str) else dict(messages=inputs_)
 
         stream_listeners = settings.stream_listeners or []
+        caller_predict = settings.caller_predict
         stream = settings.send_stream is not None
         if stream and len(stream_listeners) > 0:
-            stream = any(stream_listener.predict == predict for stream_listener in stream_listeners)
+            stream = any(stream_listener.predict == caller_predict for stream_listener in stream_listeners)
 
         if stream:
-            with settings.context(stream_predict=predict):
-                outputs = lm(**inputs_, **lm_kwargs)
+            outputs = lm(**inputs_, **lm_kwargs)
         else:
+            # Explicilty disable streaming if streaming is not enabled globally or the caller predict shouldn't be
+            # streamed.
             with settings.context(send_stream=None):
                 outputs = lm(**inputs_, **lm_kwargs)
 

dspy/adapters/json_adapter.py

@@ -32,13 +32,13 @@ class JSONAdapter(Adapter):
     def __init__(self):
         pass
 
-    def __call__(self, lm, lm_kwargs, signature, demos, inputs, predict=None):
+    def __call__(self, lm, lm_kwargs, signature, demos, inputs):
         inputs = self.format(signature, demos, inputs)
         inputs = dict(prompt=inputs) if isinstance(inputs, str) else dict(messages=inputs)
 
         stream_listeners = settings.stream_listeners or []
         if len(stream_listeners) > 0:
-            raise ValueError("Stream listener is not supported for JsonAdapter, please use ChatAdapter instead.")
+            raise ValueError("Stream listener is not yet supported for JsonAdapter, please use ChatAdapter instead.")
 
         try:
             provider = lm.model.split("/", 1)[0] or "openai"

dspy/clients/lm.py

@@ -341,7 +341,7 @@ def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cac
     )
 
     stream = dspy.settings.send_stream
-    stream_predict = dspy.settings.stream_predict
+    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.
@@ -353,7 +353,7 @@ def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cac
 
     # The stream is already opened, and will be closed by the caller.
     stream = cast(MemoryObjectSendStream, stream)
-    stream_predict_id = id(stream_predict) if stream_predict else None
+    caller_predict_id = id(caller_predict) if caller_predict else None
 
     @syncify
     async def stream_completion():
@@ -365,9 +365,9 @@ async def stream_completion():
         )
         chunks = []
         async for chunk in response:
-            if stream_predict_id:
+            if caller_predict_id:
                 # Add the predict id to the chunk so that the stream listener can identify which predict produces it.
-                chunk.predict_id = stream_predict_id
+                chunk.predict_id = caller_predict_id
             chunks.append(chunk)
             await stream.send(chunk)
         return litellm.stream_chunk_builder(chunks)

dspy/dsp/utils/settings.py

@@ -20,7 +20,7 @@
     async_max_workers=8,
     send_stream=None,
     disable_history=False,
-    stream_predict=None,
+    caller_predict=None,
     stream_listeners=[],
 )
 

dspy/predict/predict.py

@@ -2,12 +2,13 @@
 
 from pydantic import BaseModel
 
+from dspy.clients.lm import LM
+from dspy.dsp.utils.settings import settings
 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
-from dspy.clients.lm import LM
 
 
 class Predict(Module, Parameter):
@@ -98,7 +99,8 @@ def forward(self, **kwargs):
         import dspy
 
         adapter = dspy.settings.adapter or dspy.ChatAdapter()
-        completions = adapter(lm, lm_kwargs=config, signature=signature, demos=demos, inputs=kwargs, predict=self)
+        with settings.context(caller_predict=self):
+            completions = adapter(lm, lm_kwargs=config, signature=signature, demos=demos, inputs=kwargs)
 
         pred = Prediction.from_completions(completions, signature=signature)