genesis-flow 1.0.0__py3-none-any.whl

mlflow/genai/scorers/scorer_utils.py

@@ -0,0 +1,138 @@
+# This file contains utility functions for scorer functionality.
+
+import ast
+import inspect
+import logging
+import re
+from textwrap import dedent
+from typing import Any, Callable, Optional
+
+_logger = logging.getLogger(__name__)
+
+
+# FunctionBodyExtractor class is forked from https://github.com/unitycatalog/unitycatalog/blob/20dd3820be332ac04deec4e063099fb863eb3392/ai/core/src/unitycatalog/ai/core/utils/callable_utils.py
+class FunctionBodyExtractor(ast.NodeVisitor):
+    """
+    AST NodeVisitor class to extract the body of a function.
+    """
+
+    def __init__(self, func_name: str, source_code: str):
+        self.func_name = func_name
+        self.source_code = source_code
+        self.function_body = ""
+        self.indent_unit = 4
+        self.found = False
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        if not self.found and node.name == self.func_name:
+            self.found = True
+            self.extract_body(node)
+
+    def extract_body(self, node: ast.FunctionDef):
+        body = node.body
+        # Skip the docstring
+        if (
+            body
+            and isinstance(body[0], ast.Expr)
+            and isinstance(body[0].value, ast.Constant)
+            and isinstance(body[0].value.value, str)
+        ):
+            body = body[1:]
+
+        if not body:
+            return
+
+        start_lineno = body[0].lineno
+        end_lineno = body[-1].end_lineno
+
+        source_lines = self.source_code.splitlines(keepends=True)
+        function_body_lines = source_lines[start_lineno - 1 : end_lineno]
+
+        self.function_body = dedent("".join(function_body_lines)).rstrip("\n")
+
+        indents = [stmt.col_offset for stmt in body if stmt.col_offset is not None]
+        if indents:
+            self.indent_unit = min(indents)
+
+
+# extract_function_body function is forked from https://github.com/unitycatalog/unitycatalog/blob/20dd3820be332ac04deec4e063099fb863eb3392/ai/core/src/unitycatalog/ai/core/utils/callable_utils.py
+def extract_function_body(func: Callable[..., Any]) -> tuple[str, int]:
+    """
+    Extracts the body of a function as a string without the signature or docstring,
+    dedents the code, and returns the indentation unit used in the function (e.g., 2 or 4 spaces).
+    """
+    source_lines, _ = inspect.getsourcelines(func)
+    dedented_source = dedent("".join(source_lines))
+    func_name = func.__name__
+
+    extractor = FunctionBodyExtractor(func_name, dedented_source)
+    parsed_source = ast.parse(dedented_source)
+    extractor.visit(parsed_source)
+
+    return extractor.function_body, extractor.indent_unit
+
+
+def recreate_function(source: str, signature: str, func_name: str) -> Optional[Callable[..., Any]]:
+    """
+    Recreate a function from its source code, signature, and name.
+
+    Args:
+        source: The function body source code.
+        signature: The function signature string (e.g., "(inputs, outputs)").
+        func_name: The name of the function.
+
+    Returns:
+        The recreated function or None if recreation failed.
+    """
+    try:
+        # Parse the signature to build the function definition
+        sig_match = re.match(r"\((.*?)\)", signature)
+        if not sig_match:
+            return None
+
+        params_str = sig_match.group(1).strip()
+
+        # Build the function definition
+        func_def = f"def {func_name}({params_str}):\n"
+        # Indent the source code
+        indented_source = "\n".join(f"    {line}" for line in source.split("\n"))
+        func_def += indented_source
+
+        # Create a namespace with common MLflow imports that scorer functions might use
+        import_namespace = {}
+
+        # Import commonly used MLflow classes
+        try:
+            from mlflow.entities import (
+                Assessment,
+                AssessmentError,
+                AssessmentSource,
+                AssessmentSourceType,
+                Feedback,
+            )
+            from mlflow.genai.judges import CategoricalRating
+
+            import_namespace.update(
+                {
+                    "Feedback": Feedback,
+                    "Assessment": Assessment,
+                    "AssessmentSource": AssessmentSource,
+                    "AssessmentError": AssessmentError,
+                    "AssessmentSourceType": AssessmentSourceType,
+                    "CategoricalRating": CategoricalRating,
+                }
+            )
+        except ImportError:
+            pass  # Some imports might not be available in all contexts
+
+        local_namespace = {}
+
+        # Execute the function definition with MLflow imports available
+        exec(func_def, import_namespace, local_namespace)
+
+        # Return the recreated function
+        return local_namespace[func_name]
+
+    except Exception:
+        _logger.warning(f"Failed to recreate function '{func_name}' from serialized source code")
+        return None

