fix(eval): raise the simulator's output token limit and make its tool schema work with Gemini

packages/uipath/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.15.0"
+version = "2.15.1"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

packages/uipath/src/uipath/_cli/cli_debug.py

@@ -12,6 +12,7 @@
 from uipath._cli._utils._tracing import create_trace_manager
 from uipath.eval.mocks import SimulationConfig, UiPathMockRuntime, build_mocking_context
 from uipath.eval.mocks._mock_runtime import load_simulation_config
+from uipath.eval.mocks._simulation_model import get_agent_max_tokens
 from uipath.platform.common import (
     ExecutionSourceContext,
     ResourceOverwritesContext,
@@ -249,13 +250,14 @@ async def execute_debug_runtime():
                             agent_model = None
                             if schema.metadata and "settings" in schema.metadata:
                                 agent_model = schema.metadata["settings"].get("model")
+                            agent_max_tokens = get_agent_max_tokens(schema)
 
                             delegate_runtime: UiPathDebugRuntime | UiPathMockRuntime = (
                                 debug_runtime
                             )
                             if simulation_config:
                                 mocking_context = build_mocking_context(
-                                    simulation_config, agent_model
+                                    simulation_config, agent_model, agent_max_tokens
                                 )
                                 if mocking_context:
                                     delegate_runtime = UiPathMockRuntime(
@@ -264,7 +266,8 @@ async def execute_debug_runtime():
                                     )
                             else:
                                 mocking_context = load_simulation_config(
-                                    agent_model=agent_model
+                                    agent_model=agent_model,
+                                    agent_max_tokens=agent_max_tokens,
                                 )
                                 delegate_runtime = UiPathMockRuntime(
                                     delegate=debug_runtime,

packages/uipath/src/uipath/_cli/cli_run.py

@@ -10,6 +10,7 @@
 from uipath._cli._utils._debug import setup_debugging
 from uipath._cli._utils._tracing import create_trace_manager
 from uipath.eval.mocks import SimulationConfig, UiPathMockRuntime, build_mocking_context
+from uipath.eval.mocks._simulation_model import get_agent_max_tokens
 from uipath.platform.common import (
     ExecutionSourceContext,
     ResourceOverwritesContext,
@@ -316,7 +317,9 @@ async def execute() -> None:
                                         "model"
                                     )
                                 mocking_context = build_mocking_context(
-                                    simulation_config, agent_model
+                                    simulation_config,
+                                    agent_model,
+                                    get_agent_max_tokens(schema),
                                 )
                                 if mocking_context:
                                     runtime = UiPathMockRuntime(

packages/uipath/src/uipath/eval/mocks/_mock_runtime.py

@@ -23,6 +23,7 @@
 )
 from ._mock_context import mocker_context, mocking_context
 from ._mocker_factory import MockerFactory
+from ._simulation_model import simulation_model_settings
 from ._types import (
     LLMMockingStrategy,
     MockingContext,
@@ -35,9 +36,15 @@
 
 
 def build_mocking_context(
-    config: SimulationConfig, agent_model: str | None = None
+    config: SimulationConfig,
+    agent_model: str | None = None,
+    agent_max_tokens: int | None = None,
 ) -> MockingContext | None:
-    """Build a MockingContext from a validated SimulationConfig."""
+    """Build a MockingContext from a validated SimulationConfig.
+
+    Without a model in the config, the simulation uses the agent's model and
+    its max tokens (capped, see ``simulation_model_settings``).
+    """
     if not config.enabled:
         return None
 
@@ -68,9 +75,9 @@
     model = (
         ModelSettings(model=config.model)
         if config.model
-        else ModelSettings(model=agent_model)
+        else simulation_model_settings(agent_model, agent_max_tokens)
         if agent_model
         else None
     )
 
     mocking_strategy = LLMMockingStrategy(
@@ -91,7 +98,9 @@
 
 
 def build_mocking_context_from_dict(
-    simulation_data: dict[str, Any], agent_model: str | None = None
+    simulation_data: dict[str, Any],
+    agent_model: str | None = None,
+    agent_max_tokens: int | None = None,
 ) -> MockingContext | None:
     """Build a MockingContext from a simulation config dictionary.
 
@@ -100,15 +109,18 @@
     Args:
         simulation_data: Parsed simulation config (same schema as simulation.json).
         agent_model: Optional agent model name to use as fallback.
+        agent_max_tokens: Optional max tokens of the agent, used with agent_model.
 
     Returns:
         MockingContext if valid and enabled, None otherwise.
     """
     config = SimulationConfig.model_validate(simulation_data)
-    return build_mocking_context(config, agent_model)
+    return build_mocking_context(config, agent_model, agent_max_tokens)
 
 
-def load_simulation_config(agent_model: str | None = None) -> MockingContext | None:
+def load_simulation_config(
+    agent_model: str | None = None, agent_max_tokens: int | None = None
+) -> MockingContext | None:
     """Load simulation.json from current directory and convert to MockingContext.
 
     Returns:
@@ -124,7 +136,7 @@
         config = SimulationConfig.model_validate_json(
             simulation_path.read_text(encoding="utf-8")
         )
-        return build_mocking_context(config, agent_model)
+        return build_mocking_context(config, agent_model, agent_max_tokens)
 
     except Exception as e:
         logger.warning(f"Failed to load simulation.json: {e}")

packages/uipath/src/uipath/eval/mocks/_simulation_model.py

@@ -4,10 +4,47 @@
 from typing import Any
 
 from uipath.platform.chat._llm_gateway_service import ChatModels
+from uipath.runtime.schema import UiPathRuntimeSchema
 
 from ._types import ModelSettings
 
 SIMULATION_MODEL_ENV = "UIPATH_SIMULATION_MODEL"
+# Output token limit of a simulation call. Reasoning models count their
+# reasoning against it, so the chat service default (4096) is too low for long
+# simulated responses; the cap bounds a runaway reply.
+SIMULATION_MAX_TOKENS = 16384
+
+
+def get_agent_max_tokens(schema: UiPathRuntimeSchema) -> int | None:
+    """Get the agent's max tokens from the runtime schema metadata.
+
+    Read from ``schema.metadata["settings"]``, where the low-code agents runtime
+    puts the agent's settings (the model is read from there too).
+
+    Returns:
+        The agent's max tokens, or None if not found or not a positive integer.
+    """
+    settings = (schema.metadata or {}).get("settings")
+    if not isinstance(settings, dict):
+        return None
+    max_tokens = settings.get("max_tokens", settings.get("maxTokens"))
+    if isinstance(max_tokens, int) and not isinstance(max_tokens, bool):
+        return max_tokens if max_tokens > 0 else None
+    return None
+
+
+def simulation_model_settings(
+    model: str, agent_max_tokens: int | None = None
+) -> ModelSettings:
+    """Build the settings for a simulation that uses the agent's model.
+
+    The output limit is the agent's own max tokens, a value the model is known
+    to accept, capped at ``SIMULATION_MAX_TOKENS``.
+    """
+    max_tokens = SIMULATION_MAX_TOKENS
+    if agent_max_tokens and agent_max_tokens > 0:
+        max_tokens = min(agent_max_tokens, SIMULATION_MAX_TOKENS)
+    return ModelSettings(model=model, maxTokens=max_tokens)
 
 
 def simulation_completion_kwargs(
@@ -24,4 +61,5 @@ def simulation_completion_kwargs(
         "model",
         os.environ.get(SIMULATION_MODEL_ENV) or ChatModels.gpt_4_1_mini_2025_04_14,
     )
+    completion_kwargs.setdefault("max_tokens", SIMULATION_MAX_TOKENS)
     return completion_kwargs

packages/uipath/src/uipath/eval/mocks/_structured_output.py

@@ -24,6 +24,11 @@
 RESPONSE_TOOL_NAME = "submit_tool_response"
 RESPONSE_KEY = "response"
 _DEFS_PREFIX = "#/$defs/"
+# Keys whose value maps names to schemas, and keys whose value is data.
+_SCHEMA_MAP_KEYS = frozenset(
+    {"properties", "patternProperties", "$defs", "definitions"}
+)
+_DATA_KEYS = frozenset({"default", "examples", "const", "enum", "required"})
 
 logger = logging.getLogger(__name__)
 
@@ -81,6 +86,81 @@
     return inlined, leftover
 
 
+def _narrowed_type(branch_type: Any, types: list[str]) -> Any:
+    """Return the part of ``branch_type`` that ``types`` allows, or None."""
+    if branch_type in types or (branch_type == "integer" and "number" in types):
+        return branch_type
+    if branch_type == "number" and "integer" in types:
+        return "integer"
+    return None
+
+
+def _restrict_to_types(branches: list[Any], types: list[str]) -> list[Any]:
+    """Narrow ``anyOf`` branches to ``types``, for a node that has both.
+
+    A branch with its own type is kept, narrowed, only when ``types`` allows
+    it; a branch without one is copied once per type in ``types``.
+    """
+    restricted: list[Any] = []
+    for branch in branches:
+        if not isinstance(branch, dict):
+            restricted.append(branch)
+        elif "type" not in branch:
+            restricted.extend({**branch, "type": t} for t in types)
+        elif (narrowed := _narrowed_type(branch["type"], types)) is not None:
+            restricted.append({**branch, "type": narrowed})
+    # No branch fits the types: keep the branches rather than send an empty anyOf.
+    return restricted or branches
+
+
+def _to_gemini_schema(schema: Any) -> Any:
+    """Rewrite JSON-schema type lists into the form Gemini accepts.
+
+    Gemini's function declarations take a single ``type`` and mark nullable
+    fields with ``nullable: true``; a type list such as ``["string", "null"]``
+    (common in MCP tool schemas) fails the request with a 400. So
+    ``{"type": ["string", "null"]}`` becomes
+    ``{"type": "string", "nullable": true}``, and a list with several non-null
+    types becomes an ``anyOf`` of single types. When the node already has an
+    ``anyOf``, the types are applied to its branches instead, so both
+    constraints still hold. Values that are data rather than schemas
+    (``default``, ``enum``, ...) are left untouched.
+    """
+    if isinstance(schema, list):
+        return [_to_gemini_schema(item) for item in schema]
+    if not isinstance(schema, dict):
+        return schema
+
+    converted: dict[str, Any] = {}
+    for key, value in schema.items():
+        if key in _SCHEMA_MAP_KEYS and isinstance(value, dict):
+            converted[key] = {
+                name: _to_gemini_schema(sub) for name, sub in value.items()
+            }
+        elif key in _DATA_KEYS:
+            converted[key] = value
+        else:
+            converted[key] = _to_gemini_schema(value)
+
+    types = converted.get("type")
+    if isinstance(types, list):
+        non_null = [t for t in types if t != "null"]
+        if len(non_null) == 1:
+            converted["type"] = non_null[0]
+        else:
+            del converted["type"]
+        if len(non_null) > 1:
+            branches = converted.get("anyOf")
+            converted["anyOf"] = (
+                _restrict_to_types(branches, non_null)
+                if isinstance(branches, list)
+                else [{"type": t} for t in non_null]
+            )
+        if "null" in types:
+            converted["nullable"] = True
+    return converted
+
+
 def build_response_tool(schema: dict[str, Any], description: str) -> dict[str, Any]:
     """Build a normalized-API function tool that wraps ``schema`` under ``response``.
 
@@ -223,7 +303,30 @@
 
 
 class GeminiStructuredOutput(ToolCallStructuredOutput):
-    """Gemini returns empty content for ``response_format``; go straight to tools."""
+    """Gemini returns empty content for ``response_format``; go straight to tools.
+
+    The tool schema is rewritten into the form Gemini accepts first.
+    """
+
+    async def generate(
+        self,
+        llm: Any,
+        messages: list[dict[str, str]],
+        *,
+        schema: dict[str, Any],
+        response_format_name: str,
+        description: str,
+        completion_kwargs: dict[str, Any],
+    ) -> Any:
+        """Force a tool call with a Gemini-compatible schema."""
+        return await super().generate(
+            llm,
+            messages,
+            schema=_to_gemini_schema(schema),
+            response_format_name=response_format_name,
+            description=description,
+            completion_kwargs=completion_kwargs,
+        )
 
 
 def _strategy_for_model(model: str | None) -> ToolCallStructuredOutput:

packages/uipath/src/uipath/eval/runtime/runtime.py

@@ -65,7 +65,11 @@
 from ..mocks._mock_context import cache_manager_context
 from ..mocks._mock_runtime import UiPathMockRuntime
 from ..mocks._mocker import UiPathInputMockingError, format_exception_message
-from ..mocks._types import LLMMockingStrategy, MockingContext, ModelSettings
+from ..mocks._simulation_model import (
+    get_agent_max_tokens,
+    simulation_model_settings,
+)
+from ..mocks._types import LLMMockingStrategy, MockingContext
 from ..models import EvaluationResult
 from ..models.evaluation_set import (
     EvaluationItem,
@@ -662,7 +666,14 @@ async def _execute_eval(
                         mocking_model = get_agent_model(self.context.runtime_schema)
                         if mocking_model:
                             mocking_strategy = mocking_strategy.model_copy(
-                                update={"model": ModelSettings(model=mocking_model)}
+                                update={
+                                    "model": simulation_model_settings(
+                                        mocking_model,
+                                        get_agent_max_tokens(
+                                            self.context.runtime_schema
+                                        ),
+                                    )
+                                }
                             )
 
                     agent_execution_output = await self.execute_runtime(
@@ -956,7 +967,12 @@ async def _generate_input_for_eval(
             input_generation_model = get_agent_model(self.context.runtime_schema)
             if input_generation_model:
                 input_strategy = input_strategy.model_copy(
-                    update={"model": ModelSettings(model=input_generation_model)}
+                    update={
+                        "model": simulation_model_settings(
+                            input_generation_model,
+                            get_agent_max_tokens(self.context.runtime_schema),
+                        )
+                    }
                 )
 
         generated_input = await generate_llm_input(