Only throw if in union

This commit is contained in:
William Fu-Hinthorn
2025-03-18 09:58:07 -07:00
parent 987b9da4ab
commit ee650ab85f
2 changed files with 46 additions and 21 deletions
+26 -12
View File
@@ -62,18 +62,22 @@ class SchemaCoercionMapper:
processed = {}
if self._field_coercers is None:
self._field_coercers = {
n: self._build_coercer(t) for n, t in self._fields.items()
n: self._build_coercer(t, depth - 1) for n, t in self._fields.items()
}
for k, v in input_data.items():
fn = self._field_coercers.get(k)
processed[k] = fn(v, depth - 1) if fn else v
return self._construct(**processed)
def _build_coercer(self, field_type: Any) -> Callable[[Any, Any], Any]:
def _build_coercer(
self, field_type: Any, depth: int, throw: bool = False
) -> Callable[[Any, Any], Any]:
if depth == 0:
return self._passthrough
origin = get_origin(field_type)
if origin is Annotated:
real_type, *_ = get_args(field_type)
sub = self._build_coercer(real_type)
sub = self._build_coercer(real_type, depth - 1)
return lambda v, d: sub(v, d)
if isclass(field_type):
is_class_ = True
@@ -84,16 +88,16 @@ class SchemaCoercionMapper:
is_base_model = False
if is_base_model:
mapper = SchemaCoercionMapper(field_type, self.max_depth)
mapper = SchemaCoercionMapper(field_type, depth - 1)
return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v
if is_class_ and issubclass(field_type, BaseModelV1):
mapper = SchemaCoercionMapper(field_type, self.max_depth)
mapper = SchemaCoercionMapper(field_type, depth - 1)
return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v
if origin is list or field_type is list:
args = get_args(field_type)
if len(args) != 1:
return lambda v, d: v
sub = self._build_coercer(args[0])
sub = self._build_coercer(args[0], depth - 1)
def list_coercer(v: Any, d: Any) -> Any:
if not isinstance(v, (list, tuple)):
@@ -104,9 +108,17 @@ class SchemaCoercionMapper:
if origin is dict or field_type is dict:
args = get_args(field_type)
if len(args) != 2:
return self._passthrough
k_sub = self._build_coercer(args[0])
v_sub = self._build_coercer(args[1])
if not throw:
return self._passthrough
def dict_coercer(v: Any, d: Any) -> Any:
if not isinstance(v, dict):
raise TypeError("Expected dict, got %s" % type(v))
return {k_sub(k, d - 1): v_sub(val, d - 1) for k, val in v.items()}
return dict_coercer
k_sub = self._build_coercer(args[0], depth - 1)
v_sub = self._build_coercer(args[1], depth - 1)
def dict_coercer(v: Any, d: Any) -> Any:
if not isinstance(v, dict):
@@ -119,7 +131,7 @@ class SchemaCoercionMapper:
targs = get_args(field_type)
if not targs:
return lambda v, d: v
subs = [self._build_coercer(a) for a in targs]
subs = [self._build_coercer(a, depth - 1) for a in targs]
def tuple_coercer(v: Any, d: Any) -> Any:
if not isinstance(v, (list, tuple)):
@@ -133,11 +145,13 @@ class SchemaCoercionMapper:
if origin is Union:
uargs = get_args(field_type)
subs, none_in_union = [], False
for arg in uargs:
for ix, arg in enumerate(uargs):
if arg is type(None):
none_in_union = True
else:
subs.append(self._build_coercer(arg))
subs.append(
self._build_coercer(arg, depth - 1, throw=ix < len(uargs) - 1)
)
def union_coercer(v: Any, d: Any) -> Any:
if v is None and none_in_union:
+20 -9
View File
@@ -27,11 +27,7 @@ from uuid import UUID
import httpx
import pytest
from langchain_core.language_models import GenericFakeChatModel
from langchain_core.runnables import (
RunnableConfig,
RunnableLambda,
RunnablePassthrough,
)
from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough
from langchain_core.utils.aiter import aclosing
from pytest_mock import MockerFixture
from syrupy import SnapshotAssertion
@@ -4524,6 +4520,7 @@ async def test_nested_pydantic_models(version: str) -> None:
class NestedModel(BaseModel):
value: int
name: str
something: Optional[str] = None
# Forward reference model
class RecursiveModel(BaseModel):
@@ -4545,18 +4542,27 @@ async def test_nested_pydantic_models(version: str) -> None:
name: str
friends: list[str] = Field(default_factory=list) # IDs of friends
class MyTypedDict(TypedDict):
x: int
class State(BaseModel):
# Basic nested model tests
top_level: str
nested: NestedModel
optional_nested: Optional[NestedModel] = None
dict_nested: dict[str, NestedModel]
my_set: set[int]
list_nested: Annotated[
Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
]
list_nested_reversed: Annotated[
Union[list[dict[str, NestedModel]], NestedModel, dict, list],
lambda x, y: (x or []) + [y],
]
tuple_nested: tuple[str, NestedModel]
tuple_list_nested: list[tuple[int, NestedModel]]
complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]]
my_typed_dict: MyTypedDict
# Forward reference test
recursive: RecursiveModel
@@ -4572,8 +4578,11 @@ async def test_nested_pydantic_models(version: str) -> None:
"top_level": "initial",
"nested": {"value": 42, "name": "test"},
"optional_nested": {"value": 10, "name": "optional"},
"my_set": [1, 2, 4.5],
"my_typed_dict": {"x": 1},
"dict_nested": {"a": {"value": 5, "name": "a"}},
"list_nested": [{"a": {"value": 6, "name": "b"}}],
"list_nested_reversed": ["foo", "bar"],
"tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}],
"tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]],
"complex_tuple": [
@@ -5882,10 +5891,12 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
):
assert isinstance(store, BaseStore)
await store.aput(
namespace
if self.i is not None
and config["configurable"]["thread_id"] in (thread_1, thread_2)
else (f"foo_{self.i}", "bar"),
(
namespace
if self.i is not None
and config["configurable"]["thread_id"] in (thread_1, thread_2)
else (f"foo_{self.i}", "bar")
),
doc_id,
{
**doc,