mirror of
https://github.com/langchain-ai/langchain.git
synced 2026-10-05 09:25:14 +03:00
fix(core): dump args_schema as its own JSON schema
A Pydantic model class has no JSON form, so the field was stringified. Ask the schema for its own JSON schema instead -- `model_json_schema()` for a v2 class, `schema()` for a v1 one -- so every accepted `args_schema` form dumps to the shape the dict form already has, and the dumped schema validates back into a tool. A dict schema is still returned unchanged. Cached per schema class: pydantic does not memoize schema generation, and it costs ~500x the rest of the dump, which tracing would pay on every run. A schema holding an arbitrary type has no JSON schema at all, so those still fall back to the repr rather than raising. 12-tool payload: 0.140 ms on master, 0.023 ms here.
This commit is contained in:
1 parent
1d993ca683
commit
0ed59a38e9
2 files changed
+51
-6
No files matched your search
@@ -31,10 +31,11 @@ from langchain_core.tools.base import (
|
||||
_is_injected_arg_type,
|
||||
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:
|
||||
@@ -46,16 +47,29 @@ def _serialize_as_str(value: Any) -> str:
|
||||
return str(value)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=256)
|
||||
def _model_json_schema(args_schema: TypeBaseModel) -> 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 costs ~500x
|
||||
the rest of the dump, which tracing pays on every run.
|
||||
"""
|
||||
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 as a string when dumping to JSON.
|
||||
"""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`. A dict schema is already
|
||||
JSON-compatible and is returned unchanged.
|
||||
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 str(args_schema)
|
||||
return _model_json_schema(args_schema)
|
||||
|
||||
|
||||
# Attached to the fields via `Annotated` rather than declared with
|
||||
|
||||
@@ -4381,7 +4381,8 @@ def test_structured_tool_json_dump() -> None:
|
||||
dumped = write_file.model_dump(mode="json")
|
||||
# Round-trips through the stdlib encoder, i.e. it really is JSON-native.
|
||||
json.dumps(dumped)
|
||||
assert dumped["args_schema"].startswith("<class ")
|
||||
assert dumped["args_schema"] == write_file.args_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())
|
||||
|
||||
@@ -4422,6 +4423,36 @@ def test_structured_tool_json_dump_keeps_dict_args_schema() -> None:
|
||||
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_subclass_can_override_json_serializers() -> None:
|
||||
"""The JSON fallbacks must not occupy the fields' single serializer slot.
|
||||
|
||||
|
||||
Reference in new issue
Block a user