aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc2__py3-none-any.whl

aiq/data_models/interactive.py

@@ -33,6 +33,7 @@ class HumanPromptModelType(str, Enum):
     RADIO = "radio"
     CHECKBOX = "checkbox"
     DROPDOWN = "dropdown"
+    OAUTH_CONSENT = "oauth_consent"
 
 
 class BinaryChoiceOptionsType(str, Enum):
@@ -145,6 +146,14 @@ class HumanPromptNotification(HumanPromptBase):
     input_type: typing.Literal[HumanPromptModelType.NOTIFICATION] = HumanPromptModelType.NOTIFICATION
 
 
+class _HumanPromptOAuthConsent(HumanPromptBase):
+    """
+    Represents an OAuth consent prompt interaction used to notify the UI to open the authentication page for completing
+    the consent flow.
+    """
+    input_type: typing.Literal[HumanPromptModelType.OAUTH_CONSENT] = HumanPromptModelType.OAUTH_CONSENT
+
+
 class HumanPromptBinary(HumanPromptBase):
     """
     Represents a binary interaction.
@@ -190,7 +199,7 @@ class HumanPromptDropdown(HumanPromptMultipleChoiceBase):
 
 
 HumanPrompt = typing.Annotated[HumanPromptText | HumanPromptNotification | HumanPromptBinary | HumanPromptRadio
-                               | HumanPromptCheckbox | HumanPromptDropdown,
+                               | HumanPromptCheckbox | HumanPromptDropdown | _HumanPromptOAuthConsent,
                                Discriminator("input_type")]
 
 

aiq/data_models/intermediate_step.py

@@ -17,6 +17,7 @@ import time
 import typing
 import uuid
 from enum import Enum
+from typing import Literal
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
@@ -82,6 +83,26 @@ class UsageInfo(BaseModel):
     seconds_between_calls: int = 0
 
 
+class ToolParameters(BaseModel):
+    properties: dict[str, typing.Any] = Field(..., description="The properties of the function parameters.")
+    required: list[str] = Field(default_factory=list, description="The required properties of the function parameters.")
+    type_: Literal["object"] = Field(default="object", description="The type of the function parameters.", alias="type")
+    additionalProperties: bool = Field(default=False,
+                                       description="Enable function parameters allow additional properties.")
+    strict: bool = Field(default=True, description="Ensure function calls reliably adhere to the function schema.")
+
+
+class ToolDetails(BaseModel):
+    name: str = Field(..., description="The name of the function.")
+    description: str = Field(..., description="The description of the function.")
+    parameters: ToolParameters = Field(..., description="The parameters of the function.")
+
+
+class ToolSchema(BaseModel):
+    type: Literal["function"] = Field(..., description="The type of the tool.")
+    function: ToolDetails = Field(..., description="The function details.")
+
+
 class TraceMetadata(BaseModel):
     chat_responses: typing.Any | None = None
     chat_inputs: typing.Any | None = None
@@ -91,6 +112,8 @@ class TraceMetadata(BaseModel):
     span_inputs: typing.Any | None = None
     span_outputs: typing.Any | None = None
     provided_metadata: typing.Any | None = None
+    tools_schema: list[ToolSchema] = Field(default_factory=list,
+                                           description="The schema of tools used in a tool calling request.")
 
     # Allow extra fields in the model_config to support derived models
     model_config = ConfigDict(extra="allow")
@@ -211,9 +234,23 @@ class IntermediateStep(BaseModel):
     # Allow extra fields in the model_config to support derived models
     model_config = ConfigDict(extra="forbid")
 
-    function_ancestry: InvocationNode | None = InvocationNode(function_name="N/A", function_id="N/A")
+    parent_id: str
+    """
+    The parent step ID for the current step. The parent ID is the ID of the last START step which has a different UUID
+    than the current step. This value is different from the function_ancestry.parent_id value which tracks the last
+    parent FUNCTION step. For the first START step, the parent_id is 'root'.
+    """
+
+    function_ancestry: InvocationNode
+    """
+    The function ancestry for the current step showing the current AIQ function that was being executed when the step
+    was created.
+    """
 
     payload: IntermediateStepPayload
+    """
+    The payload for the current step.
+    """
 
     # ===== Payload Properties =====
     @property
@@ -263,7 +300,3 @@ class IntermediateStep(BaseModel):
     @property
     def event_state(self) -> IntermediateStepState:
         return self.payload.event_state
-
-    @property
-    def parent_id(self) -> str | None:
-        return self.function_ancestry.function_id if self.function_ancestry else None

aiq/data_models/its_strategy.py

@@ -0,0 +1,30 @@
+# 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 ITSStrategyBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    Base configuration class for Inference Time Scaling (ITS) strategy.
+    This class is used to define the structure of ITS strategy configurations.
+    """
+    pass
+
+
+ITSStrategyBaseConfigT = typing.TypeVar("ITSStrategyBaseConfigT", bound=ITSStrategyBaseConfig)

