From 326f40660575a48e90b5f4f127d02a0f672d1569 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Fri, 21 Feb 2025 11:40:37 +0000
Subject: [PATCH] feat: enhance pydantic output to include field descriptions

- Update generate_model_description to include field descriptions
- Add tests for field description handling
- Maintain backward compatibility for fields without descriptions

Fixes #2188

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/utilities/converter.py | 33 ++++++++++++++++++++-----------
 tests/utilities/test_converter.py | 22 ++++++++++++++++++++-
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/src/crewai/utilities/converter.py b/src/crewai/utilities/converter.py
index 991185f4a..46597323b 100644
--- a/src/crewai/utilities/converter.py
+++ b/src/crewai/utilities/converter.py
@@ -263,32 +263,41 @@ def generate_model_description(model: Type[BaseModel]) -> str:
     models.
     """
 
-    def describe_field(field_type):
+    def describe_field(field_type, field_info=None):
         origin = get_origin(field_type)
         args = get_args(field_type)
 
+        type_desc = ""
         if origin is Union or (origin is None and len(args) > 0):
             # Handle both Union and the new '|' syntax
             non_none_args = [arg for arg in args if arg is not type(None)]
             if len(non_none_args) == 1:
-                return f"Optional[{describe_field(non_none_args[0])}]"
+                type_desc = f"Optional[{describe_field(non_none_args[0])}]"
             else:
-                return f"Optional[Union[{', '.join(describe_field(arg) for arg in non_none_args)}]]"
+                type_desc = f"Optional[Union[{', '.join(describe_field(arg) for arg in non_none_args)}]]"
         elif origin is list:
-            return f"List[{describe_field(args[0])}]"
+            type_desc = f"List[{describe_field(args[0])}]"
         elif origin is dict:
             key_type = describe_field(args[0])
             value_type = describe_field(args[1])
-            return f"Dict[{key_type}, {value_type}]"
+            type_desc = f"Dict[{key_type}, {value_type}]"
         elif isinstance(field_type, type) and issubclass(field_type, BaseModel):
-            return generate_model_description(field_type)
+            type_desc = generate_model_description(field_type)
         elif hasattr(field_type, "__name__"):
-            return field_type.__name__
+            type_desc = field_type.__name__
         else:
-            return str(field_type)
+            type_desc = str(field_type)
 
-    fields = model.__annotations__
-    field_descriptions = [
-        f'"{name}": {describe_field(type_)}' for name, type_ in fields.items()
-    ]
+        if field_info and field_info.description:
+            return {"type": type_desc, "description": field_info.description}
+        return type_desc
+
+    fields = model.model_fields
+    field_descriptions = []
+    for name, field in fields.items():
+        field_desc = describe_field(field.annotation, field)
+        if isinstance(field_desc, dict):
+            field_descriptions.append(f'"{name}": {json.dumps(field_desc)}')
+        else:
+            field_descriptions.append(f'"{name}": {field_desc}')
     return "{\n  " + ",\n  ".join(field_descriptions) + "\n}"
diff --git a/tests/utilities/test_converter.py b/tests/utilities/test_converter.py
index 3f4a4d07b..edfee928a 100644
--- a/tests/utilities/test_converter.py
+++ b/tests/utilities/test_converter.py
@@ -4,7 +4,7 @@ from typing import Dict, List, Optional
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from crewai.llm import LLM
 from crewai.utilities.converter import (
@@ -328,6 +328,26 @@ def test_generate_model_description_dict_field():
     assert description == expected_description
 
 
+def test_generate_model_description_with_field_descriptions():
+    class ModelWithDescriptions(BaseModel):
+        name: str = Field(..., description="The user's full name")
+        age: int = Field(..., description="The user's age in years")
+        
+    description = generate_model_description(ModelWithDescriptions)
+    expected = '{\n  "name": {"type": "str", "description": "The user\'s full name"},\n  "age": {"type": "int", "description": "The user\'s age in years"}\n}'
+    assert description == expected
+
+
+def test_generate_model_description_mixed_fields():
+    class MixedModel(BaseModel):
+        name: str = Field(..., description="The user's name")
+        age: int  # No description
+        
+    description = generate_model_description(MixedModel)
+    expected = '{\n  "name": {"type": "str", "description": "The user\'s name"},\n  "age": int\n}'
+    assert description == expected
+
+
 @pytest.mark.vcr(filter_headers=["authorization"])
 def test_convert_with_instructions():
     llm = LLM(model="gpt-4o-mini")