Add a langsmith_pyo3 extra to the LangSmith SDK. (#1263)

If `langsmith` is installed with the `langsmith_pyo3` extra (e.g. `pip
install langsmith[langsmith_pyo3]`) and the `LANGSMITH_USE_PYO3_CLIENT`
env var is set to any value, the tracing client will use Rust (via PyO3
bindings) to submit runs to LangSmith servers.

python/langsmith/client.py

@@ -369,6 +369,7 @@ class Client:
         "_write_api_urls",
         "_settings",
         "_manual_cleanup",
+        "_pyo3_client",
     ]
 
     def __init__(
@@ -516,6 +517,42 @@ def __init__(
             else ls_utils.get_env_var("HIDE_OUTPUTS") == "true"
         )
 
+        # To trigger this code, set the `LANGSMITH_USE_PYO3_CLIENT` env var to any value.
+        self._pyo3_client = None
+        if ls_utils.get_env_var("USE_PYO3_CLIENT") is not None:
+            langsmith_pyo3 = None
+            try:
+                import langsmith_pyo3  # type: ignore[import-not-found, no-redef]
+            except ImportError as e:
+                logger.warning(
+                    "Failed to import `langsmith_pyo3` when PyO3 client was requested, "
+                    "falling back to Python impl: %s",
+                    repr(e),
+                )
+
+            if langsmith_pyo3:
+                # TODO: tweak these constants as needed
+                queue_capacity = 1_000_000
+                batch_size = 100
+                batch_timeout_millis = 1000
+                worker_threads = 1
+
+                try:
+                    self._pyo3_client = langsmith_pyo3.BlockingTracingClient(
+                        self.api_url,
+                        self.api_key,
+                        queue_capacity,
+                        batch_size,
+                        batch_timeout_millis,
+                        worker_threads,
+                    )
+                except Exception as e:
+                    logger.warning(
+                        "Failed to instantiate `langsmith_pyo3.BlockingTracingClient` "
+                        "when PyO3 client was requested, falling back to Python impl: %s",
+                        repr(e),
+                    )
+
         self._settings: Union[ls_schemas.LangSmithSettings, None] = None
 
         self._manual_cleanup = False
@@ -1226,16 +1263,26 @@ def create_run(
             copy=False,
         )
         self._insert_runtime_env([run_create])
+
         if (
-            self.tracing_queue is not None
             # batch ingest requires trace_id and dotted_order to be set
-            and run_create.get("trace_id") is not None
+            run_create.get("trace_id") is not None
             and run_create.get("dotted_order") is not None
         ):
-            serialized_op = serialize_run_dict("post", run_create)
-            self.tracing_queue.put(
-                TracingQueueItem(run_create["dotted_order"], serialized_op)
-            )
+            if self._pyo3_client is not None:
+                # `self._run_transform()` above turns the `id` key into a `UUID` object.
+                # We need to pass a string since `orjson` doesn't seem to serialize `UUID` objects.
+                run_create["id"] = str(run_create["id"])
+                self._pyo3_client.create_run(run_create)
+            elif self.tracing_queue is not None:
+                serialized_op = serialize_run_dict("post", run_create)
+                self.tracing_queue.put(
+                    TracingQueueItem(run_create["dotted_order"], serialized_op)
+                )
+            else:
+                # Neither Rust nor Python batch ingestion is configured,
+                # fall back to the non-batch approach.
+                self._create_run(run_create)
         else:
             self._create_run(run_create)
 

python/pyproject.toml

@@ -35,6 +35,9 @@ orjson = { version = "^3.9.14", markers = "platform_python_implementation != 'Py
 httpx = ">=0.23.0,<1"
 requests-toolbelt = "^1.0.0"
 
+# Enabled via `langsmith_pyo3` extra: `pip install langsmith[langsmith_pyo3]`.
+langsmith-pyo3 = { version = "^0.1.0rc2", optional = true }
+
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.3.1"
 black = ">=23.3,<25.0"
@@ -71,6 +74,7 @@ pytest-socket = "^0.7.0"
 
 [tool.poetry.extras]
 vcr = ["vcrpy"]
+langsmith_pyo3 = ["langsmith-pyo3"]
 
 [build-system]
 requires = ["poetry-core"]