aiq/data_models/llm.py

@@ -20,6 +20,7 @@ from .common import TypedBaseModel
 
 
 class LLMBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """Base configuration for LLM providers."""
     pass
 
 

aiq/data_models/memory.py

@@ -20,6 +20,7 @@ from .common import TypedBaseModel
 
 
 class MemoryBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """ The base level config object for a memory object. Memories provide an interface for storing and retrieving. """
     pass
 
 

aiq/data_models/object_store.py

@@ -0,0 +1,44 @@
+# 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 ObjectStoreBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+ObjectStoreBaseConfigT = typing.TypeVar("ObjectStoreBaseConfigT", bound=ObjectStoreBaseConfig)
+
+
+class KeyAlreadyExistsError(Exception):
+
+    def __init__(self, key: str, additional_message: str | None = None):
+        parts = [f"Key already exists: {key}."]
+        if additional_message:
+            parts.append(additional_message)
+        super().__init__(" ".join(parts))
+
+
+class NoSuchKeyError(Exception):
+
+    def __init__(self, key: str, additional_message: str | None = None):
+        parts = [f"No object found with key: {key}."]
+        if additional_message:
+            parts.append(additional_message)
+        super().__init__(" ".join(parts))

aiq/data_models/profiler.py

@@ -42,6 +42,7 @@ class PrefixSpanConfig(BaseModel):
 
 class ProfilerConfig(BaseModel):
 
+    base_metrics: bool = False
     token_usage_forecast: bool = False
     token_uniqueness_forecast: bool = False
     workflow_runtime_forecast: bool = False

aiq/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)

aiq/data_models/span.py

@@ -0,0 +1,187 @@
+# 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 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__)
+
+
+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):
+    AIQ_SPAN_KIND = "aiq.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"
+    AIQ_USAGE_NUM_LLM_CALLS = "aiq.usage.num_llm_calls"
+    AIQ_USAGE_SECONDS_BETWEEN_CALLS = "aiq.usage.seconds_between_calls"
+    AIQ_USAGE_TOKEN_COUNT_PROMPT = "aiq.usage.token_count.prompt"
+    AIQ_USAGE_TOKEN_COUNT_COMPLETION = "aiq.usage.token_count.completion"
+    AIQ_USAGE_TOKEN_COUNT_TOTAL = "aiq.usage.token_count.total"
+    AIQ_EVENT_TYPE = "aiq.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.")
+
+
+class SpanContext(BaseModel):
+    trace_id: int = Field(default_factory=lambda: uuid.uuid4().int, description="The 128-bit trace ID of the span.")
+    span_id: int = Field(default_factory=lambda: uuid.uuid4().int & ((1 << 64) - 1),
+                         description="The 64-bit span ID of the span.")
+
+
+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

aiq/data_models/telemetry_exporter.py

@@ -15,8 +15,8 @@
 
 import typing
 
-from .common import BaseModelRegistryTag
-from .common import TypedBaseModel
+from aiq.data_models.common import BaseModelRegistryTag
+from aiq.data_models.common import TypedBaseModel
 
 
 class TelemetryExporterBaseConfig(TypedBaseModel, BaseModelRegistryTag):

aiq/embedder/nim_embedder.py

@@ -24,6 +24,7 @@ from aiq.builder.builder import Builder
 from aiq.builder.embedder import EmbedderProviderInfo
 from aiq.cli.register_workflow import register_embedder_provider
 from aiq.data_models.embedder import EmbedderBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
 
 allowed_truncate_values = ["NONE", "START", "END"]
 
@@ -37,7 +38,7 @@ def option_in_allowed_values(v):
 TruncationOption = typing.Annotated[str, AfterValidator(option_in_allowed_values)]
 
 
-class NIMEmbedderModelConfig(EmbedderBaseConfig, name="nim"):
+class NIMEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="nim"):
     """A NVIDIA Inference Microservice (NIM) embedder provider to be used with an embedder client."""
 
     api_key: str | None = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")

aiq/embedder/openai_embedder.py

@@ -21,9 +21,10 @@ from aiq.builder.builder import Builder
 from aiq.builder.embedder import EmbedderProviderInfo
 from aiq.cli.register_workflow import register_embedder_provider
 from aiq.data_models.embedder import EmbedderBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
 
 
-class OpenAIEmbedderModelConfig(EmbedderBaseConfig, name="openai"):
+class OpenAIEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())

aiq/eval/config.py

@@ -17,13 +17,18 @@ from pathlib import Path
 
 from pydantic import BaseModel
 
+from aiq.eval.evaluator.evaluator_model import EvalInput
+from aiq.eval.evaluator.evaluator_model import EvalOutput
+from aiq.eval.usage_stats import UsageStats
+from aiq.profiler.data_models import ProfilerResults
+
 
 class EvaluationRunConfig(BaseModel):
     """
     Parameters used for a single evaluation run.
     """
     config_file: Path
-    dataset: str | None  # dataset file path can be specified in the config file
+    dataset: str | None = None  # dataset file path can be specified in the config file
     result_json_path: str = "$"
     skip_workflow: bool = False
     skip_completed_entries: bool = False
@@ -31,6 +36,14 @@ class EvaluationRunConfig(BaseModel):
     endpoint_timeout: int = 300
     reps: int = 1
     override: tuple[tuple[str, str], ...] = ()
+    # If false, the output will not be written to the output directory. This is
+    # useful when running evaluation via another tool.
+    write_output: bool = True
+    # if true, the dataset is adjusted to a multiple of the concurrency
+    adjust_dataset_size: bool = False
+    # number of passes at each concurrency, if 0 the dataset is adjusted to a multiple of the
+    # concurrency. The is only used if adjust_dataset_size is true
+    num_passes: int = 0
 
 
 class EvaluationRunOutput(BaseModel):
@@ -40,3 +53,8 @@ class EvaluationRunOutput(BaseModel):
     workflow_output_file: Path | None
     evaluator_output_files: list[Path]
     workflow_interrupted: bool
+
+    eval_input: EvalInput
+    evaluation_results: list[tuple[str, EvalOutput]]
+    usage_stats: UsageStats | None = None
+    profiler_results: ProfilerResults

aiq/eval/dataset_handler/dataset_handler.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import json
+import math
 
 import pandas as pd
 
