azure-ai-evaluation 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

azure/ai/evaluation/_evaluators/_tool_call_accuracy/tool_call_accuracy.prompty

@@ -0,0 +1,71 @@
+---
+name: Tool Call Accuracy
+description: Evaluates Tool Call Accuracy for tool used by agent
+model:
+  api: chat
+  parameters:
+    temperature: 0.0
+    max_tokens: 800
+    top_p: 1.0
+    presence_penalty: 0
+    frequency_penalty: 0
+    response_format:
+      type: text
+
+inputs:
+  query:
+    type: List
+  tool_call:
+    type: Dict
+  tool_definition:
+    type: Dict
+
+---
+system:
+# Instruction
+## Goal
+### Your are an expert in evaluating the accuracy of a tool call considering relevance and potential usefulness including syntactic and semantic correctness of a proposed tool call from an intelligent system based on provided definition and data. Your goal will involve answering the questions below using the information provided.
+- **Definition**: You are given a definition of the communication trait that is being evaluated to help guide your Score.
+- **Data**: Your input data include CONVERSATION , TOOL CALL and TOOL DEFINITION.
+- **Tasks**: To complete your evaluation you will be asked to evaluate the Data in different ways.
+
+user:
+# Definition
+**Tool Call Accuracy** refers to the relevance and potential usefulness of a TOOL CALL in the context of an ongoing CONVERSATION and EXTRACTION of RIGHT PARAMETER VALUES from the CONVERSATION.It assesses how likely the TOOL CALL is to contribute meaningfully to the CONVERSATION and help address the user's needs. Focus on evaluating the potential value of the TOOL CALL within the specific context of the given CONVERSATION, without making assumptions beyond the provided information.
+  Consider the following factors in your evaluation:
+
+  1. Relevance: How well does the proposed tool call align with the current topic and flow of the conversation?
+  2. Parameter Appropriateness: Do the parameters used in the TOOL CALL match the TOOL DEFINITION and are the parameters relevant to the latest user's query?
+  3. Parameter Value Correctness: Are the parameters values used in the TOOL CALL present or inferred by CONVERSATION and relevant to the latest user's query?
+  4. Potential Value: Is the information this tool call might provide likely to be useful in advancing the conversation or addressing the user expressed or implied needs?
+  5. Context Appropriateness: Does the tool call make sense at this point in the conversation, given what has been discussed so far?
+
+
+# Ratings
+## [Tool Call Accuracy: 0] (Irrelevant)
+**Definition:**
+ 1. The TOOL CALL is not relevant and will not help resolve the user's need.
+ 2. TOOL CALL include parameters values that are not present or inferred from CONVERSATION.
+ 3. TOOL CALL has parameters that is not present in TOOL DEFINITION.
+
+## [Tool Call Accuracy: 1] (Relevant)
+**Definition:**
+ 1. The TOOL CALL is directly relevant and very likely to help resolve the user's need.
+ 2. TOOL CALL include parameters values that are present or inferred from CONVERSATION.
+ 3. TOOL CALL has parameters that is present in TOOL DEFINITION.
+
+# Data
+CONVERSATION : {{query}}
+TOOL CALL: {{tool_call}}
+TOOL DEFINITION: {{tool_definition}}
+
+
+# Tasks
+## Please provide your assessment Score for the previous CONVERSATION , TOOL CALL and TOOL DEFINITION based on the Definitions above. Your output should include the following information:
+- **ThoughtChain**: To improve the reasoning process, think step by step and include a step-by-step explanation of your thought process as you analyze the data based on the definitions. Keep it brief and start your ThoughtChain with "Let's think step by step:".
+- **Explanation**: a very short explanation of why you think the input Data should get that Score.
+- **Score**: based on your previous analysis, provide your Score. The Score you give MUST be a integer score (i.e., "0", "1") based on the levels of the definitions.
+
+
+## Please provide your answers between the tags: <S0>your chain of thoughts</S0>, <S1>your explanation</S1>, <S2>your Score</S2>.
+# Output

azure/ai/evaluation/_evaluators/_ungrounded_attributes/__init__.py

@@ -0,0 +1,5 @@
+from ._ungrounded_attributes import UngroundedAttributesEvaluator
+
+__all__ = [
+    "UngroundedAttributesEvaluator",
+]

azure/ai/evaluation/_evaluators/_ungrounded_attributes/_ungrounded_attributes.py

@@ -0,0 +1,103 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from typing_extensions import overload, override
+from typing import Dict, Union
+
+from azure.ai.evaluation._common._experimental import experimental
+from azure.ai.evaluation._common.constants import EvaluationMetrics
+from azure.ai.evaluation._evaluators._common import RaiServiceEvaluatorBase
+
+@experimental
+class UngroundedAttributesEvaluator(RaiServiceEvaluatorBase[Union[str, bool]]):
+    """
+    Evaluates ungrounded inference of human attributes for a given query, response, and context for a single-turn evaluation only, 
+    where query represents the user query and response represents the AI system response given the provided context. 
+ 
+    Ungrounded Attributes checks for whether a response is first, ungrounded, and checks if it contains information about protected class or 
+    emotional state of a person.
+
+
+    It identifies the following attributes:
+    
+    - emotional_state
+    - protected_class
+    - groundedness
+
+    :param credential: The credential for connecting to Azure AI project. Required
+    :type credential: ~azure.core.credentials.TokenCredential
+    :param azure_ai_project: The scope of the Azure AI project.
+        It contains subscription id, resource group, and project name.
+    :type azure_ai_project: ~azure.ai.evaluation.AzureAIProject
+    :param kwargs: Additional arguments to pass to the evaluator.
+    :type kwargs: Any
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/evaluation_samples_evaluate.py
+            :start-after: [START ungrounded_attributes_evaluator]
+            :end-before: [END ungrounded_attributes_evaluator]
+            :language: python
+            :dedent: 8
+            :caption: Initialize and call a UngroundedAttributesEvaluator with a query, response and context.
+
+    .. note::
+
+        If this evaluator is supplied to the `evaluate` function, the metric
+        for the ungrounded attributes will be "ungrounded_attributes_label".
+    """
+
+    id = "ungrounded_attributes"
+    """Evaluator identifier, experimental and to be used only with evaluation in cloud."""
+
+    @override
+    def __init__(
+        self,
+        credential,
+        azure_ai_project,
+    ):
+        super().__init__(
+            eval_metric=EvaluationMetrics.UNGROUNDED_ATTRIBUTES,
+            azure_ai_project=azure_ai_project,
+            credential=credential,
+        )
+
+    @overload
+    def __call__(
+        self,
+        *,
+        query: str,
+        response: str,
+        context: str,
+    ) -> Dict[str, Union[str, float]]:
+        """Evaluate a given query/response pair and context for ungrounded attributes
+
+        :keyword query: The query to be evaluated.
+        :paramtype query: str
+        :keyword response: The response to be evaluated.
+        :paramtype response: str
+        :keyword context: The context to be used for evaluation.
+        :paramtype context: str
+        :return: The ungrounded attributes label.
+        :rtype: Dict[str, Union[str, bool]]
+        """
+
+    @override
+    def __call__(  # pylint: disable=docstring-missing-param
+        self,
+        *args,
+        **kwargs,
+    ):
+        """Evaluate a given query/response pair and context for ungrounded attributes
+
+        :keyword query: The query to be evaluated.
+        :paramtype query: str
+        :keyword response: The response to be evaluated.
+        :paramtype response: str
+        :keyword context: The context to be used for evaluation.
+        :paramtype context: str
+        :return: The ungrounded attributes label.
+        :rtype: Dict[str, Union[str, bool]]
+        """
+        
+        return super().__call__(*args, **kwargs)

azure/ai/evaluation/_evaluators/_xpia/xpia.py

@@ -43,6 +43,8 @@ class IndirectAttackEvaluator(RaiServiceEvaluatorBase[Union[str, bool]]):
     :param azure_ai_project: The scope of the Azure AI project. It contains subscription id, resource group, and project
         name.
     :type azure_ai_project: ~azure.ai.evaluation.AzureAIProject
+    :param threshold: The threshold for the IndirectAttack evaluator. Default is 0.
+    :type threshold: int
 
     .. admonition:: Example:
 

azure/ai/evaluation/_exceptions.py

@@ -62,15 +62,18 @@ class ErrorTarget(Enum):
     CODE_CLIENT = "CodeClient"
     RAI_CLIENT = "RAIClient"
     COHERENCE_EVALUATOR = "CoherenceEvaluator"
