fix: decouple mocking from evals

src/uipath/_cli/_evals/_models/_evaluation_set.py

@@ -1,9 +1,15 @@
-from enum import Enum, IntEnum
-from typing import Annotated, Any, Literal, Union
+from enum import IntEnum
+from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic.alias_generators import to_camel
 
+from uipath._cli._evals.mocks.types import (
+    InputMockingStrategy,
+    MockingStrategy,
+    ToolSimulation,
+)
+
 
 class EvaluatorReference(BaseModel):
     """Reference to an evaluator with optional weight.
@@ -52,32 +58,6 @@ def serialize(instance: "EvaluatorReference") -> Any:
         )
 
 
-class EvaluationSimulationTool(BaseModel):
-    name: str = Field(..., alias="name")
-
-
-class MockingStrategyType(str, Enum):
-    LLM = "llm"
-    MOCKITO = "mockito"
-    UNKNOWN = "unknown"
-
-
-class BaseMockingStrategy(BaseModel):
-    pass
-
-
-class ModelSettings(BaseModel):
-    """Model Generation Parameters."""
-
-    model: str = Field(..., alias="model")
-    temperature: float | str | None = Field(default=None, alias="temperature")
-    top_p: float | None = Field(default=None, alias="topP")
-    top_k: int | None = Field(default=None, alias="topK")
-    frequency_penalty: float | None = Field(default=None, alias="frequencyPenalty")
-    presence_penalty: float | None = Field(default=None, alias="presencePenalty")
-    max_tokens: int | None = Field(default=None, alias="maxTokens")
-
-
 class EvaluationSetModelSettings(BaseModel):
     """Model setting overrides within evaluation sets with ID."""
 
@@ -88,75 +68,6 @@ class EvaluationSetModelSettings(BaseModel):
     temperature: float | str | None = Field(default=None, alias="temperature")
 
 
-class LLMMockingStrategy(BaseMockingStrategy):
-    type: Literal[MockingStrategyType.LLM] = MockingStrategyType.LLM
-    prompt: str = Field(..., alias="prompt")
-    tools_to_simulate: list[EvaluationSimulationTool] = Field(
-        ..., alias="toolsToSimulate"
-    )
-    model: ModelSettings | None = Field(None, alias="model")
-
-    model_config = ConfigDict(
-        validate_by_name=True, validate_by_alias=True, extra="allow"
-    )
-
-
-class InputMockingStrategy(BaseModel):
-    prompt: str = Field(..., alias="prompt")
-    model: ModelSettings | None = Field(None, alias="model")
-
-    model_config = ConfigDict(
-        validate_by_name=True, validate_by_alias=True, extra="allow"
-    )
-
-
-class MockingArgument(BaseModel):
-    args: list[Any] = Field(default_factory=lambda: [], alias="args")
-    kwargs: dict[str, Any] = Field(default_factory=lambda: {}, alias="kwargs")
-
-
-class MockingAnswerType(str, Enum):
-    RETURN = "return"
-    RAISE = "raise"
-
-
-class MockingAnswer(BaseModel):
-    type: MockingAnswerType
-    value: Any = Field(..., alias="value")
-
-
-class MockingBehavior(BaseModel):
-    function: str = Field(..., alias="function")
-    arguments: MockingArgument = Field(..., alias="arguments")
-    then: list[MockingAnswer] = Field(..., alias="then")
-
-
-class MockitoMockingStrategy(BaseMockingStrategy):
-    type: Literal[MockingStrategyType.MOCKITO] = MockingStrategyType.MOCKITO
-    behaviors: list[MockingBehavior] = Field(..., alias="config")
-
-    model_config = ConfigDict(
-        validate_by_name=True, validate_by_alias=True, extra="allow"
-    )
-
-
-KnownMockingStrategy = Annotated[
-    Union[LLMMockingStrategy, MockitoMockingStrategy],
-    Field(discriminator="type"),
-]
-
-
-class UnknownMockingStrategy(BaseMockingStrategy):
-    type: str = Field(..., alias="type")
-
-    model_config = ConfigDict(
-        validate_by_name=True, validate_by_alias=True, extra="allow"
-    )
-
-
-MockingStrategy = Union[KnownMockingStrategy, UnknownMockingStrategy]
-
-
 class EvaluationItem(BaseModel):
     """Individual evaluation item within an evaluation set."""
 
@@ -201,7 +112,7 @@ class LegacyEvaluationItem(BaseModel):
     simulation_instructions: str | None = Field(
         default=None, alias="simulationInstructions"
     )
-    tools_to_simulate: list[EvaluationSimulationTool] = Field(
+    tools_to_simulate: list[ToolSimulation] = Field(
         default_factory=list, alias="toolsToSimulate"
     )
 

src/uipath/_cli/_evals/_runtime.py

@@ -88,6 +88,7 @@
     clear_execution_context,
     set_execution_context,
 )
+from .mocks.types import MockingContext
 
 logger = logging.getLogger(__name__)
 
@@ -507,7 +508,15 @@ async def _execute_eval(
                             eval_item, runtime
                         )
 
-                    set_execution_context(eval_item, self.span_collector, execution_id)
+                    set_execution_context(
+                        MockingContext(
+                            strategy=eval_item.mocking_strategy,
+                            name=eval_item.name,
+                            inputs=eval_item.inputs,
+                        ),
+                        span_collector=self.span_collector,
+                        execution_id=execution_id,
+                    )
 
                     await self.event_bus.publish(
                         EvaluationEvents.CREATE_EVAL_RUN,
@@ -761,8 +770,16 @@ async def _generate_input_for_eval(
         self, eval_item: EvaluationItem, runtime: UiPathRuntimeProtocol
     ) -> EvaluationItem:
         """Use LLM to generate a mock input for an evaluation item."""
+        expected_output = (
+            getattr(eval_item, "evaluation_criterias", None)
+            or getattr(eval_item, "expected_output", None)
+            or {}
+        )
         generated_input = await generate_llm_input(
-            eval_item, (await self.get_schema(runtime)).input
+            eval_item.input_mocking_strategy,
+            (await self.get_schema(runtime)).input,
+            expected_behavior=eval_item.expected_agent_behavior or "",
+            expected_output=expected_output,
         )
         updated_eval_item = eval_item.model_copy(update={"inputs": generated_input})
         return updated_eval_item

src/uipath/_cli/_evals/mocks/input_mocker.py

@@ -4,7 +4,9 @@
 from datetime import datetime
 from typing import Any
 
-from uipath._cli._evals._models._evaluation_set import EvaluationItem
+from uipath._cli._evals.mocks.types import (
+    InputMockingStrategy,
+)
 from uipath.platform import UiPath
 from uipath.tracing import traced
 
@@ -54,8 +56,10 @@ def get_input_mocking_prompt(
 
 @traced(name="__mocker__", recording=False)
 async def generate_llm_input(
-    evaluation_item: EvaluationItem,
+    mocking_strategy: InputMockingStrategy,
     input_schema: dict[str, Any],
+    expected_behavior: str,
+    expected_output: dict[str, Any],
 ) -> dict[str, Any]:
     """Generate synthetic input using an LLM based on the evaluation context."""
     from .mocks import cache_manager_context
@@ -68,18 +72,12 @@ async def generate_llm_input(
         if "additionalProperties" not in input_schema:
             input_schema["additionalProperties"] = False
 
-        expected_output = (
-            getattr(evaluation_item, "evaluation_criterias", None)
-            or getattr(evaluation_item, "expected_output", None)
-            or {}
-        )
-
         prompt_generation_args = {
             "input_schema": json.dumps(input_schema),
-            "input_generation_instructions": evaluation_item.input_mocking_strategy.prompt
-            if evaluation_item.input_mocking_strategy
+            "input_generation_instructions": mocking_strategy.prompt
+            if mocking_strategy
             else "",
-            "expected_behavior": evaluation_item.expected_agent_behavior or "",
+            "expected_behavior": expected_behavior or "",
             "expected_output": json.dumps(expected_output),
         }
 
@@ -94,11 +92,7 @@ async def generate_llm_input(
             },
         }
 
-        model_parameters = (
-            evaluation_item.input_mocking_strategy.model
-            if evaluation_item.input_mocking_strategy
-            else None
-        )
+        model_parameters = mocking_strategy.model if mocking_strategy else None
         completion_kwargs = (
             model_parameters.model_dump(by_alias=False, exclude_none=True)
             if model_parameters

src/uipath/_cli/_evals/mocks/llm_mocker.py

@@ -6,10 +6,13 @@
 
 from pydantic import BaseModel, TypeAdapter
 
+from uipath._cli._evals.mocks.types import (
+    LLMMockingStrategy,
+    MockingContext,
+)
 from uipath.tracing import traced
 from uipath.tracing._utils import _SpanUtils
 
-from .._models._evaluation_set import EvaluationItem, LLMMockingStrategy
 from .._models._mocks import ExampleCall
 from .mocker import (
     Mocker,
@@ -74,29 +77,27 @@ def pydantic_to_dict_safe(obj: Any) -> Any:
 class LLMMocker(Mocker):
     """LLM Based Mocker."""
 
-    def __init__(self, evaluation_item: EvaluationItem):
+    def __init__(self, context: MockingContext):
         """LLM Mocker constructor."""
-        self.evaluation_item = evaluation_item
-        assert isinstance(self.evaluation_item.mocking_strategy, LLMMockingStrategy)
+        self.context = context
+        assert isinstance(self.context.strategy, LLMMockingStrategy)
 
     @traced(name="__mocker__", recording=False)
     async def response(
         self, func: Callable[[T], R], params: dict[str, Any], *args: T, **kwargs
     ) -> R:
         """Respond with mocked response generated by an LLM."""
-        assert isinstance(self.evaluation_item.mocking_strategy, LLMMockingStrategy)
+        assert isinstance(self.context.strategy, LLMMockingStrategy)
 
         function_name = params.get("name") or func.__name__
-        if function_name in [
-            x.name for x in self.evaluation_item.mocking_strategy.tools_to_simulate
-        ]:
+        if function_name in [x.name for x in self.context.strategy.tools_to_simulate]:
             from uipath.platform import UiPath
             from uipath.platform.chat._llm_gateway_service import _cleanup_schema
 
             from .mocks import (
                 cache_manager_context,
-                evaluation_context,
                 execution_id_context,
+                mocking_context,
                 span_collector_context,
             )
 
@@ -127,10 +128,10 @@ async def response(
                 ]
 
                 test_run_history = "(empty)"
-                eval_item = evaluation_context.get()
+                ctx = mocking_context.get()
                 span_collector = span_collector_context.get()
                 execution_id = execution_id_context.get()
-                if eval_item and span_collector and execution_id:
+                if ctx and span_collector and execution_id:
                     spans = span_collector.get_spans(execution_id)
                     test_run_history = _SpanUtils.spans_to_llm_context(spans)
 
@@ -154,17 +155,17 @@ async def response(
                         "kwargs": kwargs,
                     },
                     "agentInfo": {  # This is incomplete
-                        # "agentName": self.evaluation_item.name, # to be obtained.
-                        "actionName": self.evaluation_item.name,  # Not sure if this is right?
-                        "userInput": self.evaluation_item.inputs,
+                        "agentName": self.context.name,
+                        "actionName": self.context.name,  # Not sure if this is right?
+                        "userInput": self.context.inputs,
                     },
-                    "testRunProctorInstructions": self.evaluation_item.mocking_strategy.prompt,
+                    "testRunProctorInstructions": self.context.strategy.prompt,
                 }
                 prompt_generation_args = {
                     k: json.dumps(pydantic_to_dict_safe(v))
                     for k, v in prompt_input.items()
                 }
-                model_parameters = self.evaluation_item.mocking_strategy.model
+                model_parameters = self.context.strategy.model
                 completion_kwargs = (
                     model_parameters.model_dump(by_alias=False, exclude_none=True)
                     if model_parameters

src/uipath/_cli/_evals/mocks/mocker_factory.py

@@ -1,25 +1,25 @@
 """Mocker Factory."""
 
-from uipath._cli._evals._models._evaluation_set import (
-    EvaluationItem,
-    LLMMockingStrategy,
-    MockitoMockingStrategy,
-)
 from uipath._cli._evals.mocks.llm_mocker import LLMMocker
 from uipath._cli._evals.mocks.mocker import Mocker
 from uipath._cli._evals.mocks.mockito_mocker import MockitoMocker
+from uipath._cli._evals.mocks.types import (
+    LLMMockingStrategy,
+    MockingContext,
+    MockitoMockingStrategy,
+)
 
 
 class MockerFactory:
     """Mocker factory."""
 
     @staticmethod
-    def create(evaluation_item: EvaluationItem) -> Mocker:
+    def create(context: MockingContext) -> Mocker:
         """Create a mocker instance."""
-        match evaluation_item.mocking_strategy:
+        match context.strategy:
             case LLMMockingStrategy():
-                return LLMMocker(evaluation_item)
+                return LLMMocker(context)
             case MockitoMockingStrategy():
-                return MockitoMocker(evaluation_item)
+                return MockitoMocker(context)
             case _:
                 raise ValueError("Unknown mocking strategy")