bigquery-agent-analytics 0.2.0__tar.gz → 0.2.1__tar.gz

{bigquery_agent_analytics-0.2.0 → bigquery_agent_analytics-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bigquery-agent-analytics
-Version: 0.2.0
+Version: 0.2.1
 Summary: SDK for analyzing and evaluating agent traces stored in BigQuery.
 Author: Google LLC
 License-Expression: Apache-2.0

{bigquery_agent_analytics-0.2.0 → bigquery_agent_analytics-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "bigquery-agent-analytics"
-version = "0.2.0"
+version = "0.2.1"
 description = "SDK for analyzing and evaluating agent traces stored in BigQuery."
 readme = "README.md"
 license = "Apache-2.0"

{bigquery_agent_analytics-0.2.0 → bigquery_agent_analytics-0.2.1}/src/bigquery_agent_analytics/trace.py

@@ -42,6 +42,42 @@ from typing import Any, Optional
 logger = logging.getLogger("bigquery_agent_analytics." + __name__)
 
 
+_ANSI_RED = "\x1b[31m"
+_ANSI_YELLOW = "\x1b[33m"
+_ANSI_RESET = "\x1b[0m"
+
+
+def _colorize(text: str, ansi_code: str, enabled: bool) -> str:
+  """Wraps text in an ANSI color code when enabled, else returns unchanged."""
+  if not enabled:
+    return text
+  return f"{ansi_code}{text}{_ANSI_RESET}"
+
+
+_TEXT_WRAPPER_PREFIX = "text: "
+
+
+def _unwrap_text_field(value: str) -> str:
+  """Strip a leading ``text: '...'`` wrapper if present.
+
+  Some ADK plugin versions serialize an LLM response with a literal
+  ``text: '...'`` Python-repr-style prefix in ``content.response``.
+  Strip it so ``Span.summary`` and ``Trace.render`` surface a clean
+  human-readable string.
+  """
+  if not value.startswith(_TEXT_WRAPPER_PREFIX):
+    return value
+  inner = value[len(_TEXT_WRAPPER_PREFIX) :]
+  if not inner or inner[0] not in ("'", '"'):
+    return inner
+  quote = inner[0]
+  # Match trailing quote only if the string is not truncated.
+  if len(inner) >= 2 and inner.endswith(quote):
+    return inner[1:-1]
+  # Truncated — drop the opening quote and leave the rest.
+  return inner[1:]
+
+
 class EventType(Enum):
   """Standard event types logged by the analytics plugin."""
 
@@ -232,13 +268,32 @@ class Span:
     if not self.is_error:
       return None
     parts = [self.event_type]
-    tool = self.content.get("tool")
-    if tool:
-      parts.append(f"tool={tool}")
+    if self.tool_name:
+      parts.append(f"tool={self.tool_name}")
     if self.error_message:
       parts.append(self.error_message[:200])
     return " | ".join(parts)
 
+  @property
+  def tool_name(self) -> Optional[str]:
+    """Returns the tool name for tool-related events.
+
+    Populated only for ``TOOL_STARTING``, ``TOOL_COMPLETED``,
+    ``TOOL_ERROR``, and ``HITL_*`` event types where the plugin
+    writes the tool name into ``content.tool``. Returns ``None``
+    for any other event type, even if ``content`` happens to
+    carry a ``"tool"`` key — callers rely on this attribute
+    meaning "this span invoked a tool."
+    """
+    if self.event_type not in (
+        "TOOL_STARTING",
+        "TOOL_COMPLETED",
+        "TOOL_ERROR",
+    ) and not self.event_type.startswith("HITL_"):
+      return None
+    tool = self.content.get("tool")
+    return tool if tool else None
+
   @property
   def label(self) -> str:
     """Returns a human-readable label for this span."""
@@ -311,6 +366,7 @@ class Span:
       text = self.content.get("text") or ""
     if not text:
       text = self.content.get("raw") or ""
+    text = _unwrap_text_field(text) if isinstance(text, str) else text
     if not text and self.content_parts:
       for p in self.content_parts:
         if p.text:
@@ -631,7 +687,7 @@ class Trace:
 
     return roots
 
-  def render(self, format: str = "tree") -> str:
+  def render(self, format: str = "tree", color: bool = False) -> str:
     """Renders the trace as a hierarchical DAG view.
 
     This generates a tree representation of the agent's
@@ -642,6 +698,12 @@ class Trace:
 
     Args:
         format: Render format. Currently supports "tree".
+        color: When ``True``, wrap error markers and warning
+            markers in ANSI color codes (red and yellow
+            respectively). Default ``False`` emits plain text
+            suitable for any output target. Enable this in TTY
+            contexts (terminal sessions) for faster visual
+            scanning of failures in large traces.
 
     Returns:
         A string containing the rendered trace. Also printed
@@ -661,10 +723,10 @@ class Trace:
     if not roots:
       # Flat rendering when no span IDs exist
       for span in self.spans:
-        self._render_flat_span(span, lines)
+        self._render_flat_span(span, lines, color=color)
     else:
       for root in roots:
-        self._render_span(root, lines, prefix="", is_last=True)
+        self._render_span(root, lines, prefix="", is_last=True, color=color)
 
     output = "\n".join(lines)
     print(output)
@@ -676,16 +738,17 @@ class Trace:
       lines: list[str],
       prefix: str,
       is_last: bool,
+      color: bool = False,
   ) -> None:
     """Recursively renders a span and its children as a tree."""
     connector = "\u2514\u2500 " if is_last else "\u251c\u2500 "
 
     if span.is_error:
-      status_icon = "\u2717"
+      status_icon = _colorize("\u2717", _ANSI_RED, color)
     elif span.subtree_has_error:
       # Propagate error visibility: mark parents whose subtree
       # contains an error so the failure is visible at every level.
-      status_icon = "\u26a0"
+      status_icon = _colorize("\u26a0", _ANSI_YELLOW, color)
     else:
       status_icon = "\u2713"
 
@@ -719,11 +782,20 @@ class Trace:
           lines,
           child_prefix,
           is_last=(i == len(span.children) - 1),
+          color=color,
       )
 
-  def _render_flat_span(self, span: Span, lines: list[str]) -> None:
+  def _render_flat_span(
+      self,
+      span: Span,
+      lines: list[str],
+      color: bool = False,
+  ) -> None:
     """Renders a single span without tree structure."""
-    status_icon = "\u2717" if span.is_error else "\u2713"
+    if span.is_error:
+      status_icon = _colorize("\u2717", _ANSI_RED, color)
+    else:
+      status_icon = "\u2713"
     latency = ""
     if span.latency_ms is not None:
       latency = f" ({span.latency_ms:.0f}ms)"
@@ -776,9 +848,13 @@ class Trace:
         if isinstance(c, dict):
           result = c.get("response")
           if result:
-            return result
+            return (
+                _unwrap_text_field(result)
+                if isinstance(result, str)
+                else result
+            )
         elif c:
-          return str(c)
+          return _unwrap_text_field(str(c))
 
     for span in reversed(self.spans):
       if span.event_type == "AGENT_COMPLETED":
@@ -786,9 +862,13 @@ class Trace:
         if isinstance(c, dict):
           result = c.get("response") or c.get("text_summary")
           if result:
-            return result
+            return (
+                _unwrap_text_field(result)
+                if isinstance(result, str)
+                else result
+            )
         elif c:
-          return str(c)
+          return _unwrap_text_field(str(c))
     return None
 
   @property

{bigquery_agent_analytics-0.2.0 → bigquery_agent_analytics-0.2.1}/tests/test_sdk_trace.py

@@ -120,6 +120,73 @@ class TestSpan:
     )
     assert "ERROR" in span.label
 
+  def test_tool_name_for_tool_event(self):
+    span = Span(
+        event_type="TOOL_STARTING",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"tool": "search_contacts"},
+    )
+    assert span.tool_name == "search_contacts"
+
+  def test_tool_name_for_hitl_event(self):
+    span = Span(
+        event_type="HITL_APPROVAL_REQUESTED",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"tool": "book_meeting"},
+    )
+    assert span.tool_name == "book_meeting"
+
+  def test_tool_name_none_for_non_tool_event(self):
+    span = Span(
+        event_type="USER_MESSAGE_RECEIVED",
+        agent=None,
+        timestamp=datetime.now(timezone.utc),
+        content={"text_summary": "Hello"},
+    )
+    assert span.tool_name is None
+
+  def test_tool_name_none_when_missing(self):
+    span = Span(
+        event_type="TOOL_STARTING",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={},
+    )
+    assert span.tool_name is None
+
+  def test_tool_name_none_when_empty(self):
+    span = Span(
+        event_type="TOOL_STARTING",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"tool": ""},
+    )
+    assert span.tool_name is None
+
+  def test_tool_name_does_not_leak_for_non_tool_event_with_tool_key(self):
+    """Non-tool events carrying an incidental content['tool'] key must
+    NOT surface it as tool_name. The attribute means "this span invoked
+    a tool" — arbitrary payloads that happen to use the same key name
+    should return None."""
+    span = Span(
+        event_type="AGENT_THOUGHT",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"tool": "looks_like_a_tool_but_isnt", "note": "reasoning"},
+    )
+    assert span.tool_name is None
+
+  def test_tool_name_does_not_leak_for_llm_response_with_tool_key(self):
+    span = Span(
+        event_type="LLM_RESPONSE",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"response": "I'll use a tool", "tool": "mentioned_in_text"},
+    )
+    assert span.tool_name is None
+
   def test_summary_with_text(self):
     span = Span(
         event_type="USER_MESSAGE_RECEIVED",
@@ -227,6 +294,57 @@ class TestSpan:
     )
     assert span.summary == "You are a helpful assistant"
 
+  def test_summary_unwraps_text_single_quoted(self):
+    """`text: 'hello'` wrapper should be stripped from response."""
+    span = Span(
+        event_type="LLM_RESPONSE",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"response": "text: 'hello world'"},
+    )
+    assert span.summary == "hello world"
+
+  def test_summary_unwraps_text_double_quoted(self):
+    """`text: \"hello\"` wrapper should be stripped too."""
+    span = Span(
+        event_type="LLM_RESPONSE",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"response": 'text: "hello world"'},
+    )
+    assert span.summary == "hello world"
+
+  def test_summary_unwraps_truncated_text_field(self):
+    """Truncated `text: 'abc...` (no closing quote) strips opening quote only."""
+    span = Span(
+        event_type="LLM_RESPONSE",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"response": "text: 'I found three Priyas"},
+    )
+    assert span.summary == "I found three Priyas"
+
+  def test_summary_leaves_plain_text_alone(self):
+    """A response without the wrapper must be returned unchanged."""
+    span = Span(
+        event_type="LLM_RESPONSE",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"response": "Booking confirmed for Priya Patel."},
+    )
+    assert span.summary == "Booking confirmed for Priya Patel."
+
+  def test_summary_leaves_non_wrapper_text_prefix_alone(self):
+    """`text:` without a following space-quoted value is not our wrapper."""
+    span = Span(
+        event_type="LLM_RESPONSE",
+        agent="agent",
+        timestamp=datetime.now(timezone.utc),
+        content={"response": "text:the-tool-id-is-x"},
+    )
+    # No leading `text: ` (with space), so the unwrapper leaves it alone.
+    assert span.summary == "text:the-tool-id-is-x"
+
 
 class TestTrace:
   """Tests for Trace class."""
@@ -300,6 +418,116 @@ class TestTrace:
     assert "USER_MESSAGE_RECEIVED" in output
     assert "TOOL_STARTING" in output
 
+  def test_render_color_default_off(self):
+    """Default render() must not emit ANSI escape codes."""
+    ts = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
+    spans = [
+        Span(
+            event_type="TOOL_ERROR",
+            agent="agent",
+            timestamp=ts,
+            content={"tool": "search"},
+            error_message="boom",
+            status="ERROR",
+            span_id="s1",
+        ),
+    ]
+    trace = Trace(trace_id="t1", session_id="sess-1", spans=spans)
+    output = trace.render()
+    assert "\x1b[" not in output
+
+  def test_render_color_true_wraps_error_icon(self):
+    """With color=True, error spans get red ANSI wrap on the cross icon."""
+    ts = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
+    spans = [
+        Span(
+            event_type="TOOL_ERROR",
+            agent="agent",
+            timestamp=ts,
+            content={"tool": "search"},
+            error_message="boom",
+            status="ERROR",
+            span_id="s1",
+        ),
+    ]
+    trace = Trace(trace_id="t1", session_id="sess-1", spans=spans)
+    output = trace.render(color=True)
+    assert "\x1b[31m\u2717\x1b[0m" in output
+
+  def test_render_color_true_wraps_subtree_warning(self):
+    """Parent of an error child gets yellow warning icon wrap."""
+    ts = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
+    spans = [
+        Span(
+            event_type="AGENT_STARTING",
+            agent="agent",
+            timestamp=ts,
+            span_id="s1",
+        ),
+        Span(
+            event_type="TOOL_ERROR",
+            agent="agent",
+            timestamp=ts,
+            content={"tool": "search"},
+            error_message="boom",
+            status="ERROR",
+            span_id="s2",
+            parent_span_id="s1",
+        ),
+    ]
+    trace = Trace(trace_id="t1", session_id="sess-1", spans=spans)
+    output = trace.render(color=True)
+    assert "\x1b[33m\u26a0\x1b[0m" in output
+    assert "\x1b[31m\u2717\x1b[0m" in output
+
+  def test_render_color_true_no_wrap_on_success(self):
+    """Successful spans stay plain even with color=True."""
+    ts = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
+    spans = [
+        Span(
+            event_type="TOOL_COMPLETED",
+            agent="agent",
+            timestamp=ts,
+            content={"tool": "search", "result": "ok"},
+            status="OK",
+            span_id="s1",
+        ),
+    ]
+    trace = Trace(trace_id="t1", session_id="sess-1", spans=spans)
+    output = trace.render(color=True)
+    assert "\x1b[" not in output
+
+  def test_render_handles_unicode_tool_names(self):
+    """Non-ASCII tool names must not crash render() and must preserve
+    the tree-connector structure on each line. Display-width alignment
+    in monospace fonts is a known limitation (tracked separately) —
+    this test pins the minimum: no exceptions, connectors emitted."""
+    ts = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
+    spans = [
+        Span(
+            event_type="TOOL_STARTING",
+            agent="agent",
+            timestamp=ts,
+            content={"tool": "搜索联系人"},
+            span_id="s1",
+        ),
+        Span(
+            event_type="TOOL_COMPLETED",
+            agent="agent",
+            timestamp=ts,
+            content={"tool": "🔍_fuzzy_search", "result": "ok"},
+            span_id="s2",
+            parent_span_id="s1",
+            latency_ms=50,
+        ),
+    ]
+    trace = Trace(trace_id="t1", session_id="sess-1", spans=spans)
+    output = trace.render()
+    assert "搜索联系人" in output
+    assert "🔍_fuzzy_search" in output
+    # tree connectors present on every non-header line
+    assert "\u2514\u2500" in output or "\u251c\u2500" in output
+
   def test_render_flat_no_span_ids(self):
     ts = datetime.now(timezone.utc)
     spans = [
@@ -363,6 +591,33 @@ class TestTrace:
     trace = Trace(trace_id="t", session_id="s", spans=spans)
     assert trace.final_response == "LLM said this"
 
+  def test_final_response_unwraps_text_prefix_from_llm_response(self):
+    """text: '...' wrapper must be stripped at the trace-level accessor too."""
+    ts = datetime.now(timezone.utc)
+    spans = [
+        Span(
+            event_type="LLM_RESPONSE",
+            agent="agent",
+            timestamp=ts,
+            content={"response": "text: 'I found three people named Priya.'"},
+        ),
+    ]
+    trace = Trace(trace_id="t", session_id="s", spans=spans)
+    assert trace.final_response == "I found three people named Priya."
+
+  def test_final_response_unwraps_text_prefix_from_agent_completed(self):
+    ts = datetime.now(timezone.utc)
+    spans = [
+        Span(
+            event_type="AGENT_COMPLETED",
+            agent="agent",
+            timestamp=ts,
+            content={"response": 'text: "Booking confirmed."'},
+        ),
+    ]
+    trace = Trace(trace_id="t", session_id="s", spans=spans)
+    assert trace.final_response == "Booking confirmed."
+
   def test_final_response_null_agent_completed(self):
     """Handles null AGENT_COMPLETED content (ADK plugin behavior)."""
     ts = datetime.now(timezone.utc)