arize-phoenix 10.1.0__py3-none-any.whl → 10.2.1__py3-none-any.whl

phoenix/session/evaluation.py

@@ -23,13 +23,18 @@ from google.protobuf.wrappers_pb2 import DoubleValue, StringValue
 import phoenix.trace.v1 as pb
 from phoenix.config import get_env_collector_endpoint, get_env_host, get_env_port
 from phoenix.session.client import Client
-from phoenix.trace.dsl.helpers import get_qa_with_reference, get_retrieved_documents
+from phoenix.trace.dsl.helpers import (
+    get_called_tools,
+    get_qa_with_reference,
+    get_retrieved_documents,
+)
 from phoenix.trace.exporter import HttpExporter
 from phoenix.trace.span_evaluations import Evaluations
 
 __all__ = [
     "get_retrieved_documents",
     "get_qa_with_reference",
+    "get_called_tools",
     "add_evaluations",
 ]
 

phoenix/trace/dsl/helpers.py

@@ -1,6 +1,7 @@
+import json
 import warnings
 from datetime import datetime
-from typing import Optional, Protocol, Union, cast
+from typing import Any, Iterable, Mapping, Optional, Protocol, Union, cast
 
 import pandas as pd
 from openinference.semconv.trace import DocumentAttributes, SpanAttributes
@@ -13,11 +14,16 @@ DOCUMENT_SCORE = DocumentAttributes.DOCUMENT_SCORE
 INPUT_VALUE = SpanAttributes.INPUT_VALUE
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
 RETRIEVAL_DOCUMENTS = SpanAttributes.RETRIEVAL_DOCUMENTS
+LLM_FUNCTION_CALL = SpanAttributes.LLM_FUNCTION_CALL
+LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
+LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
+
 
 INPUT = {"input": INPUT_VALUE}
 OUTPUT = {"output": OUTPUT_VALUE}
 IO = {**INPUT, **OUTPUT}
 
+
 IS_ROOT = "parent_id is None"
 IS_LLM = "span_kind == 'LLM'"
 IS_RETRIEVER = "span_kind == 'RETRIEVER'"
@@ -125,3 +131,86 @@ def get_qa_with_reference(
     df_ref = pd.DataFrame({"reference": ref})
     df_qa_ref = pd.concat([df_qa, df_ref], axis=1, join="inner").set_index("context.span_id")
     return df_qa_ref
+
+
+def get_called_tools(
+    obj: CanQuerySpans,
+    *,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+    project_name: Optional[str] = None,
+    timeout: Optional[int] = DEFAULT_TIMEOUT_IN_SECONDS,
+    function_name_only: bool = False,
+) -> Optional[pd.DataFrame]:
+    """Retrieve tool calls made by LLM spans within a specified time range.
+
+    This function queries LLM spans and extracts tool calls from their output messages.
+    It can return either just the function names or full function calls with arguments.
+
+    Args:
+        obj: An object that implements the CanQuerySpans protocol for querying spans.
+        start_time: Optional start time to filter spans. If None, no start time filter is applied.
+        end_time: Optional end time to filter spans. If None, no end time filter is applied.
+        project_name: Optional project name to filter spans. If None, uses the environment project name.
+        timeout: Optional timeout in seconds for the query. Defaults to DEFAULT_TIMEOUT_IN_SECONDS.
+        function_name_only: If True, returns only function names. If False, returns full function calls
+            with arguments. Defaults to False.
+
+    Returns:
+        A pandas DataFrame containing the tool calls, or None if no spans are found.
+        The DataFrame includes columns for input messages, output messages, and tool calls.
+    """  # noqa: E501
+    project_name = project_name or get_env_project_name()
+
+    def extract_tool_calls(outputs: list[dict[str, Any]]) -> Optional[list[str]]:
+        if not isinstance(outputs, list) or not outputs:
+            return None
+        ans = []
+        if isinstance(message := outputs[0].get("message"), Mapping) and isinstance(
+            tool_calls := message.get("tool_calls"), Iterable
+        ):
+            for tool_call in tool_calls:
+                if not isinstance(tool_call, Mapping):
+                    continue
+                if not isinstance(tc := tool_call.get("tool_call"), Mapping):
+                    continue
+                if not isinstance(function := tc.get("function"), Mapping):
+                    continue
+                if not isinstance(name := function.get("name"), str):
+                    continue
+                if function_name_only:
+                    ans.append(name)
+                    continue
+                kwargs = {}
+                if isinstance(arguments := function.get("arguments"), str):
+                    try:
+                        kwargs = json.loads(arguments)
+                    except Exception:
+                        pass
+                kwargs_str = "" if not kwargs else ", ".join(f"{k}={v}" for k, v in kwargs.items())
+                ans.append(f"{name}({kwargs_str})")
+        return ans or None
+
+    df_qa = cast(
+        pd.DataFrame,
+        obj.query_spans(
+            SpanQuery()
+            .where(IS_LLM)
+            .select(
+                input=LLM_INPUT_MESSAGES,
+                output=LLM_OUTPUT_MESSAGES,
+            ),
+            start_time=start_time,
+            end_time=end_time,
+            project_name=project_name,
+            timeout=timeout,
+        ),
+    )
+
+    if df_qa is None:
+        print("No spans found.")
+        return None
+
+    df_qa["tool_call"] = df_qa["output"].apply(extract_tool_calls)
+
+    return df_qa

phoenix/version.py

@@ -1 +1 @@
-__version__ = "10.1.0"
+__version__ = "10.2.1"