Update tracing name

This commit is contained in:
Nuno Campos
2024-01-13 10:55:18 -08:00
parent 3d58aff93e
commit d68faa6d06
3 changed files with 28 additions and 26 deletions
+4 -2
View File
@@ -162,6 +162,8 @@ class Pregel(
saver: Optional[BaseCheckpointSaver] = None
name: str = "LangGraph"
class Config:
arbitrary_types_allowed = True
@@ -191,7 +193,7 @@ class Pregel(
return super().get_input_schema(config)
else:
return create_model( # type: ignore[call-overload]
"PregelInput",
self.get_name("Input"),
**{
k: (self.channels[k].UpdateType, None)
for k in self.input or self.channels.keys()
@@ -210,7 +212,7 @@ class Pregel(
return super().get_output_schema(config)
else:
return create_model( # type: ignore[call-overload]
"PregelOutput",
self.get_name("Output"),
**{k: (self.channels[k].ValueType, None) for k in self.output},
)
+12 -12
View File
@@ -41,8 +41,8 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
graph.set_finish_point("add_one")
gapp = graph.compile()
assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"}
assert app.output_schema.schema() == {"title": "LangGraphOutput", "type": "integer"}
assert app.invoke(2) == 3
assert app.invoke(2, output=["output"]) == {"output": 3}
assert repr(app), "does not raise recursion error"
@@ -56,8 +56,8 @@ def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {"title": "PregelOutput"}
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {"title": "LangGraphOutput"}
assert app.invoke(2) == 3
@@ -71,9 +71,9 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"])
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {
"title": "PregelOutput",
"title": "LangGraphOutput",
"type": "object",
"properties": {
"output": {"title": "Output"},
@@ -95,8 +95,8 @@ def test_invoke_single_process_in_out_reserved_is_last(mocker: MockerFixture) ->
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {"title": "PregelOutput"}
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {"title": "LangGraphOutput"}
assert app.invoke(2) == {"input": 3, "is_last_step": False}
assert app.invoke(2, {"recursion_limit": 1}) == {"input": 3, "is_last_step": True}
@@ -112,9 +112,9 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
output=["output"],
)
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {
"title": "PregelOutput",
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output"}},
}
@@ -134,12 +134,12 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
)
assert app.input_schema.schema() == {
"title": "PregelInput",
"title": "LangGraphInput",
"type": "object",
"properties": {"input": {"title": "Input"}},
}
assert app.output_schema.schema() == {
"title": "PregelOutput",
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output"}},
}
+12 -12
View File
@@ -39,8 +39,8 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
graph.set_finish_point("add_one")
gapp = graph.compile()
assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"}
assert app.output_schema.schema() == {"title": "LangGraphOutput", "type": "integer"}
assert await app.ainvoke(2) == 3
assert await app.ainvoke(2, output=["output"]) == {"output": 3}
@@ -55,8 +55,8 @@ async def test_invoke_single_process_in_out_implicit_channels(
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {"title": "PregelOutput"}
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {"title": "LangGraphOutput"}
assert await app.ainvoke(2) == 3
@@ -70,9 +70,9 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N
app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"])
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {
"title": "PregelOutput",
"title": "LangGraphOutput",
"type": "object",
"properties": {
"output": {"title": "Output"},
@@ -96,8 +96,8 @@ async def test_invoke_single_process_in_out_reserved_is_last(
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {"title": "PregelOutput"}
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {"title": "LangGraphOutput"}
assert await app.ainvoke(2) == {"input": 3, "is_last_step": False}
assert await app.ainvoke(2, {"recursion_limit": 1}) == {
"input": 3,
@@ -114,9 +114,9 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
output=["output"],
)
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {
"title": "PregelOutput",
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output"}},
}
@@ -136,12 +136,12 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
)
assert app.input_schema.schema() == {
"title": "PregelInput",
"title": "LangGraphInput",
"type": "object",
"properties": {"input": {"title": "Input"}},
}
assert app.output_schema.schema() == {
"title": "PregelOutput",
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output"}},
}