Add tags to decorator (#109)

python/langsmith/run_helpers.py

@@ -7,7 +7,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from functools import wraps
-from typing import Any, Callable, Dict, Generator, Optional, Union
+from typing import Any, Callable, Dict, Generator, List, Mapping, Optional, Union
 from uuid import UUID
 
 from langsmith.run_trees import RunTree
@@ -55,6 +55,8 @@ def traceable(
     name: Optional[str] = None,
     extra: Optional[Dict] = None,
     executor: Optional[ThreadPoolExecutor] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    tags: Optional[List[str]] = None,
 ) -> Callable:
     """Decorator for creating or adding a run to a run tree."""
     _warn_once()
@@ -86,6 +88,8 @@ async def async_wrapper(
                 extra_inner = {**extra_outer, **run_extra}
             else:
                 extra_inner = extra_outer
+            if metadata:
+                extra_inner["metadata"] = metadata
             inputs = _get_inputs(signature, *args, **kwargs)
             if parent_run_ is not None:
                 new_run = parent_run_.create_child(
@@ -97,6 +101,7 @@ async def async_wrapper(
                         "doc": docstring,
                     },
                     inputs=inputs,
+                    tags=tags,
                     extra=extra_inner,
                 )
             else:
@@ -112,6 +117,7 @@ async def async_wrapper(
                     reference_example_id=reference_example_id,
                     project_name=project_name_,
                     extra=extra_inner,
+                    tags=tags,
                     executor=executor,
                 )
             new_run.post()
@@ -162,6 +168,8 @@ def wrapper(
                 extra_inner = {**extra_outer, **run_extra}
             else:
                 extra_inner = extra_outer
+            if metadata:
+                extra_inner["metadata"] = metadata
             inputs = _get_inputs(signature, *args, **kwargs)
             if parent_run_ is not None:
                 new_run = parent_run_.create_child(
@@ -173,6 +181,7 @@ def wrapper(
                         "doc": docstring,
                     },
                     inputs=inputs,
+                    tags=tags,
                     extra=extra_inner,
                 )
             else:
@@ -188,6 +197,7 @@ def wrapper(
                     reference_example_id=reference_example_id,
                     project_name=project_name_,
                     extra=extra_inner,
+                    tags=tags,
                     executor=executor,
                 )
             new_run.post()
@@ -231,10 +241,14 @@ def trace(
     executor: Optional[ThreadPoolExecutor] = None,
     project_name: Optional[str] = None,
     run_tree: Optional[RunTree] = None,
+    tags: Optional[List[str]] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
 ) -> Generator[RunTree, None, None]:
     """Context manager for creating a run tree."""
     _warn_cm_once()
     extra_outer = extra or {}
+    if metadata:
+        extra_outer["metadata"] = metadata
     parent_run_ = _PARENT_RUN_TREE.get() if run_tree is None else run_tree
     outer_project = _PROJECT_NAME.get()
     project_name_ = project_name or outer_project
@@ -244,6 +258,7 @@ def trace(
             run_type=run_type,
             extra=extra_outer,
             inputs=inputs,
+            tags=tags,
         )
     else:
         new_run = RunTree(
@@ -253,6 +268,7 @@ def trace(
             executor=executor,
             project_name=project_name_,
             inputs=inputs or {},
+            tags=tags,
         )
     new_run.post()
     _PARENT_RUN_TREE.set(new_run)

python/langsmith/run_trees.py

@@ -112,6 +112,7 @@ def create_child(
         reference_example_id: Optional[UUID] = None,
         start_time: Optional[datetime] = None,
         end_time: Optional[datetime] = None,
+        tags: Optional[List[str]] = None,
         extra: Optional[Dict] = None,
     ) -> RunTree:
         """Add a child run to the run tree."""
@@ -135,6 +136,7 @@ def create_child(
             session_name=self.session_name,
             client=self.client,
             executor=self.executor,
+            tags=tags,
         )
         self.child_runs.append(run)
         return run

python/tests/integration_tests/test_runs.py

@@ -42,7 +42,7 @@ def my_llm_run(text: str):
         # The function needn't accept a run
         return f"Completed: {text}"
 
-    @traceable(run_type="chain", executor=executor)
+    @traceable(run_type="chain", executor=executor, tags=["foo", "bar"])
     def my_chain_run(text: str):
         return my_run(text)
 
@@ -53,6 +53,7 @@ def my_chain_run(text: str):
     runs_dict = {run.name: run for run in runs}
     assert runs_dict["my_chain_run"].parent_run_id is None
     assert runs_dict["my_chain_run"].run_type == "chain"
+    assert runs_dict["my_chain_run"].tags == ["foo", "bar"]
     assert runs_dict["my_run"].parent_run_id == runs_dict["my_chain_run"].id
     assert runs_dict["my_run"].run_type == "chain"
     assert runs_dict["my_llm_run"].parent_run_id == runs_dict["my_run"].id