Update typehints (#448)

So you get linter warnings if you try to do a name. It doesn't actually
alter the behavior, so if we choose to extend the number of types, it
would just be a linting issue rather than a runtime issue

python/langsmith/client.py

@@ -29,6 +29,7 @@ from typing import (
     Iterable,
     Iterator,
     List,
+    Literal,
     Mapping,
     Optional,
     Sequence,
@@ -117,6 +118,9 @@ def _is_langchain_hosted(url: str) -> bool:
 
 
 ID_TYPE = Union[uuid.UUID, str]
+RUN_TYPE_T = Literal[
+    "tool", "chain", "llm", "retriever", "embedding", "prompt", "parser"
+]
 
 
 def _default_retry_config() -> Retry:
@@ -922,7 +926,7 @@ class Client:
         self,
         name: str,
         inputs: Dict[str, Any],
-        run_type: str,
+        run_type: RUN_TYPE_T,
         *,
         project_name: Optional[str] = None,
         revision_id: Optional[str] = None,

python/langsmith/run_helpers.py

@@ -32,7 +32,8 @@ from typing import (
     runtime_checkable,
 )
 
-from langsmith import client, run_trees, utils
+from langsmith import client as ls_client
+from langsmith import run_trees, utils
 
 if TYPE_CHECKING:
     from langchain.schema.runnable import Runnable
@@ -94,14 +95,14 @@ def _get_inputs(
 class LangSmithExtra(TypedDict, total=False):
     """Any additional info to be injected into the run dynamically."""
 
-    reference_example_id: Optional[client.ID_TYPE]
+    reference_example_id: Optional[ls_client.ID_TYPE]
     run_extra: Optional[Dict]
     run_tree: Optional[run_trees.RunTree]
     project_name: Optional[str]
     metadata: Optional[Dict[str, Any]]
     tags: Optional[List[str]]
-    run_id: Optional[client.ID_TYPE]
-    client: Optional[client.Client]
+    run_id: Optional[ls_client.ID_TYPE]
+    client: Optional[ls_client.Client]
 
 
 class _TraceableContainer(TypedDict, total=False):
@@ -141,13 +142,13 @@ def _collect_extra(extra_outer: dict, langsmith_extra: LangSmithExtra) -> dict:
 
 def _setup_run(
     func: Callable,
-    run_type: str,
+    run_type: ls_client.RUN_TYPE_T,
     extra_outer: dict,
     langsmith_extra: Optional[LangSmithExtra] = None,
     name: Optional[str] = None,
     metadata: Optional[Mapping[str, Any]] = None,
     tags: Optional[List[str]] = None,
-    client: Optional[client.Client] = None,
+    client: Optional[ls_client.Client] = None,
     args: Any = None,
     kwargs: Any = None,
 ) -> _TraceableContainer:
@@ -272,12 +273,12 @@ def traceable(
 
 @overload
 def traceable(
-    run_type: str = "chain",
+    run_type: ls_client.RUN_TYPE_T = "chain",
     *,
     name: Optional[str] = None,
     metadata: Optional[Mapping[str, Any]] = None,
     tags: Optional[List[str]] = None,
-    client: Optional[client.Client] = None,
+    client: Optional[ls_client.Client] = None,
     extra: Optional[Dict] = None,
     reduce_fn: Optional[Callable] = None,
 ) -> Callable[[Callable[..., R]], SupportsLangsmithExtra[R]]:
@@ -305,7 +306,7 @@ def traceable(
                 called, and the run itself will be stuck in a pending state.
 
     """
-    run_type = (
+    run_type: ls_client.RUN_TYPE_T = (
         args[0]
         if args and isinstance(args[0], str)
         else (kwargs.get("run_type") or "chain")
@@ -582,7 +583,7 @@ def traceable(
 @contextlib.contextmanager
 def trace(
     name: str,
-    run_type: str,
+    run_type: ls_client.RUN_TYPE_T = "chain",
     *,
     inputs: Optional[Dict] = None,
     extra: Optional[Dict] = None,

python/langsmith/run_trees.py

@@ -17,7 +17,7 @@ except ImportError:
     from pydantic import Field, root_validator, validator
 
 from langsmith import utils
-from langsmith.client import ID_TYPE, Client
+from langsmith.client import ID_TYPE, RUN_TYPE_T, Client
 from langsmith.schemas import RunBase
 
 logger = logging.getLogger(__name__)
@@ -112,7 +112,7 @@ class RunTree(RunBase):
     def create_child(
         self,
         name: str,
-        run_type: str,
+        run_type: RUN_TYPE_T = "chain",
         *,
         run_id: Optional[ID_TYPE] = None,
         serialized: Optional[Dict] = None,

python/tests/unit_tests/test_run_helpers.py

@@ -355,7 +355,7 @@ def test_traceable_warning() -> None:
     with warnings.catch_warnings(record=True) as warning_records:
         warnings.simplefilter("always")
 
-        @traceable(run_type="invalid_run_type")
+        @traceable(run_type="invalid_run_type")  # type: ignore
         def my_function() -> None:
             pass
 
@@ -373,7 +373,7 @@ def test_traceable_wrong_run_type_pos_arg() -> None:
     with warnings.catch_warnings(record=True) as warning_records:
         warnings.simplefilter("always")
 
-        @traceable("my_run_type")
+        @traceable("my_run_type")  # type: ignore
         def my_function() -> None:
             pass