data-designer 0.1.5__py3-none-any.whl → 0.2.1__py3-none-any.whl

data_designer/engine/column_generators/registry.py

@@ -3,6 +3,7 @@
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import (
+    EmbeddingColumnConfig,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,
     LLMJudgeColumnConfig,
@@ -12,8 +13,9 @@ from data_designer.config.column_configs import (
 )
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.generators.embedding import EmbeddingCellGenerator
 from data_designer.engine.column_generators.generators.expression import ExpressionColumnGenerator
-from data_designer.engine.column_generators.generators.llm_generators import (
+from data_designer.engine.column_generators.generators.llm_completion import (
     LLMCodeCellGenerator,
     LLMJudgeCellGenerator,
     LLMStructuredCellGenerator,
@@ -40,11 +42,11 @@ def create_default_column_generator_registry(with_plugins: bool = True) -> Colum
     registry.register(DataDesignerColumnType.LLM_CODE, LLMCodeCellGenerator, LLMCodeColumnConfig)
     registry.register(DataDesignerColumnType.LLM_JUDGE, LLMJudgeCellGenerator, LLMJudgeColumnConfig)
     registry.register(DataDesignerColumnType.EXPRESSION, ExpressionColumnGenerator, ExpressionColumnConfig)
+    registry.register(DataDesignerColumnType.EMBEDDING, EmbeddingCellGenerator, EmbeddingColumnConfig)
     registry.register(DataDesignerColumnType.SAMPLER, SamplerColumnGenerator, SamplerMultiColumnConfig)
     registry.register(DataDesignerColumnType.SEED_DATASET, SeedDatasetColumnGenerator, SeedDatasetMultiColumnConfig)
     registry.register(DataDesignerColumnType.VALIDATION, ValidationColumnGenerator, ValidationColumnConfig)
     registry.register(DataDesignerColumnType.LLM_STRUCTURED, LLMStructuredCellGenerator, LLMStructuredColumnConfig)
-
     if with_plugins:
         for plugin in PluginRegistry().get_plugins(PluginType.COLUMN_GENERATOR):
             registry.register(

data_designer/engine/column_generators/utils/judge_score_factory.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import Enum
-from typing import Type
 
 from pydantic import BaseModel, ConfigDict, Field, create_model
 
@@ -19,7 +18,7 @@ class BaseJudgeResponse(BaseModel):
     reasoning: str = Field(..., description="Reasoning for the assigned score.")
 
 
-def _stringify_scoring(options: dict, enum_type: Type[Enum]) -> str:
+def _stringify_scoring(options: dict, enum_type: type[Enum]) -> str:
     """Convert score descriptions into a single text block."""
     list_block = "\n".join(
         [SCORING_FORMAT.format(score=score, description=description) for score, description in options.items()]
@@ -27,7 +26,7 @@ def _stringify_scoring(options: dict, enum_type: Type[Enum]) -> str:
     return SCORE_FIELD_DESCRIPTION_FORMAT.format(enum_name=enum_type.__name__, scoring=list_block)
 
 
-def create_judge_response_model(score: Score) -> Type[BaseJudgeResponse]:
+def create_judge_response_model(score: Score) -> type[BaseJudgeResponse]:
     """Create a JudgeResponse data type."""
     enum_members = {}
     for option in score.options.keys():
@@ -46,12 +45,12 @@ def create_judge_response_model(score: Score) -> Type[BaseJudgeResponse]:
 
 
 def create_judge_structured_output_model(
-    judge_responses: list[Type[BaseJudgeResponse]],
-) -> Type[BaseModel]:
+    judge_responses: list[type[BaseJudgeResponse]],
+) -> type[BaseModel]:
     """Create a JudgeStructuredOutput class dynamically."""
     return create_model(
         "JudgeStructuredOutput",
         __doc__=f"Response schema for scores with the following names: {[response.__name__ for response in judge_responses]}.",
         __base__=BaseModel,
-        **{response.__name__.lower(): (response, ...) for response in judge_responses},
+        **{response.__name__: (response, ...) for response in judge_responses},
     )

data_designer/engine/configurable_task.py

@@ -3,7 +3,7 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Generic, Type, TypeVar, get_origin
+from typing import Generic, TypeVar, get_origin
 
 import pandas as pd
 
@@ -30,7 +30,7 @@ class ConfigurableTask(ABC, Generic[TaskConfigT]):
         self._initialize()
 
     @classmethod
-    def get_config_type(cls) -> Type[TaskConfigT]:
+    def get_config_type(cls) -> type[TaskConfigT]:
         for base in cls.__orig_bases__:
             if hasattr(base, "__args__") and len(base.__args__) == 1:
                 arg = base.__args__[0]

data_designer/engine/dataset_builders/artifact_storage.py

@@ -7,7 +7,6 @@ import shutil
 from datetime import datetime
 from functools import cached_property
 from pathlib import Path
-from typing import Union
 
 import pandas as pd
 from pydantic import BaseModel, field_validator, model_validator
@@ -77,7 +76,7 @@ class ArtifactStorage(BaseModel):
         return self.base_dataset_path / self.processors_outputs_folder_name
 
     @field_validator("artifact_path")
-    def validate_artifact_path(cls, v: Union[Path, str]) -> Path:
+    def validate_artifact_path(cls, v: Path | str) -> Path:
         v = Path(v)
         if not v.is_dir():
             raise ArtifactStorageError("Artifact path must exist and be a directory")

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -1,24 +1,30 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import functools
+import importlib.metadata
 import json
 import logging
 import time
+import uuid
 from pathlib import Path
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import pandas as pd
 
-from data_designer.config.column_types import ColumnConfigT, column_type_is_llm_generated
+from data_designer.config.column_types import ColumnConfigT, column_type_is_model_generated
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
     ProcessorConfig,
     ProcessorType,
 )
-from data_designer.engine.column_generators.generators.base import ColumnGenerator, GenerationStrategy
-from data_designer.engine.column_generators.generators.llm_generators import WithLLMGeneration
+from data_designer.engine.column_generators.generators.base import (
+    ColumnGenerator,
+    GenerationStrategy,
+    WithModelGeneration,
+)
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError, DatasetProcessingError
 from data_designer.engine.dataset_builders.multi_column_configs import (
@@ -32,14 +38,21 @@ from data_designer.engine.dataset_builders.utils.concurrency import (
 from data_designer.engine.dataset_builders.utils.dataset_batch_manager import (
     DatasetBatchManager,
 )
+from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
 from data_designer.engine.registry.data_designer_registry import DataDesignerRegistry
 from data_designer.engine.resources.resource_provider import ResourceProvider
 
+if TYPE_CHECKING:
+    from data_designer.engine.models.usage import ModelUsageStats
+
 logger = logging.getLogger(__name__)
 
 
+_CLIENT_VERSION: str = importlib.metadata.version("data_designer")
+
+
 class ColumnWiseDatasetBuilder:
     def __init__(
         self,
@@ -72,7 +85,7 @@ class ColumnWiseDatasetBuilder:
 
     @functools.cached_property
     def llm_generated_column_configs(self) -> list[ColumnConfigT]:
-        return [config for config in self.single_column_configs if column_type_is_llm_generated(config.column_type)]
+        return [config for config in self.single_column_configs if column_type_is_model_generated(config.column_type)]
 
     def build(
         self,
@@ -86,11 +99,12 @@ class ColumnWiseDatasetBuilder:
 
         generators = self._initialize_generators()
         start_time = time.perf_counter()
+        group_id = uuid.uuid4().hex
 
         self.batch_manager.start(num_records=num_records, buffer_size=buffer_size)
         for batch_idx in range(self.batch_manager.num_batches):
             logger.info(f"⏳ Processing batch {batch_idx + 1} of {self.batch_manager.num_batches}")
-            self._run_batch(generators)
+            self._run_batch(generators, batch_mode="batch", group_id=group_id)
             df_batch = self._run_processors(
                 stage=BuildStage.POST_BATCH,
                 dataframe=self.batch_manager.get_current_batch(as_dataframe=True),
@@ -111,10 +125,10 @@ class ColumnWiseDatasetBuilder:
         self._run_model_health_check_if_needed()
 
         generators = self._initialize_generators()
-
+        group_id = uuid.uuid4().hex
         start_time = time.perf_counter()
         self.batch_manager.start(num_records=num_records, buffer_size=num_records)
-        self._run_batch(generators, save_partial_results=False)
+        self._run_batch(generators, batch_mode="preview", save_partial_results=False, group_id=group_id)
         dataset = self.batch_manager.get_current_batch(as_dataframe=True)
         self.batch_manager.reset()
 
@@ -140,7 +154,10 @@ class ColumnWiseDatasetBuilder:
             for config in self._column_configs
         ]
 
-    def _run_batch(self, generators: list[ColumnGenerator], *, save_partial_results: bool = True) -> None:
+    def _run_batch(
+        self, generators: list[ColumnGenerator], *, batch_mode: str, save_partial_results: bool = True, group_id: str
+    ) -> None:
+        pre_batch_snapshot = self._resource_provider.model_registry.get_model_usage_snapshot()
         for generator in generators:
             generator.log_pre_generation()
             try:
@@ -163,16 +180,20 @@ class ColumnWiseDatasetBuilder:
                 )
                 raise DatasetGenerationError(f"🛑 Failed to process {column_error_str}:\n{e}")
 
+        try:
+            usage_deltas = self._resource_provider.model_registry.get_usage_deltas(pre_batch_snapshot)
+            self._emit_batch_inference_events(batch_mode, usage_deltas, group_id)
+        except Exception:
+            pass
+
     def _run_from_scratch_column_generator(self, generator: ColumnGenerator) -> None:
         df = generator.generate_from_scratch(self.batch_manager.num_records_batch)
         self.batch_manager.add_records(df.to_dict(orient="records"))
 
     def _run_cell_by_cell_generator(self, generator: ColumnGenerator) -> None:
         max_workers = MAX_CONCURRENCY_PER_NON_LLM_GENERATOR
-        if isinstance(generator, WithLLMGeneration):
+        if isinstance(generator, WithModelGeneration):
             max_workers = generator.inference_parameters.max_parallel_requests
-        elif hasattr(generator.config, "max_parallel_requests"):
-            max_workers = generator.config.max_parallel_requests
         self._fan_out_with_threads(generator, max_workers=max_workers)
 
     def _run_full_column_generator(self, generator: ColumnGenerator) -> None:
@@ -180,12 +201,12 @@ class ColumnWiseDatasetBuilder:
         self.batch_manager.update_records(df.to_dict(orient="records"))
 
     def _run_model_health_check_if_needed(self) -> bool:
-        if any(column_type_is_llm_generated(config.column_type) for config in self.single_column_configs):
+        if any(column_type_is_model_generated(config.column_type) for config in self.single_column_configs):
             self._resource_provider.model_registry.run_health_check(
-                set(config.model_alias for config in self.llm_generated_column_configs)
+                list(set(config.model_alias for config in self.llm_generated_column_configs))
             )
 
-    def _fan_out_with_threads(self, generator: WithLLMGeneration, max_workers: int) -> None:
+    def _fan_out_with_threads(self, generator: WithModelGeneration, max_workers: int) -> None:
         if generator.generation_strategy != GenerationStrategy.CELL_BY_CELL:
             raise DatasetGenerationError(
                 f"Generator {generator.metadata().name} is not a {GenerationStrategy.CELL_BY_CELL} "
@@ -288,3 +309,25 @@ class ColumnWiseDatasetBuilder:
             json_file_name="model_configs.json",
             configs=self._resource_provider.model_registry.model_configs.values(),
         )
+
+    def _emit_batch_inference_events(
+        self, batch_mode: str, usage_deltas: dict[str, ModelUsageStats], group_id: str
+    ) -> None:
+        if not usage_deltas:
+            return
+
+        events = [
+            InferenceEvent(
+                nemo_source=NemoSourceEnum.DATADESIGNER,
+                task=batch_mode,
+                task_status=TaskStatusEnum.SUCCESS,
+                model=model_name,
+                input_tokens=delta.token_usage.input_tokens,
+                output_tokens=delta.token_usage.output_tokens,
+            )
+            for model_name, delta in usage_deltas.items()
+        ]
+
+        with TelemetryHandler(source_client_version=_CLIENT_VERSION, session_id=group_id) as telemetry_handler:
+            for event in events:
+                telemetry_handler.enqueue(event)

data_designer/engine/dataset_builders/utils/concurrency.py

@@ -8,7 +8,7 @@ import json
 import logging
 from concurrent.futures import Future, ThreadPoolExecutor
 from threading import Lock, Semaphore
-from typing import Any, Optional, Protocol
+from typing import Any, Protocol
 
 from pydantic import BaseModel, Field
 
@@ -46,13 +46,13 @@ class ExecutorResults(BaseModel):
 class CallbackWithContext(Protocol):
     """Executor callback functions must accept a context kw argument."""
 
-    def __call__(self, result: Any, *, context: Optional[dict] = None) -> Any: ...
+    def __call__(self, result: Any, *, context: dict | None = None) -> Any: ...
 
 
 class ErrorCallbackWithContext(Protocol):
     """Error callbacks take the Exception instance and context."""
 
-    def __call__(self, exc: Exception, *, context: Optional[dict] = None) -> Any: ...
+    def __call__(self, exc: Exception, *, context: dict | None = None) -> Any: ...
 
 
 class ConcurrentThreadExecutor:
@@ -92,8 +92,8 @@ class ConcurrentThreadExecutor:
         *,
         max_workers: int,
         column_name: str,
-        result_callback: Optional[CallbackWithContext] = None,
-        error_callback: Optional[ErrorCallbackWithContext] = None,
+        result_callback: CallbackWithContext | None = None,
+        error_callback: ErrorCallbackWithContext | None = None,
         shutdown_error_rate: float = 0.50,
         shutdown_error_window: int = 10,
     ):
@@ -136,7 +136,7 @@ class ConcurrentThreadExecutor:
             )
         )
 
-    def submit(self, fn, *args, context: Optional[dict] = None, **kwargs) -> None:
+    def submit(self, fn, *args, context: dict | None = None, **kwargs) -> None:
         if self._executor is None:
             raise RuntimeError("Executor is not initialized, this class should be used as a context manager.")
 

data_designer/engine/models/facade.py

@@ -9,9 +9,9 @@ from copy import deepcopy
 from typing import Any
 
 from litellm.types.router import DeploymentTypedDict, LiteLLM_Params
-from litellm.types.utils import ModelResponse
+from litellm.types.utils import EmbeddingResponse, ModelResponse
 
-from data_designer.config.models import ModelConfig, ModelProvider
+from data_designer.config.models import GenerationType, ModelConfig, ModelProvider
 from data_designer.engine.model_provider import ModelProviderRegistry
 from data_designer.engine.models.errors import (
     GenerationValidationFailureError,
@@ -49,6 +49,10 @@ class ModelFacade:
     def model_provider(self) -> ModelProvider:
         return self._model_provider_registry.get_provider(self._model_config.provider)
 
+    @property
+    def model_generation_type(self) -> GenerationType:
+        return self._model_config.generation_type
+
     @property
     def model_provider_name(self) -> str:
         return self.model_provider.name
@@ -64,13 +68,12 @@ class ModelFacade:
     def completion(self, messages: list[dict[str, str]], skip_usage_tracking: bool = False, **kwargs) -> ModelResponse:
         logger.debug(
             f"Prompting model {self.model_name!r}...",
-            extra={"model": self.model_name, "messages": messages, "sensitive": True},
+            extra={"model": self.model_name, "messages": messages},
         )
         response = None
-        if self.model_provider.extra_body:
-            kwargs["extra_body"] = {**kwargs.get("extra_body", {}), **self.model_provider.extra_body}
+        kwargs = self.consolidate_kwargs(**kwargs)
         try:
-            response = self._router.completion(self.model_name, messages, **kwargs)
+            response = self._router.completion(model=self.model_name, messages=messages, **kwargs)
             logger.debug(
                 f"Received completion from model {self.model_name!r}",
                 extra={
@@ -84,9 +87,50 @@ class ModelFacade:
         except Exception as e:
             raise e
         finally:
-            if not skip_usage_tracking:
+            if not skip_usage_tracking and response is not None:
                 self._track_usage(response)
 
+    def consolidate_kwargs(self, **kwargs) -> dict[str, Any]:
+        # Remove purpose from kwargs to avoid passing it to the model
+        kwargs.pop("purpose", None)
+        kwargs = {**self._model_config.inference_parameters.generate_kwargs, **kwargs}
+        if self.model_provider.extra_body:
+            kwargs["extra_body"] = {**kwargs.get("extra_body", {}), **self.model_provider.extra_body}
+        return kwargs
+
+    @catch_llm_exceptions
+    def generate_text_embeddings(
+        self, input_texts: list[str], skip_usage_tracking: bool = False, **kwargs
+    ) -> list[list[float]]:
+        logger.debug(
+            f"Generating embeddings with model {self.model_name!r}...",
+            extra={
+                "model": self.model_name,
+                "input_count": len(input_texts),
+            },
+        )
+        kwargs = self.consolidate_kwargs(**kwargs)
+        response = None
+        try:
+            response = self._router.embedding(model=self.model_name, input=input_texts, **kwargs)
+            logger.debug(
+                f"Received embeddings from model {self.model_name!r}",
+                extra={
+                    "model": self.model_name,
+                    "embedding_count": len(response.data) if response.data else 0,
+                    "usage": self._usage_stats.model_dump(),
+                },
+            )
+            if response.data and len(response.data) == len(input_texts):
+                return [data["embedding"] for data in response.data]
+            else:
+                raise ValueError(f"Expected {len(input_texts)} embeddings, but received {len(response.data)}")
+        except Exception as e:
+            raise e
+        finally:
+            if not skip_usage_tracking and response is not None:
+                self._track_usage_from_embedding(response)
+
     @catch_llm_exceptions
     def generate(
         self,
@@ -218,8 +262,21 @@ class ModelFacade:
         ):
             self._usage_stats.extend(
                 token_usage=TokenUsageStats(
-                    prompt_tokens=response.usage.prompt_tokens,
-                    completion_tokens=response.usage.completion_tokens,
+                    input_tokens=response.usage.prompt_tokens,
+                    output_tokens=response.usage.completion_tokens,
+                ),
+                request_usage=RequestUsageStats(successful_requests=1, failed_requests=0),
+            )
+
+    def _track_usage_from_embedding(self, response: EmbeddingResponse | None) -> None:
+        if response is None:
+            self._usage_stats.extend(request_usage=RequestUsageStats(successful_requests=0, failed_requests=1))
+            return
+        if response.usage is not None and response.usage.prompt_tokens is not None:
+            self._usage_stats.extend(
+                token_usage=TokenUsageStats(
+                    input_tokens=response.usage.prompt_tokens,
+                    output_tokens=0,
                 ),
                 request_usage=RequestUsageStats(successful_requests=1, failed_requests=0),
             )

data_designer/engine/models/litellm_overrides.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 
 import random
 import threading
-from typing import Optional, Union
 
 import httpx
 import litellm
@@ -90,7 +89,7 @@ class CustomRouter(Router):
         self._initial_retry_after_s = initial_retry_after_s
         self._jitter_pct = jitter_pct
 
-    def _extract_retry_delay_from_headers(self, e: Exception) -> Optional[Union[int, float]]:
+    def _extract_retry_delay_from_headers(self, e: Exception) -> int | float | None:
         """
         Most of this code logic was extracted directly from the parent
         `Router`'s `_time_to_sleep_before_retry` function. Our override
@@ -99,7 +98,7 @@ class CustomRouter(Router):
         return this info, we'll simply use that retry value returned here.
         """
 
-        response_headers: Optional[httpx.Headers] = None
+        response_headers: httpx.Headers | None = None
         if hasattr(e, "response") and hasattr(e.response, "headers"):  # type: ignore
             response_headers = e.response.headers  # type: ignore
         if hasattr(e, "litellm_response_headers"):
@@ -119,9 +118,9 @@ class CustomRouter(Router):
         e: Exception,
         remaining_retries: int,
         num_retries: int,
-        healthy_deployments: Optional[list] = None,
-        all_deployments: Optional[list] = None,
-    ) -> Union[int, float]:
+        healthy_deployments: list | None = None,
+        all_deployments: list | None = None,
+    ) -> int | float:
         """
         Implements exponential backoff for retries.
 

data_designer/engine/models/parsers/errors.py

@@ -1,8 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional
-
 
 class ParserException(Exception):
     """Identifies errors resulting from generic parser errors.
@@ -12,7 +10,7 @@ class ParserException(Exception):
             attempted to parse.
     """
 
-    source: Optional[str]
+    source: str | None
 
     @staticmethod
     def _log_format(source: str) -> str:
@@ -24,7 +22,7 @@ class ParserException(Exception):
         # return f"<source>{source}</source>"
         return ""
 
-    def __init__(self, msg: Optional[str] = None, source: Optional[str] = None):
+    def __init__(self, msg: str | None = None, source: str | None = None):
         msg = "" if msg is None else msg.strip()
 
         if source is not None:

data_designer/engine/models/parsers/parser.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from functools import reduce
-from typing import Optional
 
 import marko
 from lxml import etree
@@ -105,8 +104,8 @@ class LLMResponseParser:
 
     def __init__(
         self,
-        tag_parsers: Optional[dict[str, TagParser]] = None,
-        postprocessors: Optional[list[PostProcessor]] = None,
+        tag_parsers: dict[str, TagParser] | None = None,
+        postprocessors: list[PostProcessor] | None = None,
     ):
         """
         Initializes the LLMResponseParser with optional tag parsers and post-processors.

data_designer/engine/models/parsers/postprocessors.py

@@ -1,7 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional, Type
 
 import json_repair
 from pydantic import BaseModel, ValidationError
@@ -60,12 +59,12 @@ def deserialize_json_code(
 
 
 class RealizePydanticTypes:
-    types: list[Type[BaseModel]]
+    types: list[type[BaseModel]]
 
-    def __init__(self, types: list[Type[BaseModel]]):
+    def __init__(self, types: list[type[BaseModel]]):
         self.types = types
 
-    def _fit_types(self, obj: dict) -> Optional[BaseModel]:
+    def _fit_types(self, obj: dict) -> BaseModel | None:
         final_obj = None
 
         for t in self.types:

data_designer/engine/models/parsers/types.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Optional, Protocol, Type, runtime_checkable
+from typing import Any, Protocol, runtime_checkable
 
 from lxml.etree import _Element
 from pydantic import BaseModel, Field
@@ -30,7 +30,7 @@ class LLMStructuredResponse(BaseModel):
         out.parsed = out.parsed[-n:]
         return out
 
-    def filter(self, block_types: list[Type[BaseModel]]) -> Self:
+    def filter(self, block_types: list[type[BaseModel]]) -> Self:
         out = self.model_copy()
         out.parsed = [b for b in out.parsed if isinstance(b, tuple(block_types))]
         return out
@@ -44,7 +44,7 @@ class TagParser(Protocol):
     element, do some computation, and return some kind of structured
     output, represented as a subclass of Pydantic `BaseModel`.
     This protocol implementation can cover both classes as well
-    as curried fuctions as parsers (e.g. `partial`).
+    as curried functions as parsers (e.g. `partial`).
     """
 
     def __call__(self, element: _Element) -> BaseModel: ...
@@ -69,7 +69,7 @@ class TextBlock(BaseModel):
 
 class CodeBlock(BaseModel):
     code: str
-    code_lang: Optional[str] = None
+    code_lang: str | None = None
 
 
 class StructuredDataBlock(BaseModel):