nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/data_models/optimizer.py

@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+from pydantic import BaseModel
+from pydantic import Field
+
+
+class OptimizerMetric(BaseModel):
+    """
+    Parameters used by the workflow optimizer to define a metric to optimize.
+    """
+    evaluator_name: str = Field(description="Name of the metric to optimize.")
+    direction: str = Field(description="Direction of the optimization. Can be 'maximize' or 'minimize'.")
+    weight: float = Field(description="Weight of the metric in the optimization process.", default=1.0)
+
+
+class NumericOptimizationConfig(BaseModel):
+    """
+    Configuration for numeric/enum optimization (Optuna).
+    """
+    enabled: bool = Field(default=True, description="Enable numeric optimization")
+    n_trials: int = Field(description="Number of trials for numeric optimization.", default=20)
+
+
+class PromptGAOptimizationConfig(BaseModel):
+    """
+    Configuration for prompt optimization using a Genetic Algorithm.
+    """
+    enabled: bool = Field(default=False, description="Enable GA-based prompt optimization")
+
+    # Prompt optimization function hooks
+    prompt_population_init_function: str | None = Field(
+        default=None,
+        description="Optional function name to initialize/mutate candidate prompts.",
+    )
+    prompt_recombination_function: str | None = Field(
+        default=None,
+        description="Optional function name to recombine two parent prompts into a child.",
+    )
+
+    # Genetic algorithm configuration
+    ga_population_size: int = Field(
+        description="Population size for genetic algorithm prompt optimization.",
+        default=24,
+    )
+    ga_generations: int = Field(
+        description="Number of generations to evolve in GA prompt optimization.",
+        default=15,
+    )
+    ga_offspring_size: int | None = Field(
+        description="Number of offspring to produce per generation. Defaults to population_size - elitism.",
+        default=None,
+    )
+    ga_crossover_rate: float = Field(
+        description="Probability of applying crossover during reproduction.",
+        default=0.8,
+        ge=0.0,
+        le=1.0,
+    )
+    ga_mutation_rate: float = Field(
+        description="Probability of mutating a child after crossover.",
+        default=0.3,
+        ge=0.0,
+        le=1.0,
+    )
+    ga_elitism: int = Field(
+        description="Number of top individuals carried over unchanged each generation.",
+        default=2,
+    )
+    ga_selection_method: str = Field(
+        description="Parent selection strategy: 'tournament' or 'roulette'.",
+        default="tournament",
+    )
+    ga_tournament_size: int = Field(
+        description="Tournament size when using tournament selection.",
+        default=3,
+    )
+    ga_parallel_evaluations: int = Field(
+        description="Max number of individuals to evaluate concurrently per generation.",
+        default=8,
+    )
+    ga_diversity_lambda: float = Field(
+        description="Strength of diversity penalty (0 disables). Penalizes identical/near-identical prompts.",
+        default=0.0,
+        ge=0.0,
+    )
+
+
+class OptimizerConfig(BaseModel):
+    """
+    Parameters used by the workflow optimizer.
+    """
+    output_path: Path | None = Field(
+        default=None,
+        description="Path to the output directory where the results will be saved.",
+    )
+
+    eval_metrics: dict[str, OptimizerMetric] | None = Field(
+        description="List of evaluation metrics to optimize.",
+        default=None,
+    )
+
+    reps_per_param_set: int = Field(
+        default=3,
+        description="Number of repetitions per parameter set for the optimization.",
+    )
+
+    target: float | None = Field(
+        description=(
+            "Target value for the optimization. If set, the optimization will stop when this value is reached."),
+        default=None,
+    )
+
+    multi_objective_combination_mode: str = Field(
+        description="Method to combine multiple objectives into a single score.",
+        default="harmonic",
+    )
+
+    # Nested configs
+    numeric: NumericOptimizationConfig = NumericOptimizationConfig()
+    prompt: PromptGAOptimizationConfig = PromptGAOptimizationConfig()
+
+
+class OptimizerRunConfig(BaseModel):
+    """
+    Parameters used for an Optimizer R=run
+    """
+    # Eval parameters
+
+    config_file: Path | BaseModel  # allow for instantiated configs to be passed in
+    dataset: str | Path | None  # dataset file path can be specified in the config file
+    result_json_path: str = "$"
+    endpoint: str | None = None  # only used when running the workflow remotely
+    endpoint_timeout: int = 300
+    override: tuple[tuple[str, str], ...] = ()

