Repository: stanfordnlp/dspy
Pull request #1053 (2024-05-22 07:29:56)

feat(dspy): try from json before applying formatting logic in typed predictor
Adds an attempt to use the function supplied to `from_json` before running string parsing. This allows users to modify the pydantic's `model_validate_json` to customise the deserialization of the generation.

Other changes were to satisfy ruff.

This PR gives users a route to inject code which can solve issues like https://github.com/stanfordnlp/dspy/issues/1024 and https://github.com/stanfordnlp/dspy/issues/1001

After this PR, the user could implement something like the following:

```python
class TaxField(BaseModel):
    key: str = Field()
    value: str = Field()
    truthful: bool = Field(
        description="is a tuple contains the extracted value and a boolean indicating if the inFieldsation is truthful and correct."
    )
    confidence: float = Field(ge=0, le=1, description="The confidence score for the answer")

    @classmethod
    def model_validate_json(
        cls, json_data: str, *, strict: bool | None = None, context: dict[str, Any] | None = None
    ) -> "list[ModelField]":
        try:
            __tracebackhide__ = True
            return cls.__pydantic_validator__.validate_json(
                json_data, strict=strict, context=context
            )
        except ValidationError:
            # custom parsing logic here
        raise ValueError("Could not find valid json") from last_exc


class QAExtractionSignature(Signature):
    """Your task is to extract the key-value pairs from the document that will follow the instructions on the field to extract.\nplease keep the content you extract truthful and correct based on the document text provided"""

    document: str = InputField()
    instruction: str = InputField(
        description="The instruction on which field to extact and how to extract it."
    )
    tax_from_fields: TaxField = OutputField(
        description="The list of fields extracted from the document text. IMPORTANT!! This must follow a semicolon separated list of values!"
    )


class TypedQAExtraction(dspy.Module):
    def __init__(self):
        self.qa_extraction = dspy.functional.TypedPredictor(QAExtractionSignature)

    def forward(self, document: str, instruction: str) -> TaxFromFields:
        return self.qa_extraction(document=document, instruction=instruction).tax_from_fields

```


Patch:
diff --git a/dspy/functional/functional.py b/dspy/functional/functional.py
index 667c86d4..8b0025b7 100644
--- a/dspy/functional/functional.py
+++ b/dspy/functional/functional.py
@@ -1,8 +1,7 @@
 import inspect
 import json
-import textwrap
 import typing
-from typing import Annotated, List, Tuple, Union  # noqa: UP035
+from typing import Annotated, Callable, List, Tuple, Union  # noqa: UP035
 
 import pydantic
 import ujson
@@ -14,7 +13,7 @@ from dspy.primitives.prediction import Prediction
 from dspy.signatures.signature import ensure_signature, make_signature
 
 
-def predictor(*args, **kwargs):
+def predictor(*args: tuple, **kwargs) -> Callable[..., dspy.Module]:
     def _predictor(func) -> dspy.Module:
         """Decorator that creates a predictor module based on the provided function."""
         signature = _func_to_signature(func)
@@ -28,7 +27,7 @@ def predictor(*args, **kwargs):
     return _predictor
 
 
-def cot(*args, **kwargs):
+def cot(*args: tuple, **kwargs) -> Callable[..., dspy.Module]:
     def _cot(func) -> dspy.Module:
         """Decorator that creates a chain of thought module based on the provided function."""
         signature = _func_to_signature(func)