mlflow/genai/scorers/validation.py

@@ -0,0 +1,165 @@
+import logging
+from collections import defaultdict
+from typing import Any, Callable, Optional
+
+from mlflow.exceptions import MlflowException
+from mlflow.genai.scorers.base import Scorer
+from mlflow.genai.scorers.builtin_scorers import (
+    BuiltInScorer,
+    MissingColumnsException,
+    get_all_scorers,
+)
+
+try:
+    # `pandas` is not required for `mlflow-skinny`.
+    import pandas as pd
+except ImportError:
+    pass
+
+_logger = logging.getLogger(__name__)
+
+
+def validate_scorers(scorers: list[Any]) -> list[Scorer]:
+    """
+    Validate a list of specified scorers.
+
+    Args:
+        scorers: A list of scorers to validate.
+
+    Returns:
+        A list of valid scorers.
+    """
+    from databricks.rag_eval.evaluation.metrics import Metric
+
+    if not isinstance(scorers, list) or len(scorers) == 0:
+        raise MlflowException.invalid_parameter_value(
+            "The `scorers` argument must be a list of scorers with at least one scorer. "
+            "If you are unsure about which scorer to use, you can specify "
+            "`scorers=mlflow.genai.scorers.get_all_scorers()` to jump start with all "
+            "available built-in scorers."
+        )
+
+    valid_scorers = []
+    legacy_metrics = []
+
+    for scorer in scorers:
+        if isinstance(scorer, Scorer):
+            valid_scorers.append(scorer)
+        elif isinstance(scorer, Metric):
+            legacy_metrics.append(scorer)
+            valid_scorers.append(scorer)
+        else:
+            # Show helpful error message for common mistakes
+            if isinstance(scorer, list) and (scorer == get_all_scorers()):
+                # Common mistake 1: scorers=[get_all_scorers()]
+                if len(scorers) == 1:
+                    hint = (
+                        "\nHint: Use `scorers=get_all_scorers()` to pass all "
+                        "builtin scorers at once."
+                    )
+                # Common mistake 2: scorers=[get_all_scorers(), scorer1, scorer2]
+                elif len(scorer) > 1:
+                    hint = (
+                        "\nHint: Use `scorers=[*get_all_scorers(), scorer1, scorer2]` to pass "
+                        "all builtin scorers at once along with your custom scorers."
+                    )
+            # Common mistake 3: scorers=[RetrievalRelevance, Correctness]
+            elif isinstance(scorer, type) and issubclass(scorer, BuiltInScorer):
+                hint = (
+                    "\nHint: It looks like you passed a scorer class instead of an instance. "
+                    f"Correct way to pass scorers is `scorers=[{scorer.__name__}()]`."
+                )
+            else:
+                hint = ""
+
+            raise MlflowException.invalid_parameter_value(
+                f"The `scorers` argument must be a list of scorers. The specified "
+                f"list contains an invalid item with type: {type(scorer).__name__}."
+                f"{hint}"
+            )
+
+    if legacy_metrics:
+        legacy_metric_names = [metric.name for metric in legacy_metrics]
+        _logger.warning(
+            f"Scorers {legacy_metric_names} are legacy metrics and will soon be deprecated "
+            "in future releases. Please use the builtin scorers defined in `mlflow.genai.scorers` "
+            "or custom scorers defined with the @scorer decorator instead."
+        )
+
+    return valid_scorers
+
+
+def valid_data_for_builtin_scorers(
+    data: "pd.DataFrame",
+    builtin_scorers: list[BuiltInScorer],
+    predict_fn: Optional[Callable[..., Any]] = None,
+) -> None:
+    """
+    Validate that the required columns are present in the data for running the builtin scorers.
+
+    Args:
+        data: The data to validate. This must be a pandas DataFrame converted to
+            the legacy evaluation set schema via `_convert_to_legacy_eval_set`.
+        builtin_scorers: The list of builtin scorers to validate the data for.
+        predict_fn: The predict function to validate the data for.
+    """
+    input_columns = set(data.columns.tolist())
+
+    # Revert the replacement of "inputs"->"request" and "outputs"->"response"
+    # in the upstream processing.
+    if "request" in input_columns:
+        input_columns.remove("request")
+        input_columns.add("inputs")
+    if "response" in input_columns:
+        input_columns.remove("response")
+        input_columns.add("outputs")
+
+    if predict_fn is not None:
+        # If the predict function is provided, the data doesn't need to
+        # contain the "outputs" column.
+        input_columns.add("outputs")
+
+    if "trace" in input_columns:
+        # Inputs and outputs are inferred from the trace.
+        input_columns |= {"inputs", "outputs"}
+
+    if predict_fn is not None:
+        input_columns |= {"trace"}
+
+    # Explode keys in the "expectations" column for easier processing.
+    if "expectations" in input_columns:
+        for value in data["expectations"].values:
+            if pd.isna(value):
+                continue
+            if not isinstance(value, dict):
+                raise MlflowException.invalid_parameter_value(
+                    "The 'expectations' column must be a dictionary of each expectation name "
+                    "to its value. For example, `{'expected_response': 'answer to the question'}`."
+                )
+            for k in value:
+                input_columns.add(f"expectations/{k}")
+
+    # Missing column -> list of scorers that require the column.
+    missing_col_to_scorers = defaultdict(list)
+    for scorer in builtin_scorers:
+        try:
+            scorer.validate_columns(input_columns)
+        except MissingColumnsException as e:
+            for col in e.missing_columns:
+                missing_col_to_scorers[col].append(scorer.name)
+
+    if missing_col_to_scorers:
+        msg = (
+            "The input data is missing following columns that are required by the specified "
+            "scorers. The results will be null for those scorers."
+        )
+        for col, scorers in missing_col_to_scorers.items():
+            if col.startswith("expectations/"):
+                col = col.replace("expectations/", "")
+                msg += (
+                    f"\n - `{col}` field in `expectations` column "
+                    f"is required by [{', '.join(scorers)}]."
+                )
+            else:
+                msg += f"\n - `{col}` column is required by [{', '.join(scorers)}]."
+        _logger.info(msg)