nat/data_models/profiler.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import BaseModel
+
+
+class PromptCachingConfig(BaseModel):
+    enable: bool = False
+    min_frequency: float = 0.5
+
+
+class BottleneckConfig(BaseModel):
+    enable_simple_stack: bool = False
+    enable_nested_stack: bool = False
+
+
+class ConcurrencySpikeConfig(BaseModel):
+    enable: bool = False
+    spike_threshold: int = 1
+
+
+class PrefixSpanConfig(BaseModel):
+    enable: bool = False
+    min_support: float = 2
+    min_coverage: float = 0
+    max_text_len: int = 1000
+    top_k: int = 10
+    chain_with_common_prefixes: bool = False
+
+
+class ProfilerConfig(BaseModel):
+
+    base_metrics: bool = False
+    token_usage_forecast: bool = False
+    token_uniqueness_forecast: bool = False
+    workflow_runtime_forecast: bool = False
+    compute_llm_metrics: bool = False
+    csv_exclude_io_text: bool = False
+    prompt_caching_prefixes: PromptCachingConfig = PromptCachingConfig()
+    bottleneck_analysis: BottleneckConfig = BottleneckConfig()
+    concurrency_spike_analysis: ConcurrencySpikeConfig = ConcurrencySpikeConfig()
+    prefix_span_analysis: PrefixSpanConfig = PrefixSpanConfig()

nat/data_models/registry_handler.py

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class RegistryHandlerBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+RegistryHandlerBaseConfigT = typing.TypeVar("RegistryHandlerBaseConfigT", bound=RegistryHandlerBaseConfig)

