Skip to content

Commit

Permalink
Add example usage of "extract_fields" for mlflow.search_traces in doc…
Browse files Browse the repository at this point in the history
…umentation (mlflow#12319)

Signed-off-by: xq-yin <xiaoqian.yin@databricks.com>
  • Loading branch information
xq-yin committed Jun 12, 2024
1 parent c5ae48e commit f97c515
Showing 1 changed file with 31 additions and 1 deletion.
32 changes: 31 additions & 1 deletion docs/source/llms/tracing/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -464,9 +464,39 @@ filter strings, and other parameters.
# Search for traces in specific experiments
traces = client.search_traces(
experiment_ids=["1", "2"], filter_string="attributes.status = 'OK'", max_results=5
experiment_ids=["1", "2"],
filter_string="attributes.status = 'OK'",
max_results=5,
)
Alternatively, you can use fluent API :py:func:`mlflow.search_traces` to search for traces, which returns a pandas DataFrame with each row containing a trace.
This method allows you to specify fields to extract from traces using the format ``"span_name.[inputs|outputs]"`` or ``"span_name.[inputs|outputs].field_name"``.
The extracted fields are included as extra columns in the pandas DataFrame. This feature can be used to build evaluation datasets to further improve model and agent performance.

.. code-block:: python
import mlflow
with mlflow.start_span(name="span1") as span:
span.set_inputs({"a": 1, "b": 2})
span.set_outputs({"c": 3, "d": 4})
# Search for traces with specific fields extracted
traces = mlflow.search_traces(
extract_fields=["span1.inputs", "span1.outputs.c"],
)
print(traces)
This outputs:

.. code-block:: text
request_id ... span1.inputs span1.outputs.c
0 tr-97c4ef97c21f4348a5698f069c1320f1 ... {'a': 1, 'b': 2} 3.0
1 tr-4dc3cd5567764499b5532e3af61b9f78 ... {'a': 1, 'b': 2} 3.0
Retrieving a Specific Trace
^^^^^^^^^^^^^^^^^^^^^^^^^^^

Expand Down

0 comments on commit f97c515

Please sign in to comment.