# Psyche Tracing

Distributed tracing for Psyche AI Virtual Assistant.

## Overview

`psyche-tracing` provides OpenTelemetry-based distributed tracing aligned with
the `@oshun/tracing` TypeScript patterns for cross-language compatibility. It
includes context propagation, instrumentation decorators, and AWS X-Ray support.

## Features

- **OpenTelemetry Integration**: Full OpenTelemetry SDK support with OTLP export
- **W3C Trace Context**: Standard context propagation via traceparent/tracestate
- **Instrumentation Decorators**: Pre-built tracing for common operations
- **AWS X-Ray Support**: Native X-Ray header format and ECS metadata
- **Chainable API**: Fluent span methods for clean code
- **Type Safety**: Branded types for trace/span IDs

## Installation

```bash
# Basic installation
poetry add psyche-tracing

# With X-Ray support
poetry add psyche-tracing[xray]
```

## Quick Start

### Basic Tracing

```python
from psyche_tracing import OshunTracer, TracerConfig, SpanStatus

# Create tracer
config = TracerConfig(
    service_name="my-service",
    service_version="1.0.0",
    exporter_endpoint="http://localhost:4318/v1/traces",
)
tracer = OshunTracer(config)

# Use context manager for automatic lifecycle
with tracer.span("process_request") as span:
    span.set_attribute("user_id", "user-123")
    span.add_event("processing_started")

    # Do work...

    span.set_status(SpanStatus.OK)
```

### Async Tracing

```python
async with tracer.async_span("async_operation") as span:
    span.set_attribute("operation", "fetch_data")
    result = await fetch_data()
    span.add_event("data_fetched", {"count": len(result)})
```

### Decorator Style

```python
@tracer.trace("my_function")
async def my_function(arg1: str) -> str:
    return f"processed: {arg1}"

# Span automatically created with name "my_function"
result = await my_function("test")
```

### Global Tracer

```python
from psyche_tracing import initialize_tracer, get_tracer

# Initialize once at startup
initialize_tracer(TracerConfig(
    service_name="my-service",
    exporter_endpoint="http://localhost:4318/v1/traces",
))

# Use anywhere
tracer = get_tracer()
with tracer.span("operation"):
    pass
```

## Context Propagation

### W3C Trace Context

```python
from psyche_tracing import (
    extract_trace_context,
    inject_trace_context,
    parse_traceparent,
    format_traceparent,
)

# Extract from incoming request
headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}
context = extract_trace_context(headers)

# Start span with extracted context
span = tracer.start_span_with_context("handle_request", context)

# Inject into outgoing request
outgoing_headers: dict[str, str] = {}
inject_trace_context(span.context, outgoing_headers)
```

### Correlation IDs

```python
from psyche_tracing import extract_correlation_id, inject_correlation_id

# Extract from headers (checks multiple common header names)
correlation_id = extract_correlation_id(headers)

# Inject into outgoing request
inject_correlation_id(correlation_id, outgoing_headers)
```

## Instrumentation

### Service Calls

```python
from psyche_tracing import ServiceCallConfig, trace_service_call_async

async with trace_service_call_async(tracer, ServiceCallConfig(
    service="auth-service",
    method="validateToken",
)) as span:
    result = await auth_client.validate_token(token)
    span.set_attribute("token_valid", result.valid)
```

### Database Calls

```python
from psyche_tracing import DatabaseCallConfig, trace_database_call_async

async with trace_database_call_async(tracer, DatabaseCallConfig(
    system="postgresql",
    operation="query",
    database="users",
    table="accounts",
    statement="SELECT * FROM accounts WHERE id = $1",
)) as span:
    result = await db.fetch_one(query, account_id)
```

### External APIs

```python
from psyche_tracing import ExternalApiCallConfig, trace_external_api_async

async with trace_external_api_async(tracer, ExternalApiCallConfig(
    provider="stripe",
    operation="charges.create",
)) as span:
    charge = await stripe.charges.create(amount=1000)
    span.set_attribute("charge_id", charge.id)
```

### AI/LLM Calls

```python
from psyche_tracing import AiCallConfig, trace_ai_call_async, record_ai_response

async with trace_ai_call_async(tracer, AiCallConfig(
    model="gpt-4",
    provider="openai",
    operation="chat",
    input_tokens=100,
)) as span:
    response = await openai.chat.completions.create(...)
    record_ai_response(
        span,
        output_tokens=response.usage.completion_tokens,
        total_tokens=response.usage.total_tokens,
    )
```

### Messaging

```python
from psyche_tracing import MessageConfig, trace_message_async

async with trace_message_async(tracer, MessageConfig(
    system="kafka",
    destination="events-topic",
    operation="publish",
    message_id="msg-123",
)) as span:
    await producer.send("events-topic", message)
```

### Jobs

```python
from psyche_tracing import JobConfig, trace_job_async

async with trace_job_async(tracer, JobConfig(
    job_type="process_upload",
    job_id="job-456",
    queue="uploads",
    attempt=1,
)) as span:
    await process_upload(upload_id)
```

### Batch Operations

