Let you specify run id in trace as chain group (#9484)

I think we'll deprecate this soon anyway but still nice to be able to
fetch the run id
This commit is contained in:
William FH 2023-08-18 17:21:53 -07:00 committed by GitHub
parent b58d492e05
commit 292ae8468e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -190,6 +190,7 @@ def trace_as_chain_group(
*,
project_name: Optional[str] = None,
example_id: Optional[Union[str, UUID]] = None,
run_id: Optional[UUID] = None,
tags: Optional[List[str]] = None,
) -> Generator[CallbackManager, None, None]:
"""Get a callback manager for a chain group in a context manager.
@ -202,6 +203,7 @@ def trace_as_chain_group(
Defaults to None.
example_id (str or UUID, optional): The ID of the example.
Defaults to None.
run_id (UUID, optional): The ID of the run.
tags (List[str], optional): The inheritable tags to apply to all runs.
Defaults to None.
@ -229,7 +231,7 @@ def trace_as_chain_group(
inheritable_tags=tags,
)
run_manager = cm.on_chain_start({"name": group_name}, {})
run_manager = cm.on_chain_start({"name": group_name}, {}, run_id=run_id)
yield run_manager.get_child()
run_manager.on_chain_end({})
@ -241,6 +243,7 @@ async def atrace_as_chain_group(
*,
project_name: Optional[str] = None,
example_id: Optional[Union[str, UUID]] = None,
run_id: Optional[UUID] = None,
tags: Optional[List[str]] = None,
) -> AsyncGenerator[AsyncCallbackManager, None]:
"""Get an async callback manager for a chain group in a context manager.
@ -253,6 +256,7 @@ async def atrace_as_chain_group(
Defaults to None.
example_id (str or UUID, optional): The ID of the example.
Defaults to None.
run_id (UUID, optional): The ID of the run.
tags (List[str], optional): The inheritable tags to apply to all runs.
Defaults to None.
Returns:
@ -276,7 +280,7 @@ async def atrace_as_chain_group(
)
cm = AsyncCallbackManager.configure(inheritable_callbacks=cb, inheritable_tags=tags)
run_manager = await cm.on_chain_start({"name": group_name}, {})
run_manager = await cm.on_chain_start({"name": group_name}, {}, run_id=run_id)
try:
yield run_manager.get_child()
finally: