@cloudbase/agent-observability 1.0.1-alpha.9

package/src/index.ts

@@ -0,0 +1,775 @@
+/**
+ * Observability - OpenTelemetry-based tracing with OpenInference semantic conventions
+ *
+ * @packageDocumentation
+ */
+
+import { trace, context, TimeInput, SpanStatusCode, Span, SpanContext } from "@opentelemetry/api";
+
+import {
+  createObservationAttributes,
+  createTraceAttributes,
+} from "./core/attributes.js";
+import {
+  ObservationSpan,
+  ObservationLLM,
+  ObservationEmbedding,
+  ObservationAgent,
+  ObservationTool,
+  ObservationChain,
+  ObservationRetriever,
+  ObservationReranker,
+  ObservationEvaluator,
+  ObservationGuardrail,
+  type Observation,
+} from "./core/spanWrapper.js";
+import { getTracer } from "./core/tracerProvider.js";
+import {
+  ObservationType,
+  ObservationLevel,
+  BaseSpanAttributes,
+  LLMAttributes,
+  ToolAttributes,
+  AgentAttributes,
+  ChainAttributes,
+  RetrieverAttributes,
+  RerankerAttributes,
+  EvaluatorAttributes,
+  GuardrailAttributes,
+  EmbeddingAttributes,
+  ObservationAttributes,
+  TraceAttributes,
+} from "./types.js";
+
+// Export types
+export type {
+  ObservationType,
+  ObservationLevel,
+  BaseSpanAttributes,
+  LLMAttributes,
+  ToolAttributes,
+  AgentAttributes,
+  ChainAttributes,
+  RetrieverAttributes,
+  RerankerAttributes,
+  EvaluatorAttributes,
+  GuardrailAttributes,
+  EmbeddingAttributes,
+  ObservationAttributes,
+  TraceAttributes,
+};
+
+// Export observation classes
+export {
+  ObservationSpan,
+  ObservationLLM,
+  ObservationEmbedding,
+  ObservationAgent,
+  ObservationTool,
+  ObservationChain,
+  ObservationRetriever,
+  ObservationReranker,
+  ObservationEvaluator,
+  ObservationGuardrail,
+};
+
+// Export observation union type
+export type { Observation };
+
+// Export core functions
+export {
+  createTraceAttributes,
+  createObservationAttributes,
+} from "./core/attributes.js";
+export {
+  setTracerProvider,
+  getTracerProvider,
+  getTracer,
+} from "./core/tracerProvider.js";
+
+/**
+ * Options for starting observations (spans).
+ *
+ * @public
+ */
+export type StartObservationOptions = {
+  /** Custom start time for the observation */
+  startTime?: Date;
+  /** Parent span context to attach this observation to */
+  parentSpanContext?: SpanContext;
+};
+
+/**
+ * Options for startObservation function.
+ *
+ * @public
+ */
+export type StartObservationOpts = StartObservationOptions & {
+  /** Type of observation to create. Defaults to 'span' */
+  asType?: ObservationType;
+};
+
+/**
+ * Creates an OpenTelemetry span with the AG-Kit tracer.
+ *
+ * @param params - Parameters for span creation
+ * @returns The created OpenTelemetry span
+ * @internal
+ */
+function createOtelSpan(params: {
+  name: string;
+  startTime?: TimeInput;
+  parentSpanContext?: SpanContext;
+}): Span {
+  return getTracer().startSpan(
+    params.name,
+    { startTime: params.startTime },
+    createParentContext(params.parentSpanContext),
+  );
+}
+
+/**
+ * Creates a parent context from a span context.
+ *
+ * @param parentSpanContext - The span context to use as parent
+ * @returns The created context or undefined if no parent provided
+ * @internal
+ */
+function createParentContext(
+  parentSpanContext?: SpanContext,
+): ReturnType<typeof trace.setSpanContext> | undefined {
+  if (!parentSpanContext) return;
+  return trace.setSpanContext(context.active(), parentSpanContext);
+}
+
+// Function overloads for proper type inference
+// Generic overload for dynamic asType (returns Observation union)
+export function startObservation(
+  name: string,
+  attributes?: BaseSpanAttributes,
+  options?: StartObservationOpts & { asType?: ObservationType },
+): Observation;
+
+// Type-specific overloads for precise type inference
+export function startObservation(
+  name: string,
+  attributes: LLMAttributes,
+  options: StartObservationOpts & { asType: "llm" },
+): ObservationLLM;
+export function startObservation(
+  name: string,
+  attributes: EmbeddingAttributes,
+  options: StartObservationOpts & { asType: "embedding" },
+): ObservationEmbedding;
+export function startObservation(
+  name: string,
+  attributes: AgentAttributes,
+  options: StartObservationOpts & { asType: "agent" },
+): ObservationAgent;
+export function startObservation(
+  name: string,
+  attributes: ToolAttributes,
+  options: StartObservationOpts & { asType: "tool" },
+): ObservationTool;
+export function startObservation(
+  name: string,
+  attributes: ChainAttributes,
+  options: StartObservationOpts & { asType: "chain" },
+): ObservationChain;
+export function startObservation(
+  name: string,
+  attributes: RetrieverAttributes,
+  options: StartObservationOpts & { asType: "retriever" },
+): ObservationRetriever;
+export function startObservation(
+  name: string,
+  attributes: RerankerAttributes,
+  options: StartObservationOpts & { asType: "reranker" },
+): ObservationReranker;
+export function startObservation(
+  name: string,
+  attributes: EvaluatorAttributes,
+  options: StartObservationOpts & { asType: "evaluator" },
+): ObservationEvaluator;
+export function startObservation(
+  name: string,
+  attributes: GuardrailAttributes,
+  options: StartObservationOpts & { asType: "guardrail" },
+): ObservationGuardrail;
+export function startObservation(
+  name: string,
+  attributes?: BaseSpanAttributes,
+  options?: StartObservationOpts & { asType?: "span" },
+): ObservationSpan;
+
+/**
+ * Creates and starts a new AG-Kit observation.
+ *
+ * Supports multiple observation types with full TypeScript type safety:
+ * - **span**: General-purpose operations (default)
+ * - **llm**: LLM calls and AI model interactions
+ * - **embedding**: Text embedding and vector operations
+ * - **agent**: AI agent workflows
+ * - **tool**: Individual tool calls
+ * - **chain**: Multi-step processes
+ * - **retriever**: Document retrieval
+ * - **reranker**: Result reranking
+ * - **evaluator**: Quality assessment
+ * - **guardrail**: Safety checks
+ *
+ * @param name - Descriptive name for the observation
+ * @param attributes - Type-specific attributes
+ * @param options - Configuration options
+ * @returns Strongly-typed observation object
+ *
+ * @example
+ * ```typescript
+ * import { startObservation } from './observability';
+ *
+ * // LLM observation
+ * const llm = startObservation('openai-gpt-4', {
+ *   input: [{ role: 'user', content: 'Hello' }],
+ *   model: 'gpt-4',
+ *   modelParameters: { temperature: 0.7 }
+ * }, { asType: 'llm' });
+ *
+ * // Tool observation
+ * const tool = startObservation('weather-api', {
+ *   input: { location: 'SF' }
+ * }, { asType: 'tool' });
+ *
+ * // Chain observation
+ * const chain = startObservation('rag-pipeline', {
+ *   input: { question: 'What is AI?' }
+ * }, { asType: 'chain' });
+ * ```
+ *
+ * @public
+ */
+export function startObservation(
+  name: string,
+  attributes?:
+    | BaseSpanAttributes
+    | LLMAttributes
+    | EmbeddingAttributes
+    | AgentAttributes
+    | ToolAttributes
+    | ChainAttributes
+    | RetrieverAttributes
+    | RerankerAttributes
+    | EvaluatorAttributes
+    | GuardrailAttributes,
+  options?: StartObservationOpts,
+): Observation {
+  const { asType = "span", ...observationOptions } = options || {};
+
+  const otelSpan = createOtelSpan({
+    name,
+    ...observationOptions,
+  });
+
+  switch (asType) {
+    case "llm":
+      return new ObservationLLM({
+        otelSpan,
+        attributes: attributes as LLMAttributes,
+      });
+
+    case "embedding":
+      return new ObservationEmbedding({
+        otelSpan,
+        attributes: attributes as EmbeddingAttributes,
+      });
+
+    case "agent":
+      return new ObservationAgent({
+        otelSpan,
+        attributes: attributes as AgentAttributes,
+      });
+
+    case "tool":
+      return new ObservationTool({
+        otelSpan,
+        attributes: attributes as ToolAttributes,
+      });
+
+    case "chain":
+      return new ObservationChain({
+        otelSpan,
+        attributes: attributes as ChainAttributes,
+      });
+
+    case "retriever":
+      return new ObservationRetriever({
+        otelSpan,
+        attributes: attributes as RetrieverAttributes,
+      });
+
+    case "reranker":
+      return new ObservationReranker({
+        otelSpan,
+        attributes: attributes as RerankerAttributes,
+      });
+
+    case "evaluator":
+      return new ObservationEvaluator({
+        otelSpan,
+        attributes: attributes as EvaluatorAttributes,
+      });
+
+    case "guardrail":
+      return new ObservationGuardrail({
+        otelSpan,
+        attributes: attributes as GuardrailAttributes,
+      });
+
+    case "span":
+    default:
+      return new ObservationSpan({
+        otelSpan,
+        attributes: attributes as BaseSpanAttributes,
+      });
+  }
+}
+
+/**
+ * Updates the currently active trace with new attributes.
+ *
+ * @param attributes - Trace attributes to set
+ *
+ * @example
+ * ```typescript
+ * import { updateActiveTrace } from './observability';
+ *
+ * updateActiveTrace({
+ *   name: 'user-workflow',
+ *   userId: 'user-123',
+ *   tags: ['production']
+ * });
+ * ```
+ *
+ * @public
+ */
+export function updateActiveTrace(attributes: TraceAttributes) {
+  const span = trace.getActiveSpan();
+
+  if (!span) {
+    console.warn(
+      "[Observability] No active OTEL span in context. Skipping trace update.",
+    );
+    return;
+  }
+
+  span.setAttributes(createTraceAttributes(attributes));
+}
+
+/**
+ * Gets the current active trace ID.
+ *
+ * @returns The trace ID of the currently active span, or undefined
+ *
+ * @public
+ */
+export function getActiveTraceId(): string | undefined {
+  return trace.getActiveSpan()?.spanContext().traceId;
+}
+
+/**
+ * Gets the current active observation ID.
+ *
+ * @returns The span ID of the currently active span, or undefined
+ *
+ * @public
+ */
+export function getActiveSpanId(): string | undefined {
+  return trace.getActiveSpan()?.spanContext().spanId;
+}
+
+// ============================================================================
+// Active Observation Functions
+// ============================================================================
+
+/**
+ * Options for startActiveObservation.
+ *
+ * @public
+ */
+export type StartActiveObservationOpts = StartObservationOpts & {
+  /** Whether to automatically end the observation when the function exits. Default: true */
+  endOnExit?: boolean;
+};
+
+/**
+ * Wraps a Promise to automatically end the span when it resolves/rejects.
+ *
+ * @param promise - The promise to wrap
+ * @param span - The OpenTelemetry span
+ * @param endOnExit - Whether to end the span on exit
+ * @returns The wrapped promise
+ * @internal
+ */
+function wrapPromise<T>(
+  promise: Promise<T>,
+  span: Span,
+  endOnExit: boolean | undefined,
+): Promise<T> {
+  return promise.then(
+    (value) => {
+      if (endOnExit !== false) {
+        span.end();
+      }
+      return value;
+    },
+    (err: unknown) => {
+      span.setStatus({
+        code: SpanStatusCode.ERROR,
+        message: err instanceof Error ? err.message : "Unknown error",
+      });
+      if (endOnExit !== false) {
+        span.end();
+      }
+      throw err;
+    },
+  );
+}
+
+// Function overloads for startActiveObservation
+export function startActiveObservation<
+  F extends (observation: ObservationSpan) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+export function startActiveObservation<
+  F extends (observation: ObservationLLM) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+export function startActiveObservation<
+  F extends (observation: ObservationEmbedding) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+export function startActiveObservation<
+  F extends (observation: ObservationAgent) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+export function startActiveObservation<
+  F extends (observation: ObservationTool) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+export function startActiveObservation<
+  F extends (observation: ObservationChain) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+export function startActiveObservation<
+  F extends (observation: ObservationRetriever) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+export function startActiveObservation<
+  F extends (observation: ObservationReranker) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+export function startActiveObservation<
+  F extends (observation: ObservationEvaluator) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+export function startActiveObservation<
+  F extends (observation: ObservationGuardrail) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F>;
+
+/**
+ * Creates an observation with automatic lifecycle management.
+ *
+ * This function creates an observation and executes a function with that observation
+ * as a parameter. The observation is automatically ended when the function completes
+ * (unless `endOnExit` is set to `false`).
+ *
+ * Supports both synchronous and asynchronous functions, with automatic error handling.
+ *
+ * @param name - Descriptive name for the observation
+ * @param fn - Function to execute with the observation
+ * @param options - Configuration options
+ * @returns The result of the function
+ *
+ * @example
+ * ```typescript
+ * import { startActiveObservation } from './observability';
+ *
+ * // Synchronous function
+ * const result = startActiveObservation('data-processing', (span) => {
+ *   span.update({ input: { data: [1, 2, 3] } });
+ *   const processed = data.map(x => x * 2);
+ *   span.update({ output: { result: processed } });
+ *   return processed;
+ * }, { asType: 'span' });
+ *
+ * // Asynchronous function
+ * const embeddings = await startActiveObservation(
+ *   'text-embeddings',
+ *   async (embedding) => {
+ *     embedding.update({
+ *       input: { texts: ['Hello', 'World'] },
+ *       model: 'text-embedding-ada-002'
+ *     });
+ *
+ *     const vectors = await generateEmbeddings(texts);
+ *
+ *     embedding.update({ output: { embeddings: vectors } });
+ *     return vectors;
+ *   },
+ *   { asType: 'embedding' }
+ * );
+ *
+ * // Disable automatic ending (for long-running operations)
+ * startActiveObservation(
+ *   'background-task',
+ *   (span) => {
+ *     span.update({ input: { taskId: '123' } });
+ *     startBackgroundProcess(span);
+ *     return 'started';
+ *   },
+ *   { asType: 'span', endOnExit: false }
+ * );
+ * ```
+ *
+ * @see {@link startObservation} for manual observation lifecycle management
+ * @see {@link observe} for decorator-style function wrapping
+ *
+ * @public
+ */
+export function startActiveObservation<
+  F extends (observation: Observation) => unknown,
+>(name: string, fn: F, options?: StartActiveObservationOpts): ReturnType<F> {
+  const { asType = "span", endOnExit, ...observationOptions } = options || {};
+
+  return getTracer().startActiveSpan(
+    name,
+    { startTime: observationOptions?.startTime },
+    createParentContext(observationOptions?.parentSpanContext) ??
+      context.active(),
+    (span) => {
+      try {
+        let observation: Observation;
+
+        switch (asType) {
+          case "llm":
+            observation = new ObservationLLM({ otelSpan: span });
+            break;
+          case "embedding":
+            observation = new ObservationEmbedding({ otelSpan: span });
+            break;
+          case "agent":
+            observation = new ObservationAgent({ otelSpan: span });
+            break;
+          case "tool":
+            observation = new ObservationTool({ otelSpan: span });
+            break;
+          case "chain":
+            observation = new ObservationChain({ otelSpan: span });
+            break;
+          case "retriever":
+            observation = new ObservationRetriever({ otelSpan: span });
+            break;
+          case "reranker":
+            observation = new ObservationReranker({ otelSpan: span });
+            break;
+          case "evaluator":
+            observation = new ObservationEvaluator({ otelSpan: span });
+            break;
+          case "guardrail":
+            observation = new ObservationGuardrail({ otelSpan: span });
+            break;
+          case "span":
+          default:
+            observation = new ObservationSpan({ otelSpan: span });
+        }
+
+        const result = fn(observation as Parameters<F>[0]);
+
+        if (result instanceof Promise) {
+          return wrapPromise(
+            result,
+            span,
+            endOnExit,
+          ) as ReturnType<F>;
+        } else {
+          if (endOnExit !== false) {
+            span.end();
+          }
+          return result as ReturnType<F>;
+        }
+      } catch (err) {
+        span.setStatus({
+          code: SpanStatusCode.ERROR,
+          message: err instanceof Error ? err.message : "Unknown error",
+        });
+        if (endOnExit !== false) {
+          span.end();
+        }
+        throw err;
+      }
+    },
+  );
+}
+
+/**
+ * Updates the currently active observation with new attributes.
+ *
+ * @param attributes - Observation attributes to set
+ *
+ * @example
+ * ```typescript
+ * import { updateActiveObservation } from './observability';
+ *
+ * // Within an active observation context
+ * updateActiveObservation({
+ *   metadata: { stage: 'processing' }
+ * });
+ * ```
+ *
+ * @public
+ */
+export function updateActiveObservation(
+  attributes: BaseSpanAttributes,
+): void {
+  const span = trace.getActiveSpan();
+
+  if (!span) {
+    console.warn(
+      "[Observability] No active OTEL span in context. Skipping observation update.",
+    );
+    return;
+  }
+
+  span.setAttributes(createObservationAttributes("span", attributes));
+}
+
+// ============================================================================
+// Decorator Function
+// ============================================================================
+
+/**
+ * Options for the observe decorator.
+ *
+ * @public
+ */
+export type ObserveOptions = Omit<StartObservationOpts, "name"> & {
+  /** Whether to capture function arguments as input. Default: true */
+  captureInput?: boolean;
+  /** Whether to capture return value as output. Default: true */
+  captureOutput?: boolean;
+};
+
+/**
+ * Captures function arguments for observability input.
+ *
+ * @param args - Function arguments to capture
+ * @returns Serialized arguments
+ * @internal
+ */
+function _captureArguments(args: unknown[]): Record<string, unknown> {
+  if (args.length === 0) return {};
+  if (args.length === 1) return { arg: args[0] };
+  return { args };
+}
+
+/**
+ * Decorator function to add observability to any function.
+ *
+ * Wraps a function with automatic observation creation, input/output capture,
+ * and lifecycle management. The observation is automatically ended when the
+ * function completes.
+ *
+ * @param fn - Function to wrap
+ * @param options - Configuration options
+ * @returns Wrapped function with observability
+ *
+ * @example
+ * ```typescript
+ * import { observe } from './observability';
+ *
+ * // Wrap an existing function
+ * const fetchData = observe(async (url: string) => {
+ *   const response = await fetch(url);
+ *   return response.json();
+ * }, { asType: 'tool' });
+ *
+ * // Wrap with custom name
+ * const processPayment = observe(
+ *   async (amount: number, currency: string) => {
+ *     return await paymentGateway.charge(amount, currency);
+ *   },
+ *   { name: 'payment-gateway-call', asType: 'tool' }
+ * );
+ *
+ * // Class method decoration
+ * class UserService {
+ *   @observe({ asType: 'chain' })
+ *   async getUser(id: string) {
+ *     return await db.users.find(id);
+ *   }
+ * }
+ * ```
+ *
+ * @public
+ */
+export function observe<T extends (...args: any[]) => any>(
+  fn: T,
+  options: ObserveOptions = {},
+): T {
+  const {
+    asType = "span",
+    captureInput = true,
+    captureOutput = true,
+    ...observationOptions
+  } = options;
+
+  const wrappedFunction = function (
+    this: any,
+    ...args: Parameters<T>
+  ): ReturnType<T> {
+    const name = fn.name || "anonymous-function";
+
+    // Prepare input data
+    const inputData = captureInput ? _captureArguments(args) : undefined;
+
+    // Create the observation
+    const observation = startObservation(
+      name,
+      inputData ? { input: inputData } : {},
+      {
+        ...observationOptions,
+        asType: asType as "span",
+      },
+    );
+
+    // Set the observation span as active in the context
+    const activeContext = trace.setSpan(context.active(), observation.otelSpan);
+
+    // Execute the function within the observation context
+    const result = context.with(activeContext, () => fn.apply(this, args));
+
+    // Handle promises
+    if (result instanceof Promise) {
+      return result.then(
+        (value) => {
+          if (captureOutput) {
+            observation.update({ output: value });
+          }
+          observation.end();
+          return value;
+        },
+        (err: unknown) => {
+          observation.update({
+            level: "ERROR",
+            statusMessage: err instanceof Error ? err.message : "Unknown error",
+          });
+          observation.end();
+          throw err;
+        },
+      ) as ReturnType<T>;
+    }
+
+    // Handle synchronous functions
+    if (captureOutput) {
+      observation.update({ output: result });
+    }
+    observation.end();
+
+    return result as ReturnType<T>;
+  };
+
+  // Preserve function properties
+  Object.defineProperty(wrappedFunction, "name", { value: fn.name });
+  Object.defineProperty(wrappedFunction, "length", { value: fn.length });
+
+  return wrappedFunction as T;
+}