diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 169538d8..b2676397 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -155,6 +155,9 @@ jobs: AZURE_API_BASE: https://example.openai.azure.com/ PHOENIX_DISABLE_AUTO_INSTRUMENT: '1' run: python -c "import ${{ matrix.import }}" + - name: Test Phoenix collector against installed client + if: matrix.example == 'phoenix-auto-trace-openai' + run: python source/tests/test_phoenix_collector.py -v - name: Test Bank Manager model routing if: matrix.example == 'bank-manager-agent-control' working-directory: source @@ -166,6 +169,45 @@ jobs: working-directory: source run: python -m unittest discover -s examples/science_research_agent/tests -p 'test_tools.py' -v + test-phoenix-versions: + name: "Phoenix collector: ${{ matrix.name }}" + needs: build + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - name: minimum supported + phoenix: '15.0.0' + phoenix_client: '2.1.0' + - name: maximum supported + phoenix: '19.17.0' + phoenix_client: '' + steps: + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - uses: actions/checkout@v4 + with: + path: source + - uses: actions/download-artifact@v4 + with: + name: dist + path: dist + - name: Install Phoenix compatibility stack + shell: bash + run: | + python -m pip install --upgrade pip + WHEEL="$(ls dist/*.whl)[phoenix]" + PACKAGES=("$WHEEL" "arize-phoenix==${{ matrix.phoenix }}") + if [ -n "${{ matrix.phoenix_client }}" ]; then + PACKAGES+=("arize-phoenix-client==${{ matrix.phoenix_client }}") + fi + python -m pip install "${PACKAGES[@]}" + python -m pip check + - name: Run Phoenix collector compatibility tests + run: python source/tests/test_phoenix_collector.py -v + test-foundry-host-install: name: "Example install: langgraph-foundry-hosted server" runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index b7f54637..805cbf36 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Restore `PhoenixCollector` compatibility with the supported Phoenix client API, including bounded turn-time queries and conversion of Phoenix DataFrame timestamps and indexed span IDs. - `assert-ai --version` now reads the installed distribution metadata instead of reporting a hard-coded stale version. - Declare `aiohttp`, which is directly used by the HTTP endpoint target, instead of relying on LiteLLM to install it transitively. - Keep Bank Manager's GPT and non-GPT Azure routes compatible with ASSERT's OpenAI dependency range, and verify its documented installation with `pip check`. diff --git a/assert_ai/core/collector.py b/assert_ai/core/collector.py index a522ed35..f639c3c5 100644 --- a/assert_ai/core/collector.py +++ b/assert_ai/core/collector.py @@ -12,6 +12,13 @@ from __future__ import annotations +import heapq +import json +from collections import defaultdict +from collections.abc import Mapping, Sequence +from datetime import UTC, datetime +from numbers import Integral, Real +from sys import maxsize as _PHOENIX_RAW_QUERY_LIMIT from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: @@ -29,6 +36,10 @@ "output.value", }) +# Legacy DataFrame clients cannot express an unbounded query; use the largest +# portable GraphQL Int while keeping the supported raw API genuinely cursor-complete. +_PHOENIX_LEGACY_QUERY_LIMIT = 2_147_483_647 + @runtime_checkable class SpanCollector(Protocol): @@ -120,7 +131,7 @@ class PhoenixCollector: Phoenix is an OPTIONAL dependency — only imported when instantiated. Install: pip install 'assert-ai[phoenix]' - Queries Phoenix for DataFrame, then converts to list[OTelSpan] internally. + Uses Phoenix's event-bearing span API and converts to list[OTelSpan]. """ def __init__( @@ -130,13 +141,15 @@ def __init__( project_name: str | None = None, ) -> None: try: - import phoenix as px - self._client = px.Client(endpoint=endpoint) + from phoenix.client import Client # type: ignore[import-not-found] + + self._client = Client(base_url=endpoint) except ImportError as e: raise ImportError( "PhoenixCollector requires arize-phoenix. " "Install with: pip install 'assert-ai[phoenix]'" ) from e + self._endpoint = endpoint self._default_project = project_name def get_spans( @@ -147,41 +160,123 @@ def get_spans( end_time: str | None = None, trace_ids: list[str] | None = None, ) -> list[Any]: - import pandas as pd - name = project_name or self._default_project if name is None: raise ValueError("project_name required") + start_datetime = _parse_datetime(start_time, field_name="start_time") + end_datetime = _parse_datetime(end_time, field_name="end_time") try: - df: pd.DataFrame = self._client.get_spans_dataframe( - project_name=name, - start_time=start_time, - end_time=end_time, + span_client = self._client.spans + raw_get_spans = getattr(span_client, "get_spans", None) + if callable(raw_get_spans): + query: dict[str, Any] = { + "project_identifier": name, + "start_time": start_datetime, + "end_time": end_datetime, + "limit": _PHOENIX_RAW_QUERY_LIMIT, + } + if trace_ids: + query["trace_ids"] = list(trace_ids) + raw_spans = raw_get_spans(**query) + if not isinstance(raw_spans, Sequence) or isinstance(raw_spans, (str, bytes)): + raise TypeError("Phoenix get_spans() returned a non-sequence value") + return _order_spans([ + _phoenix_span_to_otel_span(raw_span) + for raw_span in raw_spans + ]) + + # Compatibility fallback for clients predating the raw span API. + # It cannot recover span events, so supported clients are tested on + # the raw path and this remains only a bounded legacy escape hatch. + import pandas as pd # type: ignore[import-not-found] + + df: pd.DataFrame = span_client.get_spans_dataframe( + project_identifier=name, + start_time=start_datetime, + end_time=end_datetime, + limit=_PHOENIX_LEGACY_QUERY_LIMIT, ) + if trace_ids: + if "context.trace_id" not in df.columns: + raise RuntimeError( + "Phoenix DataFrame missing 'context.trace_id' column. " + f"Available columns: {list(df.columns)}" + ) + df = df[df["context.trace_id"].isin(trace_ids)] + return _order_spans(_dataframe_to_otel_spans(df)) except ConnectionError as exc: raise RuntimeError( - f"Cannot connect to Phoenix at {self._client._base_url if hasattr(self._client, '_base_url') else 'unknown'} " + f"Cannot connect to Phoenix at {self._endpoint} " f"for project '{name}': {exc}" ) from exc except Exception as exc: raise RuntimeError( - f"Failed to fetch spans from Phoenix for project '{name}': {type(exc).__name__}: {exc}" + f"Failed to fetch spans from Phoenix for project '{name}': " + f"{type(exc).__name__}: {exc}" ) from exc - if trace_ids: - if "context.trace_id" not in df.columns: - raise RuntimeError( - "Phoenix DataFrame missing 'context.trace_id' column. " - f"Available columns: {list(df.columns)}" - ) - df = df[df["context.trace_id"].isin(trace_ids)] - - return _dataframe_to_otel_spans(df) def validate(self, spans: list[Any]) -> list[str]: return _validate_otel_spans(spans) +def _parse_datetime(value: str | None, *, field_name: str) -> datetime | None: + """Convert the collector protocol's ISO timestamp to Phoenix's datetime API.""" + if value is None: + return None + if not isinstance(value, str): + raise TypeError(f"{field_name} must be an ISO-8601 string") + normalized = value[:-1] + "+00:00" if value.endswith("Z") else value + try: + parsed = datetime.fromisoformat(normalized) + except ValueError as exc: + raise ValueError(f"{field_name} must be a valid ISO-8601 timestamp") from exc + if parsed.tzinfo is None or parsed.tzinfo.utcoffset(parsed) is None: + raise ValueError(f"{field_name} must include a timezone offset") + return parsed.astimezone(UTC) + + +def _is_missing(value: Any) -> bool: + """Return whether a scalar DataFrame value represents missing data.""" + if value is None: + return True + try: + missing = value != value + return bool(missing) + except (TypeError, ValueError): + try: + import pandas as pd # type: ignore[import-not-found] + + return bool(pd.isna(value)) + except (TypeError, ValueError): + return False + + +def _timestamp_to_ns(value: Any, *, field_name: str) -> int: + """Normalize Phoenix numeric or timezone-aware timestamp values to nanoseconds.""" + if _is_missing(value): + return 0 + native_ns = getattr(value, "value", None) + if isinstance(native_ns, Integral) and not isinstance(native_ns, bool): + return int(native_ns) + if isinstance(value, datetime): + if value.tzinfo is None or value.tzinfo.utcoffset(value) is None: + raise ValueError(f"{field_name} must be timezone-aware") + utc_value = value.astimezone(UTC) + delta = utc_value - datetime(1970, 1, 1, tzinfo=UTC) + return ( + (delta.days * 86_400 + delta.seconds) * 1_000_000_000 + + delta.microseconds * 1_000 + ) + if isinstance(value, Real) and not isinstance(value, bool): + return int(value) # type: ignore[arg-type] + if isinstance(value, str): + parsed = _parse_datetime(value, field_name=field_name) + assert parsed is not None + return _timestamp_to_ns(parsed, field_name=field_name) + raise TypeError(f"{field_name} has unsupported type {type(value).__name__}") + + def _dataframe_to_otel_spans(df: Any) -> list[Any]: """Convert an OpenInference-format DataFrame to list[OTelSpan]. @@ -190,23 +285,158 @@ def _dataframe_to_otel_spans(df: Any) -> list[Any]: from assert_ai.core.otel import OTelSpan spans = [] - for _, row in df.iterrows(): + for index, row in df.iterrows(): attrs: dict[str, Any] = {} for col in df.columns: if col.startswith("attributes."): key = col[len("attributes."):] val = row[col] - if val is not None and not (isinstance(val, float) and val != val): + if not _is_missing(val): attrs[key] = val + trace_id = row.get("context.trace_id", "") + span_id = row.get("context.span_id", row.get("span_id", "")) + if _is_missing(span_id) or span_id == "": + span_id = index if isinstance(index, str) else "" + parent_span_id = row.get("parent_id") + if _is_missing(parent_span_id) or parent_span_id == "": + parent_span_id = None + kind = attrs.get("openinference.span.kind") + if _is_missing(kind) or not kind: + kind = row.get("span_kind") + if _is_missing(kind) or not kind: + kind = "UNKNOWN" + spans.append(OTelSpan( - trace_id=str(row.get("context.trace_id", "")), - span_id=str(row.get("context.span_id", "")), - parent_span_id=str(row["parent_id"]) if row.get("parent_id") else None, + trace_id="" if _is_missing(trace_id) else str(trace_id), + span_id=str(span_id), + parent_span_id=str(parent_span_id) if parent_span_id is not None else None, name=str(row.get("name", "")), - kind=attrs.get("openinference.span.kind", "UNKNOWN"), - start_time_ns=int(row.get("start_time", 0)) if row.get("start_time") else 0, - end_time_ns=int(row.get("end_time", 0)) if row.get("end_time") else 0, + kind=str(kind), + start_time_ns=_timestamp_to_ns( + row.get("start_time"), + field_name="start_time", + ), + end_time_ns=_timestamp_to_ns( + row.get("end_time"), + field_name="end_time", + ), attributes=attrs, )) return spans + + +def _phoenix_span_to_otel_span(raw: Mapping[str, Any]) -> "OTelSpan": + """Convert one event-bearing Phoenix API span to ASSERT's neutral shape.""" + from assert_ai.core.otel import OTelSpan + + context = raw.get("context") or {} + if not isinstance(context, Mapping): + raise ValueError("Phoenix span context must be an object") + raw_attributes = raw.get("attributes") or {} + if not isinstance(raw_attributes, Mapping): + raise ValueError("Phoenix span attributes must be an object") + attributes = dict(raw_attributes) + return OTelSpan( + trace_id=str(context.get("trace_id") or ""), + span_id=str(context.get("span_id") or ""), + parent_span_id=(str(raw["parent_id"]) if raw.get("parent_id") else None), + name=str(raw.get("name") or ""), + kind=str( + attributes.get("openinference.span.kind") + or raw.get("span_kind") + or "UNKNOWN" + ), + start_time_ns=_timestamp_to_ns(raw.get("start_time"), field_name="start_time"), + end_time_ns=_timestamp_to_ns(raw.get("end_time"), field_name="end_time"), + attributes=attributes, + status=str(raw.get("status_code") or "OK"), + events=_phoenix_events_to_otlp(raw.get("events") or []), + ) + + +def _phoenix_events_to_otlp(raw_events: Any) -> list[dict[str, Any]]: + """Preserve Phoenix event content in the OTLP-JSON shape ASSERT consumes.""" + if not isinstance(raw_events, Sequence) or isinstance(raw_events, (str, bytes)): + raise ValueError("Phoenix span events must be a sequence") + events: list[dict[str, Any]] = [] + for raw_event in raw_events: + if not isinstance(raw_event, Mapping): + raise ValueError("Phoenix span event must be an object") + raw_attributes = raw_event.get("attributes") or {} + if not isinstance(raw_attributes, Mapping): + raise ValueError("Phoenix span event attributes must be an object") + event: dict[str, Any] = { + "name": str(raw_event.get("name") or ""), + "attributes": [ + {"key": str(key), "value": _to_otlp_value(value)} + for key, value in raw_attributes.items() + ], + } + if raw_event.get("timestamp") is not None: + event["timeUnixNano"] = str( + _timestamp_to_ns(raw_event["timestamp"], field_name="event timestamp") + ) + events.append(event) + return events + + +def _to_otlp_value(value: Any) -> dict[str, Any]: + """Encode a Phoenix event attribute using OTLP's JSON value shape.""" + if isinstance(value, bool): + return {"boolValue": value} + if isinstance(value, int): + return {"intValue": str(value)} + if isinstance(value, float): + return {"doubleValue": value} + if isinstance(value, (list, tuple)): + return {"arrayValue": {"values": [_to_otlp_value(item) for item in value]}} + if isinstance(value, Mapping): + value = json.dumps(value, ensure_ascii=False, sort_keys=True, default=str) + return {"stringValue": "" if value is None else str(value)} + + +def _order_spans(spans: list["OTelSpan"]) -> list["OTelSpan"]: + """Produce deterministic chronology with parent-first ordering on time ties.""" + groups: defaultdict[tuple[int, str], list["OTelSpan"]] = defaultdict(list) + for span in spans: + groups[(span.start_time_ns, span.trace_id)].append(span) + + ordered: list["OTelSpan"] = [] + for group_key in sorted(groups): + group = groups[group_key] + by_id = {span.span_id: span for span in group} + if len(by_id) != len(group): + raise ValueError( + f"Phoenix returned duplicate span IDs for trace {group_key[1]!r}" + ) + + children: defaultdict[str, list[str]] = defaultdict(list) + indegree = {span_id: 0 for span_id in by_id} + for span in group: + parent_id = span.parent_span_id + if parent_id in by_id and parent_id != span.span_id: + children[parent_id].append(span.span_id) + indegree[span.span_id] += 1 + + pending = set(by_id) + ready = [span_id for span_id, count in indegree.items() if count == 0] + heapq.heapify(ready) + while pending: + if not ready: + # Break malformed parent cycles by the stable span ID rather + # than inheriting Phoenix's page/input order. + heapq.heappush(ready, min(pending)) + span_id = heapq.heappop(ready) + if span_id not in pending: + continue + pending.remove(span_id) + ordered.append(by_id[span_id]) + for child_id in sorted(children.get(span_id, [])): + if child_id not in pending: + continue + indegree[child_id] -= 1 + if indegree[child_id] == 0: + heapq.heappush(ready, child_id) + + return ordered diff --git a/tests/test_exception_handling.py b/tests/test_exception_handling.py index e7664518..24d6c80b 100644 --- a/tests/test_exception_handling.py +++ b/tests/test_exception_handling.py @@ -361,9 +361,10 @@ def test_connection_error_raises_runtime_error(self) -> None: from assert_ai.core.collector import PhoenixCollector collector = PhoenixCollector.__new__(PhoenixCollector) + collector._endpoint = "http://localhost:6006" collector._default_project = "test-project" collector._client = MagicMock() - collector._client.get_spans_dataframe.side_effect = ConnectionError("refused") + collector._client.spans.get_spans.side_effect = ConnectionError("refused") with self.assertRaises(RuntimeError) as ctx: collector.get_spans(project_name="test-project") @@ -379,9 +380,10 @@ def test_generic_error_raises_runtime_error(self) -> None: from assert_ai.core.collector import PhoenixCollector collector = PhoenixCollector.__new__(PhoenixCollector) + collector._endpoint = "http://localhost:6006" collector._default_project = "test-project" collector._client = MagicMock() - collector._client.get_spans_dataframe.side_effect = RuntimeError("unexpected") + collector._client.spans.get_spans_dataframe.side_effect = RuntimeError("unexpected") with self.assertRaises(RuntimeError) as ctx: collector.get_spans(project_name="test-project") diff --git a/tests/test_framework_agnostic.py b/tests/test_framework_agnostic.py index 73d514dc..32269c95 100644 --- a/tests/test_framework_agnostic.py +++ b/tests/test_framework_agnostic.py @@ -954,14 +954,6 @@ def test_dataframe_collector_validate_non_dataframe(self): warnings = collector.validate([]) self.assertEqual(warnings, []) # empty list → no warnings - @unittest.skip("Pre-existing: conflicting phoenix module lacks Client attribute") - def test_phoenix_collector_import_error(self): - from assert_ai.core.collector import PhoenixCollector - - with self.assertRaises(ImportError) as ctx: - PhoenixCollector() - self.assertIn("arize-phoenix", str(ctx.exception)) - def test_custom_collector_satisfies_protocol(self): """A plain class with get_spans/validate should satisfy SpanCollector.""" from assert_ai.core.collector import SpanCollector diff --git a/tests/test_phoenix_collector.py b/tests/test_phoenix_collector.py new file mode 100644 index 00000000..2872a651 --- /dev/null +++ b/tests/test_phoenix_collector.py @@ -0,0 +1,357 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Compatibility tests for the optional Phoenix span collector.""" +from __future__ import annotations + +import asyncio +import importlib.util +import json +import sys +import types +import unittest +from datetime import UTC, datetime +from unittest.mock import patch + +PHOENIX_AVAILABLE = importlib.util.find_spec("phoenix") is not None + + +def _raw_span( + span_id: str, + *, + trace_id: str = "trace-1", + parent_id: str | None = None, + start_time: str = "2026-09-01T00:00:00Z", + end_time: str = "2026-09-01T00:00:01Z", + attributes: dict | None = None, + events: list[dict] | None = None, +) -> dict: + return { + "id": f"relay-{span_id}", + "name": "model call", + "context": {"trace_id": trace_id, "span_id": span_id}, + "parent_id": parent_id, + "span_kind": "LLM", + "start_time": start_time, + "end_time": end_time, + "status_code": "OK", + "status_message": "", + "attributes": attributes or {"openinference.span.kind": "LLM"}, + "events": events or [], + } + + +class PhoenixCollectorMissingDependencyTest(unittest.TestCase): + def test_missing_dependency_has_actionable_install_error(self) -> None: + from assert_ai.core.collector import PhoenixCollector + + with patch.dict(sys.modules, {"phoenix": None, "phoenix.client": None}): + with self.assertRaisesRegex(ImportError, r"assert-ai\[phoenix\]"): + PhoenixCollector() + + +@unittest.skipUnless(PHOENIX_AVAILABLE, "install assert-ai[phoenix] to test the adapter") +class PhoenixCollectorCompatibilityTest(unittest.TestCase): + """Exercise the real Phoenix client surface installed by the public extra.""" + + def test_constructs_with_supported_phoenix_client(self) -> None: + from assert_ai.core.collector import PhoenixCollector + + collector = PhoenixCollector( + endpoint="http://localhost:6006", + project_name="review-project", + ) + + self.assertTrue(hasattr(collector._client, "spans")) + self.assertTrue(hasattr(collector._client.spans, "get_spans_dataframe")) + + def test_raw_query_preserves_events_and_pushes_trace_filter_to_phoenix(self) -> None: + from assert_ai.core.collector import PhoenixCollector + from assert_ai.core.otel import _spans_to_events + + message = json.dumps({ + "role": "assistant", + "content": "event answer", + "tool_calls": [{ + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q":"x"}'}, + }], + }) + raw_span = _raw_span( + "span-event", + start_time="2026-09-01T00:00:00.000001Z", + attributes={"gen_ai.operation.name": "chat"}, + events=[{ + "name": "gen_ai.choice", + "timestamp": "2026-09-01T00:00:00.500000Z", + "attributes": {"message": message}, + }, { + "name": "gen_ai.tool.message", + "timestamp": "2026-09-01T00:00:00.750000Z", + "attributes": {"id": "call-1", "content": '{"answer":42}'}, + }], + ) + collector = PhoenixCollector(project_name="review-project") + + with patch.object( + collector._client.spans, + "get_spans", + return_value=[raw_span], + ) as get_spans: + spans = collector.get_spans(trace_ids=["trace-1"]) + + query = get_spans.call_args.kwargs + self.assertEqual(query["project_identifier"], "review-project") + self.assertEqual(query["trace_ids"], ["trace-1"]) + self.assertEqual(query["limit"], sys.maxsize) + self.assertEqual(spans[0].start_time_ns, 1_788_220_800_000_001_000) + self.assertEqual(len(spans[0].events), 2) + events, aggregate = _spans_to_events(spans) + self.assertEqual(events[0]["edit"]["message"]["content"], "event answer") + self.assertEqual(events[1]["edit"]["tool_name"], "lookup") + self.assertEqual(events[1]["edit"]["tool_args"], {"q": "x"}) + self.assertEqual( + json.loads(events[1]["edit"]["tool_result"]), + {"answer": 42}, + ) + self.assertEqual(aggregate["llm_call_count"], 1) + + def test_raw_query_follows_phoenix_cursor_pages(self) -> None: + from assert_ai.core.collector import PhoenixCollector + + class Response: + def __init__(self, payload: dict): + self._payload = payload + + def raise_for_status(self) -> None: + return None + + def json(self) -> dict: + return self._payload + + collector = PhoenixCollector(project_name="review-project") + responses = [ + Response({"data": [_raw_span("span-page-1")], "next_cursor": "page-2"}), + Response({"data": [_raw_span("span-page-2")]}), + ] + + with patch.object( + collector._client.spans._client, + "get", + side_effect=responses, + ) as request: + spans = collector.get_spans() + + self.assertEqual([span.span_id for span in spans], ["span-page-1", "span-page-2"]) + self.assertEqual(request.call_count, 2) + self.assertNotIn("cursor", request.call_args_list[0].kwargs["params"]) + self.assertEqual(request.call_args_list[1].kwargs["params"]["cursor"], "page-2") + + def test_legacy_dataframe_fallback_uses_explicit_limit_and_converts_spans(self) -> None: + import pandas as pd # type: ignore[import-not-found] + + from assert_ai.core.collector import PhoenixCollector + + start = pd.Timestamp("2026-09-01T00:00:00Z") + end = pd.Timestamp("2026-09-01T00:00:01Z") + dataframe = pd.DataFrame( + { + "context.trace_id": ["trace-1", "trace-2"], + "parent_id": [None, "span-parent"], + "name": ["model call", "tool call"], + "span_kind": ["LLM", "TOOL"], + "start_time": [start, start], + "end_time": [end, end], + "attributes.session.id": ["session-1", "session-1"], + "attributes.output.value": ["answer", "tool result"], + "attributes.llm.model_name": ["gpt-5.4", None], + }, + index=pd.Index(["span-1", "span-2"], name="context.span_id"), + ) + collector = PhoenixCollector( + endpoint="http://localhost:6006", + project_name="review-project", + ) + + with ( + patch.object(collector._client.spans, "get_spans", None), + patch.object( + collector._client.spans, + "get_spans_dataframe", + return_value=dataframe, + ) as get_spans, + ): + spans = collector.get_spans( + start_time="2026-09-01T00:00:00Z", + end_time="2026-09-01T00:00:01+00:00", + trace_ids=["trace-1"], + ) + + query = get_spans.call_args.kwargs + self.assertEqual(query["project_identifier"], "review-project") + self.assertIsInstance(query["start_time"], datetime) + self.assertIsInstance(query["end_time"], datetime) + self.assertEqual(query["start_time"].tzinfo, UTC) + self.assertEqual(query["end_time"].tzinfo, UTC) + self.assertGreater(query["limit"], 1000) + + self.assertEqual(len(spans), 1) + span = spans[0] + self.assertEqual(span.trace_id, "trace-1") + self.assertEqual(span.span_id, "span-1") + self.assertIsNone(span.parent_span_id) + self.assertEqual(span.kind, "LLM") + self.assertEqual(span.start_time_ns, start.value) + self.assertEqual(span.end_time_ns, end.value) + self.assertEqual(span.attributes["session.id"], "session-1") + self.assertEqual(span.attributes["output.value"], "answer") + + def test_orders_newest_first_phoenix_results_chronologically(self) -> None: + from assert_ai.core.collector import PhoenixCollector + + initial = _raw_span( + "span-initial", + start_time="2026-09-01T00:00:00Z", + end_time="2026-09-01T00:00:01Z", + ) + final = _raw_span( + "span-final", + start_time="2026-09-01T00:00:02Z", + end_time="2026-09-01T00:00:03Z", + ) + collector = PhoenixCollector(project_name="review-project") + + with patch.object( + collector._client.spans, + "get_spans", + return_value=[final, initial], + ): + spans = collector.get_spans() + + self.assertEqual([span.span_id for span in spans], ["span-initial", "span-final"]) + + def test_orders_parent_before_child_when_timestamps_tie(self) -> None: + from assert_ai.core.collector import PhoenixCollector + + parent = _raw_span("span-parent") + child = _raw_span("span-child", parent_id="span-parent") + collector = PhoenixCollector(project_name="review-project") + + with patch.object( + collector._client.spans, + "get_spans", + return_value=[child, parent], + ): + spans = collector.get_spans() + + self.assertEqual([span.span_id for span in spans], ["span-parent", "span-child"]) + + def test_orders_deep_parent_chain_without_recursion(self) -> None: + from assert_ai.core.collector import PhoenixCollector + + raw_spans = [ + _raw_span( + f"span-{index:04d}", + parent_id=(f"span-{index - 1:04d}" if index else None), + ) + for index in range(1100) + ] + collector = PhoenixCollector(project_name="review-project") + + with patch.object( + collector._client.spans, + "get_spans", + return_value=list(reversed(raw_spans)), + ): + spans = collector.get_spans() + + self.assertEqual( + [span.span_id for span in spans], + [f"span-{index:04d}" for index in range(1100)], + ) + + def test_orders_parent_cycle_independently_of_phoenix_input_order(self) -> None: + from assert_ai.core.collector import PhoenixCollector + + span_a = _raw_span("span-a", parent_id="span-b") + span_b = _raw_span("span-b", parent_id="span-a") + collector = PhoenixCollector(project_name="review-project") + orders = [] + for raw_order in ([span_a, span_b], [span_b, span_a]): + with patch.object( + collector._client.spans, + "get_spans", + return_value=raw_order, + ): + orders.append([span.span_id for span in collector.get_spans()]) + + self.assertEqual(orders, [["span-a", "span-b"], ["span-a", "span-b"]]) + + def test_rejects_invalid_or_timezone_naive_bounds_before_query(self) -> None: + from assert_ai.core.collector import PhoenixCollector + + collector = PhoenixCollector(project_name="review-project") + with patch.object(collector._client.spans, "get_spans") as get_spans: + with self.assertRaisesRegex(ValueError, "valid ISO-8601"): + collector.get_spans(start_time="not-a-timestamp") + with self.assertRaisesRegex(ValueError, "timezone offset"): + collector.get_spans(start_time="2026-09-01T00:00:00") + + get_spans.assert_not_called() + + def test_otel_session_consumes_bounded_phoenix_spans(self) -> None: + from assert_ai.core.collector import PhoenixCollector + from assert_ai.core.model_client import Message + from assert_ai.core.otel_session import OTelTracedSession + + raw_span = _raw_span( + "span-session", + trace_id="trace-session", + attributes={ + "openinference.span.kind": "LLM", + "session.id": "session-1", + "output.value": "answer", + "llm.model_name": "gpt-5.4", + "llm.token_count.prompt": 5, + "llm.token_count.completion": 2, + }, + ) + collector = PhoenixCollector(project_name="review-project") + target_module = types.ModuleType("_phoenix_collector_target") + setattr(target_module, "target", lambda message: f"response to {message}") + sys.modules[target_module.__name__] = target_module + + async def run_session(): + session = OTelTracedSession( + callable_ref="_phoenix_collector_target:target", + collector=collector, + ) + await session.open() + try: + return await session.run_turn([Message(role="user", content="test")]) + finally: + await session.close() + + try: + with patch.object( + collector._client.spans, + "get_spans", + return_value=[raw_span], + ) as get_spans: + result = asyncio.run(run_session()) + finally: + sys.modules.pop(target_module.__name__, None) + + query = get_spans.call_args.kwargs + self.assertIsInstance(query["start_time"], datetime) + self.assertIsInstance(query["end_time"], datetime) + self.assertLessEqual(query["start_time"], query["end_time"]) + assert result.raw is not None + self.assertTrue(result.raw["span_validation"]["valid"]) + self.assertEqual(len(result.raw["trace_events"]), 1) + self.assertEqual(result.raw["trace_metadata"]["llm_call_count"], 1) + + +if __name__ == "__main__": + unittest.main()