Add support for Axiom backend

[nirga]

Description

Add support for Axiom as an observability backend, enabling users to query OpenTelemetry traces from Axiom's serverless data platform.

Motivation

Axiom is a serverless observability platform with excellent OpenTelemetry support, designed for high-volume data at low cost. It's particularly popular with serverless and edge computing workloads. Adding Axiom support would enable users running modern cloud-native applications to leverage this MCP server.

Implementation Requirements

1. Backend Implementation

• Create src/openllmetry_mcp/backends/axiom.py following the abstract interface in backends/base.py
• Implement the following methods:
  • search_traces() - Query traces using Axiom Query API (APL)
  • get_trace() - Retrieve full trace details
  • list_services() - List available services
  • get_aggregated_usage() - Aggregate token usage metrics

2. Configuration

• Add axiom as a backend type option
• Required configuration:
  • BACKEND_URL: Axiom API URL (default: https://api.axiom.co or https://cloud.axiom.co)
  • BACKEND_API_KEY: Axiom API token with query permissions
  • AXIOM_DATASET: Dataset name where traces are stored
  • AXIOM_ORG_ID: Organization ID (optional, for multi-org scenarios)
• Update .env.example with Axiom configuration example

3. API Integration

Reference Axiom APIs:

• Axiom Query API
• APL (Axiom Processing Language)
• OpenTelemetry with Axiom
• Authentication: Bearer token in Authorization header

4. Query Language (APL)

Axiom uses APL (Axiom Processing Language), similar to KQL:

['otel-traces']
| where _time >= datetime(2024-01-01T00:00:00Z)
| where _time <= datetime(2024-01-01T23:59:59Z)
| where serviceName == "api"
| where ['gen_ai.system'] == "openai"
| project traceId, spanId, serviceName, duration, 
          gen_ai_system=['gen_ai.system'],
          gen_ai_model=['gen_ai.request.model'],
          prompt_tokens=tolong(['gen_ai.usage.prompt_tokens']),
          completion_tokens=tolong(['gen_ai.usage.completion_tokens'])
| limit 100

5. Trace Data Structure

Axiom stores OTel traces with flattened attribute structure:

• Span attributes become top-level fields with original names
• Resource attributes prefixed or stored separately
• _time field for timestamp
• Support for nested JSON attributes

6. OpenLLMetry Support

Ensure proper parsing of OpenLLMetry semantic conventions (gen_ai.* attributes) from Axiom's flattened trace data.

7. Documentation

• Add Axiom backend configuration to README.md
• Include Claude Desktop integration example
• Document dataset setup and OTel integration
• Add troubleshooting section for common Axiom connection issues
• Include APL query examples

8. Testing

• Add unit tests in tests/backends/test_axiom.py
• Test APL query generation
• Test trace querying, filtering, and aggregation
• Test error handling and authentication
• Mock Axiom API responses

Example Configuration

BACKEND_TYPE=axiom
BACKEND_URL=https://api.axiom.co
BACKEND_API_KEY=xaat-your-api-token-here
AXIOM_DATASET=otel-traces

APL Query Examples

Search traces with filters

['otel-traces']
| where _time between (datetime(2024-01-01) .. datetime(2024-01-02))
| where serviceName == "api"
| where statusCode == "ERROR"

Aggregate token usage

['otel-traces']
| where isnotempty(['gen_ai.system'])
| summarize 
    total_prompt_tokens=sum(tolong(['gen_ai.usage.prompt_tokens'])),
    total_completion_tokens=sum(tolong(['gen_ai.usage.completion_tokens'])),
    request_count=count()
  by model=['gen_ai.request.model'], service=serviceName

Axiom-Specific Features

1. Serverless Scale: No infrastructure management
2. Fast Ingestion: Real-time data availability
3. APL Power: Advanced query capabilities
4. Virtual Fields: Can create computed fields
5. Streaming: Real-time query results

Dataset Setup

Document how to send OTel traces to Axiom:

# OTel Collector config
exporters:
  otlphttp:
    endpoint: https://api.axiom.co/v1/traces
    headers:
      authorization: Bearer ${AXIOM_TOKEN}
      x-axiom-dataset: otel-traces

References

• Axiom Official Docs
• Axiom OpenTelemetry Guide
• Axiom APL Reference
• Axiom REST API
• Existing backend implementations: backends/jaeger.py, backends/tempo.py, backends/traceloop.py

Considerations

• APL syntax differs significantly from other backends
• Field names may be flattened differently than expected
• Rate limiting on API calls (document limits)
• Dataset must be pre-configured with OTel data
• Virtual fields could be used for computed attributes

Acceptance Criteria

• Axiom backend implementation complete
• All abstract methods implemented
• APL query builder for search filters
• Configuration documented
• Dataset setup guide included
• Tests passing with mocked Axiom API
• README updated with Axiom examples
• Works with Claude Desktop integration
• Handle APL-specific query limitations
• Document rate limits and best practices