From 28cc124cd347322444699e12717819a63e14f01c Mon Sep 17 00:00:00 2001 From: Predrag Gruevski <2348618+obi1kenobi@users.noreply.github.com> Date: Wed, 27 Nov 2024 00:06:46 -0500 Subject: [PATCH] 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 | 59 ++++++++++++++++++++++++++++++++++---- python/poetry.lock | 50 ++++++++++++++++++++++++++++++-- python/pyproject.toml | 4 +++ 3 files changed, 105 insertions(+), 8 deletions(-) diff --git a/python/langsmith/client.py b/python/langsmith/client.py index b3b5fc993..d097867f4 100644 --- a/python/langsmith/client.py +++ b/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) diff --git a/python/poetry.lock b/python/poetry.lock index 2b362f986..598ceaaf5 100644 --- a/python/poetry.lock +++ b/python/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand. [[package]] name = "annotated-types" @@ -590,6 +590,51 @@ files = [ {file = "jiter-0.7.0.tar.gz", hash = "sha256:c061d9738535497b5509f8970584f20de1e900806b239a39a9994fc191dad630"}, ] +[[package]] +name = "langsmith-pyo3" +version = "0.1.0rc2" +description = "" +optional = true +python-versions = ">=3.8" +files = [ + {file = "langsmith_pyo3-0.1.0rc2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6a77834cf7225b863615456b4110fcc7df3ebd41a2d6ea0b8359f7ff8a785f21"}, + {file = "langsmith_pyo3-0.1.0rc2-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:57a3f4a777f601305bdd8a40d103cc4e24f06fdfddb26d9e2713991e636ed26d"}, + {file = "langsmith_pyo3-0.1.0rc2-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:bbc771e40e78dfd02f55b81c9ad6dda94922f7feebb9193963bbb83bd8af3eae"}, + {file = "langsmith_pyo3-0.1.0rc2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:108739044d80909ac60069b2dc4f0b6c4ba46ce4bf6a2cfbded2b25b67524f7c"}, + {file = "langsmith_pyo3-0.1.0rc2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:9df02ad699243b2e54dea4309918dbe6923bc366a89e9b5f7ad857f9ae910f0d"}, + {file = "langsmith_pyo3-0.1.0rc2-cp310-none-win_amd64.whl", hash = "sha256:f4f79a3b6e8d58c2123c022a3e314064e5b170b94bde966fd352253631fa4857"}, + {file = "langsmith_pyo3-0.1.0rc2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e93a45b4d08fa3f188c6d76a98ab9b3fba7d0d604b0aa5e6370ce65334c0af6a"}, + {file = "langsmith_pyo3-0.1.0rc2-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:c98ea4804d6a5d9213c7833b6d36fa967f8201bfbc57ac9e743f9b15f455d389"}, + {file = "langsmith_pyo3-0.1.0rc2-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:15f32257d5182324541aa2f370acf81b797afcb14238187b50244255676570e3"}, + {file = "langsmith_pyo3-0.1.0rc2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:7676baa04abce61298118b8790d0743246f8617e47b97036bd734a4623160c9a"}, + {file = "langsmith_pyo3-0.1.0rc2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:725c687b036b333334c394dee110e40c73db2d86551c11821f1b089e61487407"}, + {file = "langsmith_pyo3-0.1.0rc2-cp311-none-win_amd64.whl", hash = "sha256:3eb3ad8804d215b9670ef6c135714ced1e6db6d5f53c335fa3c1da9cbc24fef8"}, + {file = "langsmith_pyo3-0.1.0rc2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b4dd2f456686b400bb47a400ceea571fb6c6cc6757cf6a9a4d5174ffa7c188a4"}, + {file = "langsmith_pyo3-0.1.0rc2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:248b1aaab324f8f535b888d6ea1fff0f5e639b21686fe772010ae2cf360b2327"}, + {file = "langsmith_pyo3-0.1.0rc2-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:c3dcd1f8bb6951f0ef4181d74f713fcf864b86f49132228acdf8f8c877605daa"}, + {file = "langsmith_pyo3-0.1.0rc2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:62ee1c0ac5079809d8fb746d4522f573e37457197aebb71965eb2672a75bff38"}, + {file = "langsmith_pyo3-0.1.0rc2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:17ffbbe8a62f8d87b4d9096bddfa2c4421cb29d45043b0b09d78bed8a9b7741f"}, + {file = "langsmith_pyo3-0.1.0rc2-cp312-none-win_amd64.whl", hash = "sha256:3a06377c9ac390ed76a65f62e29b88480be32739b96ed9f51b6e6f6210551202"}, + {file = "langsmith_pyo3-0.1.0rc2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d66e40f56495a84d11b9d47d69421750214e2de3ba683cdbc9eb04ded11a3f66"}, + {file = "langsmith_pyo3-0.1.0rc2-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:12b27b6441666a3123a6016fcf78288d9194f54e48f021b5172fe8fc58994eba"}, + {file = "langsmith_pyo3-0.1.0rc2-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6e124f0aa1142087b7a6b0d2b9f6dd82415fa64899f12e9650174957918300f4"}, + {file = "langsmith_pyo3-0.1.0rc2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:46bd1c411ebbda514954020c46eb65b3a8a9378cfc153fc35a09e375fc5feead"}, + {file = "langsmith_pyo3-0.1.0rc2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4d08430ba4e93ec9ac704c1b0116130d9af7cee86b7d4b3a74829b239d5d557a"}, + {file = "langsmith_pyo3-0.1.0rc2-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:7ac4c3a138cea449d5ed5621425daf148b0ed000df4a490cfb304099e0770004"}, + {file = "langsmith_pyo3-0.1.0rc2-cp38-cp38-manylinux_2_28_aarch64.whl", hash = "sha256:0d9d2bef3a0098e2ff28d7135320660abdf342b857c00f5ca17c0b03870193c8"}, + {file = "langsmith_pyo3-0.1.0rc2-cp38-cp38-manylinux_2_28_x86_64.whl", hash = "sha256:4b82c9b6ba9bb6fd464aaca50b2f8094aba92f2948df0e6301b8b0fc2bb46baf"}, + {file = "langsmith_pyo3-0.1.0rc2-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:19e70971e432661314b6196357eb92633049e2dd0bc6fba61b86aa113a09aedf"}, + {file = "langsmith_pyo3-0.1.0rc2-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:372d80979cd1b7d59e290ab80f9cad3d7059f5aa66c9757d522898f0d399bbed"}, + {file = "langsmith_pyo3-0.1.0rc2-cp38-none-win_amd64.whl", hash = "sha256:7fbb73b1c448ac4964358c9ca1be3107bb2c0c38715343c5da7f92ed0e3ee490"}, + {file = "langsmith_pyo3-0.1.0rc2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:ddeca8b3ae912d090ec9fd5589e0b80cd5475c80dbc439976d3f92bcbe93da81"}, + {file = "langsmith_pyo3-0.1.0rc2-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:6daa12978c18f4560858eac2c84d60090cd5ba7e55e657e052ba7b558f23c1d8"}, + {file = "langsmith_pyo3-0.1.0rc2-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:60ec1e51f674141ab96f8c2d814d42410f8163f9323f1e98bde8d26cf4676513"}, + {file = "langsmith_pyo3-0.1.0rc2-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:2c4813d11d515386b34a827c958edabecd9ef32306800a5e7d2f12ea2d1d0943"}, + {file = "langsmith_pyo3-0.1.0rc2-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:11c8d34d8583d5bb009a081fbfebab8b73c1730069626940ee05644c40f77625"}, + {file = "langsmith_pyo3-0.1.0rc2-cp39-none-win_amd64.whl", hash = "sha256:f341dff48be2c289c23733489e60adf7e1f005eea95ebb6275b20314fd7fb5a6"}, + {file = "langsmith_pyo3-0.1.0rc2.tar.gz", hash = "sha256:30eb26aa33deca44eb9210b77d478ec2157a0cb51f96da30f87072dd5912e3ed"}, +] + [[package]] name = "marshmallow" version = "3.22.0" @@ -2065,9 +2110,10 @@ multidict = ">=4.0" propcache = ">=0.2.0" [extras] +langsmith-pyo3 = ["langsmith-pyo3"] vcr = [] [metadata] lock-version = "2.0" python-versions = ">=3.8.1,<4.0" -content-hash = "a5a6c61cba1b5ce9cf739700a780c2df63ff7aaa482c29de9910418263318586" +content-hash = "5eb7654ca6ae40c61164f6da76b8b4bd6baf0ef0967c77251bd01efad9d7d320" diff --git a/python/pyproject.toml b/python/pyproject.toml index 51dfb7e3f..a8ab4b445 100644 --- a/python/pyproject.toml +++ b/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"]