From 292ae8468ea8ae3983cf7d71d5ec8450e0b19be1 Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Fri, 18 Aug 2023 17:21:53 -0700
Subject: [PATCH] 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
---
 libs/langchain/langchain/callbacks/manager.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/libs/langchain/langchain/callbacks/manager.py b/libs/langchain/langchain/callbacks/manager.py
index 7016a13ed95..fc05b547773 100644
--- a/libs/langchain/langchain/callbacks/manager.py
+++ b/libs/langchain/langchain/callbacks/manager.py
@@ -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: