@lmnr-ai/lmnr 0.4.5 → 0.4.7-alpha

package/src/sdk/tracing/index.ts

@@ -0,0 +1,357 @@
+/* eslint-disable @typescript-eslint/no-var-requires */
+import { NodeSDK } from "@opentelemetry/sdk-node";
+import {
+  SimpleSpanProcessor,
+  BatchSpanProcessor,
+  SpanProcessor,
+} from "@opentelemetry/sdk-trace-node";
+import { baggageUtils } from "@opentelemetry/core";
+import { Span, context, diag } from "@opentelemetry/api";
+import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-proto";
+import { Resource } from "@opentelemetry/resources";
+import { SEMRESATTRS_SERVICE_NAME } from "@opentelemetry/semantic-conventions";
+import { Instrumentation } from "@opentelemetry/instrumentation";
+import { InitializeOptions } from "../interfaces";
+import {
+  Detector,
+  DetectorSync,
+  IResource,
+  ResourceDetectionConfig,
+  envDetectorSync,
+  hostDetectorSync,
+  processDetectorSync,
+} from "@opentelemetry/resources"
+import {
+  ASSOCIATION_PROPERTIES_KEY,
+} from "./tracing";
+import { Telemetry } from "../telemetry/telemetry";
+import { _configuration } from "../configuration";
+import {
+  SpanAttributes,
+} from "@traceloop/ai-semantic-conventions";
+import { AnthropicInstrumentation } from "@traceloop/instrumentation-anthropic";
+import { OpenAIInstrumentation } from "@traceloop/instrumentation-openai";
+import { AzureOpenAIInstrumentation } from "@traceloop/instrumentation-azure";
+import { LlamaIndexInstrumentation } from "@traceloop/instrumentation-llamaindex";
+import {
+  AIPlatformInstrumentation,
+  VertexAIInstrumentation,
+} from "@traceloop/instrumentation-vertexai";
+import { BedrockInstrumentation } from "@traceloop/instrumentation-bedrock";
+import { CohereInstrumentation } from "@traceloop/instrumentation-cohere";
+import { PineconeInstrumentation } from "@traceloop/instrumentation-pinecone";
+import { LangChainInstrumentation } from "@traceloop/instrumentation-langchain";
+import { ChromaDBInstrumentation } from "@traceloop/instrumentation-chromadb";
+import { QdrantInstrumentation } from "@traceloop/instrumentation-qdrant";
+
+let _sdk: NodeSDK;
+let _spanProcessor: SimpleSpanProcessor | BatchSpanProcessor;
+let openAIInstrumentation: OpenAIInstrumentation | undefined;
+let anthropicInstrumentation: AnthropicInstrumentation | undefined;
+let azureOpenAIInstrumentation: AzureOpenAIInstrumentation | undefined;
+let cohereInstrumentation: CohereInstrumentation | undefined;
+let vertexaiInstrumentation: VertexAIInstrumentation | undefined;
+let aiplatformInstrumentation: AIPlatformInstrumentation | undefined;
+let bedrockInstrumentation: BedrockInstrumentation | undefined;
+let langchainInstrumentation: LangChainInstrumentation | undefined;
+let llamaIndexInstrumentation: LlamaIndexInstrumentation | undefined;
+let pineconeInstrumentation: PineconeInstrumentation | undefined;
+let chromadbInstrumentation: ChromaDBInstrumentation | undefined;
+let qdrantInstrumentation: QdrantInstrumentation | undefined;
+
+
+const instrumentations: Instrumentation[] = [];
+
+const initInstrumentations = () => {
+  const exceptionLogger = (e: Error) => Telemetry.getInstance().logException(e);
+  const enrichTokens =
+    (process.env.TRACELOOP_ENRICH_TOKENS || "true").toLowerCase() === "true";
+
+  openAIInstrumentation = new OpenAIInstrumentation({
+    enrichTokens,
+    exceptionLogger,
+  });
+  instrumentations.push(openAIInstrumentation);
+
+  anthropicInstrumentation = new AnthropicInstrumentation({ exceptionLogger });
+  instrumentations.push(anthropicInstrumentation);
+
+  azureOpenAIInstrumentation = new AzureOpenAIInstrumentation({
+    exceptionLogger,
+  });
+  instrumentations.push(azureOpenAIInstrumentation);
+
+  cohereInstrumentation = new CohereInstrumentation({ exceptionLogger });
+  instrumentations.push(cohereInstrumentation);
+
+  vertexaiInstrumentation = new VertexAIInstrumentation({
+    exceptionLogger,
+  });
+  instrumentations.push(vertexaiInstrumentation);
+
+  aiplatformInstrumentation = new AIPlatformInstrumentation({
+    exceptionLogger,
+  });
+  instrumentations.push(aiplatformInstrumentation);
+
+  bedrockInstrumentation = new BedrockInstrumentation({ exceptionLogger });
+  instrumentations.push(bedrockInstrumentation);
+
+  pineconeInstrumentation = new PineconeInstrumentation({ exceptionLogger });
+  instrumentations.push(pineconeInstrumentation);
+
+  langchainInstrumentation = new LangChainInstrumentation({ exceptionLogger });
+  instrumentations.push(langchainInstrumentation);
+
+  llamaIndexInstrumentation = new LlamaIndexInstrumentation({
+    exceptionLogger,
+  });
+  instrumentations.push(llamaIndexInstrumentation);
+
+  chromadbInstrumentation = new ChromaDBInstrumentation({ exceptionLogger });
+  instrumentations.push(chromadbInstrumentation);
+
+  qdrantInstrumentation = new QdrantInstrumentation({ exceptionLogger });
+  instrumentations.push(qdrantInstrumentation);
+};
+
+const manuallyInitInstrumentations = (
+  instrumentModules: InitializeOptions["instrumentModules"],
+) => {
+  const exceptionLogger = (e: Error) => Telemetry.getInstance().logException(e);
+  const enrichTokens =
+    (process.env.TRACELOOP_ENRICH_TOKENS || "true").toLowerCase() === "true";
+
+  instrumentations.length = 0;
+
+  if (instrumentModules?.openAI) {
+    openAIInstrumentation = new OpenAIInstrumentation({
+      enrichTokens,
+      exceptionLogger,
+    });
+    instrumentations.push(openAIInstrumentation);
+    openAIInstrumentation.manuallyInstrument(instrumentModules.openAI);
+  }
+
+  if (instrumentModules?.anthropic) {
+    anthropicInstrumentation = new AnthropicInstrumentation({
+      exceptionLogger,
+    });
+    instrumentations.push(anthropicInstrumentation);
+    anthropicInstrumentation.manuallyInstrument(instrumentModules.anthropic);
+  }
+
+  if (instrumentModules?.azureOpenAI) {
+    const instrumentation = new AzureOpenAIInstrumentation({ exceptionLogger });
+    instrumentations.push(instrumentation as Instrumentation);
+    azureOpenAIInstrumentation = instrumentation;
+    instrumentation.manuallyInstrument(instrumentModules.azureOpenAI);
+  }
+
+  if (instrumentModules?.cohere) {
+    cohereInstrumentation = new CohereInstrumentation({ exceptionLogger });
+    instrumentations.push(cohereInstrumentation);
+    cohereInstrumentation.manuallyInstrument(instrumentModules.cohere);
+  }
+
+  if (instrumentModules?.google_vertexai) {
+    vertexaiInstrumentation = new VertexAIInstrumentation({
+      exceptionLogger,
+    });
+    instrumentations.push(vertexaiInstrumentation);
+    vertexaiInstrumentation.manuallyInstrument(
+      instrumentModules.google_vertexai,
+    );
+  }
+
+  if (instrumentModules?.google_aiplatform) {
+    aiplatformInstrumentation = new AIPlatformInstrumentation({
+      exceptionLogger,
+    });
+    instrumentations.push(aiplatformInstrumentation);
+    aiplatformInstrumentation.manuallyInstrument(
+      instrumentModules.google_aiplatform,
+    );
+  }
+
+  if (instrumentModules?.bedrock) {
+    bedrockInstrumentation = new BedrockInstrumentation({ exceptionLogger });
+    instrumentations.push(bedrockInstrumentation);
+    bedrockInstrumentation.manuallyInstrument(instrumentModules.bedrock);
+  }
+
+  if (instrumentModules?.pinecone) {
+    const instrumentation = new PineconeInstrumentation({ exceptionLogger });
+    instrumentations.push(instrumentation as Instrumentation);
+    instrumentation.manuallyInstrument(instrumentModules.pinecone);
+  }
+
+  if (instrumentModules?.langchain) {
+    langchainInstrumentation = new LangChainInstrumentation({
+      exceptionLogger,
+    });
+    instrumentations.push(langchainInstrumentation);
+    langchainInstrumentation.manuallyInstrument(instrumentModules.langchain);
+  }
+
+  if (instrumentModules?.llamaIndex) {
+    llamaIndexInstrumentation = new LlamaIndexInstrumentation({
+      exceptionLogger,
+    });
+    instrumentations.push(llamaIndexInstrumentation);
+    llamaIndexInstrumentation.manuallyInstrument(instrumentModules.llamaIndex);
+  }
+
+  if (instrumentModules?.chromadb) {
+    chromadbInstrumentation = new ChromaDBInstrumentation({ exceptionLogger });
+    instrumentations.push(chromadbInstrumentation);
+    chromadbInstrumentation.manuallyInstrument(instrumentModules.chromadb);
+  }
+
+  if (instrumentModules?.qdrant) {
+    qdrantInstrumentation = new QdrantInstrumentation({ exceptionLogger });
+    instrumentations.push(qdrantInstrumentation);
+    qdrantInstrumentation.manuallyInstrument(instrumentModules.qdrant);
+  }
+};
+
+function awaitAttributes(detector: DetectorSync): Detector {
+  return {
+    async detect(config?: ResourceDetectionConfig): Promise<IResource> {
+      const resource = detector.detect(config)
+      await resource.waitForAsyncAttributes?.()
+
+      return resource
+    },
+  }
+}
+
+/**
+ * Initializes the Traceloop SDK.
+ * Must be called once before any other SDK methods.
+ *
+ * @param options - The options to initialize the SDK. See the {@link InitializeOptions} for details.
+ */
+export const startTracing = (options: InitializeOptions) => {
+  if (options.instrumentModules !== undefined) {
+    // If options.instrumentModules is empty, it will not initialize anything,
+    // so empty dict can essentially be passed to disable any kind of automatic instrumentation.
+    manuallyInitInstrumentations(options.instrumentModules);
+  } else {
+    initInstrumentations();
+  }
+
+  if (!shouldSendTraces()) {
+    openAIInstrumentation?.setConfig({
+      traceContent: false,
+    });
+    azureOpenAIInstrumentation?.setConfig({
+      traceContent: false,
+    });
+    llamaIndexInstrumentation?.setConfig({
+      traceContent: false,
+    });
+    vertexaiInstrumentation?.setConfig({
+      traceContent: false,
+    });
+    aiplatformInstrumentation?.setConfig({
+      traceContent: false,
+    });
+    bedrockInstrumentation?.setConfig({
+      traceContent: false,
+    });
+    cohereInstrumentation?.setConfig({
+      traceContent: false,
+    });
+    chromadbInstrumentation?.setConfig({
+      traceContent: false,
+    });
+  }
+
+  const headers = process.env.TRACELOOP_HEADERS
+    ? baggageUtils.parseKeyPairsIntoRecord(process.env.TRACELOOP_HEADERS)
+    : { Authorization: `Bearer ${options.apiKey}` };
+
+  const traceExporter =
+    options.exporter ??
+    new OTLPTraceExporter({
+      url: `${options.baseUrl}/v1/traces`,
+      headers,
+    });
+  _spanProcessor = options.disableBatch
+    ? new SimpleSpanProcessor(traceExporter)
+    : new BatchSpanProcessor(traceExporter);
+
+  _spanProcessor.onStart = (span: Span) => {
+    // This sets the properties only if the context has them
+    const associationProperties = context
+      .active()
+      .getValue(ASSOCIATION_PROPERTIES_KEY);
+    if (associationProperties) {
+      for (const [key, value] of Object.entries(associationProperties)) {
+        span.setAttribute(
+          `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.${key}`,
+          value,
+        );
+      }
+    }
+  };
+
+  if (options.exporter) {
+    Telemetry.getInstance().capture("tracer:init", {
+      exporter: "custom",
+      processor: options.disableBatch ? "simple" : "batch",
+    });
+  } else {
+    Telemetry.getInstance().capture("tracer:init", {
+      exporter: options.baseUrl ?? "",
+      processor: options.disableBatch ? "simple" : "batch",
+    });
+  }
+
+  const spanProcessors: SpanProcessor[] = [_spanProcessor];
+  if (options.processor) {
+    spanProcessors.push(options.processor);
+  }
+
+  _sdk = new NodeSDK({
+    resource: new Resource({
+      [SEMRESATTRS_SERVICE_NAME]:
+        options.appName || process.env.npm_package_name,
+    }),
+    spanProcessors,
+    contextManager: options.contextManager,
+    textMapPropagator: options.propagator,
+    traceExporter,
+    instrumentations,
+    // We should re-consider removing irrelevant spans here in the future
+    // sampler: new TraceloopSampler(),
+    resourceDetectors: [
+      awaitAttributes(envDetectorSync),
+      awaitAttributes(processDetectorSync),
+      awaitAttributes(hostDetectorSync),
+    ],
+  });
+
+  _sdk.start();
+};
+
+export const shouldSendTraces = () => {
+  if (!_configuration) {
+    diag.warn("Traceloop not initialized");
+    return false;
+  }
+
+  if (
+    _configuration.traceContent === false ||
+    (process.env.TRACELOOP_TRACE_CONTENT || "true").toLowerCase() === "false"
+  ) {
+    return false;
+  }
+
+  return true;
+};
+
+export const forceFlush = async () => {
+  await _spanProcessor.forceFlush();
+};

package/src/sdk/tracing/manual.ts

@@ -0,0 +1,179 @@
+import { Span, context, trace } from "@opentelemetry/api";
+import { getTracer } from "./tracing";
+import {
+  Events,
+  EventAttributes,
+  SpanAttributes,
+} from "@traceloop/ai-semantic-conventions";
+import { shouldSendTraces } from ".";
+
+type VectorDBCallConfig = {
+  vendor: string;
+  type: "query" | "upsert" | "delete";
+};
+
+type LLMCallConfig = {
+  vendor: string;
+  type: "chat" | "completion";
+};
+
+export class VectorSpan {
+  private span: Span;
+
+  constructor(span: Span) {
+    this.span = span;
+  }
+
+  reportQuery({ queryVector }: { queryVector: number[] }) {
+    if (!shouldSendTraces()) {
+      this.span.addEvent(Events.DB_QUERY_EMBEDDINGS);
+    }
+
+    this.span.addEvent(Events.DB_QUERY_EMBEDDINGS, {
+      [EventAttributes.DB_QUERY_EMBEDDINGS_VECTOR]: JSON.stringify(queryVector),
+    });
+  }
+
+  reportResults({
+    results,
+  }: {
+    results: {
+      ids?: string;
+      scores?: number;
+      distances?: number;
+      metadata?: Record<string, unknown>;
+      vectors?: number[];
+      documents?: string;
+    }[];
+  }) {
+    for (let i = 0; i < results.length; i++) {
+      this.span.addEvent(Events.DB_QUERY_RESULT, {
+        [EventAttributes.DB_QUERY_RESULT_ID]: results[i].ids,
+        [EventAttributes.DB_QUERY_RESULT_SCORE]: results[i].scores,
+        [EventAttributes.DB_QUERY_RESULT_DISTANCE]: results[i].distances,
+        [EventAttributes.DB_QUERY_RESULT_METADATA]: JSON.stringify(
+          results[i].metadata,
+        ),
+        [EventAttributes.DB_QUERY_RESULT_VECTOR]: results[i].vectors,
+        [EventAttributes.DB_QUERY_RESULT_DOCUMENT]: results[i].documents,
+      });
+    }
+  }
+}
+
+export class LLMSpan {
+  private span: Span;
+
+  constructor(span: Span) {
+    this.span = span;
+  }
+
+  reportRequest({
+    model,
+    messages,
+  }: {
+    model: string;
+    messages: {
+      role: string;
+      content?: string | unknown;
+    }[];
+  }) {
+    this.span.setAttributes({
+      [SpanAttributes.LLM_REQUEST_MODEL]: model,
+    });
+
+    messages.forEach((message, index) => {
+      this.span.setAttributes({
+        [`${SpanAttributes.LLM_PROMPTS}.${index}.role`]: message.role,
+        [`${SpanAttributes.LLM_PROMPTS}.${index}.content`]:
+          typeof message.content === "string"
+            ? message.content
+            : JSON.stringify(message.content),
+      });
+    });
+  }
+
+  reportResponse({
+    model,
+    usage,
+    completions,
+  }: {
+    model: string;
+    usage?: {
+      prompt_tokens: number;
+      completion_tokens: number;
+      total_tokens: number;
+    };
+    completions?: {
+      finish_reason: string;
+      message: {
+        role: "system" | "user" | "assistant";
+        content: string | null;
+      };
+    }[];
+  }) {
+    this.span.setAttribute(SpanAttributes.LLM_RESPONSE_MODEL, model);
+
+    if (usage) {
+      this.span.setAttributes({
+        [SpanAttributes.LLM_USAGE_PROMPT_TOKENS]: usage.prompt_tokens,
+        [SpanAttributes.LLM_USAGE_COMPLETION_TOKENS]: usage.completion_tokens,
+        [SpanAttributes.LLM_USAGE_TOTAL_TOKENS]: usage.total_tokens,
+      });
+    }
+
+    completions?.forEach((completion, index) => {
+      this.span.setAttributes({
+        [`${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`]:
+          completion.finish_reason,
+        [`${SpanAttributes.LLM_COMPLETIONS}.${index}.role`]:
+          completion.message.role,
+        [`${SpanAttributes.LLM_COMPLETIONS}.${index}.content`]:
+          completion.message.content || "",
+      });
+    });
+  }
+}
+
+export function withVectorDBCall<
+  F extends ({ span }: { span: VectorSpan }) => ReturnType<F>,
+>({ vendor, type }: VectorDBCallConfig, fn: F, thisArg?: ThisParameterType<F>) {
+  const entityContext = context.active();
+
+  return getTracer().startActiveSpan(
+    `${vendor}.${type}`,
+    { [SpanAttributes.LLM_REQUEST_TYPE]: type },
+    entityContext,
+    (span: Span) => {
+      const res = fn.apply(thisArg, [{ span: new VectorSpan(span) }]);
+      if (res instanceof Promise) {
+        return res.then((resolvedRes) => {
+          span.end();
+          return resolvedRes;
+        });
+      }
+
+      span.end();
+      return res;
+    },
+  );
+}
+
+export function withLLMCall<
+  F extends ({ span }: { span: LLMSpan }) => ReturnType<F>,
+>({ vendor, type }: LLMCallConfig, fn: F, thisArg?: ThisParameterType<F>) {
+  const span = getTracer().startSpan(`${vendor}.${type}`, {}, context.active());
+  span.setAttribute(SpanAttributes.LLM_REQUEST_TYPE, type);
+  trace.setSpan(context.active(), span);
+
+  const res = fn.apply(thisArg, [{ span: new LLMSpan(span) }]);
+  if (res instanceof Promise) {
+    return res.then((resolvedRes) => {
+      span.end();
+      return resolvedRes;
+    });
+  }
+
+  span.end();
+  return res;
+}

