fix(sdk): manually report VectorDB spans (#174)

Author: nirga

packages/ai-semantic-conventions/src/SemanticAttributes.ts

@@ -33,8 +33,8 @@ export const SpanAttributes = {
   LLM_CHAT_STOP_SEQUENCES: "llm.chat.stop_sequences",
 
   // Vector DB
-  VECTOR_DB_VENDOR: "vector_db.vendor",
-  VECTOR_DB_QUERY_TOP_K: "vector_db.query.top_k",
+  VECTOR_DB_VENDOR: "db.system",
+  VECTOR_DB_QUERY_TOP_K: "db.vector.query.top_k",
 
   // LLM Workflows
   TRACELOOP_SPAN_KIND: "traceloop.span.kind",
@@ -45,31 +45,47 @@ export const SpanAttributes = {
   TRACELOOP_ENTITY_OUTPUT: "traceloop.entity.output",
 };
 
+export const Events = {
+  DB_QUERY_EMBEDDINGS: "db.query.embeddings",
+  DB_QUERY_RESULT: "db.query.result",
+};
+
 export const EventAttributes = {
-  // Vector DB Query Request
-  VECTOR_DB_QUERY_TOP_K: "vector_db.query.top_k",
-  VECTOR_DB_QUERY_INCLUDE_VALUES: "vector_db.query.include_values",
-  VECTOR_DB_QUERY_INCLUDE_METADATA: "vector_db.query.include_metadata",
-  VECTOR_DB_QUERY_ID: "vector_db.query.id",
-  VECTOR_DB_QUERY_EMBEDDINGS_VECTOR: "vector_db.query.embeddings.vector",
-  VECTOR_DB_QUERY_METADATA_FILTER: "vector_db.query.metadata_filter",
+  // Query Embeddings
+  DB_QUERY_EMBEDDINGS_VECTOR: "db.query.embeddings.vector",
+
+  // Query Result (canonical format)
+  DB_QUERY_RESULT_ID: "db.query.result.id",
+  DB_QUERY_RESULT_SCORE: "db.query.result.score",
+  DB_QUERY_RESULT_DISTANCE: "db.query.result.distance",
+  DB_QUERY_RESULT_METADATA: "db.query.result.metadata",
+  DB_QUERY_RESULT_VECTOR: "db.query.result.vector",
+  DB_QUERY_RESULT_DOCUMENT: "db.query.result.document",
+
+  // DEPRECATED: Vector DB Query Request
+  VECTOR_DB_QUERY_TOP_K: "db.vector.query.top_k",
+  VECTOR_DB_QUERY_INCLUDE_VALUES: "db.vector.query.include_values",
+  VECTOR_DB_QUERY_INCLUDE_METADATA: "db.vector.query.include_metadata",
+  VECTOR_DB_QUERY_ID: "db.vector.query.id",
+  VECTOR_DB_QUERY_EMBEDDINGS_VECTOR: "db.vector.query.embeddings.vector",
+  VECTOR_DB_QUERY_METADATA_FILTER: "db.vector.query.metadata_filter",
 
-  // Vector DB Query Response
-  VECTOR_DB_QUERY_RESULT_NAMESPACE: "vector_db.query.result.namespace",
+  // DEPRECATED: Vector DB Query Response
+  VECTOR_DB_QUERY_RESULT_NAMESPACE: "db.vector.query.result.namespace",
   VECTOR_DB_QUERY_RESULT_READ_UNITS_CONSUMED:
-    "vector_db.query.result.read_units",
+    "db.vector.query.result.read_units",
   VECTOR_DB_QUERY_RESULT_MATCHES_LENGTH:
-    "vector_db.query.result.matches_length",
+    "db.vector.query.result.matches_length",
 
-  // Vector DB Query Response of each result
-  VECTOR_DB_QUERY_RESULT_SCORE: "vector_db.query.result.{i}.score",
-  VECTOR_DB_QUERY_RESULT_ID: "vector_db.query.result.{i}.id",
-  VECTOR_DB_QUERY_RESULT_VALUES: "vector_db.query.result.{i}.values",
+  // DEPRECATED: Vector DB Query Response of each result
+  VECTOR_DB_QUERY_RESULT_SCORE: "db.vector.query.result.{i}.score",
+  VECTOR_DB_QUERY_RESULT_ID: "db.vector.query.result.{i}.id",
+  VECTOR_DB_QUERY_RESULT_VALUES: "db.vector.query.result.{i}.values",
   VECTOR_DB_QUERY_RESULT_SPARSE_INDICES:
-    "vector_db.query.result.{i}.sparse.indices",
+    "db.vector.query.result.{i}.sparse.indices",
   VECTOR_DB_QUERY_RESULT_SPARSE_VALUES:
-    "vector_db.query.result.{i}.sparse.values",
-  VECTOR_DB_QUERY_RESULT_METADATA: "vector_db.query.result.{i}.metadata",
+    "db.vector.query.result.{i}.sparse.values",
+  VECTOR_DB_QUERY_RESULT_METADATA: "db.vector.query.result.{i}.metadata",
 };
 
 export enum LLMRequestTypeValues {

packages/traceloop-sdk/src/lib/node-server-sdk.ts

@@ -5,6 +5,7 @@ export { InitializeOptions } from "./interfaces";
 export { initialize } from "./configuration";
 export { forceFlush } from "./tracing";
 export * from "./tracing/decorators";
+export * from "./tracing/manual";
 export * from "./tracing/association";
 export * from "./tracing/score";
 export * from "./prompts";

packages/traceloop-sdk/src/lib/tracing/manual.ts

@@ -0,0 +1,75 @@
+import { Span, context } from "@opentelemetry/api";
+import { getTracer } from "./tracing";
+import { Events, EventAttributes } from "@traceloop/ai-semantic-conventions";
+import { shouldSendTraces } from ".";
+
+type VectorDBCallConfig = {
+  vendor: string;
+  type: "query" | "upsert" | "delete";
+};
+
+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 function withVectorDBCall<
+  F extends ({ span }: { span: VectorSpan }) => ReturnType<F>,
+>({ vendor, type }: VectorDBCallConfig, fn: F, thisArg?: ThisParameterType<F>) {
+  const entityContext = context.active();
+
+  getTracer().startActiveSpan(
+    `${vendor}.${type}`,
+    {},
+    entityContext,
+    async (span: Span) => {
+      const res = fn.apply(thisArg, [{ span: new VectorSpan(span) }]);
+      if (res instanceof Promise) {
+        return res.then(() => {
+          span.end();
+        });
+      }
+
+      span.end();
+    },
+  );
+}