Check function signature instead of output for auto-tracing retrievers (unitycatalog#861)

**PR Checklist**

- [x] A description of the changes is added to the description of this
PR.
- [x] If there is a related issue, make sure it is linked to this PR.
- [x] If you've fixed a bug or added code that should be tested, add
tests!
- [ ] If you've added or modified a feature, documentation in `docs` is
updated
-> documenting this in docs for Agent Tools
https://github.com/databricks-eng/docs/pull/19358

**Description of changes**
Following up from this
[comment](unitycatalog#821 (comment)),
we want to modify auto-tracing for retrievers to check the function
signature instead of the output format so that we can capture errors
that occur for retrievers as well. Our recommended way for querying
vector search in the
[documentation](https://docs.databricks.com/en/generative-ai/agent-framework/unstructured-retrieval-tools.html#vector-search-retriever-tool-with-unity-catalog-functions)
has data type `TABLE_TYPE` and we can match on this signature.

Manual testing:
- Langchain
<img width="1307" alt="Screenshot 2025-01-27 at 4 20 28 PM"
src="https://github.com/user-attachments/assets/2765a580-67c5-4de0-9526-03097ec35cca"
/>
<img width="1309" alt="Screenshot 2025-01-27 at 4 21 16 PM"
src="https://github.com/user-attachments/assets/cef0bf0b-6af7-4a96-b721-f5cdf2611daf"
/>

- CrewAI
<img width="1301" alt="Screenshot 2025-01-27 at 4 53 10 PM"
src="https://github.com/user-attachments/assets/3b60b494-3dd0-40e4-9cac-5674b7790b55"
/>
<img width="1306" alt="Screenshot 2025-01-27 at 5 21 02 PM"
src="https://github.com/user-attachments/assets/533d0e12-361e-4e01-99c0-095087b555cb"
/>


- LlamaIndex
<img width="1303" alt="Screenshot 2025-01-27 at 5 10 42 PM"
src="https://github.com/user-attachments/assets/81325a91-083d-4eef-badc-908eb892a53f"
/>
<img width="1312" alt="Screenshot 2025-01-27 at 5 11 50 PM"
src="https://github.com/user-attachments/assets/499bad04-3ee2-4966-b211-84cd5f3c908e"
/>

---------

Signed-off-by: Ann Zhang <[email protected]>

Author: annzhang-db

.github/workflows/ucai-integration-tests.yml

@@ -194,6 +194,7 @@ jobs:
           pip install core/.
           echo "PYTHONPATH=$(pwd)/core/src:\$PYTHONPATH" >> $GITHUB_ENV
           pip install integrations/gemini[dev]
+          pip install mlflow
       - name: Run tests
         run: |
           pytest integrations/gemini/tests --ignore=integrations/gemini/tests/test_gemini_toolkit_oss.py

ai/core/src/unitycatalog/ai/core/base.py

@@ -1,14 +1,16 @@
 import json
 import logging
 import threading
-import time
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, Literal, Optional
 
 from unitycatalog.ai.core.paged_list import PagedList
-from unitycatalog.ai.core.utils.function_processing_utils import auto_trace_retriever
+from unitycatalog.ai.core.utils.function_processing_utils import (
+    _execute_uc_function_with_retriever_tracing,
+)
+from unitycatalog.ai.core.utils.validation_utils import has_retriever_signature
 
 _logger = logging.getLogger(__name__)
 
@@ -157,13 +159,14 @@ def execute_function(
             parameters = parameters or {}
             self.validate_input_params(function_info.input_params, parameters)
 
-            start_time_ns = time.time_ns()
-            result = self._execute_uc_function(function_info, parameters, **kwargs)
-            end_time_ns = time.time_ns()
+            if kwargs.get("enable_retriever_tracing", False) and has_retriever_signature(
+                function_info
+            ):
+                return _execute_uc_function_with_retriever_tracing(
+                    self._execute_uc_function, function_info, parameters, **kwargs
+                )
 
-            if kwargs.get("enable_retriever_tracing", False):
-                auto_trace_retriever(function_name, parameters, result, start_time_ns, end_time_ns)
-            return result
+            return self._execute_uc_function(function_info, parameters, **kwargs)
 
     @abstractmethod
     def _execute_uc_function(

ai/core/src/unitycatalog/ai/core/utils/function_processing_utils.py

@@ -20,10 +20,7 @@
     PydanticType,
 )
 from unitycatalog.ai.core.utils.type_utils import UC_TYPE_JSON_MAPPING
-from unitycatalog.ai.core.utils.validation_utils import (
-    FullFunctionName,
-    is_valid_retriever_output,
-)
+from unitycatalog.ai.core.utils.validation_utils import FullFunctionName
 
 _logger = logging.getLogger(__name__)
 
@@ -321,67 +318,76 @@ def supported_function_info_types():
     return types
 
 
-def auto_trace_retriever(
-    function_name: str,
+def process_retriever_output(result: "FunctionExecutionResult") -> List[Dict[str, Any]]:
+    """
+    Process retriever output from result into mlflow.entities.Document format for tracing.
+
+    Args:
+        result: The result of the function execution to be processed.
+
+    Returns:
+        Retriever output formatted into a list of Documents.
+    """
+    if result.format == "CSV":
+        df = pd.read_csv(StringIO(result.value))
+        if "metadata" in df.columns:
+            df["metadata"] = df["metadata"].apply(ast.literal_eval)
+        output = df.to_dict(orient="records")
+    else:
+        value = result.value
+        output = ast.literal_eval(value) if isinstance(value, str) else value
+
+    return output
+
+
+def _execute_uc_function_with_retriever_tracing(
+    _execute_uc_function: Callable,
+    function_info: "FunctionInfo",
     parameters: Dict[str, Any],
-    result: "FunctionExecutionResult",
-    start_time_ns: int,
-    end_time_ns: int,
-):
+    **kwargs: Any,
+) -> "FunctionExecutionResult":
     """
-    If the given function is a retriever, trace the function given the provided start and end time.
-    A function is considered a retriever if the result is of valid retriever output format.
+    Executes a UC function with MLflow tracing with span type RETRIEVER enabled. If MLflow cannot
+    be imported, the function executes without tracing and logs a warning.
 
     Args:
-        function_name: The function name.
-        parameters: The input parameters to the function.
-        result: The output result of the function.
-        start_time_ns: The start time of the function in nanoseconds.
-        end_time_ns: The end time of the function in nanoseconds.
+        _execute_uc_function (Callable): A function that executes the given UC function.
+        function_info (FunctionInfo): Metadata about the UC function to be executed.
+        parameters (Dict[str, Any]): Parameters to be passed to the function during execution.
+        **kwargs (Any): Additional keyword arguments to be passed to the function.
+
+    Returns:
+        Any: The output of the function execution.
     """
     try:
-        if result.format == "CSV":
-            df = pd.read_csv(StringIO(result.value))
-            if "metadata" in df.columns:
-                df["metadata"] = df["metadata"].apply(ast.literal_eval)
-            output = df.to_dict(orient="records")
-        else:
-            value = result.value
-            output = ast.literal_eval(value) if isinstance(value, str) else value
-
-        if is_valid_retriever_output(output):
-            import mlflow
-            from mlflow import MlflowClient
-            from mlflow.entities import SpanType
-
-            client = MlflowClient()
-            common_params = dict(
-                name=function_name,
-                span_type=SpanType.RETRIEVER,
-                inputs=parameters,
-                start_time_ns=start_time_ns,
-            )
+        import mlflow
+        from mlflow.entities import SpanType
 
-            if parent_span := mlflow.get_current_active_span():
-                span = client.start_span(
-                    request_id=parent_span.request_id,
-                    parent_id=parent_span.span_id,
-                    **common_params,
-                )
-                client.end_span(
-                    request_id=span.request_id,
-                    span_id=span.span_id,
-                    outputs=output,
-                    end_time_ns=end_time_ns,
-                )
-            else:
-                span = client.start_trace(**common_params)
-                client.end_trace(
-                    request_id=span.request_id, outputs=output, end_time_ns=end_time_ns
-                )
-    except Exception as e:
-        # Ignoring exceptions because auto-tracing retriever is not essential functionality
-        _logger.debug(
-            f"Skipping tracing {function_name} as a retriever because of the following error:\n {e}"
+        result = None
+
+        @mlflow.trace(name=function_info.full_name, span_type=SpanType.RETRIEVER)
+        def execute_retriever(parameters):
+            # Set inputs manually so we log {"query": "..."} instead of {"parameters": {"query": "..."}}
+            if span := mlflow.get_current_active_span():
+                span.set_inputs(parameters)
+
+            nonlocal result
+            result = _execute_uc_function(function_info, parameters, **kwargs)
+
+            # Re-raise errors so they can get traced
+            if result.error:
+                raise Exception(result.error)
+
+            return process_retriever_output(result)
+
+        try:
+            execute_retriever(parameters)
+        except Exception:  # Catch all errors that are re-raised
+            pass
+
+        return result
+    except ImportError as e:
+        _logger.warn(
+            f"Skipping tracing {function_info.full_name} as a retriever because of the following error:\n {e}"
         )
-        pass
+        return _execute_uc_function(function_info, parameters, **kwargs)

ai/core/src/unitycatalog/ai/core/utils/validation_utils.py

@@ -142,39 +142,32 @@ def validate_function_name_length(function_name: str) -> None:
         )
 
 
-def is_valid_retriever_output(output: Any) -> bool:
+def has_retriever_signature(function_info: "FunctionInfo") -> bool:
     """
-    Checks if the given output follows the retriever format for MLflow, which is a list of Documents
-    or dictionaries that follow the Document format.
+    Checks if the given function signature follows the retriever format for MLflow, which is a
+    list of Documents.
 
     Args:
-        output: The value to determine if it is a valid retriever output.
+        function_info: The function to determine if it has a valid retriever signature.
 
     Returns:
-        bool: If the provided output is a valid retriever output.
+        bool: If the provided function has a valid retriever signature.
     """
-    if not isinstance(output, list):
+    if "TABLE_TYPE" not in str(function_info.data_type):
         return False
 
-    if len(output) < 1:
+    return_params = function_info.return_params
+
+    if not (return_params and return_params.parameters):
         return False
 
-    def is_valid_retriever_item(item: Any) -> bool:
-        from mlflow.entities import Document
+    for param in return_params.parameters:
+        param_dict = param.as_dict() if hasattr(param, "as_dict") else dict(param)
 
-        if isinstance(item, Document):
+        if param_dict.get("name") == "page_content" and param_dict.get("type_name") == "STRING":
             return True
 
-        if isinstance(item, dict):
-            try:
-                Document(**item)
-                return True
-            except TypeError:
-                return False
-
-        return False
-
-    return all(is_valid_retriever_item(item) for item in output)
+    return False
 
 
 def mlflow_tracing_enabled(integration_name: str) -> bool:

ai/core/src/unitycatalog/ai/test_utils/function_utils.py

@@ -3,14 +3,68 @@
 from contextlib import contextmanager
 from typing import Any, Callable, Generator, NamedTuple, Optional
 
+from databricks.sdk.service.catalog import (
+    ColumnTypeName,
+    FunctionParameterInfo,
+    FunctionParameterInfos,
+)
+
 from unitycatalog.ai.core.databricks import DatabricksFunctionClient
 from unitycatalog.ai.core.utils.function_processing_utils import get_tool_name
+from unitycatalog.client import (
+    ColumnTypeName as OSSColumnTypeName,
+)
+from unitycatalog.client import (
+    FunctionParameterInfo as OSSFunctionParameterInfo,
+)
+from unitycatalog.client import (
+    FunctionParameterInfos as OSSFunctionParameterInfos,
+)
 
 CATALOG = "integration_testing"
 
 RETRIEVER_OUTPUT_SCALAR = '[{"page_content": "# Technology partners\\n## What is Databricks Partner Connect?\\n", "metadata": {"similarity_score": 0.010178182, "chunk_id": "0217a07ba2fec61865ce408043acf1cf"}}, {"page_content": "# Technology partners\\n## What is Databricks?\\n", "metadata": {"similarity_score": 0.010178183, "chunk_id": "0217a07ba2fec61865ce408043acf1cd"}}]'
 RETRIEVER_OUTPUT_CSV = "page_content,metadata\n\"# Technology partners\n## What is Databricks Partner Connect?\n\",\"{'similarity_score': 0.010178182, 'chunk_id': '0217a07ba2fec61865ce408043acf1cf'}\"\n\"# Technology partners\n## What is Databricks?\n\",\"{'similarity_score': 0.010178183, 'chunk_id': '0217a07ba2fec61865ce408043acf1cd'}\"\n"
 
+RETRIEVER_TABLE_FULL_DATA_TYPE = "(page_content STRING, metadata MAP<STRING, STRING>)"
+RETRIEVER_TABLE_RETURN_PARAMS = FunctionParameterInfos(
+    parameters=[
+        FunctionParameterInfo(
+            name="page_content",
+            type_text="string",
+            type_name=ColumnTypeName.STRING,
+            type_json='{"name":"page_content","type":"string","nullable":true,"metadata":{}}',
+            position=0,
+        ),
+        FunctionParameterInfo(
+            name="metadata",
+            type_text="map<string,string>",
+            type_name=ColumnTypeName.MAP,
+            type_json='{"name":"metadata","type":{"type":"map","keyType":"string","valueType":"string","valueContainsNull":true},"nullable":true,"metadata":{}}',
+            position=1,
+        ),
+    ]
+)
+RETRIEVER_TABLE_RETURN_PARAMS_OSS = OSSFunctionParameterInfos(
+    parameters=[
+        OSSFunctionParameterInfo(
+            name="page_content",
+            type_text="string",
+            type_name=OSSColumnTypeName.STRING,
+            type_json='{"name":"page_content","type":"string","nullable":true,"metadata":{}}',
+            position=0,
+        ),
+        OSSFunctionParameterInfo(
+            name="metadata",
+            type_text="map<string,string>",
+            type_name=OSSColumnTypeName.MAP,
+            type_json='{"name":"metadata","type":{"type":"map","keyType":"string","valueType":"string","valueContainsNull":true},"nullable":true,"metadata":{}}',
+            position=1,
+        ),
+    ]
+)
+
+
 _logger = logging.getLogger(__name__)