fix: serialize args_schema as JSON schema for checkpoint roundtrip

Use BeforeValidator/PlainSerializer with create_model_from_schema to
serialize type[BaseModel] as its JSON schema dict and reconstruct a
dynamic model on deserialization.
This commit is contained in:
Greyson LaLonde
2026-04-07 00:04:45 +08:00
parent 07aa7c0558
commit c9c3ac8dcc

View File

@@ -5,12 +5,14 @@ from collections.abc import Callable
import inspect
import json
import textwrap
from typing import TYPE_CHECKING, Any, get_type_hints
from typing import TYPE_CHECKING, Annotated, Any, get_type_hints
from pydantic import (
BaseModel,
BeforeValidator,
ConfigDict,
Field,
PlainSerializer,
PrivateAttr,
create_model,
model_validator,
@@ -18,9 +20,22 @@ from pydantic import (
from typing_extensions import Self
from crewai.utilities.logger import Logger
from crewai.utilities.pydantic_schema_utils import create_model_from_schema
from crewai.utilities.string_utils import sanitize_tool_name
def _serialize_schema(v: type[BaseModel] | None) -> dict[str, Any] | None:
return v.model_json_schema() if v else None
def _deserialize_schema(v: Any) -> type[BaseModel] | None:
if v is None or isinstance(v, type):
return v
if isinstance(v, dict):
return create_model_from_schema(v)
return None
if TYPE_CHECKING:
pass
@@ -61,7 +76,11 @@ class CrewStructuredTool(BaseModel):
name: str = Field(default="")
description: str = Field(default="")
args_schema: type[BaseModel] | None = Field(default=None)
args_schema: Annotated[
type[BaseModel] | None,
BeforeValidator(_deserialize_schema),
PlainSerializer(_serialize_schema),
] = Field(default=None)
func: Any = Field(default=None, exclude=True)
result_as_answer: bool = Field(default=False)
max_usage_count: int | None = Field(default=None)