nat/data_models/retriever.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+
+from nat.data_models.common import BaseModelRegistryTag
+from nat.data_models.common import TypedBaseModel
+
+
+class RetrieverBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    The base level config object for a retriever object. Retrievers use different provider clients (e.g., Milvus) to
+    provide an interface for searching for and retrieving documents from the configured data store.
+    """
+    pass
+
+
+RetrieverBaseConfigT = typing.TypeVar("RetrieverBaseConfigT", bound=RetrieverBaseConfig)

nat/data_models/retry_mixin.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import BaseModel
+from pydantic import Field
+
+
+class RetryMixin(BaseModel):
+    """Mixin class for retry configuration."""
+    do_auto_retry: bool = Field(default=True,
+                                description="Whether to automatically retry method calls"
+                                " that fail with a retryable error.",
+                                exclude=True)
+    num_retries: int = Field(default=5,
+                             description="Number of times to retry a method call that fails"
+                             " with a retryable error.",
+                             exclude=True)
+    retry_on_status_codes: list[int | str] = Field(default_factory=lambda: [429, 500, 502, 503, 504],
+                                                   description="List of HTTP status codes that should trigger a retry.",
+                                                   exclude=True)
+    retry_on_errors: list[str] | None = Field(default_factory=lambda: ["Too Many Requests"],
+                                              description="List of error substrings that should trigger a retry.",
+                                              exclude=True)

nat/data_models/span.py

@@ -0,0 +1,228 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import os
+import time
+import uuid
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import field_validator
+
+logger = logging.getLogger(__name__)
+
+_SPAN_PREFIX = os.getenv("NAT_SPAN_PREFIX", "nat").strip() or "nat"
+
+
+class SpanKind(Enum):
+    LLM = "LLM"
+    TOOL = "TOOL"
+    WORKFLOW = "WORKFLOW"
+    TASK = "TASK"
+    FUNCTION = "FUNCTION"
+    CUSTOM = "CUSTOM"
+    SPAN = "SPAN"
+    EMBEDDER = "EMBEDDER"
+    RETRIEVER = "RETRIEVER"
+    AGENT = "AGENT"
+    RERANKER = "RERANKER"
+    GUARDRAIL = "GUARDRAIL"
+    EVALUATOR = "EVALUATOR"
+    UNKNOWN = "UNKNOWN"
+
+
+EVENT_TYPE_TO_SPAN_KIND_MAP = {
+    "LLM_START": SpanKind.LLM,
+    "LLM_END": SpanKind.LLM,
+    "LLM_NEW_TOKEN": SpanKind.LLM,
+    "TOOL_START": SpanKind.TOOL,
+    "TOOL_END": SpanKind.TOOL,
+    "WORKFLOW_START": SpanKind.WORKFLOW,
+    "WORKFLOW_END": SpanKind.WORKFLOW,
+    "TASK_START": SpanKind.TASK,
+    "TASK_END": SpanKind.TASK,
+    "FUNCTION_START": SpanKind.FUNCTION,
+    "FUNCTION_END": SpanKind.FUNCTION,
+    "CUSTOM_START": SpanKind.CUSTOM,
+    "CUSTOM_END": SpanKind.CUSTOM,
+    "SPAN_START": SpanKind.SPAN,
+    "SPAN_END": SpanKind.SPAN,
+    "EMBEDDER_START": SpanKind.EMBEDDER,
+    "EMBEDDER_END": SpanKind.EMBEDDER,
+    "RETRIEVER_START": SpanKind.RETRIEVER,
+    "RETRIEVER_END": SpanKind.RETRIEVER,
+    "AGENT_START": SpanKind.AGENT,
+    "AGENT_END": SpanKind.AGENT,
+    "RERANKER_START": SpanKind.RERANKER,
+    "RERANKER_END": SpanKind.RERANKER,
+    "GUARDRAIL_START": SpanKind.GUARDRAIL,
+    "GUARDRAIL_END": SpanKind.GUARDRAIL,
+    "EVALUATOR_START": SpanKind.EVALUATOR,
+    "EVALUATOR_END": SpanKind.EVALUATOR,
+}
+
+
+def event_type_to_span_kind(event_type: str) -> SpanKind:
+    """Convert an event type to a span kind.
+
+    Args:
+        event_type (str): The event type to convert.
+
+    Returns:
+        SpanKind: The span kind.
+    """
+    return EVENT_TYPE_TO_SPAN_KIND_MAP.get(event_type, SpanKind.UNKNOWN)
+
+
+class SpanAttributes(Enum):
+    NAT_SPAN_KIND = f"{_SPAN_PREFIX}.span.kind"
+    INPUT_VALUE = "input.value"
+    INPUT_MIME_TYPE = "input.mime_type"
+    LLM_TOKEN_COUNT_PROMPT = "llm.token_count.prompt"
+    LLM_TOKEN_COUNT_COMPLETION = "llm.token_count.completion"
+    LLM_TOKEN_COUNT_TOTAL = "llm.token_count.total"
+    OUTPUT_VALUE = "output.value"
+    OUTPUT_MIME_TYPE = "output.mime_type"
+    NAT_USAGE_NUM_LLM_CALLS = f"{_SPAN_PREFIX}.usage.num_llm_calls"
+    NAT_USAGE_SECONDS_BETWEEN_CALLS = f"{_SPAN_PREFIX}.usage.seconds_between_calls"
+    NAT_USAGE_TOKEN_COUNT_PROMPT = f"{_SPAN_PREFIX}.usage.token_count.prompt"
+    NAT_USAGE_TOKEN_COUNT_COMPLETION = f"{_SPAN_PREFIX}.usage.token_count.completion"
+    NAT_USAGE_TOKEN_COUNT_TOTAL = f"{_SPAN_PREFIX}.usage.token_count.total"
+    NAT_EVENT_TYPE = f"{_SPAN_PREFIX}.event_type"
+
+
+class MimeTypes(Enum):
+    TEXT = "text/plain"
+    JSON = "application/json"
+
+
+class SpanStatusCode(Enum):
+    OK = "OK"
+    ERROR = "ERROR"
+    UNSET = "UNSET"
+
+
+class SpanEvent(BaseModel):
+    timestamp: float = Field(default_factory=lambda: int(time.time() * 1e9), description="The timestamp of the event.")
+    name: str = Field(description="The name of the event.")
+    attributes: dict[str, Any] = Field(default_factory=dict, description="The attributes of the event.")
+
+
+class SpanStatus(BaseModel):
+    code: SpanStatusCode = Field(default=SpanStatusCode.OK, description="The status code of the span.")
+    message: str | None = Field(default=None, description="The status message of the span.")
+
+
+def _generate_nonzero_trace_id() -> int:
+    """Generate a non-zero 128-bit trace ID."""
+    return uuid.uuid4().int
+
+
+def _generate_nonzero_span_id() -> int:
+    """Generate a non-zero 64-bit span ID."""
+    return uuid.uuid4().int >> 64
+
+
+class SpanContext(BaseModel):
+    trace_id: int = Field(default_factory=_generate_nonzero_trace_id,
+                          description="The OTel-syle 128-bit trace ID of the span.")
+    span_id: int = Field(default_factory=_generate_nonzero_span_id,
+                         description="The OTel-syle 64-bit span ID of the span.")
+
+    @field_validator("trace_id", mode="before")
+    @classmethod
+    def _validate_trace_id(cls, v: int | str | None) -> int:
+        """Regenerate if trace_id is None; raise an exception if trace_id is invalid;"""
+        if isinstance(v, str):
+            v = uuid.UUID(v).int
+        if isinstance(v, type(None)):
+            v = _generate_nonzero_trace_id()
+        if v <= 0 or v >> 128:
+            raise ValueError(f"Invalid trace_id: must be a non-zero 128-bit integer, got {v}")
+        return v
+
+    @field_validator("span_id", mode="before")
+    @classmethod
+    def _validate_span_id(cls, v: int | str | None) -> int:
+        """Regenerate if span_id is None; raise an exception if span_id is invalid;"""
+        if isinstance(v, str):
+            try:
+                v = int(v, 16)
+            except ValueError:
+                raise ValueError(f"span_id unable to be parsed: {v}")
+        if isinstance(v, type(None)):
+            v = _generate_nonzero_span_id()
+        if v <= 0 or v >> 64:
+            raise ValueError(f"Invalid span_id: must be a non-zero 64-bit integer, got {v}")
+        return v
+
+
+class Span(BaseModel):
+    name: str = Field(description="The name of the span.")
+    context: SpanContext | None = Field(default=None, description="The context of the span.")
+    parent: "Span | None" = Field(default=None, description="The parent span of the span.")
+    start_time: int = Field(default_factory=lambda: int(time.time() * 1e9), description="The start time of the span.")
+    end_time: int | None = Field(default=None, description="The end time of the span.")
+    attributes: dict[str, Any] = Field(default_factory=dict, description="The attributes of the span.")
+    events: list[SpanEvent] = Field(default_factory=list, description="The events of the span.")
+    status: SpanStatus = Field(default_factory=SpanStatus, description="The status of the span.")
+
+    @field_validator('context', mode='before')
+    @classmethod
+    def set_default_context(cls, v: SpanContext | None) -> SpanContext:
+        """Set the default context if the context is not provided.
+
+        Args:
+            v (SpanContext | None): The context to set.
+
+        Returns:
+            SpanContext: The context.
+        """
+        if v is None:
+            return SpanContext()
+        return v
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        """Set the attribute of the span.
+
+        Args:
+            key (str): The key of the attribute.
+            value (Any): The value of the attribute.
+        """
+        self.attributes[key] = value
+
+    def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None:
+        """Add an event to the span.
+
+        Args:
+            name (str): The name of the event.
+            attributes (dict[str, Any] | None): The attributes of the event.
+        """
+        if attributes is None:
+            attributes = {}
+        self.events = self.events + [SpanEvent(name=name, attributes=attributes)]
+
+    def end(self, end_time: int | None = None) -> None:
+        """End the span.
+
+        Args:
+            end_time (int | None): The end time of the span.
+        """
+        if end_time is None:
+            end_time = int(time.time() * 1e9)
+        self.end_time = end_time

nat/data_models/step_adaptor.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from enum import Enum
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import model_validator
+
+from nat.data_models.intermediate_step import IntermediateStepType
+
+logger = logging.getLogger(__name__)
+
+
+class StepAdaptorMode(str, Enum):
+    DEFAULT = "default"
+    CUSTOM = "custom"
+    OFF = "off"
+
+
+class StepAdaptorConfig(BaseModel):
+    """
+    Configures how intermediate steps are filtered and normalized by the StepAdaptor.
+
+    Args:
+        mode (StepAdaptorMode): One of:
+            - 'current' => pass only LLM (all LLM_* events) + TOOL_END
+            - 'end_events_only' => pass only LLM_END and TOOL_END
+            - 'custom' => pass only the events in custom_event_types
+        custom_event_types (list[IntermediateStepType]):
+            If mode == 'custom', we only pass events whose event_type is in this list.
+            Otherwise, this field is ignored.
+    """
+    mode: StepAdaptorMode = StepAdaptorMode.DEFAULT
+    custom_event_types: list[IntermediateStepType] = Field(default_factory=list)
+
+    @model_validator(mode="after")
+    def check_custom_event_types(self) -> "StepAdaptorConfig":
+        """
+        Validates custom configurations
+        """
+        if self.mode != StepAdaptorMode.CUSTOM and self.custom_event_types:
+            logger.warning("Ignoring custom_event_types because mode is not 'custom'")
+            self.custom_event_types = []
+        elif self.mode == StepAdaptorMode.CUSTOM and not self.custom_event_types:
+            logger.warning("No custom_event_types provided for custom mode. Defaulting to CUSTOM_START and CUSTOM_END")
+            self.custom_event_types = [IntermediateStepType.CUSTOM_START, IntermediateStepType.CUSTOM_END]
+        elif self.mode == StepAdaptorMode.OFF:
+            logger.warning("StepAdaptor is disabled. Ignoring all intermediate event types")
+            self.custom_event_types = []
+        return self

nat/data_models/streaming.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import typing
+from collections.abc import Callable
+
+_T = typing.TypeVar("_T")
+
+
+class Streaming(typing.Generic[_T]):
+
+    def __init__(self, *, convert: Callable[[typing.Any], _T] = None):
+        self.convert = convert
+
+        sig = inspect.signature(self.convert)
+
+        if (sig.return_annotation == inspect.Parameter.empty):
+            raise ValueError("convert must have a return annotation")
+
+        self.single_output_type = sig.return_annotation

nat/data_models/swe_bench_model.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import field_validator
+
+
+class SWEBenchInput(BaseModel):
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    repo: str
+    instance_id: str
+    base_commit: str
+    patch: str
+    test_patch: str
+    problem_statement: str
+    hints_text: str
+    created_at: str | int
+    version: float
+    FAIL_TO_PASS: list[str]
+    PASS_TO_PASS: list[str]
+    environment_setup_commit: str
+
+    # Handle improperly formatted JSON strings for list fields
+    @field_validator("FAIL_TO_PASS", "PASS_TO_PASS", mode="before")
+    def parse_list_fields(cls, value):
+        if isinstance(value, str):
+            # Attempt to parse the string as a list
+            return json.loads(value)
+        return value
+
+
+class SWEBenchOutput(BaseModel):
+    model_config = ConfigDict(protected_namespaces=(), populate_by_name=True)
+
+    instance_id: str
+    model_name_or_path: str
+    model_patch: str