@@ -33,12 +34,23 @@ class DatasetHandler:
     One DatasetHandler object is needed for each dataset to be evaluated.
     """
 
-    def __init__(self, dataset_config: EvalDatasetConfig, reps: int):
+    def __init__(self,
+                 dataset_config: EvalDatasetConfig,
+                 reps: int,
+                 concurrency: int,
+                 num_passes: int | None = None,
+                 adjust_dataset_size: bool = False):
         from aiq.eval.intermediate_step_adapter import IntermediateStepAdapter
 
         self.dataset_config = dataset_config
         self.dataset_filter = DatasetFilter(dataset_config.filter)
         self.reps = reps
+
+        # number of passes at specific concurrency
+        self.concurrency = concurrency
+        self.num_passes = num_passes
+        self.adjust_dataset_size = adjust_dataset_size
+
         # Helpers
         self.intermediate_step_adapter = IntermediateStepAdapter()
 
@@ -81,6 +93,7 @@ class DatasetHandler:
                 output_obj=row.get(self.generated_answer_key, "") if structured else "",
                 trajectory=row.get(self.trajectory_key, []) if structured else [],
                 expected_trajectory=row.get(self.expected_trajectory_key, []) if structured else [],
+                full_dataset_entry=row.to_dict(),
             )
 
         # if input dataframe is empty return an empty list
@@ -108,6 +121,63 @@ class DatasetHandler:
 
         return input_df
 
+    def adjust_dataset(self, input_df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Adjust the dataset so its length is a multiple of concurrency.
+
+        If num_passes > 0:
+            dataset size is adjusted to concurrency * num_passes
+        else:
+            dataset size is adjusted to the largest multiple of concurrency
+            that is less than or equal to the current dataset size
+        """
+        if self.concurrency <= 0:
+            raise ValueError("Concurrency must be > 0")
+
+        if self.num_passes < 0:
+            raise ValueError("num_passes must be >= 0")
+
+        original_size = input_df.shape[0]
+
+        # Calculate target size
+        if self.num_passes > 0:
+            # When num_passes is specified, always use concurrency * num_passes
+            # This respects the user's intent for exact number of passes
+            target_size = self.concurrency * self.num_passes
+        else:
+            # When num_passes = 0, use the largest multiple of concurrency <= original_size
+            # If original_size < concurrency, we need at least concurrency rows
+            if original_size >= self.concurrency:
+                target_size = (original_size // self.concurrency) * self.concurrency
+            else:
+                target_size = self.concurrency
+
+        if target_size == 0:
+            raise ValueError("Input dataset too small for even one batch at given concurrency.")
+
+        id_col = self.dataset_config.id_key
+
+        # If we need more rows than we have, replicate the dataset
+        if original_size < target_size:
+            # Clean existing _rep suffix if present
+            input_df[id_col] = input_df[id_col].astype(str).str.replace(r"_rep\d+$", "", regex=True)
+
+            # Calculate how many complete copies we need
+            copies_needed = math.ceil(target_size / original_size)
+
+            # Create the replicated dataframe
+            replicated_dfs = []
+            for i in range(copies_needed):
+                df_copy = input_df.copy()
+                if i > 0:  # Add suffix to all but the first copy
+                    df_copy[id_col] = df_copy[id_col].astype(str) + f"_rep{i}"
+                replicated_dfs.append(df_copy)
+
+            input_df = pd.concat(replicated_dfs, ignore_index=True)
+
+        # Return exactly the target size
+        return input_df.head(target_size)
+
     def get_eval_input_from_dataset(self, dataset: str) -> EvalInput:
         # read the dataset and convert it to EvalInput
 
@@ -126,9 +196,14 @@ class DatasetHandler:
         input_df = self.dataset_filter.apply_filters(input_df)
         input_df.drop_duplicates(subset=[self.dataset_config.id_key], inplace=True)
 
+        if self.reps > 1 and self.adjust_dataset_size:
+            raise ValueError("reps and adjust_dataset_size are mutually exclusive")
+
         # If more than one repetition is needed, replicate the rows
         if self.reps > 1:
             input_df = self.setup_reps(input_df)
+        elif self.adjust_dataset_size:
+            input_df = self.adjust_dataset(input_df)
 
         # Convert the DataFrame to a list of EvalInput objects
         return self.get_eval_input_from_df(input_df)
@@ -151,6 +226,16 @@ class DatasetHandler:
         allow re-running evaluation using the orignal config file and '--skip_workflow' option.
         """
 
+        def parse_if_json_string(value):
+            if isinstance(value, str):
+                try:
+                    return json.loads(value)
+                except json.JSONDecodeError:
+                    return value
+            if hasattr(value, "model_dump"):
+                return value.model_dump()
+            return value
+
         indent = 2
         if self.is_structured_input():
             # Extract structured data from EvalInputItems
@@ -164,6 +249,6 @@ class DatasetHandler:
             } for item in eval_input.eval_input_items]
         else:
             # Unstructured case: return only raw output objects as a JSON array
-            data = [json.loads(item.output_obj) for item in eval_input.eval_input_items]
+            data = [parse_if_json_string(item.output_obj) for item in eval_input.eval_input_items]
 
         return json.dumps(data, indent=indent, ensure_ascii=False, default=str)