Add Patronus OTEL backend to AIQToolkit

Git diff paste
 avatar
unknown
python
12 days ago
1.5 kB
7
Indexable
diff --git a/src/aiq/observability/register.py b/src/aiq/observability/register.py
index d48249a..0894d9c 100644
--- a/src/aiq/observability/register.py
+++ b/src/aiq/observability/register.py
@@ -155,3 +155,32 @@ async def file_logging_method(config: FileLoggingMethod, builder: Builder):
     handler = logging.FileHandler(filename=config.path, mode="a", encoding="utf-8")
     handler.setLevel(level)
     yield handler
+
+
+class PatronusTelemetryExporter(TelemetryExporterBaseConfig, name="patronus"):
+    """A telemetry exporter to transmit traces to Patronus service."""
+
+    endpoint: str = Field(description="The Patronus OTEL endpoint")
+    api_key: str = Field(description="The Patronus API key", default="")
+    project: str = Field(description="The project name to group the telemetry traces.")
+
+
+@register_telemetry_exporter(config_type=PatronusTelemetryExporter)
+async def patronus_telemetry_exporter(config: PatronusTelemetryExporter, builder: Builder):
+    """Create a Patronus telemetry exporter."""
+    trace_exporter = telemetry_optional_import("opentelemetry.exporter.otlp.proto.grpc.trace_exporter")
+
+    api_key = config.api_key or os.environ.get("PATRONUS_API_KEY")
+    if not api_key:
+        raise ValueError("API key is required for Patronus")
+
+    headers = {
+        "x-api-key": api_key,
+        "pat-project-name": config.project,
+    }
+    yield trace_exporter.OTLPSpanExporter(endpoint=config.endpoint, headers=headers)
Editor is loading...
Leave a Comment