Support async path for ChatAdapter and JSONAdapter (#8419)

* init

* Add async path to chatadapter and jsonadapter call

* fix tests

Author: chenmoneygithub

dspy/adapters/chat_adapter.py

@@ -50,6 +50,26 @@ def __call__(
                 raise e
             return JSONAdapter()(lm, lm_kwargs, signature, demos, inputs)
 
+    async def acall(
+        self,
+        lm: LM,
+        lm_kwargs: dict[str, Any],
+        signature: Type[Signature],
+        demos: list[dict[str, Any]],
+        inputs: dict[str, Any],
+    ) -> list[dict[str, Any]]:
+        try:
+            return await super().acall(lm, lm_kwargs, signature, demos, inputs)
+        except Exception as e:
+            # fallback to JSONAdapter
+            from dspy.adapters.json_adapter import JSONAdapter
+
+            if isinstance(e, ContextWindowExceededError) or isinstance(self, JSONAdapter):
+                # On context window exceeded error or already using JSONAdapter, we don't want to retry with a different
+                # adapter.
+                raise e
+            return await JSONAdapter().acall(lm, lm_kwargs, signature, demos, inputs)
+
     def format_field_description(self, signature: Type[Signature]) -> str:
         return (
             f"Your input fields are:\n{get_field_description_string(signature.input_fields)}\n"

dspy/adapters/json_adapter.py

@@ -22,6 +22,11 @@
 
 logger = logging.getLogger(__name__)
 
+ERROR_MESSAGE_ON_JSON_ADAPTER_FAILURE = (
+    "Both structured output format and JSON mode failed. Please choose a model that supports "
+    "`response_format` argument. Original error: {}"
+)
+
 
 def _has_open_ended_mapping(signature: SignatureMeta) -> bool:
     """
@@ -37,6 +42,18 @@ def _has_open_ended_mapping(signature: SignatureMeta) -> bool:
 
 
 class JSONAdapter(ChatAdapter):
+    def _json_adapter_call_common(self, lm, lm_kwargs, signature, demos, inputs, call_fn):
+        """Common call logic to be used for both sync and async calls."""
+        provider = lm.model.split("/", 1)[0] or "openai"
+        params = litellm.get_supported_openai_params(model=lm.model, custom_llm_provider=provider)
+
+        if not params or "response_format" not in params:
+            return call_fn(lm, lm_kwargs, signature, demos, inputs)
+
+        if _has_open_ended_mapping(signature):
+            lm_kwargs["response_format"] = {"type": "json_object"}
+            return call_fn(lm, lm_kwargs, signature, demos, inputs)
+
     def __call__(
         self,
         lm: LM,
@@ -45,37 +62,49 @@ def __call__(
         demos: list[dict[str, Any]],
         inputs: dict[str, Any],
     ) -> list[dict[str, Any]]:
-        provider = lm.model.split("/", 1)[0] or "openai"
-        params = litellm.get_supported_openai_params(model=lm.model, custom_llm_provider=provider)
+        result = self._json_adapter_call_common(lm, lm_kwargs, signature, demos, inputs, super().__call__)
+        if result:
+            return result
 
-        # If response_format is not supported, use basic call
-        if not params or "response_format" not in params:
+        try:
+            structured_output_model = _get_structured_outputs_response_format(signature)
+            lm_kwargs["response_format"] = structured_output_model
             return super().__call__(lm, lm_kwargs, signature, demos, inputs)
+        except Exception:
+            logger.warning("Failed to use structured output format, falling back to JSON mode.")
+            try:
+                lm_kwargs["response_format"] = {"type": "json_object"}
+                return super().__call__(lm, lm_kwargs, signature, demos, inputs)
+            except AdapterParseError as e:
+                raise e
+            except Exception as e:
+                raise RuntimeError(ERROR_MESSAGE_ON_JSON_ADAPTER_FAILURE.format(e)) from e
 
-        # Check early for open-ended mapping types before trying structured outputs.
-        if _has_open_ended_mapping(signature):
-            lm_kwargs["response_format"] = {"type": "json_object"}
-            return super().__call__(lm, lm_kwargs, signature, demos, inputs)
+    async def acall(
+        self,
+        lm: LM,
+        lm_kwargs: dict[str, Any],
+        signature: Type[Signature],
+        demos: list[dict[str, Any]],
+        inputs: dict[str, Any],
+    ) -> list[dict[str, Any]]:
+        result = self._json_adapter_call_common(lm, lm_kwargs, signature, demos, inputs, super().acall)
+        if result:
+            return await result
 
-        # Try structured output first, fall back to basic JSON if it fails.
         try:
             structured_output_model = _get_structured_outputs_response_format(signature)
             lm_kwargs["response_format"] = structured_output_model
-            return super().__call__(lm, lm_kwargs, signature, demos, inputs)
+            return await super().acall(lm, lm_kwargs, signature, demos, inputs)
         except Exception:
             logger.warning("Failed to use structured output format, falling back to JSON mode.")
             try:
                 lm_kwargs["response_format"] = {"type": "json_object"}
-                return super().__call__(lm, lm_kwargs, signature, demos, inputs)
+                return await super().acall(lm, lm_kwargs, signature, demos, inputs)
             except AdapterParseError as e:
-                # On AdapterParseError, we raise the original error.
                 raise e
             except Exception as e:
-                # On any other error, we raise a RuntimeError with the original error message.
-                raise RuntimeError(
-                    "Both structured output format and JSON mode failed. Please choose a model that supports "
-                    f"`response_format` argument. Original error: {e}"
-                ) from e
+                raise RuntimeError(ERROR_MESSAGE_ON_JSON_ADAPTER_FAILURE.format(e)) from e
 
     def _call_preprocess(
         self,

tests/adapters/test_chat_adapter.py

@@ -3,6 +3,7 @@
 
 import pydantic
 import pytest
+from litellm.utils import Choices, Message, ModelResponse
 
 import dspy
 
@@ -376,3 +377,48 @@ class MySignature(dspy.Signature):
     assert messages[2]["content"] == "[[ ## answer ## ]]\nParis\n\n[[ ## completed ## ]]\n"
     assert messages[3]["content"] == "[[ ## question ## ]]\nWhat is the capital of Germany?"
     assert messages[4]["content"] == "[[ ## answer ## ]]\nBerlin\n\n[[ ## completed ## ]]\n"
+
+
+def test_chat_adapter_fallback_to_json_adapter_on_exception():
+    signature = dspy.make_signature("question->answer")
+    adapter = dspy.ChatAdapter()
+
+    with mock.patch("litellm.completion") as mock_completion:
+        # Mock returning a response compatible with JSONAdapter but not ChatAdapter
+        mock_completion.return_value = ModelResponse(
+            choices=[Choices(message=Message(content="{'answer': 'Paris'}"))],
+            model="openai/gpt-4o-mini",
+        )
+
+        lm = dspy.LM("openai/gpt-4o-mini", cache=False)
+
+        with mock.patch("dspy.adapters.json_adapter.JSONAdapter.__call__") as mock_json_adapter_call:
+            adapter(lm, {}, signature, [], {"question": "What is the capital of France?"})
+            mock_json_adapter_call.assert_called_once()
+
+        # The parse should succeed
+        result = adapter(lm, {}, signature, [], {"question": "What is the capital of France?"})
+        assert result == [{"answer": "Paris"}]
+
+
+@pytest.mark.asyncio
+async def test_chat_adapter_fallback_to_json_adapter_on_exception_async():
+    signature = dspy.make_signature("question->answer")
+    adapter = dspy.ChatAdapter()
+
+    with mock.patch("litellm.acompletion") as mock_completion:
+        # Mock returning a response compatible with JSONAdapter but not ChatAdapter
+        mock_completion.return_value = ModelResponse(
+            choices=[Choices(message=Message(content="{'answer': 'Paris'}"))],
+            model="openai/gpt-4o-mini",
+        )
+
+        lm = dspy.LM("openai/gpt-4o-mini", cache=False)
+
+        with mock.patch("dspy.adapters.json_adapter.JSONAdapter.acall") as mock_json_adapter_acall:
+            await adapter.acall(lm, {}, signature, [], {"question": "What is the capital of France?"})
+            mock_json_adapter_acall.assert_called_once()
+
+        # The parse should succeed
+        result = await adapter.acall(lm, {}, signature, [], {"question": "What is the capital of France?"})
+        assert result == [{"answer": "Paris"}]