mlflow/genai/utils/data_validation.py

@@ -0,0 +1,146 @@
+import inspect
+import logging
+from typing import Any, Callable
+
+from mlflow.exceptions import MlflowException
+from mlflow.tracing.provider import trace_disabled
+
+_logger = logging.getLogger(__name__)
+
+
+def check_model_prediction(predict_fn: Callable[..., Any], sample_input: Any):
+    """
+    Validate if the predict function executes properly with the provided input.
+
+    Args:
+        predict_fn: The predict function to be evaluated.
+        sample_input: A sample input to the model.
+    """
+    _logger.info("Testing model prediction with the first sample in the dataset.")
+
+    # Wrap the function to add a decorator for disabling tracing
+    @trace_disabled
+    def _check():
+        predict_fn(**sample_input)
+
+    try:
+        _check()
+    except Exception as e:
+        # Check input format and raise friendly message for typical error patterns
+        _validate_function_and_input_compatibility(predict_fn, sample_input, e)
+        _logger.debug(f"Failed to run predict_fn with input: {sample_input}", exc_info=True)
+
+
+def _validate_function_and_input_compatibility(
+    predict_fn: Callable[..., Any], sample_input: dict[str, Any], e: Exception
+) -> Callable[..., Any]:
+    """
+    Validate the data format in the input column against the predict_fn.
+
+    The input column must contain a dictionary of field names and values. When the
+    predict_fn is provided, the field names must match the arguments of the predict_fn.
+    """
+    params = inspect.signature(predict_fn).parameters
+    if not params:
+        raise MlflowException.invalid_parameter_value(
+            "`predict_fn` must accept at least one argument."
+        ) from e
+
+    # Check for *args-style parameters which aren't supported
+    _validate_no_var_args(params, e)
+
+    # Check if input keys match function parameters
+    _validate_input_keys_match_function_params(params, sample_input.keys(), e)
+
+    # For other errors, show a generic error message
+    raise MlflowException.invalid_parameter_value(
+        "Failed to run the prediction function specified in the `predict_fn` "
+        f"parameter. Input: {sample_input}. Error: {e}\n\n"
+    ) from e
+
+
+def _has_variable_positional_arguments(params: inspect.Signature) -> bool:
+    """Check if the function has variable positional arguments."""
+    return any(p.kind == inspect.Parameter.VAR_POSITIONAL for p in params.values())
+
+
+def _validate_no_var_args(params: inspect.Signature, e: Exception):
+    if not any(p.kind == inspect.Parameter.VAR_POSITIONAL for p in params.values()):
+        return
+
+    """Raise an error for functions using *args which aren't supported."""
+    code_sample = """```python
+def predict_fn(param1, param2):
+    # Invoke the original predict function with positional arguments
+    return fn(param1, param2)
+
+data = [
+    {
+        "inputs": {
+            "param1": "value1",
+            "param2": "value2",
+        }
+    }
+]
+
+mlflow.genai.evaluate(predict_fn=predict_fn, data=data, ...)
+```
+"""
+
+    raise MlflowException.invalid_parameter_value(
+        "The `predict_fn` has dynamic positional arguments (e.g. `*args`), "
+        "so it cannot be used as a `predict_fn`. Please wrap it into another "
+        "function that accepts explicit keyword arguments.\n"
+        f"Example:\n\n{code_sample}\n"
+    ) from e
+
+
+def _validate_input_keys_match_function_params(
+    params: inspect.Signature,
+    input_keys: list[str],
+    e: Exception,
+):
+    if _has_required_keyword_arguments(params, input_keys):
+        return
+
+    """Raise an error when input keys don't match function parameters."""
+    param_names = list(params.keys())
+    input_example = {arg: f"value{i + 1}" for i, arg in enumerate(param_names[:3])}
+
+    if len(param_names) > 3:
+        input_example["..."] = "..."
+
+    code_sample = "\n".join(
+        [
+            "```python",
+            "data = [",
+            "    {",
+            '        "inputs": {',
+            *(f'            "{k}": "{v}",' for k, v in input_example.items()),
+            "        }",
+            "    }",
+            "]",
+            "```",
+        ]
+    )
+
+    raise MlflowException.invalid_parameter_value(
+        "The `inputs` column must be a dictionary with the parameter names of "
+        f"the `predict_fn` as keys. It seems the specified keys do not match "
+        f"with the `predict_fn`'s arguments. Correct example:\n\n{code_sample}"
+    ) from e
+
+
+def _has_required_keyword_arguments(params: inspect.Signature, required_args: list[str]) -> bool:
+    """Check if the function accepts the specified keyword arguments."""
+    func_args = []
+
+    for name, param in params.items():
+        # If the function has **kwargs, it accepts all keyword arguments
+        if param.kind == inspect.Parameter.VAR_KEYWORD:
+            return True
+
+        func_args.append(name)
+
+    # Required argument must be a subset of the function's arguments
+    return set(required_args) <= set(func_args)

