@fallom/trace 0.1.11 → 0.2.0

package/dist/index.d.ts

@@ -2,10 +2,10 @@ import { SpanExporter, ReadableSpan } from '@opentelemetry/sdk-trace-base';
 import { ExportResult } from '@opentelemetry/core';
 
 /**
- * Fallom tracing module.
- *
- * Auto-instruments all LLM calls via OTEL and groups them by session.
- * Also supports custom spans for business metrics.
+ * Type definitions for Fallom tracing module.
+ */
+/**
+ * Session context for grouping traces.
  */
 interface SessionContext {
     configKey: string;
@@ -13,27 +13,92 @@ interface SessionContext {
     customerId?: string;
 }
 /**
- * Initialize Fallom tracing. Auto-instruments all LLM calls.
+ * Trace context for linking spans together.
+ */
+interface TraceContext {
+    traceId: string;
+    parentSpanId?: string;
+}
+/**
+ * Data structure for a trace sent to the Fallom API.
+ */
+interface TraceData {
+    config_key: string;
+    session_id: string;
+    customer_id?: string;
+    trace_id: string;
+    span_id: string;
+    parent_span_id?: string;
+    name: string;
+    kind?: string;
+    model?: string;
+    start_time: string;
+    end_time: string;
+    duration_ms: number;
+    status: "OK" | "ERROR";
+    error_message?: string;
+    prompt_tokens?: number;
+    completion_tokens?: number;
+    total_tokens?: number;
+    time_to_first_token_ms?: number;
+    is_streaming?: boolean;
+    attributes?: Record<string, unknown>;
+    prompt_key?: string;
+    prompt_version?: number;
+    prompt_ab_test_key?: string;
+    prompt_variant_index?: number;
+}
+/**
+ * Options for creating a Fallom session.
+ */
+interface SessionOptions {
+    /** Your config name (e.g., "linkedin-agent") */
+    configKey: string;
+    /** Your session/conversation ID */
+    sessionId: string;
+    /** Optional customer/user identifier for analytics */
+    customerId?: string;
+}
+/**
+ * Options for wrapAISDK.
+ */
+interface WrapAISDKOptions {
+    /**
+     * Enable debug logging to see the raw Vercel AI SDK response structure.
+     * Useful for debugging token extraction issues with different providers.
+     */
+    debug?: boolean;
+}
+
+/**
+ * Core Fallom tracing functionality.
+ *
+ * Handles initialization and trace sending.
+ * Session management is now handled by FallomSession.
+ */
+
+/**
+ * Initialize Fallom tracing.
  *
  * @param options - Configuration options
  * @param options.apiKey - Your Fallom API key. Defaults to FALLOM_API_KEY env var.
- * @param options.baseUrl - API base URL. Defaults to FALLOM_BASE_URL env var, or https://spans.fallom.com
- * @param options.captureContent - Whether to capture prompt/completion content in traces.
- *                                 Set to false for privacy/compliance. Defaults to true.
- *                                 Also respects FALLOM_CAPTURE_CONTENT env var ("true"/"false").
+ * @param options.baseUrl - API base URL. Defaults to https://traces.fallom.com
+ * @param options.captureContent - Whether to capture prompt/completion content.
+ * @param options.debug - Enable debug logging.
  *
  * @example
  * ```typescript
- * import fallom from 'fallom';
+ * import fallom from '@fallom/trace';
  *
- * // Normal usage (captures everything)
- * fallom.trace.init();
+ * await fallom.init({ apiKey: process.env.FALLOM_API_KEY });
  *
- * // Privacy mode (no prompts/completions stored)
- * fallom.trace.init({ captureContent: false });
+ * const session = fallom.session({
+ *   configKey: "my-agent",
+ *   sessionId: "session-123",
+ * });
  *
- * fallom.trace.setSession("my-agent", sessionId);
- * await agent.run(message); // Automatically traced
+ * const { generateText } = session.wrapAISDK(ai);
+ * await generateText({ model: openai("gpt-4o"), prompt: "Hello!" });
  * ```
  */
 declare function init$3(options?: {
@@ -42,228 +107,122 @@ declare function init$3(options?: {
     captureContent?: boolean;
     debug?: boolean;
 }): Promise<void>;
-/**
- * Set the current session context.
- *
- * All subsequent LLM calls in this async context will be
- * automatically tagged with this configKey, sessionId, and customerId.
- *
- * @param configKey - Your config name (e.g., "linkedin-agent")
- * @param sessionId - Your session/conversation ID
- * @param customerId - Optional customer/user identifier for analytics
- *
- * @example
- * ```typescript
- * trace.setSession("linkedin-agent", sessionId, "user_123");
- * await agent.run(message); // Automatically traced with session + customer
- * ```
- */
-declare function setSession(configKey: string, sessionId: string, customerId?: string): void;
-/**
- * Run a function with session context.
- * Use this to ensure session context propagates across async boundaries.
- *
- * @param configKey - Your config name
- * @param sessionId - Your session ID
- * @param customerId - Optional customer/user identifier
- * @param fn - Function to run with session context
- *
- * @example
- * ```typescript
- * await trace.runWithSession("my-agent", sessionId, "user_123", async () => {
- *   await agent.run(message); // Has session context
- * });
- * ```
- */
-declare function runWithSession<T>(configKey: string, sessionId: string, customerIdOrFn: string | (() => T), fn?: () => T): T;
-/**
- * Get current session context, if any.
- */
-declare function getSession(): SessionContext | undefined;
-/**
- * Clear session context.
- */
-declare function clearSession(): void;
-/**
- * Record custom business metrics. Latest value per field wins.
- *
- * Use this for metrics that OTEL can't capture automatically:
- * - Outlier scores
- * - Engagement metrics
- * - Conversion rates
- * - Any business-specific outcome
- *
- * @param data - Dict of metrics to record
- * @param options - Optional session identifiers
- * @param options.configKey - Config name (optional if setSession was called)
- * @param options.sessionId - Session ID (optional if setSession was called)
- *
- * @example
- * ```typescript
- * // If session context is set:
- * trace.span({ outlier_score: 0.8, engagement: 42 });
- *
- * // Or explicitly:
- * trace.span(
- *   { outlier_score: 0.8 },
- *   { configKey: "linkedin-agent", sessionId: "user123-convo456" }
- * );
- * ```
- */
-declare function span(data: Record<string, unknown>, options?: {
-    configKey?: string;
-    sessionId?: string;
-}): void;
 /**
  * Shutdown the tracing SDK gracefully.
  */
 declare function shutdown(): Promise<void>;
+
 /**
- * Wrap an OpenAI client to automatically trace all chat completions.
- * Works with OpenAI, OpenRouter, Azure OpenAI, LiteLLM, and any OpenAI-compatible API.
- *
- * @param client - The OpenAI client instance
- * @returns The same client with tracing enabled
- *
- * @example
- * ```typescript
- * import OpenAI from "openai";
- * import { trace } from "@fallom/trace";
- *
- * const openai = trace.wrapOpenAI(new OpenAI());
- *
- * trace.setSession("my-config", sessionId);
- * const response = await openai.chat.completions.create({...}); // Automatically traced!
- * ```
+ * FallomSession - Session-scoped tracing for concurrent-safe operations.
  */
-declare function wrapOpenAI<T extends {
-    chat: {
-        completions: {
-            create: (...args: any[]) => Promise<any>;
-        };
-    };
-}>(client: T): T;
+
 /**
- * Wrap an Anthropic client to automatically trace all message creations.
+ * A session-scoped Fallom instance.
  *
- * @param client - The Anthropic client instance
- * @returns The same client with tracing enabled
+ * All wrappers created from this session automatically use the session context,
+ * making them safe for concurrent operations without global state issues.
  *
  * @example
  * ```typescript
- * import Anthropic from "@anthropic-ai/sdk";
- * import { trace } from "@fallom/trace";
+ * const session = fallom.session({
+ *   configKey: "my-app",
+ *   sessionId: "session-123",
+ *   customerId: "user-456"
+ * });
  *
- * const anthropic = trace.wrapAnthropic(new Anthropic());
+ * // All calls use the session context
+ * const { generateText } = session.wrapAISDK(ai);
+ * await generateText({ model: openai("gpt-4o"), prompt: "..." });
  *
- * trace.setSession("my-config", sessionId);
- * const response = await anthropic.messages.create({...}); // Automatically traced!
+ * // Or wrap the model directly
+ * const model = session.traceModel(openai("gpt-4o"));
+ * await generateText({ model, prompt: "..." });
  * ```
  */
-declare function wrapAnthropic<T extends {
-    messages: {
-        create: (...args: any[]) => Promise<any>;
+declare class FallomSession {
+    private ctx;
+    constructor(options: SessionOptions);
+    /** Get the session context. */
+    getContext(): SessionContext;
+    /**
+     * Get model assignment for this session (A/B testing).
+     */
+    getModel(configKeyOrOptions?: string | {
+        fallback?: string;
+        version?: number;
+    }, options?: {
+        fallback?: string;
+        version?: number;
+    }): Promise<string>;
+    /**
+     * Wrap a Vercel AI SDK model to trace all calls.
+     */
+    traceModel<T extends {
+        doGenerate?: (...args: any[]) => Promise<any>;
+        doStream?: (...args: any[]) => Promise<any>;
+    }>(model: T): T;
+    /** Wrap OpenAI client. Delegates to shared wrapper. */
+    wrapOpenAI<T extends {
+        chat: {
+            completions: {
+                create: (...args: any[]) => Promise<any>;
+            };
+        };
+    }>(client: T): T;
+    /** Wrap Anthropic client. Delegates to shared wrapper. */
+    wrapAnthropic<T extends {
+        messages: {
+            create: (...args: any[]) => Promise<any>;
+        };
+    }>(client: T): T;
+    /** Wrap Google AI model. Delegates to shared wrapper. */
+    wrapGoogleAI<T extends {
+        generateContent: (...args: any[]) => Promise<any>;
+    }>(model: T): T;
+    /** Wrap Vercel AI SDK. Delegates to shared wrapper. */
+    wrapAISDK<T extends {
+        generateText: (...args: any[]) => Promise<any>;
+        streamText: (...args: any[]) => any;
+        generateObject?: (...args: any[]) => Promise<any>;
+        streamObject?: (...args: any[]) => any;
+    }>(ai: T, options?: WrapAISDKOptions): {
+        generateText: T["generateText"];
+        streamText: T["streamText"];
+        generateObject: T["generateObject"];
+        streamObject: T["streamObject"];
     };
-}>(client: T): T;
+    /** Wrap Mastra agent. Delegates to shared wrapper. */
+    wrapMastraAgent<T extends {
+        generate: (...args: any[]) => Promise<any>;
+        name?: string;
+    }>(agent: T): T;
+}
 /**
- * Wrap a Google Generative AI client to automatically trace all content generations.
- *
- * @param client - The GoogleGenerativeAI client instance
- * @returns The same client with tracing enabled
- *
- * @example
- * ```typescript
- * import { GoogleGenerativeAI } from "@google/generative-ai";
- * import { trace } from "@fallom/trace";
- *
- * const genAI = new GoogleGenerativeAI(apiKey);
- * const model = trace.wrapGoogleAI(genAI.getGenerativeModel({ model: "gemini-pro" }));
- *
- * trace.setSession("my-config", sessionId);
- * const response = await model.generateContent("Hello!"); // Automatically traced!
- * ```
+ * Create a session-scoped Fallom instance.
  */
-declare function wrapGoogleAI<T extends {
-    generateContent: (...args: any[]) => Promise<any>;
-}>(model: T): T;
-/**
- * Wrap the Vercel AI SDK to automatically trace all LLM calls.
- * Works with generateText, streamText, generateObject, streamObject.
- *
- * @param ai - The ai module (import * as ai from "ai")
- * @returns Object with wrapped generateText, streamText, generateObject, streamObject
- *
- * @example
- * ```typescript
- * import * as ai from "ai";
- * import { createOpenAI } from "@ai-sdk/openai";
- * import { trace } from "@fallom/trace";
- *
- * await trace.init({ apiKey: process.env.FALLOM_API_KEY });
- * const { generateText, streamText } = trace.wrapAISDK(ai);
- *
- * const openrouter = createOpenAI({
- *   apiKey: process.env.OPENROUTER_API_KEY,
- *   baseURL: "https://openrouter.ai/api/v1",
- * });
- *
- * trace.setSession("my-config", sessionId);
- * const { text } = await generateText({
- *   model: openrouter("openai/gpt-4o-mini"),
- *   prompt: "Hello!",
- * }); // Automatically traced!
- * ```
- */
-declare function wrapAISDK<T extends {
-    generateText: (...args: any[]) => Promise<any>;
-    streamText: (...args: any[]) => any;
-    generateObject?: (...args: any[]) => Promise<any>;
-    streamObject?: (...args: any[]) => any;
-}>(ai: T): {
-    generateText: T["generateText"];
-    streamText: T["streamText"];
-    generateObject: T["generateObject"];
-    streamObject: T["streamObject"];
-};
+declare function session(options: SessionOptions): FallomSession;
+
 /**
- * Wrap a Mastra agent to automatically trace all generate() calls.
- *
- * @param agent - The Mastra Agent instance
- * @returns The same agent with tracing enabled
- *
- * @example
- * ```typescript
- * import { trace } from "@fallom/trace";
- * import { Agent } from "@mastra/core";
- *
- * await trace.init({ apiKey: "your-key" });
+ * Fallom tracing module.
  *
- * const agent = new Agent({ ... });
- * const tracedAgent = trace.wrapMastraAgent(agent);
+ * Auto-instruments all LLM calls via OTEL and groups them by session.
+ * Also supports custom spans for business metrics.
  *
- * trace.setSession("my-app", "session-123", "user-456");
- * const result = await tracedAgent.generate([{ role: "user", content: "Hello" }]);
- * // ^ Automatically traced!
- * ```
+ * This file re-exports from the modular trace/ directory.
+ * Each wrapper is in its own file for better maintainability.
  */
-declare function wrapMastraAgent<T extends {
-    generate: (...args: any[]) => Promise<any>;
-    name?: string;
-}>(agent: T): T;
 
-declare const trace_clearSession: typeof clearSession;
-declare const trace_getSession: typeof getSession;
-declare const trace_runWithSession: typeof runWithSession;
-declare const trace_setSession: typeof setSession;
+type trace_FallomSession = FallomSession;
+declare const trace_FallomSession: typeof FallomSession;
+type trace_SessionContext = SessionContext;
+type trace_SessionOptions = SessionOptions;
+type trace_TraceContext = TraceContext;
+type trace_TraceData = TraceData;
+type trace_WrapAISDKOptions = WrapAISDKOptions;
+declare const trace_session: typeof session;
 declare const trace_shutdown: typeof shutdown;
-declare const trace_span: typeof span;
-declare const trace_wrapAISDK: typeof wrapAISDK;
-declare const trace_wrapAnthropic: typeof wrapAnthropic;
-declare const trace_wrapGoogleAI: typeof wrapGoogleAI;
-declare const trace_wrapMastraAgent: typeof wrapMastraAgent;
-declare const trace_wrapOpenAI: typeof wrapOpenAI;
 declare namespace trace {
-  export { trace_clearSession as clearSession, trace_getSession as getSession, init$3 as init, trace_runWithSession as runWithSession, trace_setSession as setSession, trace_shutdown as shutdown, trace_span as span, trace_wrapAISDK as wrapAISDK, trace_wrapAnthropic as wrapAnthropic, trace_wrapGoogleAI as wrapGoogleAI, trace_wrapMastraAgent as wrapMastraAgent, trace_wrapOpenAI as wrapOpenAI };
+  export { trace_FallomSession as FallomSession, type trace_SessionContext as SessionContext, type trace_SessionOptions as SessionOptions, type trace_TraceContext as TraceContext, type trace_TraceData as TraceData, type trace_WrapAISDKOptions as WrapAISDKOptions, init$3 as init, trace_session as session, trace_shutdown as shutdown };
 }
 
 /**
@@ -296,9 +255,6 @@ declare function init$2(options?: {
  *
  * Same session_id always returns same model (sticky assignment).
  *
- * Also automatically sets trace context, so all subsequent LLM calls
- * are tagged with this session.
- *
  * @param configKey - Your config name (e.g., "linkedin-agent")
  * @param sessionId - Your session/conversation ID (must be consistent)
  * @param options - Optional settings
@@ -477,7 +433,7 @@ declare function init(options?: InitOptions): Promise<void>;
  * Fallom Exporter for Mastra
  *
  * Custom OpenTelemetry exporter that sends traces from Mastra agents to Fallom.
- * Reads session context from the shared trace module (set via trace.setSession()).
+ * Session context should be passed to the exporter constructor.
  *
  * Usage with Mastra:
  * ```typescript
@@ -487,7 +443,14 @@ declare function init(options?: InitOptions): Promise<void>;
  * // Initialize trace module
  * await trace.init({ apiKey: process.env.FALLOM_API_KEY });
  *
- * // Create Mastra with Fallom exporter
+ * // Create session for this request
+ * const session = trace.session({
+ *   configKey: "my-app",
+ *   sessionId: "session-123",
+ *   customerId: "user-456"
+ * });
+ *
+ * // Create Mastra with Fallom exporter (pass session context)
  * const mastra = new Mastra({
  *   agents: { myAgent },
  *   telemetry: {
@@ -495,13 +458,13 @@ declare function init(options?: InitOptions): Promise<void>;
  *     enabled: true,
  *     export: {
  *       type: "custom",
- *       exporter: new FallomExporter(),
+ *       exporter: new FallomExporter({
+ *         session: session.getContext()
+ *       }),
  *     },
  *   },
  * });
  *
- * // In your request handler:
- * trace.setSession("my-app", "session-123", "user-456");
  * const result = await mastra.getAgent("myAgent").generate("Hello!");
  * ```
  */
@@ -513,6 +476,8 @@ interface FallomExporterOptions {
     baseUrl?: string;
     /** Enable debug logging */
     debug?: boolean;
+    /** Session context for tracing */
+    session?: SessionContext;
 }
 /**
  * Set prompt tracking info.
@@ -531,13 +496,14 @@ declare function clearMastraPrompt(): void;
 /**
  * OpenTelemetry SpanExporter that sends traces to Fallom.
  *
- * Reads session context from trace.setSession() automatically.
+ * Pass session context via constructor options.
  * Compatible with Mastra's custom exporter interface.
  */
 declare class FallomExporter implements SpanExporter {
     private apiKey;
     private baseUrl;
     private debug;
+    private session?;
     private pendingExports;
     constructor(options?: FallomExporterOptions);
     private log;
@@ -584,35 +550,30 @@ declare class FallomExporter implements SpanExporter {
  *
  * @example
  * ```typescript
- * import fallom from 'fallom';
- *
- * // Initialize (call this early, before LLM imports if possible)
- * fallom.init({ apiKey: "your-api-key" });
- *
- * // Set session context for tracing
- * fallom.trace.setSession("my-agent", sessionId);
- *
- * // Get A/B tested model
- * const model = await fallom.models.get("my-config", sessionId, {
- *   fallback: "gpt-4o-mini"
+ * import fallom from '@fallom/trace';
+ * import * as ai from 'ai';
+ * import { createOpenAI } from '@ai-sdk/openai';
+ *
+ * // Initialize once
+ * await fallom.init({ apiKey: "your-api-key" });
+ *
+ * // Create a session for this conversation/request
+ * const session = fallom.session({
+ *   configKey: "my-agent",
+ *   sessionId: "session-123",
+ *   customerId: "user-456",
  * });
  *
- * // Get managed prompts (with optional A/B testing)
- * const prompt = await fallom.prompts.get("onboarding", {
- *   variables: { userName: "John" }
- * });
+ * // Option 1: Wrap the AI SDK (our style)
+ * const { generateText } = session.wrapAISDK(ai);
+ * await generateText({ model: openai("gpt-4o"), prompt: "Hello!" });
  *
- * // Use with OpenAI
- * const response = await openai.chat.completions.create({
- *   model,
- *   messages: [
- *     { role: "system", content: prompt.system },
- *     { role: "user", content: prompt.user }
- *   ]
- * });
+ * // Option 2: Wrap the model directly (PostHog style)
+ * const model = session.traceModel(openai("gpt-4o"));
+ * await ai.generateText({ model, prompt: "Hello!" });
  *
- * // Record custom metrics
- * fallom.trace.span({ user_satisfaction: 5 });
+ * // Get A/B tested model within session
+ * const modelName = await session.getModel({ fallback: "gpt-4o-mini" });
  * ```
  */
 
@@ -621,6 +582,7 @@ declare const _default: {
     trace: typeof trace;
     models: typeof models;
     prompts: typeof prompts;
+    session: typeof session;
 };
 
-export { FallomExporter, type FallomExporterOptions, type InitOptions, type PromptResult, clearMastraPrompt, _default as default, init, models, prompts, setMastraPrompt, setMastraPromptAB, trace };
+export { FallomExporter, type FallomExporterOptions, FallomSession, type InitOptions, type PromptResult, type SessionContext, type SessionOptions, clearMastraPrompt, _default as default, init, models, prompts, session, setMastraPrompt, setMastraPromptAB, trace };