diff --git a/google/genai/_common.py b/google/genai/_common.py index 64dd44274..b2a1b4d17 100644 --- a/google/genai/_common.py +++ b/google/genai/_common.py @@ -332,6 +332,14 @@ def _remove_extra_fields(model: Any, response: dict[str, object]) -> None: Mutates the response in place. """ + # Models are built on first use (`defer_build`), and an unbuilt model's field + # annotations still hold the forward references they were declared with, e.g. + # the string 'Part' rather than the `Part` class. Building the model resolves + # them, and costs nothing extra here because the caller validates against + # this model immediately afterwards. + if isinstance(model, type) and issubclass(model, pydantic.BaseModel): + model.model_rebuild(raise_errors=False) + key_values = list(response.items()) for key, value in key_values: diff --git a/google/genai/tests/common/test_common.py b/google/genai/tests/common/test_common.py index 1ead6e46d..75ca4b506 100644 --- a/google/genai/tests/common/test_common.py +++ b/google/genai/tests/common/test_common.py @@ -952,3 +952,30 @@ class TestModel(_common.BaseModel): TestModel.model_validate(data) assert len(caplog.records) == 0 + + +class _ForwardRefOuterModel(_common.BaseModel): + # Declared before the model it points at, the way the generated types + # sometimes are, so the annotation is a forward reference. + items: Optional[List["_ForwardRefInnerModel"]] = None + + +class _ForwardRefInnerModel(_common.BaseModel): + value: Optional[str] = None + + +def test_from_response_resolves_forward_references(): + """Extra fields are pruned from nested models declared as forward references. + + Models are built on first use, so until something builds the outer model its + `items` annotation holds an unresolved reference to `_ForwardRefInnerModel` + rather than the class, and the pruning walk has nothing to recurse into. + """ + assert not _ForwardRefOuterModel.__pydantic_complete__ + + result = _ForwardRefOuterModel._from_response( + response={"items": [{"value": "a", "extraField": 1}]}, kwargs={} + ) + + assert result.items is not None + assert result.items[0].value == "a"