+    COMPLETENESS_EVALUATOR = "CompletenessEvaluator"
     CONTENT_SAFETY_CHAT_EVALUATOR = "ContentSafetyEvaluator"
     ECI_EVALUATOR = "ECIEvaluator"
     F1_EVALUATOR = "F1Evaluator"
     GROUNDEDNESS_EVALUATOR = "GroundednessEvaluator"
     PROTECTED_MATERIAL_EVALUATOR = "ProtectedMaterialEvaluator"
+    INTENT_RESOLUTION_EVALUATOR = "IntentResolutionEvaluator"
     RELEVANCE_EVALUATOR = "RelevanceEvaluator"
     SIMILARITY_EVALUATOR = "SimilarityEvaluator"
     FLUENCY_EVALUATOR = "FluencyEvaluator"
     RETRIEVAL_EVALUATOR = "RetrievalEvaluator"
+    TASK_ADHERENCE_EVALUATOR = "TaskAdherenceEvaluator"
     INDIRECT_ATTACK_EVALUATOR = "IndirectAttackEvaluator"
     INDIRECT_ATTACK_SIMULATOR = "IndirectAttackSimulator"
     ADVERSARIAL_SIMULATOR = "AdversarialSimulator"
@@ -80,6 +83,8 @@ class ErrorTarget(Enum):
     MODELS = "Models"
     UNKNOWN = "Unknown"
     CONVERSATION = "Conversation"
+    TOOL_CALL_ACCURACY_EVALUATOR = "ToolCallAccuracyEvaluator"
+    RED_TEAM = "RedTeam"
 
 
 class EvaluationException(AzureError):

azure/ai/evaluation/_legacy/__init__.py

@@ -0,0 +1,3 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------

azure/ai/evaluation/_legacy/_batch_engine/__init__.py

@@ -0,0 +1,9 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# NOTE: This is a direct port of the bare minimum needed for BatchEngine functionality from
+#       the original Promptflow code. The goal here is expediency, not elegance. As such
+#       parts of this code may be a little "quirky", seem incomplete in places, or contain
+#       longer TODOs comments than usual. In a future code update, large swaths of this code
+#       will be refactored or deleted outright.

azure/ai/evaluation/_legacy/_batch_engine/_config.py

@@ -0,0 +1,45 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from dataclasses import dataclass
+from logging import Logger
+
+from ..._constants import PF_BATCH_TIMEOUT_SEC_DEFAULT
+
+
+@dataclass
+class BatchEngineConfig:
+    """Context for a batch of evaluations. This will contain the configuration,
+    logging, and other needed information."""
+
+    logger: Logger
+    """The logger to use for logging messages."""
+
+    batch_timeout_seconds: int = PF_BATCH_TIMEOUT_SEC_DEFAULT
+    """The maximum amount of time to wait for all evaluations in the batch to complete."""
+
+    run_timeout_seconds: int = 600
+    """The maximum amount of time to wait for an evaluation to run against a single entry
+    in the data input to complete."""
+
+    max_concurrency: int = 10
+    """The maximum number of evaluations to run concurrently."""
+
+    use_async: bool = True
+    """Whether to use asynchronous evaluation."""
+
+    default_num_results: int = 100
+    """The default number of results to return if you don't ask for all results."""
+
+    def __post_init__(self):
+        if self.logger is None:
+            raise ValueError("logger cannot be None")
+        if self.batch_timeout_seconds <= 0:
+            raise ValueError("batch_timeout_seconds must be greater than 0")
+        if self.run_timeout_seconds <= 0:
+            raise ValueError("run_timeout_seconds must be greater than 0")
+        if self.max_concurrency <= 0:
+            raise ValueError("max_concurrency must be greater than 0")
+        if self.default_num_results <= 0:
+            raise ValueError("default_num_results must be greater than 0")

azure/ai/evaluation/_legacy/_batch_engine/_engine.py