@@ -72,11 +71,11 @@ def TypedChainOfThought(signature, instructions=None, reasoning=None, *, max_ret
     signature = ensure_signature(signature, instructions)
     output_keys = ", ".join(signature.output_fields.keys())
 
-    DEFAULT_RATIONALE = dspy.OutputField(
+    default_rationale = dspy.OutputField(
         prefix="Reasoning: Let's think step by step in order to",
         desc="${produce the " + output_keys + "}. We ...",
     )
-    reasoning = reasoning or DEFAULT_RATIONALE
+    reasoning = reasoning or default_rationale
 
     return TypedPredictor(
         signature.prepend(
@@ -93,8 +92,10 @@ class TypedPredictor(dspy.Module):
 
         Args:
             signature: The signature of the module. Can use type annotations.
+            instructions: A description of what the model should do.
             max_retries: The number of times to retry the prediction if the output is invalid.
             wrap_json: If True, json objects in the input will be wrapped in ```json ... ```
+            explain_errors: If True, the model will try to explain the errors it encounters.
         """
         super().__init__()
         self.signature = ensure_signature(signature, instructions)
@@ -128,7 +129,7 @@ class TypedPredictor(dspy.Module):
         )(json_schema=schema).json_object
         # We use the model_validate_json method to make sure the example is valid
         try:
-            type_.model_validate_json(_unwrap_json(json_object))
+            type_.model_validate_json(_unwrap_json(json_object, type_.model_validate_json))
         except (pydantic.ValidationError, ValueError):
             return ""  # Unable to make an example
         return json_object
@@ -172,11 +173,10 @@ class TypedPredictor(dspy.Module):
 
     def _make_explanation(self, task_description: str, model_output: str, error: str) -> str:
         class Signature(dspy.Signature):
-            __doc__ = textwrap.dedent(
-                """
-                I gave my language model a task, but it failed. Figure out what went wrong,
-                and write instructions to help it avoid the error next time.""",
-            )
+            """I gave my language model a task, but it failed.
+
+            Figure out what went wrong, and write instructions to help it avoid the error next time.
+            """
 
             task_description: str = dspy.InputField(desc="What I asked the model to do")
             language_model_output: str = dspy.InputField(desc="The output of the model")
@@ -259,7 +259,7 @@ class TypedPredictor(dspy.Module):
                         desc=field.json_schema_extra.get("desc", "")
                         + (". Respond with a single JSON object. JSON Schema: " + schema),
                         format=lambda x, to_json=to_json: (x if isinstance(x, str) else to_json(x)),
-                        parser=lambda x, from_json=from_json: from_json(_unwrap_json(x)),
+                        parser=lambda x, from_json=from_json: from_json(_unwrap_json(x, from_json)),
                         type_=type_,
                     )
             else:  # If input field
@@ -405,14 +405,17 @@ def _func_to_signature(func):
     return dspy.Signature(fields, instructions)
 
 
-def _unwrap_json(output):
-    output = output.strip()
-    if output.startswith("```"):
-        if not output.startswith("```json"):
-            raise ValueError("json output should start with ```json")
-        if not output.endswith("```"):
-            raise ValueError("Don't write anything after the final json ```")
-        output = output[7:-3].strip()
-    if not output.startswith("{") or not output.endswith("}"):
-        raise ValueError("json output should start and end with { and }")
-    return ujson.dumps(ujson.loads(output))  # ujson is a bit more robust than the standard json
+def _unwrap_json(output, from_json: Callable[[str], Union[pydantic.BaseModel, str]]):
+    try:
+        return from_json(output).model_dump_json()
+    except (ValueError, pydantic.ValidationError, AttributeError):
+        output = output.strip()
+        if output.startswith("```"):
+            if not output.startswith("```json"):
+                raise ValueError("json output should start with ```json") from None
+            if not output.endswith("```"):
+                raise ValueError("Don't write anything after the final json ```") from None
+            output = output[7:-3].strip()
+        if not output.startswith("{") or not output.endswith("}"):
+            raise ValueError("json output should start and end with { and }") from None
+        return ujson.dumps(ujson.loads(output))  # ujson is a bit more robust than the standard json


Test patch:
diff --git a/tests/functional/test_functional.py b/tests/functional/test_functional.py
index bb98b9b3..3e605566 100644
--- a/tests/functional/test_functional.py
+++ b/tests/functional/test_functional.py
@@ -1,14 +1,13 @@
 import datetime
 import textwrap
-import pydantic
-from pydantic import AfterValidator, Field, BaseModel, field_validator, model_validator
-from typing import Annotated, Generic, Literal, TypeVar
-from typing import List
+from typing import Annotated, Any, Generic, List, Literal, Optional, TypeVar
 
+import pydantic
 import pytest
+from pydantic import AfterValidator, BaseModel, Field, ValidationError, field_validator, model_validator
 
 import dspy
-from dspy.functional import predictor, cot, FunctionalModule, TypedPredictor, TypedChainOfThought
+from dspy.functional import FunctionalModule, TypedChainOfThought, TypedPredictor, cot, predictor
 from dspy.predict.predict import Predict
 from dspy.primitives.example import Example
 from dspy.teleprompt.bootstrap import BootstrapFewShot
@@ -282,7 +281,7 @@ def test_regex():
 
     email = textwrap.dedent(
         """\
-        We're excited to welcome you aboard your upcoming flight from 
+        We're excited to welcome you aboard your upcoming flight from
         John F. Kennedy International Airport (JFK) to Los Angeles International Airport (LAX)
         on December 25, 2022. Here's everything you need to know before you take off: ...
     """
@@ -304,6 +303,69 @@ def test_regex():
     )
 
 
+def test_custom_model_validate_json():
+    class Airport(BaseModel):
+        code: str = Field(pattern=r"^[A-Z]{3}$")
+        lat: float
+        lon: float
+
+    class TravelInformation(BaseModel):
+        origin: Airport
+        destination: Airport
+        date: datetime.date
+
+        @classmethod
+        def model_validate_json(
+            cls, json_data: str, *, strict: Optional[bool] = None, context: Optional[dict[str, Any]] = None
+        ) -> "TravelInformation":
+            try:
+                __tracebackhide__ = True
+                return cls.__pydantic_validator__.validate_json(json_data, strict=strict, context=context)
+            except ValidationError:
+                for substring_length in range(len(json_data), 1, -1):
+                    for start in range(len(json_data) - substring_length + 1):
+                        substring = json_data[start : start + substring_length]
+                        try:
+                            __tracebackhide__ = True
+                            res = cls.__pydantic_validator__.validate_json(substring, strict=strict, context=context)
+                            return res
+                        except ValidationError as exc:
+                            last_exc = exc
+                            pass
+            raise ValueError("Could not find valid json") from last_exc
+
+    @predictor
+    def flight_information(email: str) -> TravelInformation:
+        pass
+
+    email = textwrap.dedent(
+        """\
+        We're excited to welcome you aboard your upcoming flight from
+        John F. Kennedy International Airport (JFK) to Los Angeles International Airport (LAX)
+        on December 25, 2022. Here's everything you need to know before you take off: ...
+    """
+    )
+    lm = DummyLM(
+        [
+            # Example with a bad origin code.
+            (
+                "Here is your json: "
+                "{"
+                '"origin": {"code":"JFK", "lat":40.6446, "lon":-73.7797}, '
+                '"destination": {"code":"LAX", "lat":33.942791, "lon":-118.410042}, '
+                '"date": "2022-12-25"}'
+            ),
+        ]
+    )
+    dspy.settings.configure(lm=lm)
+
+    assert flight_information(email=email) == TravelInformation(
+        origin={"code": "JFK", "lat": 40.6446, "lon": -73.7797},
+        destination={"code": "LAX", "lat": 33.942791, "lon": -118.410042},
+        date=datetime.date(2022, 12, 25),
+    )
+
+
 def test_raises():
     class TravelInformation(BaseModel):
         origin: str = Field(pattern=r"^[A-Z]{3}$")
@@ -601,7 +663,8 @@ def test_list_input2():
 
     assert output == "Output"
 
-    assert lm.get_convo(-1) == textwrap.dedent("""\
+    assert lm.get_convo(-1) == textwrap.dedent(
+        """\
         Given the fields `attempted_signatures`, produce the fields `proposed_signature`.
 
         ---
@@ -616,7 +679,8 @@ def test_list_input2():
 
         Attempted Signatures: [{"string":"string 1","score":0.5},{"string":"string 2","score":0.4},{"string":"string 3","score":0.3}]
         Reasoning: Let's think step by step in order to Thoughts
-        Proposed Signature: Output""")
+        Proposed Signature: Output"""
+    )
 
 
 def test_custom_reasoning_field():
@@ -643,7 +707,8 @@ def test_custom_reasoning_field():
     assert isinstance(output.question, Question)
     assert output.question.value == expected
 
-    assert lm.get_convo(-1) == textwrap.dedent("""\
+    assert lm.get_convo(-1) == textwrap.dedent(
+        """\
         Given the fields `topic`, produce the fields `question`.
 
         ---
@@ -658,7 +723,8 @@ def test_custom_reasoning_field():
 
         Topic: Physics
         Custom Reasoning: Let's break this down. To generate a question about Thoughts
-        Question: {"value": "What is the speed of light?"}""")
+        Question: {"value": "What is the speed of light?"}"""
+    )
 
 
 def test_generic_signature():
@@ -772,7 +838,8 @@ def test_demos():
 
     assert program(input="What is the capital of France?").output == "Paris"
 
-    assert lm.get_convo(-1) == textwrap.dedent("""\
+    assert lm.get_convo(-1) == textwrap.dedent(
+        """\
         Given the fields `input`, produce the fields `output`.
 
         ---
@@ -790,7 +857,8 @@ def test_demos():
         ---
 
         Input: What is the capital of France?
-        Output: Paris""")
+        Output: Paris"""
+    )
 
 
 def _test_demos_missing_input():
@@ -802,7 +870,8 @@ def _test_demos_missing_input():
     dspy.settings.configure(lm=DummyLM(["My thoughts", "Paris"]))
     assert program(input="What is the capital of France?").output == "Paris"
 
-    assert dspy.settings.lm.get_convo(-1) == textwrap.dedent("""\
+    assert dspy.settings.lm.get_convo(-1) == textwrap.dedent(
+        """\
         Given the fields `input`, produce the fields `output`.
 
         ---
@@ -822,7 +891,8 @@ def _test_demos_missing_input():
 
         Input: What is the capital of France?
         Thoughts: My thoughts
-        Output: Paris""")
+        Output: Paris"""
+    )
 
 
 def test_conlist():