adaptive-memory-multi-model-router 2.8.0 → 2.9.1

package/src/observability/middleware.ts

@@ -0,0 +1,164 @@
+import { Request, Response, NextFunction } from 'express';
+import { getTracer, createTracer } from './tracer';
+import { getMetrics, createMetricsCollector } from './metrics';
+
+/**
+ * Express middleware for observability
+ * - Adds trace ID to all requests
+ * - Records request/response metrics
+ * - Attaches span context
+ */
+export function observabilityMiddleware(
+  req: Request,
+  res: Response,
+  next: NextFunction
+): void {
+  const tracer = getTracer();
+  const metrics = getMetrics();
+
+  // Generate or extract trace ID
+  const traceId = (req.headers['x-trace-id'] as string) || tracer.generateTraceId();
+
+  // Attach to request object
+  (req as any).traceId = traceId;
+  (req as any).spanId = null;
+
+  // Add trace ID to response headers
+  res.setHeader('X-Trace-ID', traceId);
+
+  // Start span for this request
+  const span = tracer.startSpan(`${req.method} ${req.path}`);
+  (req as any).spanId = span.spanId;
+  (req as any).span = span;
+
+  // Record start time for latency calculation
+  const startTime = Date.now();
+
+  // Hook into response finish
+  res.on('finish', () => {
+    const duration = Date.now() - startTime;
+
+    // End the span
+    tracer.endSpan(span.spanId, {
+      method: req.method,
+      path: req.path,
+      statusCode: res.statusCode,
+      durationMs: duration,
+    });
+
+    // Record metrics
+    const model = (req as any).model || 'unknown';
+    const provider = (req as any).provider || 'unknown';
+    const cacheHit = (req as any).cacheHit || false;
+    const tier = (req as any).tier || 'standard';
+
+    // Request counter
+    metrics.incrementCounter('a3m_router_requests_total', {
+      model,
+      provider,
+      tier,
+      cache_hit: String(cacheHit),
+    });
+
+    // Latency histogram (convert to seconds)
+    metrics.recordHistogram('a3m_request_latency_seconds', duration / 1000, {
+      model,
+      provider,
+    });
+
+    // Error counter
+    if (res.statusCode >= 400) {
+      metrics.incrementCounter('a3m_router_errors_total', {
+        provider,
+        error_type: res.statusCode >= 500 ? 'server_error' : 'client_error',
+      });
+    }
+  });
+
+  next();
+}
+
+/**
+ * Fastify plugin for observability middleware
+ */
+export function observabilityPlugin(
+  instance: any,
+  options: any,
+  next: (err?: Error) => void
+): void {
+  instance.addHook('onRequest', async (request: any, reply: any) => {
+    const tracer = getTracer();
+    const metrics = getMetrics();
+
+    const traceId = request.headers['x-trace-id'] || tracer.generateTraceId();
+    request.traceId = traceId;
+
+    reply.header('X-Trace-ID', traceId);
+
+    const span = tracer.startSpan(`${request.method} ${request.url}`);
+    request.spanId = span.spanId;
+    request.span = span;
+    request.startTime = Date.now();
+  });
+
+  instance.addHook('onResponse', async (request: any, reply: any) => {
+    const tracer = getTracer();
+    const metrics = getMetrics();
+
+    const duration = Date.now() - request.startTime;
+
+    if (request.spanId) {
+      tracer.endSpan(request.spanId, {
+        method: request.method,
+        path: request.url,
+        statusCode: reply.statusCode,
+        durationMs: duration,
+      });
+    }
+
+    const model = request.model || 'unknown';
+    const provider = request.provider || 'unknown';
+    const cacheHit = request.cacheHit || false;
+    const tier = request.tier || 'standard';
+
+    metrics.incrementCounter('a3m_router_requests_total', {
+      model,
+      provider,
+      tier,
+      cache_hit: String(cacheHit),
+    });
+
+    metrics.recordHistogram('a3m_request_latency_seconds', duration / 1000, {
+      model,
+      provider,
+    });
+
+    if (reply.statusCode >= 400) {
+      metrics.incrementCounter('a3m_router_errors_total', {
+        provider,
+        error_type: reply.statusCode >= 500 ? 'server_error' : 'client_error',
+      });
+    }
+  });
+
+  next();
+}
+
+/**
+ * Middleware for budget warning alerts
+ */
+export function budgetAlertMiddleware(
+  req: Request,
+  res: Response,
+  next: NextFunction
+): void {
+  const tracer = getTracer();
+
+  // Hook into cost tracking to emit budget warnings
+  tracer.on('route_complete', (trace: any) => {
+    // This would be connected to budget enforcement in practice
+    // For now just a placeholder for the event hook
+  });
+
+  next();
+}

