✨ Add metadatas to embeddings (#141)

* ✨ Add metadatas to embeddings * fix any in Record * add missing metadatas in Ressources type * change one test to use gpt-4o because gpt-3.5 is now dumb * ✅ tests: Fix memory test to await embedding * ✅ tests: Oopsie * up patch version --------- Co-authored-by: Lancelot Owczarczak <[email protected]>
polyfact · Jul 19, 2024 · 138f1f1 · 138f1f1
1 parent d5019a5
commit 138f1f1
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 4 deletions.
diff --git a/lib/embeddings.ts b/lib/embeddings.ts
@@ -92,7 +92,9 @@ async function searchEmbeddings(
     id: string,
     input: string,
     clientOptions: InputClientOptions = {},
-): Promise<{ id: string; content: string; similarity: number }[]> {
+): Promise<
+    { id: string; content: string; similarity: number; metadatas: Record<string, unknown> }[]
+> {
     const { token, endpoint } = await defaultOptions(clientOptions);
 
     try {
@@ -166,7 +168,11 @@ class Embeddings {
         return this.memoryId;
     }
 
-    async search(input: string): Promise<{ id: string; content: string; similarity: number }[]> {
+    async search(
+        input: string,
+    ): Promise<
+        { id: string; content: string; similarity: number; metadatas: Record<string, unknown> }[]
+    > {
         const id = await this.memoryId;
         return searchEmbeddings(id, input, await this.clientOptions);
     }

diff --git a/lib/generate.ts b/lib/generate.ts
@@ -110,6 +110,7 @@ export type Ressource = {
     similarity: number;
     id: string;
     content: string;
+    metadatas?: Record<string, unknown>;
 };
 
 export type GenerationResult = {

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "polyfire-js",
-    "version": "0.2.59",
+    "version": "0.2.60",
     "main": "index.js",
     "types": "index.d.ts",
     "author": "Lancelot Owczarczak <[email protected]>",

diff --git a/tests/generation.test.js b/tests/generation.test.js
@@ -86,12 +86,13 @@ test(
 
         const embeddings = Embeddings();
 
-        embeddings.add('The word you have to remember is "banana42"');
+        await embeddings.add('The word you have to remember is "banana42"');
 
         const message = await generate("Just answer with the word you have to remember", {
             temperature: 0,
             cache: false,
             embeddings,
+            model: "gpt-3.5-turbo",
         });
 
         expect(message).toBe("banana42");