feat(js): llama-index-ts embeddings support (#761)

Co-authored-by: Parker Stafford <parker.stafford92@gmail.com> Co-authored-by: Mikyo King <mikyo@arize.com>
Arize-ai · Aug 9, 2024 · 337fb86 · 337fb86
1 parent ae7af3f
commit 337fb86
Show file tree

Hide file tree

Showing 7 changed files with 416 additions and 180 deletions.
diff --git a/js/packages/openinference-instrumentation-llama-index/package.json b/js/packages/openinference-instrumentation-llama-index/package.json
@@ -24,6 +24,8 @@
     "@opentelemetry/instrumentation": "^0.46.0"
   },
   "devDependencies": {
-    "llamaindex": "^0.3.14"
+    "jest": "^29.7.0",
+    "llamaindex": "^0.3.14",
+    "openai": "^4.24.1"
   }
 }
diff --git a/js/packages/openinference-instrumentation-llama-index/src/instrumentation.ts b/js/packages/openinference-instrumentation-llama-index/src/instrumentation.ts
@@ -7,7 +7,13 @@ import {
   InstrumentationNodeModuleDefinition,
 } from "@opentelemetry/instrumentation";
 import { diag } from "@opentelemetry/api";
-import { patchQueryMethod, patchRetrieveMethod } from "./utils";
+import {
+  patchQueryEngineQueryMethod,
+  patchRetrieveMethod,
+  patchQueryEmbeddingMethod,
+  isRetrieverPrototype,
+  isEmbeddingPrototype,
+} from "./utils";
 import { VERSION } from "./version";
 
 const MODULE_NAME = "llamaindex";
@@ -58,31 +64,52 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<
     }
 
     // TODO: Support streaming
+    // TODO: Generalize to QueryEngine interface (RetrieverQueryEngine, RouterQueryEngine)
     this._wrap(
       moduleExports.RetrieverQueryEngine.prototype,
       "query",
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       (original): any => {
-        return patchQueryMethod(original, moduleExports, this.tracer);
+        return patchQueryEngineQueryMethod(original, this.tracer);
       },
     );
 
-    this._wrap(
-      moduleExports.VectorIndexRetriever.prototype,
-      "retrieve",
-      (original) => {
-        return patchRetrieveMethod(original, moduleExports, this.tracer);
-      },
-    );
+    for (const value of Object.values(moduleExports)) {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const prototype = (value as any).prototype;
+
+      if (isRetrieverPrototype(prototype)) {
+        this._wrap(prototype, "retrieve", (original) => {
+          return patchRetrieveMethod(original, this.tracer);
+        });
+      }
 
+      if (isEmbeddingPrototype(prototype)) {
+        this._wrap(prototype, "getQueryEmbedding", (original) => {
+          return patchQueryEmbeddingMethod(original, this.tracer);
+        });
+      }
+    }
     _isOpenInferencePatched = true;
     return moduleExports;
   }
 
   private unpatch(moduleExports: typeof llamaindex, moduleVersion?: string) {
     this._diag.debug(`Un-patching ${MODULE_NAME}@${moduleVersion}`);
     this._unwrap(moduleExports.RetrieverQueryEngine.prototype, "query");
-    this._unwrap(moduleExports.VectorIndexRetriever.prototype, "retrieve");
+
+    for (const value of Object.values(moduleExports)) {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const prototype = (value as any).prototype;
+
+      if (isRetrieverPrototype(prototype)) {
+        this._unwrap(prototype, "retrieve");
+      }
+
+      if (isEmbeddingPrototype(prototype)) {
+        this._unwrap(prototype, "getQueryEmbedding");
+      }
+    }
 
     _isOpenInferencePatched = false;
   }

diff --git a/js/packages/openinference-instrumentation-llama-index/src/types.ts b/js/packages/openinference-instrumentation-llama-index/src/types.ts
@@ -1,39 +1,5 @@
-import { SemanticConventions } from "@arizeai/openinference-semantic-conventions";
-
-export type RetrievalDocument = {
-  [SemanticConventions.DOCUMENT_ID]?: string;
-  [SemanticConventions.DOCUMENT_CONTENT]?: string;
-  [SemanticConventions.DOCUMENT_SCORE]?: number | undefined;
-  [SemanticConventions.DOCUMENT_METADATA]?: string;
-};
-
-type LLMMessageToolCall = {
-  [SemanticConventions.TOOL_CALL_FUNCTION_NAME]?: string;
-  [SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON]?: string;
-};
-
-export type LLMMessageToolCalls = {
-  [SemanticConventions.MESSAGE_TOOL_CALLS]?: LLMMessageToolCall[];
-};
-
-export type LLMMessageFunctionCall = {
-  [SemanticConventions.MESSAGE_FUNCTION_CALL_NAME]?: string;
-  [SemanticConventions.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON]?: string;
-};
-
-export type LLMMessage = LLMMessageToolCalls &
-  LLMMessageFunctionCall & {
-    [SemanticConventions.MESSAGE_ROLE]?: string;
-    [SemanticConventions.MESSAGE_CONTENT]?: string;
-  };
-
-export type LLMMessagesAttributes =
-  | {
-      [SemanticConventions.LLM_INPUT_MESSAGES]: LLMMessage[];
-    }
-  | {
-      [SemanticConventions.LLM_OUTPUT_MESSAGES]: LLMMessage[];
-    };
+import * as llamaindex from "llamaindex";
+import { BaseRetriever } from "llamaindex";
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export type GenericFunction = (...args: any[]) => any;
@@ -42,22 +8,12 @@ export type SafeFunction<T extends GenericFunction> = (
   ...args: Parameters<T>
 ) => ReturnType<T> | null;
 
-export type LLMParameterAttributes = {
-  [SemanticConventions.LLM_MODEL_NAME]?: string;
-  [SemanticConventions.LLM_INVOCATION_PARAMETERS]?: string;
-};
+export type ObjectWithModel = { model: string };
+
+export type RetrieverQueryEngineQueryMethodType =
+  typeof llamaindex.RetrieverQueryEngine.prototype.query;
 
-export type PromptTemplateAttributes = {
-  [SemanticConventions.PROMPT_TEMPLATE_TEMPLATE]?: string;
-  [SemanticConventions.PROMPT_TEMPLATE_VARIABLES]?: string;
-};
-export type TokenCountAttributes = {
-  [SemanticConventions.LLM_TOKEN_COUNT_COMPLETION]?: number;
-  [SemanticConventions.LLM_TOKEN_COUNT_PROMPT]?: number;
-  [SemanticConventions.LLM_TOKEN_COUNT_TOTAL]?: number;
-};
+export type RetrieverRetrieveMethodType = BaseRetriever["retrieve"];
 
-export type ToolAttributes = {
-  [SemanticConventions.TOOL_NAME]?: string;
-  [SemanticConventions.TOOL_DESCRIPTION]?: string;
-};
+export type QueryEmbeddingMethodType =
+  typeof llamaindex.BaseEmbedding.prototype.getQueryEmbedding;