mlflow/genai/utils/enum_utils.py

@@ -0,0 +1,23 @@
+from enum import Enum, EnumMeta
+
+
+class MetaEnum(EnumMeta):
+    """Metaclass for Enum classes that allows to check if a value is a valid member of the Enum."""
+
+    def __contains__(cls, item):
+        try:
+            cls(item)
+        except ValueError:
+            return False
+        return True
+
+
+class StrEnum(str, Enum, metaclass=MetaEnum):
+    def __str__(self):
+        """Return the string representation of the enum using its value."""
+        return self.value
+
+    @classmethod
+    def values(cls) -> list[str]:
+        """Return a list of all string values of the Enum."""
+        return [str(member) for member in cls]

mlflow/genai/utils/trace_utils.py

@@ -0,0 +1,211 @@
+import logging
+from typing import Any, Callable, Optional
+
+from opentelemetry.trace import NoOpTracer
+
+import mlflow
+from mlflow.entities.span import Span, SpanType
+from mlflow.entities.trace import Trace
+from mlflow.genai.utils.data_validation import check_model_prediction
+from mlflow.tracing.constant import TraceTagKey
+from mlflow.tracing.display.display_handler import IPythonTraceDisplayHandler
+from mlflow.tracking.client import MlflowClient
+
+_logger = logging.getLogger(__name__)
+
+
+def convert_predict_fn(predict_fn: Callable[..., Any], sample_input: Any) -> Callable[..., Any]:
+    """
+    Check the predict_fn is callable and add trace decorator if it is not already traced.
+    """
+    with NoOpTracerPatcher() as counter:
+        check_model_prediction(predict_fn, sample_input)
+
+    if counter.count == 0:
+        predict_fn = mlflow.trace(predict_fn)
+
+    # Wrap the prediction function to unwrap the inputs dictionary into keyword arguments.
+    return lambda request: predict_fn(**request)
+
+
+class NoOpTracerPatcher:
+    """
+    A context manager to count the number of times NoOpTracer's start_span is called.
+
+    The check is done in the following steps so it doesn't have any side effects:
+    1. Disable tracing.
+    2. Patch the NoOpTracer.start_span method to count the number of times it is called.
+        NoOpTracer is used when tracing is disabled.
+    3. Call the predict function with the sample input.
+    4. Restore the original NoOpTracer.start_span method and re-enable tracing.
+
+
+    WARNING: This function is not thread-safe. We do not provide support for running
+        `mlflow.genai.evaluate` in multi-threaded environments.`
+    """
+
+    def __init__(self):
+        self.count = 0
+
+    def __enter__(self):
+        self.original = NoOpTracer.start_span
+
+        def _patched_start_span(_self, *args, **kwargs):
+            self.count += 1
+            return self.original(_self, *args, **kwargs)
+
+        NoOpTracer.start_span = _patched_start_span
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        NoOpTracer.start_span = self.original
+
+
+def parse_inputs_to_str(inputs: Any) -> str:
+    """Parse the inputs to a request string compatible with the judges API"""
+    from databricks.rag_eval.utils import input_output_utils
+
+    return input_output_utils.request_to_string(inputs)
+
+
+def parse_output_to_str(output: Any) -> str:
+    """Parse the output to a string compatible with the judges API"""
+    from databricks.rag_eval.utils import input_output_utils
+
+    return input_output_utils.response_to_string(output)
+
+
+def extract_retrieval_context_from_trace(trace: Optional[Trace]) -> dict[str, list[Any]]:
+    """
+    Extract the retrieval context from the trace.
+    Only consider the last retrieval span in the trace if there are multiple retrieval spans.
+    If the trace does not have a retrieval span, return None.
+    ⚠️ Warning: Please make sure to not throw exception. If fails, return None.
+    """
+    if trace is None or trace.data is None:
+        return {}
+
+    # Only consider the top-level retrieval spans
+    top_level_retrieval_spans = _get_top_level_retrieval_spans(trace)
+    if len(top_level_retrieval_spans) == 0:
+        return {}
+
+    retrieved = {}  # span_id -> list of documents
+
+    for retrieval_span in top_level_retrieval_spans:
+        try:
+            contexts = [_parse_chunk(chunk) for chunk in retrieval_span.outputs or []]
+            retrieved[retrieval_span.span_id] = [c for c in contexts if c is not None]
+        except Exception as e:
+            _logger.debug(
+                f"Fail to get retrieval context from span: {retrieval_span}. Error: {e!r}"
+            )
+
+    return retrieved
+
+
+def _get_top_level_retrieval_spans(trace: Trace) -> list[Span]:
+    """
+    Get the top-level retrieval spans in the trace.
+    Top-level retrieval spans are retrieval spans that are not children of other retrieval spans.
+    For example, given the following spans:
+    - Span A (Chain)
+      - Span B (Retriever)
+        - Span C (Retriever)
+      - Span D (Retriever)
+        - Span E (LLM)
+          - Span F (Retriever)
+    Span B and Span D are top-level retrieval spans.
+    Span C and Span F are NOT top-level because they are children of other retrieval spans.
+    """
+    top_level_retrieval_spans = []
+    # Cache span_id -> span mapping for fast lookup
+    all_spans = {span.span_id: span for span in trace.data.spans}
+    for span in trace.search_spans(span_type=SpanType.RETRIEVER):
+        # Check if this span is a child of another retrieval span
+        parent_id = span.parent_id
+        while parent_id:
+            parent_span = all_spans.get(parent_id)
+            if not parent_span:
+                # Malformed trace
+                _logger.debug(
+                    f"Malformed trace: span {span} has parent span ID {parent_id}, "
+                    "but the parent span is not found in the trace."
+                )
+                break
+
+            if parent_span.span_type == SpanType.RETRIEVER:
+                # This span is a child of another retrieval span
+                break
+
+            parent_id = parent_span.parent_id
+        else:
+            # If the loop completes without breaking, this is a top-level span
+            top_level_retrieval_spans.append(span)
+
+    return top_level_retrieval_spans
+
+
+def _parse_chunk(chunk: Any) -> Optional[dict[str, Any]]:
+    if not isinstance(chunk, dict):
+        return None
+
+    doc = {"content": chunk.get("page_content")}
+    if doc_uri := chunk.get("metadata", {}).get("doc_uri"):
+        doc["doc_uri"] = doc_uri
+    return doc
+
+
+def clean_up_extra_traces(run_id: str, start_time_ms: int):
+    """
+    Clean up noisy traces generated outside predict function.
+
+    Evaluation run should only contain traces that is being evaluated or generated by the predict
+    function. If not, the result will not show the correct list of traces.
+    Sometimes, there are extra traces generated during the evaluation, for example, custom scorer
+    code might generate traces. This function cleans up those noisy traces.
+
+    TODO: This is not a fundamental solution. Ideally, evaluation result should be able to render
+        correct result even if there are extra traces in the run.
+
+    Args:
+        run_id: The ID of the run to clean up.
+        start_time_ms: The start time of the evaluation in milliseconds.
+    """
+    from mlflow.tracking.fluent import _get_experiment_id
+
+    try:
+        # Search for all traces generated during evaluation
+        traces = mlflow.search_traces(
+            run_id=run_id,
+            # Not download spans for efficiency
+            include_spans=False,
+            # Limit to traces generated after evaluation time to ensure we will not
+            # delete traces generated before evaluation.
+            filter_string=f"trace.timestamp >= {start_time_ms}",
+            return_type="list",
+        )
+        extra_trace_ids = [
+            # Traces from predict function should always have the EVAL_REQUEST_ID tag
+            trace.info.trace_id
+            for trace in traces
+            if TraceTagKey.EVAL_REQUEST_ID not in trace.info.tags
+        ]
+        if extra_trace_ids:
+            _logger.debug(
+                f"Found {len(extra_trace_ids)} extra traces generated during evaluation run. "
+                "Deleting them."
+            )
+            MlflowClient().delete_traces(
+                experiment_id=_get_experiment_id(), trace_ids=extra_trace_ids
+            )
+            # Avoid displaying the deleted trace in notebook cell output
+            for trace_id in extra_trace_ids:
+                IPythonTraceDisplayHandler.get_instance().traces_to_display.pop(trace_id, None)
+        else:
+            _logger.debug("No extra traces found during evaluation run.")
+    except Exception as e:
+        _logger.warning(
+            f"Failed to clean up extra traces generated during evaluation. The "
+            f"result page might not show the correct list of traces. Error: {e}"
+        )