[Py] Suport yield & return from wrapped generators

Sync only (since python async generators dont' support it anyhow)
langchain-ai · Aug 26, 2024 · 879f45b · 879f45b
1 parent b846a4c
commit 879f45b
Show file tree

Hide file tree

Showing 4 changed files with 108 additions and 82 deletions.
diff --git a/python/langsmith/run_helpers.py b/python/langsmith/run_helpers.py
@@ -619,6 +619,7 @@ def generator_wrapper(
                 inspect.signature(func).parameters.get("run_tree", None) is not None
             )
             results: List[Any] = []
+            function_return: Any = None
             try:
                 if func_accepts_parent_run:
                     kwargs["run_tree"] = run_container["new_run"]
@@ -647,8 +648,10 @@ def generator_wrapper(
                             yield item
                         except GeneratorExit:
                             break
-                except StopIteration:
-                    pass
+                except StopIteration as e:
+                    function_return = e.value
+                    if function_return is not None:
+                        results.append(function_return)
 
             except BaseException as e:
                 _on_run_end(run_container, error=e)
@@ -665,6 +668,7 @@ def generator_wrapper(
             else:
                 function_result = None
             _on_run_end(run_container, outputs=function_result)
+            return function_return
 
         if inspect.isasyncgenfunction(func):
             selected_wrapper: Callable = async_generator_wrapper