package/src/sdk/tracing/sampler.ts

@@ -0,0 +1,35 @@
+import {
+  Sampler,
+  SamplingResult,
+  SamplingDecision,
+} from "@opentelemetry/sdk-trace-base";
+import { Context, SpanKind, Attributes } from "@opentelemetry/api";
+
+const FILTERED_ATTRIBUTE_KEYS = ["next.span_name"];
+
+export class TraceloopSampler implements Sampler {
+  shouldSample(
+    _context: Context,
+    _traceId: string,
+    _spanName: string,
+    _spanKind: SpanKind,
+    attributes: Attributes,
+  ): SamplingResult {
+    let filter = false;
+    FILTERED_ATTRIBUTE_KEYS.forEach((key) => {
+      if (attributes?.[key]) {
+        filter = true;
+      }
+    });
+
+    return {
+      decision: filter
+        ? SamplingDecision.NOT_RECORD
+        : SamplingDecision.RECORD_AND_SAMPLED,
+    };
+  }
+
+  toString(): string {
+    return "TraceloopSampler";
+  }
+}

package/src/sdk/tracing/tracing.ts

@@ -0,0 +1,19 @@
+import { trace, createContextKey } from "@opentelemetry/api";
+import { Context } from "@opentelemetry/api/build/src/context/types";
+
+const TRACER_NAME = "traceloop.tracer";
+export const WORKFLOW_NAME_KEY = createContextKey("workflow_name");
+export const ENTITY_NAME_KEY = createContextKey("entity_name");
+export const ASSOCIATION_PROPERTIES_KEY = createContextKey(
+  "association_properties",
+);
+
+export const getTracer = () => {
+  return trace.getTracer(TRACER_NAME);
+};
+
+export const getEntityPath = (entityContext: Context): string | undefined => {
+  const path = entityContext.getValue(ENTITY_NAME_KEY);
+
+  return path ? `${path}` : undefined;
+};