```python
from psyche_tracing import trace_batch

async def process_item(span, item, index):
    span.set_attribute(f"item.{index}.id", item.id)
    return await process(item)

results = await trace_batch(tracer, "process_items", items, process_item)
```

### Retry Operations

```python
from psyche_tracing import RetryConfig, trace_with_retry

result = await trace_with_retry(
    tracer,
    "fetch_with_retry",
    RetryConfig(max_attempts=3, delay_ms=1000, backoff=2.0),
    async def operation(span, attempt):
        span.add_event("attempting", {"attempt": attempt})
        return await fetch_data()
)
```

## AWS X-Ray Integration

### X-Ray Header Handling

```python
from psyche_tracing import (
    parse_xray_header,
    format_xray_header,
    xray_header_to_span_context,
    span_context_to_xray_header,
)

# Extract from X-Ray header
xray_header = request.headers.get("X-Amzn-Trace-Id")
context = xray_header_to_span_context(xray_header)

# Format for X-Ray
xray_header = span_context_to_xray_header(span.context)
```

### X-Ray Tracer Configuration

```python
from psyche_tracing import XRayTracerConfig, create_xray_tracer_config, OshunTracer

xray_config = XRayTracerConfig(
    service_name="my-service",
    aws_region="us-east-1",
    use_daemon=True,  # Use X-Ray daemon sidecar
    sample_rate=0.05,  # 5% sampling in production
)

config = create_xray_tracer_config(xray_config)
tracer = OshunTracer(config)
```

### ECS Metadata

```python
from psyche_tracing import fetch_ecs_metadata, get_ecs_tracing_metadata

# Get ECS container metadata
metadata = await fetch_ecs_metadata()
if metadata:
    print(f"Task ARN: {metadata.task_arn}")
    print(f"Cluster: {metadata.cluster_name}")

# Get as span attributes
attrs = await get_ecs_tracing_metadata()
span.set_attributes(attrs)
```

## Span Attributes

### HTTP Semantic Conventions

```python
from psyche_tracing import HttpSpanAttributes

span.set_attributes({
    HttpSpanAttributes.METHOD: "POST",
    HttpSpanAttributes.URL: "https://api.example.com/users",
    HttpSpanAttributes.STATUS_CODE: 200,
})
```

### Database Semantic Conventions

```python
from psyche_tracing import DbSpanAttributes

span.set_attributes({
    DbSpanAttributes.SYSTEM: "postgresql",
    DbSpanAttributes.NAME: "users",
    DbSpanAttributes.OPERATION: "query",
})
```

### AI/LLM Attributes

```python
from psyche_tracing import AiSpanAttributes

span.set_attributes({
    AiSpanAttributes.MODEL: "claude-3",
    AiSpanAttributes.PROVIDER: "anthropic",
    AiSpanAttributes.INPUT_TOKENS: 1000,
    AiSpanAttributes.OUTPUT_TOKENS: 500,
})
```

## Configuration

### TracerConfig Options

| Option                    | Type     | Default  | Description                    |
| ------------------------- | -------- | -------- | ------------------------------ |
| `service_name`            | str      | required | Name of the service            |
| `service_version`         | str      | "1.0.0"  | Service version                |
| `environment`             | str      | None     | Deployment environment         |
| `enabled`                 | bool     | True     | Whether tracing is enabled     |
| `exporter_endpoint`       | str      | None     | OTLP exporter endpoint         |
| `exporter_headers`        | dict     | None     | Headers for OTLP exporter      |
| `sample_rate`             | float    | 1.0      | Sampling rate (0.0-1.0)        |
| `use_batch_processor`     | bool     | True     | Use batch span processor       |
| `batch_flush_interval_ms` | int      | 5000     | Batch flush interval           |
| `max_batch_size`          | int      | 512      | Maximum batch size             |
| `max_queue_size`          | int      | 2048     | Maximum queue size             |
| `resource_attributes`     | dict     | None     | Additional resource attributes |
| `on_span_end`             | callable | None     | Callback when span ends        |
| `debug`                   | bool     | False    | Enable console exporter        |

### Environment Variables

```bash
# Service info
SERVICE_NAME=my-service
SERVICE_VERSION=1.0.0
ENVIRONMENT=production

# OpenTelemetry
OTEL_ENABLED=true
OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318/v1/traces
OTEL_SAMPLE_RATE=0.1
OTEL_DEBUG=false
```

## Integration with @oshun/tracing

This library is designed to be compatible with the TypeScript `@oshun/tracing`:

- Same W3C trace context format
- Same span attribute names
- Same X-Ray integration patterns
- Compatible instrumentation patterns

Traces from Python services can be correlated with TypeScript service traces.

## Dependencies

- `pydantic` >= 2.5.0 - Type validation
- `opentelemetry-api` >= 1.20.0 - OpenTelemetry API
- `opentelemetry-sdk` >= 1.20.0 - OpenTelemetry SDK
- `opentelemetry-exporter-otlp-proto-http` >= 1.20.0 - OTLP exporter
- `opentelemetry-propagator-aws-xray` >= 1.0.0 (optional) - X-Ray support

## License

Proprietary - Oshun Platform