Fix issues in BAMLAdapter (#8654)

* init

* increment

* fix various issues on bamlAdapter

* lint

Author: chenmoneygithub

dspy/adapters/baml_adapter.py

@@ -17,141 +17,117 @@
 COMMENT_SYMBOL = "#"
 
 
-def _render_type_str(annotation: Any, _depth: int = 0, indent: int = 0) -> str:
+def _render_type_str(
+    annotation: Any,
+    depth: int = 0,
+    indent: int = 0,
+    seen_models: set[type] | None = None,
+) -> str:
     """Recursively renders a type annotation into a simplified string.
 
     Args:
         annotation: The type annotation to render
-        _depth: Current recursion depth (prevents infinite recursion)
+        depth: Current recursion depth (prevents infinite recursion)
         indent: Current indentation level for nested structures
     """
-    max_depth = 10
-    if _depth > max_depth:  # Prevent excessive recursion
-        return f"<max depth of {max_depth} exceeded>"
+    # Non-nested types
+    if annotation is str:
+        return "string"
+    if annotation is int:
+        return "int"
+    if annotation is float:
+        return "float"
+    if annotation is bool:
+        return "boolean"
+    if inspect.isclass(annotation) and issubclass(annotation, BaseModel):
+        return _build_simplified_schema(annotation, indent, seen_models)
 
     try:
         origin = get_origin(annotation)
         args = get_args(annotation)
     except Exception:
         return str(annotation)
 
-    # Handle Optional[T] or T | None
+    # Optional[T] or T | None
     if origin in (types.UnionType, Union):
         non_none_args = [arg for arg in args if arg is not type(None)]
         # Render the non-None part of the union
-        type_render = " or ".join([_render_type_str(arg, _depth + 1, indent) for arg in non_none_args])
-        # Add 'or null' if None was part of the union
+        type_render = " or ".join([_render_type_str(arg, depth + 1, indent) for arg in non_none_args])
+        # Add "or null" if None was part of the union
         if len(non_none_args) < len(args):
             return f"{type_render} or null"
         return type_render
 
-    # Base types
-    if annotation is str:
-        return "string"
-    if annotation is int:
-        return "int"
-    if annotation is float:
-        return "float"
-    if annotation is bool:
-        return "boolean"
-
-    # Composite types
+    # Literal[T1, T2, ...]
     if origin is Literal:
         return " or ".join(f'"{arg}"' for arg in args)
+
+    # list[T]
     if origin is list:
         # For Pydantic models in lists, use bracket notation
         inner_type = args[0]
         if inspect.isclass(inner_type) and issubclass(inner_type, BaseModel):
             # Build inner schema - the Pydantic model inside should use indent level for array contents
-            inner_schema = _build_simplified_schema(inner_type, indent + 1)
+            inner_schema = _build_simplified_schema(inner_type, indent + 1, seen_models)
             # Format with proper bracket notation and indentation
             current_indent = "  " * indent
             return f"[\n{inner_schema}\n{current_indent}]"
         else:
-            return f"{_render_type_str(inner_type, _depth + 1, indent)}[]"
-    if origin is dict:
-        return f"dict[{_render_type_str(args[0], _depth + 1, indent)}, {_render_type_str(args[1], _depth + 1, indent)}]"
+            return f"{_render_type_str(inner_type, depth + 1, indent)}[]"
 
-    # Pydantic models (we'll recurse in the main function)
-    if inspect.isclass(annotation) and issubclass(annotation, BaseModel):
-        try:
-            return _build_simplified_schema(annotation, indent)
-        except Exception:
-            return f"<{annotation.__name__}>"
+    # dict[T1, T2]
+    if origin is dict:
+        return f"dict[{_render_type_str(args[0], depth + 1, indent)}, {_render_type_str(args[1], depth + 1, indent)}]"
 
-    # Fallback
+    # fallback
     if hasattr(annotation, "__name__"):
         return annotation.__name__
     return str(annotation)
 
 
-def _build_simplified_schema(model: type[BaseModel], indent: int = 0, _seen: set[type] | None = None) -> str:
+def _build_simplified_schema(
+    pydantic_model: type[BaseModel],
+    indent: int = 0,
+    seen_models: set[type] | None = None,
+) -> str:
     """Builds a simplified, human-readable schema from a Pydantic model.
 
     Args:
-        model: The Pydantic model to build schema for
+        pydantic_model: The Pydantic model to build schema for
         indent: Current indentation level
-        _seen: Set to track visited models (prevents infinite recursion)
+        seen_models: Set to track visited pydantic models (prevents infinite recursion)
     """
-    if _seen is None:
-        _seen = set()
+    seen_models = seen_models or set()
 
-    if model in _seen:
-        return f"<circular reference to {model.__name__}>"
+    if pydantic_model in seen_models:
+        raise ValueError("BAMLAdapter cannot handle recursive pydantic models, please use a different adapter.")
 
-    _seen.add(model)
+    # Add `pydantic_model` to `seen_models` with a placeholder value to avoid infinite recursion.
+    seen_models.add(pydantic_model)
 
-    try:
-        lines = []
-        current_indent = "  " * indent
-        next_indent = "  " * (indent + 1)
-
-        lines.append(f"{current_indent}{{")
-
-        fields = model.model_fields
-        if not fields:
-            lines.append(f"{next_indent}{COMMENT_SYMBOL} No fields defined")
-        for name, field in fields.items():
-            if field.description:
-                lines.append(f"{next_indent}{COMMENT_SYMBOL} {field.description}")
-            elif field.alias and field.alias != name:
-                # If there's an alias but no description, show the alias as a comment
-                lines.append(f"{next_indent}{COMMENT_SYMBOL} alias: {field.alias}")
-
-            # Check for a nested Pydantic model
-            field_type_to_render = field.annotation
-
-            # Unpack Optional[T] to get T
-            origin = get_origin(field_type_to_render)
-            if origin in (types.UnionType, Union):
-                non_none_args = [arg for arg in get_args(field_type_to_render) if arg is not type(None)]
-                if len(non_none_args) == 1:
-                    field_type_to_render = non_none_args[0]
-
-            # Unpack list[T] to get T
-            origin = get_origin(field_type_to_render)
-            if origin is list:
-                field_type_to_render = get_args(field_type_to_render)[0]
-
-            if inspect.isclass(field_type_to_render) and issubclass(field_type_to_render, BaseModel):
-                # Recursively build schema for nested models with circular reference protection
-                nested_schema = _build_simplified_schema(field_type_to_render, indent + 1, _seen)
-                rendered_type = _render_type_str(field.annotation, indent=indent + 1).replace(
-                    field_type_to_render.__name__, nested_schema
-                )
-            else:
-                rendered_type = _render_type_str(field.annotation, indent=indent + 1)
-
-            line = f"{next_indent}{name}: {rendered_type},"
-
-            lines.append(line)
-
-        lines.append(f"{current_indent}}}")
-        return "\n".join(lines)
-    except Exception as e:
-        return f"<error building schema for {model.__name__}: {e}>"
-    finally:
-        _seen.discard(model)
+    lines = []
+    current_indent = "  " * indent
+    next_indent = "  " * (indent + 1)
+
+    lines.append(f"{current_indent}{{")
+
+    fields = pydantic_model.model_fields
+    if not fields:
+        lines.append(f"{next_indent}{COMMENT_SYMBOL} No fields defined")
+    for name, field in fields.items():
+        if field.description:
+            lines.append(f"{next_indent}{COMMENT_SYMBOL} {field.description}")
+        elif field.alias and field.alias != name:
+            # If there's an alias but no description, show the alias as a comment
+            lines.append(f"{next_indent}{COMMENT_SYMBOL} alias: {field.alias}")
+
+        rendered_type = _render_type_str(field.annotation, indent=indent + 1, seen_models=seen_models)
+        line = f"{next_indent}{name}: {rendered_type},"
+
+        lines.append(line)
+
+    lines.append(f"{current_indent}}}")
+    return "\n".join(lines)
 
 
 class BAMLAdapter(JSONAdapter):
@@ -176,6 +152,7 @@ class PatientAddress(BaseModel):
         street: str
         city: str
         country: Literal["US", "CA"]
+
     class PatientDetails(BaseModel):
         name: str = Field(description="Full name of the patient.")
         age: int
@@ -245,41 +222,14 @@ def format_field_structure(self, signature: type[Signature]) -> str:
         if signature.output_fields:
             for name, field in signature.output_fields.items():
                 field_type = field.annotation
-                main_type = field_type
-
-                # Find the core type if it's wrapped in Optional or Union
-                origin = get_origin(field_type)
-                if origin in (types.UnionType, Union):
-                    non_none_args = [arg for arg in get_args(field_type) if arg is not type(None)]
-                    if len(non_none_args) == 1:
-                        main_type = non_none_args[0]
-
                 sections.append(f"[[ ## {name} ## ]]")
-
-                if inspect.isclass(main_type) and issubclass(main_type, BaseModel):
-                    # We have a pydantic model, so build the simplified schema for it.
-                    schema_str = _build_simplified_schema(main_type)
-                    sections.append(schema_str)
-                else:
-                    # Handle non-pydantic or primitive types simply
-                    type_str = _render_type_str(field_type, indent=0)
-                    sections.append(f"Output field `{name}` should be of type: {type_str}")
-
-                sections.append("")  # Empty line after each output
+                sections.append(f"Output field `{name}` should be of type: {_render_type_str(field_type, indent=0)}\n")
 
         # Add completed section
         sections.append("[[ ## completed ## ]]")
 
         return "\n".join(sections)
 
-    def format_task_description(self, signature: type[Signature]) -> str:
-        """Format the task description for the system message."""
-        import textwrap
-
-        instructions = textwrap.dedent(signature.instructions)
-        objective = ("\n" + " " * 8).join([""] + instructions.splitlines())
-        return f"In adhering to this structure, your objective is: {objective}"
-
     def format_user_message_content(
         self,
         signature: type[Signature],

tests/adapters/test_baml_adapter.py

@@ -40,14 +40,9 @@ class ImageWrapper(pydantic.BaseModel):
     tag: list[str]
 
 
-class CircularModelA(pydantic.BaseModel):
+class CircularModel(pydantic.BaseModel):
     name: str
-    related_b: "CircularModelB | None" = None
-
-
-class CircularModelB(pydantic.BaseModel):
-    id: int
-    related_a: CircularModelA | None = None
+    field: "CircularModel"
 
 
 def test_baml_adapter_basic_schema_generation():
@@ -143,20 +138,18 @@ class TestSignature(dspy.Signature):
     assert "metadata: dict[string, string]," in schema
 
 
-def test_baml_adapter_prevents_circular_references():
+def test_baml_adapter_raise_error_on_circular_references():
     """Test that circular references are handled gracefully."""
 
     class TestSignature(dspy.Signature):
         input: str = dspy.InputField()
-        circular: CircularModelA = dspy.OutputField()
+        circular: CircularModel = dspy.OutputField()
 
     adapter = BAMLAdapter()
-    schema = adapter.format_field_structure(TestSignature)
+    with pytest.raises(ValueError) as error:
+        adapter.format_field_structure(TestSignature)
 
-    # Should not cause infinite recursion and should handle ForwardRef gracefully
-    assert "name: string," in schema
-    # The actual output shows ForwardRef handling, which is acceptable
-    assert "ForwardRef" in schema or "<circular reference" in schema
+    assert "BAMLAdapter cannot handle recursive pydantic models" in str(error.value)
 
 
 def test_baml_adapter_formats_pydantic_inputs_as_clean_json():
@@ -227,53 +220,6 @@ class TestSignature(dspy.Signature):
         pass
 
 
-def test_baml_adapter_inherits_json_parsing_behavior():
-    """Test that BAMLAdapter maintains JSONAdapter's parsing compatibility."""
-
-    class TestSignature(dspy.Signature):
-        question: str = dspy.InputField()
-        answer: str = dspy.OutputField()
-
-    baml_adapter = BAMLAdapter()
-    json_adapter = dspy.JSONAdapter()
-
-    # Both should parse the same JSON response identically
-    completion = '{"answer": "Paris"}'
-
-    baml_result = baml_adapter.parse(TestSignature, completion)
-    json_result = json_adapter.parse(TestSignature, completion)
-
-    assert baml_result == json_result == {"answer": "Paris"}
-
-
-def test_baml_adapter_parse_complex_pydantic_models():
-    """Test parsing JSON into complex Pydantic models."""
-
-    class TestSignature(dspy.Signature):
-        input: str = dspy.InputField()
-        patient: PatientDetails = dspy.OutputField()
-
-    adapter = BAMLAdapter()
-
-    completion = """{"patient": {
-        "name": "John Smith",
-        "age": 42,
-        "address": {
-            "street": "456 Oak St",
-            "city": "Springfield",
-            "country": "US"
-        }
-    }}"""
-
-    result = adapter.parse(TestSignature, completion)
-
-    assert isinstance(result["patient"], PatientDetails)
-    assert result["patient"].name == "John Smith"
-    assert result["patient"].age == 42
-    assert isinstance(result["patient"].address, PatientAddress)
-    assert result["patient"].address.street == "456 Oak St"
-
-
 def test_baml_adapter_raises_on_missing_fields():
     """Test that missing required fields raise appropriate errors."""