From 5502cbf2334bf7610087b12f84ef2611bd3c2266 Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Tue, 13 Feb 2024 12:26:45 -0800
Subject: [PATCH] Add first token time for oai wrapper (#435)

---
 js/package.json                 |  2 +-
 python/langsmith/run_helpers.py | 34 ++++++++++++++++++++++++++++-----
 python/langsmith/run_trees.py   |  8 ++++++--
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/js/package.json b/js/package.json
index 1990f198f..91028fc28 100644
--- a/js/package.json
+++ b/js/package.json
@@ -122,4 +122,4 @@
     },
     "./package.json": "./package.json"
   }
-}
\ No newline at end of file
+}
diff --git a/python/langsmith/run_helpers.py b/python/langsmith/run_helpers.py
index e7c34d836..c369e44b4 100644
--- a/python/langsmith/run_helpers.py
+++ b/python/langsmith/run_helpers.py
@@ -4,6 +4,7 @@
 
 import contextlib
 import contextvars
+import datetime
 import functools
 import inspect
 import logging
@@ -117,6 +118,7 @@ def _container_end(
     container: _TraceableContainer,
     outputs: Optional[Any] = None,
     error: Optional[str] = None,
+    events: Optional[List[dict]] = None,
 ):
     """End the run."""
     run_tree = container.get("new_run")
@@ -124,7 +126,7 @@ def _container_end(
         # Tracing disabled
         return
     outputs_ = outputs if isinstance(outputs, dict) else {"output": outputs}
-    run_tree.end(outputs=outputs_, error=error)
+    run_tree.end(outputs=outputs_, error=error, events=events)
     run_tree.patch()
 
 
@@ -350,6 +352,7 @@ async def async_wrapper(
         async def async_generator_wrapper(
             *args: Any, langsmith_extra: Optional[LangSmithExtra] = None, **kwargs: Any
         ) -> AsyncGenerator:
+            events: List[dict] = []
             context_run = _PARENT_RUN_TREE.get()
             run_container = _setup_run(
                 func,
@@ -385,11 +388,21 @@ async def async_generator_wrapper(
                 if inspect.iscoroutine(async_gen_result):
                     async_gen_result = await async_gen_result
                 async for item in async_gen_result:
+                    if run_type == "llm":
+                        events.append(
+                            {
+                                "name": "new_token",
+                                "time": datetime.datetime.now(
+                                    datetime.timezone.utc
+                                ).isoformat(),
+                                "kwargs": {"token": item},
+                            },
+                        )
                     results.append(item)
                     yield item
             except BaseException as e:
                 stacktrace = traceback.format_exc()
-                _container_end(run_container, error=stacktrace)
+                _container_end(run_container, error=stacktrace, events=events)
                 raise e
             finally:
                 _PARENT_RUN_TREE.set(context_run)
@@ -407,7 +420,7 @@ async def async_generator_wrapper(
                     function_result = results
             else:
                 function_result = None
-            _container_end(run_container, outputs=function_result)
+            _container_end(run_container, outputs=function_result, events=events)
 
         @functools.wraps(func)
         def wrapper(
@@ -456,6 +469,7 @@ def generator_wrapper(
             *args: Any, langsmith_extra: Optional[LangSmithExtra] = None, **kwargs: Any
         ) -> Any:
             context_run = _PARENT_RUN_TREE.get()
+            events: List[dict] = []
             run_container = _setup_run(
                 func,
                 run_type=run_type,
@@ -483,6 +497,16 @@ def generator_wrapper(
                     # around this.
                     generator_result = func(*args, **kwargs)
                 for item in generator_result:
+                    if run_type == "llm":
+                        events.append(
+                            {
+                                "name": "new_token",
+                                "time": datetime.datetime.now(
+                                    datetime.timezone.utc
+                                ).isoformat(),
+                                "kwargs": {"token": item},
+                            },
+                        )
                     results.append(item)
                     try:
                         yield item
@@ -490,7 +514,7 @@ def generator_wrapper(
                         break
             except BaseException as e:
                 stacktrace = traceback.format_exc()
-                _container_end(run_container, error=stacktrace)
+                _container_end(run_container, error=stacktrace, events=events)
                 raise e
             finally:
                 _PARENT_RUN_TREE.set(context_run)
@@ -508,7 +532,7 @@ def generator_wrapper(
                     function_result = results
             else:
                 function_result = None
-            _container_end(run_container, outputs=function_result)
+            _container_end(run_container, outputs=function_result, events=events)
 
         if inspect.isasyncgenfunction(func):
             selected_wrapper: Callable = async_generator_wrapper
diff --git a/python/langsmith/run_trees.py b/python/langsmith/run_trees.py
index 7f10a820a..5d4e91277 100644
--- a/python/langsmith/run_trees.py
+++ b/python/langsmith/run_trees.py
@@ -73,8 +73,6 @@ def infer_defaults(cls, values: dict) -> dict:
                 values["trace_id"] = values["parent_run"].trace_id
             else:
                 values["trace_id"] = values["id"]
-        else:
-            print(values["trace_id"])
         cast(dict, values.setdefault("extra", {}))
         return values
 
@@ -100,6 +98,7 @@ def end(
         outputs: Optional[Dict] = None,
         error: Optional[str] = None,
         end_time: Optional[datetime] = None,
+        events: Optional[List[Dict]] = None,
     ) -> None:
         """Set the end time of the run and all child runs."""
         self.end_time = end_time or datetime.now(timezone.utc)
@@ -107,6 +106,8 @@ def end(
             self.outputs = outputs
         if error is not None:
             self.error = error
+        if events is not None:
+            self.events = events
 
     def create_child(
         self,
@@ -181,6 +182,9 @@ def patch(self) -> None:
             end_time=self.end_time,
             dotted_order=self.dotted_order,
             trace_id=self.trace_id,
+            events=self.events,
+            tags=self.tags,
+            extra=self.extra,
         )
 
     def wait(self) -> None: