added test for openai, anthropic embeding, added fastembed for anthropic

elizaOS · Feb 13, 2025 · 1a5cd64 · 1a5cd64
1 parent fb381cc
commit 1a5cd64
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 0 deletions.
diff --git a/packages/plugin-anthropic/package.json b/packages/plugin-anthropic/package.json
@@ -19,6 +19,7 @@
     ],
     "dependencies": {
         "@ai-sdk/anthropic": "^1.1.6",
+        "fastembed": "^1.0.0",
         "@elizaos/core": "workspace:*",
         "tsup": "8.3.5",
         "zod": "3.21.4"

diff --git a/packages/plugin-anthropic/src/index.ts b/packages/plugin-anthropic/src/index.ts
@@ -6,6 +6,7 @@ import {
 } from "@elizaos/core";
 import { generateText } from "ai";
 import { z } from "zod";
+import { EmbeddingModel, FlagEmbedding } from "fastembed";
 
 // Define a configuration schema for the Anthropics plugin.
 const configSchema = z.object({
@@ -82,12 +83,43 @@ export const anthropicPlugin: Plugin = {
         stopSequences,
       });
       return text;
+    },
+    [ModelClass.TEXT_EMBEDDING]: async (runtime, text: string | null) => {
+      if (!text) return new Array(1536).fill(0);
+
+      const model = await FlagEmbedding.init({ model: EmbeddingModel.BGESmallENV15 });
+      const embedding = await model.queryEmbed(text);
+
+      const finalEmbedding = Array.isArray(embedding) 
+        ? ArrayBuffer.isView(embedding[0]) 
+          ? Array.from(embedding[0] as never)
+          : embedding
+        : Array.from(embedding);
+
+      if (!Array.isArray(finalEmbedding) || finalEmbedding[0] === undefined) {
+        throw new Error("Invalid embedding format");
+      }
+
+      return finalEmbedding.map(Number);
     }
   },
   tests: [
     {
       name: "anthropic_plugin_tests",
       tests: [
+        {
+          name: 'anthropic_test_text_embedding',
+          fn: async (runtime) => {
+            try {
+              console.log("testing embedding");
+              const embedding = await runtime.useModel(ModelClass.TEXT_EMBEDDING, "Hello, world!");
+              console.log("embedding done", embedding);
+            } catch (error) {
+              console.error("Error in test_text_embedding:", error);
+              throw error;
+            }
+          }
+        },
         {
           name: 'anthropic_test_text_small',  
           fn: async (runtime) => {

diff --git a/packages/plugin-openai/src/index.ts b/packages/plugin-openai/src/index.ts
@@ -310,6 +310,18 @@ export const openaiPlugin: Plugin = {
             }
           }
         },
+        {
+          name: 'openai_test_text_embedding',
+          fn: async (runtime) => {
+            try {
+              const embedding = await runtime.useModel(ModelClass.TEXT_EMBEDDING, "Hello, world!");
+              console.log("embedding", embedding);
+            } catch (error) {
+              console.error("Error in test_text_embedding:", error);
+              throw error;
+            }
+          }
+        },
         {
           name: 'openai_test_text_large',
           fn: async (runtime) => {