@@ -0,0 +1,368 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# This contains code merged together from the following files:
+# promptflow-devkit/promptflow/batch/_batch_engine.py
+# promptflow-devkit/promptflow/_proxy/_python_executor_proxy.py
+# promptflow-core/promptflow/executor/_script_executor.py
+# TODO ralphe: The way this code does batch execution needs to be improved. For now
+#              porting over the code largely as is to remove the Promptflow dependency
+#              as quickly as possible. In phase 2 this code will be heavily refactored.
+
+import re
+import asyncio
+from math import floor
+from asyncio import Semaphore
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Any, Callable, Dict, Final, Generator, Mapping, MutableMapping, Optional, Sequence, Set, Tuple
+from uuid import uuid4
+
+from ._utils import get_int_env_var, get_value_from_path
+from ._status import BatchStatus
+from ._result import BatchResult, BatchRunDetails, BatchRunError, TokenMetrics
+from ._run_storage import AbstractRunStorage, NoOpRunStorage
+from ._logging import log_progress, NodeLogManager
+from ..._exceptions import ErrorBlame
+from ._exceptions import (
+    BatchEngineCanceledError,
+    BatchEngineError,
+    BatchEngineRunFailedError,
+    BatchEngineTimeoutError,
+    BatchEngineValidationError,
+)
+from ._utils_deprecated import (
+    async_run_allowing_running_loop,
+    convert_eager_flow_output_to_dict,
+)
+
+
+MAX_WORKER_COUNT: Final[int] = 10
+KEYWORD_PATTERN: Final = re.compile(r"^\${([^{}]+)}$")
+
+
+class BatchEngine:
+    """This class is used to execute flows in batch mode"""
+
+    def __init__(
+        self,
+        executor: Callable,
+        *,
+        storage: Optional[AbstractRunStorage] = None,
+        batch_timeout_sec: Optional[int] = None,
+        line_timeout_sec: Optional[int] = None,
+        max_worker_count: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        """Create a new batch engine instance
+
+        :param Callable executor: The executor to run the flow
+        :param Optional[AbstractRunStorage] storage: The storage to store execution results
+        :param Optional[int] batch_timeout_sec: The timeout of batch run in seconds
+        :param Optional[int] line_timeout_sec: The timeout of each line in seconds
+        :param Optional[int] max_worker_count: The concurrency limit of batch run
+        :param kwargs: The keyword arguments related to creating the executor proxy class
+        :type kwargs: Any
+        """
+
+        self._executor = executor
+        # self._working_dir = working_dir
+
+        # self._is_eager_flow = True
+        # self._is_prompty_flow = False
+        # self._program_language = FlowLanguage.Python
+        # self._message_format = MessageFormatType.BASIC
+        # self._multimedia_processor = MultimediaProcessor.create(self._message_format)
+        # self._connections = {}
+
+        self._storage: AbstractRunStorage = storage or NoOpRunStorage()
+
+        # TODO ralphe: Consume these from the batch context/config instead of from
+        #              kwargs or (even worse) environment variables
+        # self._batch_use_async = kwargs.get("batch_use_async", True)
+        self._batch_timeout_sec = batch_timeout_sec or get_int_env_var("PF_BATCH_TIMEOUT_SEC")
+        self._line_timeout_sec = line_timeout_sec or get_int_env_var("PF_LINE_TIMEOUT_SEC", 600)
+        self._max_worker_count = max_worker_count or get_int_env_var("PF_WORKER_COUNT") or MAX_WORKER_COUNT
+        # update kwargs with worker_count and line_timeout_sec
+        kwargs.update({"worker_count": self._max_worker_count, "line_timeout_sec": self._line_timeout_sec})
+
+        self._is_canceled: bool = False
+        self._kwargs: Mapping[str, Any] = kwargs
+        # self._init_kwargs: Mapping[str, Any] = init_kwargs or {}
+
+    def run(
+        self,
+        data: Sequence[Mapping[str, Any]],
+        column_mapping: Mapping[str, str],
+        *,
+        id: Optional[str] = None,
+        max_lines: Optional[int] = None,
+    ) -> BatchResult:
+        if not data:
+            raise BatchEngineValidationError("Please provide a non-empty data mapping.")
+        if not column_mapping:
+            raise BatchEngineValidationError("The column mapping is required.")
+
+        start_time = datetime.now(timezone.utc)
+
+        batch_inputs = self._apply_column_mapping(data, column_mapping, max_lines)
+        if not batch_inputs or all(len(data) == 0 for data in batch_inputs):
+            raise BatchEngineValidationError("No data to process.")
+
+        try:
+            id = id or str(uuid4())
+
+            result: BatchResult = async_run_allowing_running_loop(self._exec_in_task, id, batch_inputs, start_time)
+
+            return result
+        except Exception as ex:
+            raise BatchEngineError(
+                "Unexpected error while running the batch run.", blame=ErrorBlame.SYSTEM_ERROR
+            ) from ex
+
+    def cancel(self):
+        # TODO ralphe: Make sure this works
+        self._is_canceled = True
+
+    @staticmethod
+    def _apply_column_mapping(
+        data: Sequence[Mapping[str, Any]],
+        column_mapping: Mapping[str, str],
+        max_lines: Optional[int],
+    ) -> Sequence[Mapping[str, str]]:
+        data = data[:max_lines] if max_lines else data
+
+        inputs: Sequence[Mapping[str, Any]] = []
+        line: int = 0
+
+        for input in data:
+            line += 1
+            mapped: Dict[str, Any] = {}
+            missing_inputs: Set[str] = set()
+
+            for key, value in column_mapping.items():
+                if not isinstance(value, str):
+                    # All non-string values are literal values.
+                    mapped[key] = value
+                    continue
+
+                match: Optional[re.Match[str]] = re.search(KEYWORD_PATTERN, value)
+                if match is None:
+                    # Literal string value value
+                    mapped[key] = value
+                    continue
+
+                dict_path = match.group(1)
+                found, value = get_value_from_path(dict_path, input)
+                if found:
+                    mapped[key] = value
+                else:
+                    missing_inputs.add(dict_path)
+
+            if missing_inputs:
+                missing = ", ".join(missing_inputs)
+                raise BatchEngineValidationError(f"Missing inputs for line {line}: '{missing}'")
+
+            inputs.append(mapped)
+
+        return inputs
+
+    async def _exec_in_task(
+        self, run_id: str, batch_inputs: Sequence[Mapping[str, Any]], start_time: datetime
+    ) -> BatchResult:
+        # Since the batch execution is not guaranteed to be completed in the same order
+        # as the inputs, we keep track of these in a mapping from index to result
+        results: Dict[int, BatchRunDetails] = {}
+        status: BatchStatus = BatchStatus.Completed
+        error: Optional[Exception] = None
+
+        task = asyncio.create_task(self._exec_batch(run_id, batch_inputs, start_time, results))
+
+        while not task.done():
+            # check whether the task is completed or canceled every 1s
+            await asyncio.sleep(1)
+            if self._is_canceled:
+                task.cancel()
+                # use current completed line results and aggregation results to create a BatchResult
+                status = BatchStatus.Canceled
+                error = BatchEngineCanceledError("The batch run is canceled by user.")
+                break
+            elif self._batch_timeout_expired(start_time):
+                task.cancel()
+                status = BatchStatus.Failed
+                error = BatchEngineTimeoutError(
+                    f"The batch run failed due to timeout [{self._batch_timeout_sec}s]. "
+                    f"Please adjust the timeout to a higher value."
+                )
+                break
+
+        end_time = datetime.now(timezone.utc)
+        metrics = TokenMetrics(0, 0, 0)
+        failed_lines: int = 0
+
+        # generate the details in the same order as the inputs and fill in the missing results
+        # with a failed status
+        result_details = [
+            (
+                results[i]
+                if i in results
+                else BatchRunDetails(
+                    id=BatchRunDetails.create_id(run_id, i),
+                    status=BatchStatus.Failed,
+                    result=None,
+                    start_time=None,
+                    end_time=None,
+                    tokens=TokenMetrics(0, 0, 0),
+                    error=BatchRunError("The line run is not completed.", None),
+                )
+            )
+            for i in range(len(batch_inputs))
+        ]
+
+        for line_result in result_details:
+            # Indicate the worst status of the batch run. This works because
+            # canceled and failed have a higher value than completed.
+            status = max(status, line_result.status)
+            if BatchStatus.is_failed(line_result.status):
+                failed_lines += 1
+            if line_result.tokens:
+                metrics.prompt_tokens += line_result.tokens.prompt_tokens
+                metrics.completion_tokens += line_result.tokens.completion_tokens
+                metrics.total_tokens += line_result.tokens.total_tokens
+
+        if failed_lines and not error:
+            error = BatchEngineRunFailedError(
+                str(floor(failed_lines / len(batch_inputs) * 100)) + f"% of the batch run failed."
+            )
+
+        return BatchResult(
+            status=status,
+            total_lines=len(batch_inputs),
+            failed_lines=failed_lines,
+            start_time=start_time,
+            end_time=end_time,
+            tokens=metrics,
+            details=result_details,
+            error=error,
+        )
+
+    async def _exec_batch(
+        self,
+        run_id: str,
+        batch_inputs: Sequence[Mapping[str, Any]],
+        start_time: datetime,
+        results: MutableMapping[int, BatchRunDetails],
+    ) -> None:
+        semaphore: Semaphore = Semaphore(self._max_worker_count)
+
+        # TODO ralphe: This async code needs to refactored to use e.g. asyncio.gather, or
+        #              asyncio.as_completed.
+        # TODO ralphe: This code needs to handle cancellation better
+        async def create_under_semaphore(index: int, inputs: Mapping[str, Any]):
+            async with semaphore:
+                return await self._exec_line_async(run_id, inputs, index)
+
+        pending = [
+            asyncio.create_task(create_under_semaphore(index, inputs)) for index, inputs in enumerate(batch_inputs)
+        ]
+
+        total_lines: int = len(batch_inputs)
+        completed_lines: int = 0
+        while completed_lines < total_lines:
+            # TODO ralphe: Fix this code so it doesn't re-order the outputs
+            # wait for any task to complete
+            done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+            completed_line_results = [task.result() for task in done]
+            # persist node run infos and flow run info in line result to storage
+            self._persist_run_info([result for _, result in completed_line_results])
+            results.update({index: result for index, result in completed_line_results})
+            # update the progress log
+            completed_lines += len(completed_line_results)
+            log_progress(
+                run_start_time=start_time,
+                total_count=total_lines,
+                current_count=completed_lines,
+                # TODO ralphe: set logger to use here
+            )
+
+    async def _exec_line_async(
+        self,
+        run_id: str,
+        inputs: Mapping[str, Any],
+        index: int,
+    ) -> Tuple[int, BatchRunDetails]:
+        with self._exec_line_context(run_id, index):
+            details: BatchRunDetails = BatchRunDetails(
+                id=f"{run_id}_{index}",
+                status=BatchStatus.NotStarted,
+                result=None,
+                start_time=datetime.now(timezone.utc),
+                end_time=None,
+                tokens=TokenMetrics(0, 0, 0),
+                error=None,
+            )
+
+            try:
+                # TODO ralphe: Handle line timeouts here
+                output: Any = await self._executor(**inputs)
+                details.status = BatchStatus.Completed
+                details.result = convert_eager_flow_output_to_dict(output)
+
+                # TODO figure out how to get the token metrics here
+            except Exception as ex:
+                details.status = BatchStatus.Failed
+                details.error = BatchRunError(
+                    f"Error while evaluating single input: {ex.__class__.__name__}: {str(ex)}", ex
+                )
+            finally:
+                details.end_time = datetime.now(timezone.utc)
+
+        return index, details
+
+    def _persist_run_info(self, line_results: Sequence[BatchRunDetails]):
+        # TODO ralphe: implement?
+        pass
+
+    def _batch_timeout_expired(self, start_time: datetime) -> bool:
+        if self._batch_timeout_sec is None:
+            return False
+        return (datetime.now(timezone.utc) - start_time).total_seconds() > self._batch_timeout_sec
+
+    @contextmanager
+    def _exec_line_context(self, run_id: str, line_number: int) -> Generator[None, Any, None]:
+        # TODO ralphe: Do proper tracing and logging here
+        log_manager = NodeLogManager()
+        log_manager.set_node_context(run_id, "Flex", line_number)
+        with log_manager, self._update_operation_context(run_id, line_number):
+            yield
+
+    @contextmanager
+    def _update_operation_context(self, run_id: str, line_number: int) -> Generator[None, Any, None]:
+        # operation_context = OperationContext.get_instance()
+        # original_context = operation_context.copy()
+        # original_mode = operation_context.get("run_mode", RunMode.Test.name)
+        # values_for_context = {"flow_id": self._flow_id, "root_run_id": run_id}
+        # if original_mode == RunMode.Batch.name:
+        #     values_for_otel = {
+        #         "batch_run_id": run_id,
+        #         "line_number": line_number,
+        #     }
+        # else:
+        #     values_for_otel = {"line_run_id": run_id}
+        # try:
+        #     append_promptflow_package_ua(operation_context)
+        #     operation_context.set_execution_target(execution_target=self._execution_target)
+        #     operation_context.set_default_tracing_keys(DEFAULT_TRACING_KEYS)
+        #     operation_context.run_mode = original_mode
+        #     operation_context.update(values_for_context)
+        #     for k, v in values_for_otel.items():
+        #         operation_context._add_otel_attributes(k, v)
+        #     # Inject OpenAI API to make sure traces and headers injection works and
+        #     # update OpenAI API configs from environment variables.
+        #     inject_openai_api()
+        yield
+
+        # finally:
+        #     OperationContext.set_instance(original_context)