package/src/observability/tracer.ts

@@ -0,0 +1,229 @@
+import { EventEmitter } from 'events';
+import { Span, RouteTrace, ObservabilityEvent } from './types';
+
+// Langfuse integration interface
+interface LangfuseConfig {
+  publicKey?: string;
+  secretKey?: string;
+  baseUrl?: string;
+}
+
+interface LangfuseClient {
+  trace: (data: any) => Promise<void>;
+  flush: () => Promise<void>;
+}
+
+function generateId(): string {
+  return Math.random().toString(36).substring(2, 15) +
+         Math.random().toString(36).substring(2, 15);
+}
+
+class Tracer extends EventEmitter {
+  private traces: Map<string, Span> = new Map();
+  private routeTraces: RouteTrace[] = [];
+  private langfuseClient?: LangfuseClient;
+  private langfuseConfig?: LangfuseConfig;
+
+  constructor() {
+    super();
+  }
+
+  /**
+   * Initialize Langfuse if configured
+   */
+  initLangfuse(config: LangfuseConfig): void {
+    if (!config.publicKey || !config.secretKey) {
+      console.warn('[Tracer] Langfuse keys not provided, running without Langfuse');
+      return;
+    }
+
+    this.langfuseConfig = config;
+
+    // Simple Langfuse REST client
+    this.langfuseClient = {
+      trace: async (data: any) => {
+        try {
+          const response = await fetch(`${config.baseUrl || 'https://cloud.langfuse.com'}/api/public/ingestion`, {
+            method: 'POST',
+            headers: {
+              'Content-Type': 'application/json',
+              'Authorization': `Bearer ${config.publicKey}`,
+            },
+            body: JSON.stringify(data),
+          });
+          if (!response.ok) {
+            console.warn('[Tracer] Langfuse ingestion failed:', response.status);
+          }
+        } catch (err) {
+          console.warn('[Tracer] Langfuse error:', err);
+        }
+      },
+      flush: async () => {
+        // Langfuse uses async ingestion, noop here
+      },
+    };
+
+    console.log('[Tracer] Langfuse initialized');
+  }
+
+  /**
+   * Create a new trace span
+   */
+  startSpan(operationName: string, parentSpanId?: string): Span {
+    const traceId = parentSpanId
+      ? this.traces.get(parentSpanId)?.traceId || this.generateTraceId()
+      : this.generateTraceId();
+
+    const span: Span = {
+      traceId,
+      spanId: this.generateTraceId(),
+      parentSpanId,
+      operationName,
+      startTime: Date.now(),
+      attributes: {},
+      status: 'started',
+    };
+
+    this.traces.set(span.spanId, span);
+    this.emit('span_started', span);
+    this.emit('observability_event', { type: 'span_started', span } as ObservabilityEvent);
+
+    return span;
+  }
+
+  /**
+   * End a span
+   */
+  endSpan(spanId: string, attributes?: Record<string, any>): void {
+    const span = this.traces.get(spanId);
+    if (!span) {
+      console.warn(`[Tracer] span ${spanId} not found`);
+      return;
+    }
+
+    span.endTime = Date.now();
+    span.duration = span.endTime - span.startTime;
+    span.status = 'completed';
+
+    if (attributes) {
+      span.attributes = { ...span.attributes, ...attributes };
+    }
+
+    this.emit('span_completed', span);
+    this.emit('observability_event', { type: 'span_completed', span } as ObservabilityEvent);
+  }
+
+  /**
+   * Record an error on a span
+   */
+  errorSpan(spanId: string, error: string, attributes?: Record<string, any>): void {
+    const span = this.traces.get(spanId);
+    if (!span) {
+      console.warn(`[Tracer] span ${spanId} not found`);
+      return;
+    }
+
+    span.endTime = Date.now();
+    span.duration = span.endTime - span.startTime;
+    span.status = 'error';
+    span.error = error;
+
+    if (attributes) {
+      span.attributes = { ...span.attributes, ...attributes };
+    }
+
+    this.emit('error', spanId, error);
+    this.emit('observability_event', { type: 'error', spanId, error } as ObservabilityEvent);
+  }
+
+  /**
+   * Record a route decision
+   */
+  recordRoute(trace: RouteTrace): void {
+    this.routeTraces.push(trace);
+
+    // Emit event
+    this.emit('route_complete', trace);
+    this.emit('observability_event', { type: 'route_complete', trace } as ObservabilityEvent);
+
+    // Export to Langfuse async
+    if (this.langfuseClient) {
+      this.langfuseClient.trace({
+        name: 'route_decision',
+        traceId: trace.traceId,
+        timestamp: new Date(trace.timestamp).toISOString(),
+        input: { query: trace.query },
+        output: {
+          model: trace.model,
+          provider: trace.provider,
+          tokens: trace.queryTokens + trace.responseTokens,
+          latencyMs: trace.latencyMs,
+          cost: trace.cost,
+        },
+        metadata: {
+          queryTokens: trace.queryTokens,
+          responseTokens: trace.responseTokens,
+          cacheHit: trace.cacheHit,
+          complexity: trace.complexity,
+          tier: trace.tier,
+        },
+      }).catch(() => {}); // Fire and forget
+    }
+  }
+
+  /**
+   * Generate a unique trace ID
+   */
+  generateTraceId(): string {
+    return generateId();
+  }
+
+  /**
+   * Get all route traces
+   */
+  getTraces(limit?: number): RouteTrace[] {
+    if (limit) {
+      return this.routeTraces.slice(-limit);
+    }
+    return [...this.routeTraces];
+  }
+
+  /**
+   * Get all spans for a trace
+   */
+  getSpans(traceId: string): Span[] {
+    return Array.from(this.traces.values()).filter(s => s.traceId === traceId);
+  }
+
+  /**
+   * Export to Langfuse (async, non-blocking)
+   */
+  async flushToLangfuse(): Promise<void> {
+    if (this.langfuseClient) {
+      await this.langfuseClient.flush();
+    }
+  }
+
+  /**
+   * Clear all traces (for testing)
+   */
+  reset(): void {
+    this.traces.clear();
+    this.routeTraces = [];
+  }
+}
+
+// Singleton instance
+let tracerInstance: Tracer | null = null;
+
+export function getTracer(): Tracer {
+  if (!tracerInstance) {
+    tracerInstance = new Tracer();
+  }
+  return tracerInstance;
+}
+
+export function createTracer(): Tracer {
+  tracerInstance = new Tracer();
+  return tracerInstance;
+}

package/src/observability/types.ts

@@ -0,0 +1,45 @@
+// Observability types for A3M Router
+
+export interface Span {
+  traceId: string;
+  spanId: string;
+  parentSpanId?: string;
+  operationName: string;
+  startTime: number;
+  endTime?: number;
+  duration?: number;
+  attributes: Record<string, string | number | boolean>;
+  status: 'started' | 'completed' | 'error';
+  error?: string;
+}
+
+export interface Metric {
+  name: string;
+  value: number;
+  timestamp: number;
+  labels: Record<string, string>;
+  type: 'counter' | 'gauge' | 'histogram';
+}
+
+export interface RouteTrace {
+  traceId: string;
+  timestamp: number;
+  query: string;
+  queryTokens: number;
+  model: string;
+  provider: string;
+  responseTokens: number;
+  latencyMs: number;
+  cost: number;
+  cacheHit: boolean;
+  complexity: number;
+  tier: 'budget' | 'standard' | 'premium';
+}
+
+// Event types emitted by Tracer
+export type ObservabilityEvent =
+  | { type: 'route_complete'; trace: RouteTrace }
+  | { type: 'span_started'; span: Span }
+  | { type: 'span_completed'; span: Span }
+  | { type: 'budget_warning'; apiKey: string; usedCents: number; thresholdCents: number }
+  | { type: 'error'; spanId: string; error: string };