diff --git a/python/langsmith/_internal/_compressed_runs.py b/python/langsmith/_internal/_compressed_runs.py index 7871ebffc..beb100850 100644 --- a/python/langsmith/_internal/_compressed_runs.py +++ b/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) diff --git a/python/langsmith/_internal/_operations.py b/python/langsmith/_internal/_operations.py index 23e2d2c1d..497637db1 100644 --- a/python/langsmith/_internal/_operations.py +++ b/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: