WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-10 03:37:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

Support optional output schema

Browse Source
This commit is contained in:
William Fu-Hinthorn
2024-09-06 11:29:43 -07:00
parent a9ccd0fd79
commit 587335bfe6
1 changed files with 47 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/state.py
+47 -35
View File
@@ -474,45 +474,22 @@ class CompiledStateGraph(CompiledGraph):
def get_input_schema(
self, config: Optional[RunnableConfig] = None
) -> type[BaseModel]:
if isclass(self.builder.input) and issubclass(
self.builder.input, (BaseModel, BaseModelV1)
):
return self.builder.input
else:
keys = list(self.builder.schemas[self.builder.input].keys())
if len(keys) == 1 and keys[0] == "__root__":
return create_model( # type: ignore[call-overload]
self.get_name("Input"),
__root__=(self.channels[keys[0]].UpdateType, None),
)
else:
return create_model( # type: ignore[call-overload]
self.get_name("Input"),
**{
k: (
self.channels[k].UpdateType,
(
get_field_default(
k,
self.channels[k].UpdateType,
self.builder.input,
)
),
)
for k in self.builder.schemas[self.builder.input]
if isinstance(self.channels[k], BaseChannel)
},
)
return _get_schema(
typ=self.builder.input,
schemas=self.builder.schemas,
channels=self.builder.channels,
name=self.get_name("Input"),
)
def get_output_schema(
self, config: Optional[RunnableConfig] = None
) -> type[BaseModel]:
if isclass(self.builder.output) and issubclass(
self.builder.output, (BaseModel, BaseModelV1)
):
return self.builder.output
return super().get_output_schema(config)
return _get_schema(
typ=self.builder.output,
schemas=self.builder.schemas,
channels=self.builder.channels,
name=self.get_name("Output"),
)
def attach_node(self, key: str, node: Optional[StateNodeSpec]) -> None:
if key == START:
@@ -772,3 +749,38 @@ def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[ManagedValueS
return decoration
return None
def _get_schema(
typ: Type,
schemas: dict,
channels: dict,
name: str,
) -> type[BaseModel]:
if isclass(typ) and issubclass(typ, (BaseModel, BaseModelV1)):
return typ
else:
keys = list(schemas[typ].keys())
if len(keys) == 1 and keys[0] == "__root__":
return create_model( # type: ignore[call-overload]
name,
__root__=(channels[keys[0]].UpdateType, None),
)
else:
return create_model( # type: ignore[call-overload]
name,
**{
k: (
channels[k].UpdateType,
(
get_field_default(
k,
channels[k].UpdateType,
typ,
)
),
)
for k in schemas[typ]
if isinstance(channels[k], BaseChannel)
},
)