mirror of
https://github.com/langchain-ai/langchain.git
synced 2026-10-05 09:25:14 +03:00
fix(core): make StructuredTool JSON-serializable (#39631)
## Problem
`StructuredTool` cannot be dumped to JSON:
```python
@tool
def write_file(file_path: str, content: str) -> str:
"""Write content to the given path."""
return "ok"
write_file.model_dump(mode="json")
# PydanticSerializationError: Unable to serialize unknown type:
# <class 'pydantic._internal._model_construction.ModelMetaclass'>
```
`args_schema` holds a Pydantic model class, and `func` / `coroutine`
hold callables. None of them have a JSON form. Python-mode
`model_dump()` works; only the JSON modes raise.
This also costs tracing performance. The [LangSmith
SDK](https://github.com/langchain-ai/langsmith-sdk/blob/main/python/langsmith/_internal/_serde.py)
catches the error, dumps again in Python mode, and then sends every
class and function left in the result through its own `default` hook.
## Change
A `PlainSerializer(..., when_used="json-unless-none")` on the three
fields.
- `args_schema` dumps as its own JSON schema: `model_json_schema()` for
a Pydantic v2 class, `schema()` for a v1 one. A dict schema passes
through unchanged. A schema holding an arbitrary type has no JSON schema
at all, so it falls back to its repr instead of raising.
- `func` and `coroutine` dump as strings.
- Python-mode dumps are unchanged — still the live schema class and
callables.
- `exclude` / `include` / `exclude_none` keep working.
- Attached with `Annotated`, not `@field_serializer`. A field has only
one serializer slot, so `@field_serializer` would break any subclass
that declares its own serializer for the same field.
- Schema generation is cached per class. Pydantic does not memoize it,
and tracing would pay for it on every run.
## Result
The dump is JSON-native, and the schema it carries has the same shape a
dict `args_schema` already has, so it validates back into a working
tool.
Measured on the 8 filesystem tools of a deepagents agent, dumped through
the LangSmith serializer:
| | master | this PR |
|---|---|---|
| time per dump | 0.119 ms | **0.025 ms** |
| payload | 7,706 B | 12,213 B |
| objects reaching the SDK's `default` hook | 32 | 8 |
The payload grows because `args_schema` now carries the real schema
instead of `"<class ...>"`. Only the tool itself still enters the
`default` hook; the class and functions inside it no longer do. Without
the per-class cache the same dump takes 1.10 ms, so the cache is what
makes this a win rather than a regression.
## Why not on `BaseTool`
`args_schema` is declared there as well, so `Tool` and custom subclasses
hit the same error. But an `Annotated` serializer only applies where the
field is declared, and `StructuredTool` redeclares `args_schema` — it
would not inherit one from `BaseTool`.
## Tests
In `libs/core/tests/unit_tests/test_tools.py`: JSON round trip, Python
mode unchanged, dump options respected, dict `args_schema` preserved,
Pydantic v1 schema class, arbitrary-type fallback, and a subclass
declaring its own serializers for the same fields.
This commit is contained in:
1 parent
3b3b308e6d
commit
13b1b2feae
2 files changed
+203
-5
No files matched your search
@@ -13,7 +13,7 @@ from typing import (
|
||||
Literal,
|
||||
)
|
||||
|
||||
from pydantic import Field, SkipValidation
|
||||
from pydantic import Field, PlainSerializer, SkipValidation
|
||||
from typing_extensions import override
|
||||
|
||||
# Cannot move to TYPE_CHECKING as _run/_arun parameter annotations are needed at runtime
|
||||
@@ -31,10 +31,63 @@ from langchain_core.tools.base import (
|
||||
_get_runnable_config_param,
|
||||
create_schema_from_function,
|
||||
)
|
||||
from langchain_core.utils.pydantic import is_basemodel_subclass
|
||||
from langchain_core.utils.pydantic import is_basemodel_subclass, model_json_schema
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_core.messages import ToolCall
|
||||
from langchain_core.utils.pydantic import TypeBaseModel
|
||||
|
||||
|
||||
def _serialize_as_str(value: Any) -> str:
|
||||
"""Stringify a value that has no JSON form.
|
||||
|
||||
A plain function rather than the `str` builtin: Pydantic < 2.9 inspects the
|
||||
serializer's signature, and `inspect.signature` raises for C builtins.
|
||||
"""
|
||||
return str(value)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=256)
|
||||
def _model_json_schema(args_schema: TypeBaseModel, _rebuild_token: object) -> Any:
|
||||
"""Ask a v1 or v2 model class for its JSON schema, or fall back to its repr.
|
||||
|
||||
Cached because pydantic does not memoize this per class and it dominates the
|
||||
dump -- 1.10 ms against 0.02 ms for an eight-tool payload -- which tracing
|
||||
pays on every run.
|
||||
|
||||
`_rebuild_token` is the class's validator, which `model_rebuild()` replaces.
|
||||
Keying on it retires the entry for a rebuilt schema, and for one whose
|
||||
forward reference was still unresolved when it was first dumped. Mutating a
|
||||
class in place without rebuilding it is not detected.
|
||||
"""
|
||||
try:
|
||||
return model_json_schema(args_schema)
|
||||
except Exception: # a schema holding an arbitrary type has no JSON schema
|
||||
return str(args_schema)
|
||||
|
||||
|
||||
def _serialize_args_schema(args_schema: ArgsSchema) -> Any:
|
||||
"""Represent a schema class by its JSON schema when dumping to JSON.
|
||||
|
||||
A Pydantic model class has no JSON form, so leaving it to the default
|
||||
serializer raises `PydanticSerializationError`. Its own JSON schema is the
|
||||
shape a dict schema already has, and a dict is returned unchanged.
|
||||
"""
|
||||
if isinstance(args_schema, dict):
|
||||
return args_schema
|
||||
return _model_json_schema(
|
||||
args_schema, getattr(args_schema, "__pydantic_validator__", None)
|
||||
)
|
||||
|
||||
|
||||
# Attached to the fields via `Annotated` rather than declared with
|
||||
# `@field_serializer`, which would take the field's one serializer slot and make
|
||||
# any subclass that declares its own serializer for the same field fail at class
|
||||
# creation with `PydanticUserError: Multiple field serializer functions ...`.
|
||||
_JsonSchemaFallback = PlainSerializer(
|
||||
_serialize_args_schema, when_used="json-unless-none"
|
||||
)
|
||||
_JsonCallableFallback = PlainSerializer(_serialize_as_str, when_used="json-unless-none")
|
||||
|
||||
|
||||
class StructuredTool(BaseTool):
|
||||
@@ -42,15 +95,17 @@ class StructuredTool(BaseTool):
|
||||
|
||||
description: str = ""
|
||||
|
||||
args_schema: Annotated[ArgsSchema, SkipValidation()] = Field(
|
||||
args_schema: Annotated[ArgsSchema, SkipValidation(), _JsonSchemaFallback] = Field(
|
||||
..., description="The tool schema."
|
||||
)
|
||||
"""The input arguments' schema."""
|
||||
|
||||
func: Callable[..., Any] | None = None
|
||||
func: Annotated[Callable[..., Any] | None, _JsonCallableFallback] = None
|
||||
"""The function to run when the tool is called."""
|
||||
|
||||
coroutine: Callable[..., Awaitable[Any]] | None = None
|
||||
coroutine: Annotated[
|
||||
Callable[..., Awaitable[Any]] | None, _JsonCallableFallback
|
||||
] = None
|
||||
"""The asynchronous version of the function."""
|
||||
|
||||
# --- Runnable ---
|
||||
|
||||
@@ -31,6 +31,7 @@ from pydantic import (
|
||||
Field,
|
||||
RootModel,
|
||||
ValidationError,
|
||||
field_serializer,
|
||||
)
|
||||
from pydantic.errors import PydanticUndefinedAnnotation
|
||||
from pydantic.v1 import BaseModel as BaseModelV1
|
||||
@@ -4734,3 +4735,145 @@ def test_tool_call_schema_json_schema_cache_invalidated_on_reassignment() -> Non
|
||||
new_schema = new_cls.model_json_schema()
|
||||
assert new_schema is not old_schema
|
||||
assert new_schema["description"] == "New description for cache test."
|
||||
|
||||
|
||||
def test_structured_tool_json_dump() -> None:
|
||||
"""`mode="json"` dumps must not raise on the callable / schema-class fields."""
|
||||
|
||||
@tool
|
||||
def write_file(file_path: str, content: str) -> str:
|
||||
"""Write content to the given path."""
|
||||
return "ok"
|
||||
|
||||
dumped = write_file.model_dump(mode="json")
|
||||
# Round-trips through the stdlib encoder, i.e. it really is JSON-native.
|
||||
json.dumps(dumped)
|
||||
schema = cast("type[BaseModel]", write_file.args_schema)
|
||||
assert dumped["args_schema"] == schema.model_json_schema()
|
||||
assert set(dumped["args_schema"]["properties"]) == {"file_path", "content"}
|
||||
assert "write_file" in dumped["func"]
|
||||
json.loads(write_file.model_dump_json())
|
||||
|
||||
# Python-mode dumps still hand out the live objects.
|
||||
native = write_file.model_dump()
|
||||
assert callable(native["func"])
|
||||
assert isinstance(native["args_schema"], type)
|
||||
|
||||
|
||||
def test_structured_tool_json_dump_respects_options() -> None:
|
||||
"""The field serializers compose with the usual `model_dump` options."""
|
||||
|
||||
@tool
|
||||
def write_file(file_path: str, content: str) -> str:
|
||||
"""Write content to the given path."""
|
||||
return "ok"
|
||||
|
||||
assert "description" not in write_file.model_dump(
|
||||
mode="json", exclude={"description"}
|
||||
)
|
||||
assert set(write_file.model_dump(mode="json", include={"name", "func"})) == {
|
||||
"name",
|
||||
"func",
|
||||
}
|
||||
# `coroutine` is None here, so `exclude_none` must still drop it.
|
||||
assert "coroutine" not in write_file.model_dump(mode="json", exclude_none=True)
|
||||
|
||||
|
||||
def test_structured_tool_json_dump_keeps_dict_args_schema() -> None:
|
||||
"""A dict schema is already JSON-native and must be passed through as-is."""
|
||||
schema = {"type": "object", "properties": {"a": {"type": "string"}}}
|
||||
dict_tool = StructuredTool(
|
||||
name="d",
|
||||
description="d",
|
||||
func=lambda **kwargs: "x",
|
||||
args_schema=schema,
|
||||
)
|
||||
assert dict_tool.model_dump(mode="json")["args_schema"] == schema
|
||||
|
||||
|
||||
def test_structured_tool_json_dump_uses_v1_schema_method() -> None:
|
||||
"""A `pydantic.v1` schema class is a supported `args_schema` form too."""
|
||||
|
||||
class V1Schema(BaseModelV1):
|
||||
a: str
|
||||
|
||||
v1_tool = StructuredTool(
|
||||
name="v", description="d", func=lambda a: "x", args_schema=V1Schema
|
||||
)
|
||||
assert v1_tool.model_dump(mode="json")["args_schema"] == V1Schema.schema()
|
||||
|
||||
|
||||
def test_structured_tool_json_dump_falls_back_for_arbitrary_types() -> None:
|
||||
"""A schema with no JSON schema form must still dump instead of raising."""
|
||||
|
||||
class NotJsonSchemable:
|
||||
pass
|
||||
|
||||
@tool
|
||||
def uses_injected(
|
||||
a: int, conn: Annotated[NotJsonSchemable, InjectedToolArg]
|
||||
) -> str:
|
||||
"""Doc."""
|
||||
return "ok"
|
||||
|
||||
assert uses_injected.model_dump(mode="json")["args_schema"].startswith("<class ")
|
||||
# The injected arg never reaches the model, so the tool is still usable.
|
||||
assert set(uses_injected.args) == {"a"}
|
||||
|
||||
|
||||
def test_structured_tool_json_dump_follows_model_rebuild() -> None:
|
||||
"""A rebuilt schema must not keep serving its pre-rebuild JSON schema."""
|
||||
|
||||
class Renamable(BaseModel):
|
||||
a: str
|
||||
|
||||
renamable = StructuredTool(
|
||||
name="r", description="d", func=lambda a: "x", args_schema=Renamable
|
||||
)
|
||||
assert renamable.model_dump(mode="json")["args_schema"]["title"] == "Renamable"
|
||||
|
||||
Renamable.model_config["title"] = "Renamed"
|
||||
Renamable.model_rebuild(force=True)
|
||||
assert renamable.model_dump(mode="json")["args_schema"]["title"] == "Renamed"
|
||||
|
||||
class Deferred(BaseModel):
|
||||
a: "Resolved"
|
||||
|
||||
deferred = StructuredTool(
|
||||
name="f", description="d", func=lambda a: "x", args_schema=Deferred
|
||||
)
|
||||
# Unresolvable so far, so it dumps as a repr like any other schema-less type.
|
||||
assert deferred.model_dump(mode="json")["args_schema"].startswith("<class ")
|
||||
|
||||
class Resolved(BaseModel):
|
||||
z: int
|
||||
|
||||
Deferred.model_rebuild(force=True)
|
||||
assert deferred.model_dump(mode="json")["args_schema"]["properties"] == {
|
||||
"a": {"$ref": "#/$defs/Resolved"}
|
||||
}
|
||||
|
||||
|
||||
def test_structured_tool_subclass_can_override_json_serializers() -> None:
|
||||
"""The JSON fallbacks must not occupy the fields' single serializer slot.
|
||||
|
||||
A subclass declaring its own `@field_serializer` for the same fields used to
|
||||
fail at class creation with `PydanticUserError: Multiple field serializer
|
||||
functions were defined`.
|
||||
"""
|
||||
|
||||
class MyTool(StructuredTool):
|
||||
@field_serializer("func", when_used="json-unless-none")
|
||||
def _my_func_repr(self, func: Any) -> str:
|
||||
return "custom-func"
|
||||
|
||||
@field_serializer("args_schema", when_used="json-unless-none")
|
||||
def _my_schema_repr(self, args_schema: Any) -> Any:
|
||||
return {"custom": "schema"}
|
||||
|
||||
my_tool = MyTool.from_function(
|
||||
func=lambda file_path, content: "ok", name="w", description="d"
|
||||
)
|
||||
dumped = my_tool.model_dump(mode="json")
|
||||
assert dumped["func"] == "custom-func"
|
||||
assert dumped["args_schema"] == {"custom": "schema"}
|
||||
Reference in new issue
Block a user