231 lines
7.8 KiB
Python
231 lines
7.8 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from agent_framework.observability.code_mapper import ObservabilityCodeMapper
|
|
from agent_framework.observability.observer import AgentObserver
|
|
from agent_framework.observability.telemetry import Telemetry
|
|
from app.observability.telemetry_observer import TelemetryBackedAgentObserver
|
|
|
|
|
|
class CaptureAnalytics:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def publish(self, event_type, payload):
|
|
self.calls.append((event_type, payload))
|
|
|
|
|
|
class CaptureBus:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def publish(self, event_type, payload, **kwargs):
|
|
self.calls.append((event_type, payload, kwargs))
|
|
|
|
|
|
def test_mapper_passthrough_and_contract_fields(tmp_path: Path):
|
|
mapper = ObservabilityCodeMapper({"GRL.TOXOUT": "GRL.004"})
|
|
code, payload, metadata = mapper.normalize_payload("GRL.TOXOUT", {"x": 1}, {})
|
|
assert code == "GRL.004"
|
|
assert payload["event_code_internal"] == "GRL.TOXOUT"
|
|
assert metadata["event_code_internal"] == "GRL.TOXOUT"
|
|
assert metadata["event_code_mapped"] == "GRL.004"
|
|
assert mapper.map("GRL.NEW") == "GRL.NEW"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_observer_maps_before_analytics_and_event_bus():
|
|
analytics = CaptureAnalytics()
|
|
bus = CaptureBus()
|
|
observer = AgentObserver(
|
|
analytics=analytics,
|
|
event_bus=bus,
|
|
code_mapper=ObservabilityCodeMapper({"GRL.TOXOUT": "GRL.004"}),
|
|
)
|
|
event = await observer.emit("GRL.TOXOUT", {"reason": "tox"})
|
|
|
|
assert event["eventType"] == "GRL.004"
|
|
assert event["payload"]["tag"] == "GRL.004"
|
|
assert event["payload"]["event_code_internal"] == "GRL.TOXOUT"
|
|
assert analytics.calls[0][0] == "GRL.004"
|
|
assert bus.calls[0][0] == "GRL.004"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_telemetry_maps_direct_events(tmp_path: Path):
|
|
mapping = tmp_path / "mapping.yaml"
|
|
mapping.write_text("mappings:\n GRL.TOXOUT: GRL.004\n", encoding="utf-8")
|
|
settings = SimpleNamespace(
|
|
OBSERVABILITY_CODE_MAPPING_ENABLED=True,
|
|
OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
|
|
ENABLE_LANGFUSE=False,
|
|
ENABLE_OCI_STREAMING=False,
|
|
ENABLE_OTEL=False,
|
|
LANGFUSE_TRACE_MODE="verbose",
|
|
)
|
|
telemetry = Telemetry(settings)
|
|
capture = []
|
|
|
|
async def subscriber(evt):
|
|
capture.append(evt)
|
|
|
|
telemetry.event_bus.subscribe(subscriber)
|
|
await telemetry.event("GRL.TOXOUT", {"a": 1}, kind="grl")
|
|
|
|
assert capture
|
|
assert capture[0].name == "GRL.004"
|
|
assert capture[0].payload["event_code_internal"] == "GRL.TOXOUT"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_contas_telemetry_observer_returns_mapped_envelope(tmp_path: Path):
|
|
mapping = tmp_path / "mapping.yaml"
|
|
mapping.write_text("mappings:\n GRL.TOXOUT: GRL.004\n", encoding="utf-8")
|
|
settings = SimpleNamespace(
|
|
OBSERVABILITY_CODE_MAPPING_ENABLED=True,
|
|
OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
|
|
ENABLE_LANGFUSE=False,
|
|
ENABLE_OCI_STREAMING=False,
|
|
ENABLE_OTEL=False,
|
|
LANGFUSE_TRACE_MODE="verbose",
|
|
)
|
|
telemetry = Telemetry(settings)
|
|
observer = TelemetryBackedAgentObserver(telemetry)
|
|
event = await observer.emit("GRL.TOXOUT", {"rail": "TOXOUT"})
|
|
|
|
assert event["eventType"] == "GRL.004"
|
|
assert event["body"]["tag"] == "GRL.004"
|
|
assert event["metadata"]["event_code_internal"] == "GRL.TOXOUT"
|
|
|
|
class _FakeObservation:
|
|
def __init__(self):
|
|
self.updates = []
|
|
|
|
def update(self, **kwargs):
|
|
self.updates.append(kwargs)
|
|
|
|
|
|
class _FakeObservationCM:
|
|
def __init__(self, observation):
|
|
self.observation = observation
|
|
|
|
def __enter__(self):
|
|
return self.observation
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
|
|
class _FakeLangfuse:
|
|
def __init__(self):
|
|
self.started = []
|
|
|
|
def start_as_current_observation(self, **kwargs):
|
|
self.started.append(dict(kwargs))
|
|
return _FakeObservationCM(_FakeObservation())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generation_observation_name_is_mapped_before_langfuse_and_event_bus(tmp_path: Path):
|
|
mapping = tmp_path / "mapping.yaml"
|
|
mapping.write_text(
|
|
"mappings:\n guardrail.dlex_in: GRL.004\n guardrail.tox: GRL.005\n",
|
|
encoding="utf-8",
|
|
)
|
|
settings = SimpleNamespace(
|
|
OBSERVABILITY_CODE_MAPPING_ENABLED=True,
|
|
OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
|
|
ENABLE_LANGFUSE=False,
|
|
ENABLE_OCI_STREAMING=False,
|
|
ENABLE_OTEL=False,
|
|
LANGFUSE_TRACE_MODE="verbose",
|
|
)
|
|
telemetry = Telemetry(settings)
|
|
fake = _FakeLangfuse()
|
|
telemetry.langfuse = fake
|
|
telemetry.enabled = True
|
|
capture = []
|
|
|
|
async def subscriber(evt):
|
|
capture.append(evt)
|
|
|
|
telemetry.event_bus.subscribe(subscriber)
|
|
async with telemetry.generation_span(
|
|
"guardrail.dlex_in",
|
|
"model-x",
|
|
"input",
|
|
metadata={"component": "guardrail.dlex_in"},
|
|
) as generation:
|
|
generation.set_output("ok")
|
|
|
|
assert fake.started[0]["name"] == "GRL.004"
|
|
assert fake.started[0]["metadata"]["observability_name_internal"] == "guardrail.dlex_in"
|
|
assert fake.started[0]["metadata"]["observability_name_mapped"] == "GRL.004"
|
|
assert capture[0].name == "GRL.004"
|
|
assert capture[0].payload["observability_name_internal"] == "guardrail.dlex_in"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_span_name_is_mapped_before_otel_event_bus_and_langfuse(tmp_path: Path):
|
|
mapping = tmp_path / "mapping.yaml"
|
|
mapping.write_text("mappings:\n guardrail.tox: GRL.005\n", encoding="utf-8")
|
|
settings = SimpleNamespace(
|
|
OBSERVABILITY_CODE_MAPPING_ENABLED=True,
|
|
OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
|
|
ENABLE_LANGFUSE=False,
|
|
ENABLE_OCI_STREAMING=False,
|
|
ENABLE_OTEL=False,
|
|
LANGFUSE_TRACE_MODE="verbose",
|
|
)
|
|
telemetry = Telemetry(settings)
|
|
fake = _FakeLangfuse()
|
|
telemetry.langfuse = fake
|
|
telemetry.enabled = True
|
|
capture = []
|
|
|
|
async def subscriber(evt):
|
|
capture.append(evt)
|
|
|
|
telemetry.event_bus.subscribe(subscriber)
|
|
async with telemetry.span("guardrail.tox", input={"text": "x"}):
|
|
pass
|
|
|
|
assert fake.started[0]["name"] == "GRL.005"
|
|
assert fake.started[0]["metadata"]["observability_name_internal"] == "guardrail.tox"
|
|
assert capture[0].name == "GRL.005.started"
|
|
assert capture[-1].name == "GRL.005.completed"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_langfuse_analytics_direct_sdk_path_also_maps_observation_name(tmp_path: Path):
|
|
from agent_framework.analytics.providers.langfuse import LangfuseAnalyticsPublisher
|
|
|
|
mapping = tmp_path / "mapping.yaml"
|
|
mapping.write_text("mappings:\n guardrail.dlex_in: GRL.004\n", encoding="utf-8")
|
|
settings = SimpleNamespace(
|
|
OBSERVABILITY_CODE_MAPPING_ENABLED=True,
|
|
OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
|
|
LANGFUSE_PUBLIC_KEY="pk",
|
|
LANGFUSE_SECRET_KEY="sk",
|
|
LANGFUSE_HOST="http://localhost",
|
|
)
|
|
fake = _FakeLangfuse()
|
|
publisher = LangfuseAnalyticsPublisher(settings=settings, langfuse=fake)
|
|
|
|
# Give the technical event correlation so the provider does not intentionally skip it.
|
|
from agent_framework.observability.context import set_observability_context, clear_observability_context
|
|
clear_observability_context()
|
|
set_observability_context(request_id="req-123", trace_id="req-123")
|
|
try:
|
|
await publisher.publish("guardrail.dlex_in", {"metadata": {"request_id": "req-123"}})
|
|
finally:
|
|
clear_observability_context()
|
|
|
|
assert fake.started, "Langfuse analytics provider did not create an observation"
|
|
assert fake.started[0]["name"] == "GRL.004"
|
|
assert fake.started[0]["metadata"]["observability_name_internal"] == "guardrail.dlex_in"
|
|
assert fake.started[0]["metadata"]["observability_name_mapped"] == "GRL.004"
|