fix optional zstd

python/langsmith/_internal/_compressed_runs.py

@@ -2,7 +2,7 @@
 import threading
 
 try:
-    from zstandard import ZstdCompressor
+    import zstandard as zstd
 
     HAVE_ZSTD = True
 except ImportError:
@@ -23,7 +23,7 @@ def __init__(self):
                 "zstandard package required for compression. "
                 "Install with 'pip install langsmith[compression]'"
             )
-        self.compressor_writer = ZstdCompressor(
+        self.compressor_writer = zstd.ZstdCompressor(
             level=compression_level, threads=-1
         ).stream_writer(self.buffer, closefd=False)
 
@@ -35,6 +35,6 @@ def reset(self):
                 "zstandard package required for compression. "
                 "Install with 'pip install langsmith[compression]'"
             )
-        self.compressor_writer = ZstdCompressor(
+        self.compressor_writer = zstd.ZstdCompressor(
             level=compression_level, threads=-1
         ).stream_writer(self.buffer, closefd=False)

python/langsmith/_internal/_operations.py

@@ -3,15 +3,12 @@
 import itertools
 import logging
 import uuid
-from typing import Literal, Optional, Union, cast
+from typing import TYPE_CHECKING, Literal, Optional, Union, cast
 
 try:
-    from zstandard import ZstdCompressionWriter
+    import zstandard as zstd
 except ImportError:
-
-    class ZstdCompressionWriter:  # type: ignore[no-redef]
-        """only used for typing checks."""
-
+    pass
 
 from langsmith import schemas as ls_schemas
 from langsmith._internal import _orjson
@@ -283,7 +280,7 @@ def serialized_run_operation_to_multipart_parts_and_context(
 
 def compress_multipart_parts_and_context(
     parts_and_context: MultipartPartsAndContext,
-    compressor_writer: ZstdCompressionWriter,
+    compressor_writer: zstd.ZstdCompressionWriter,
     boundary: str,
 ) -> None:
     for part_name, (filename, data, content_type, headers) in parts_and_context.parts: