From 4d0afaebda8da6f7bee0f7fd24d86b350fa9c24f Mon Sep 17 00:00:00 2001
From: ogzhanolguncu <ogzhan11@gmail.com>
Date: Thu, 23 May 2024 14:45:02 +0300
Subject: [PATCH] fix: allow passing ids to addContext

---
 src/services/retrieval.ts | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/src/services/retrieval.ts b/src/services/retrieval.ts
index 9d6cc83..e6feb2e 100644
--- a/src/services/retrieval.ts
+++ b/src/services/retrieval.ts
@@ -6,11 +6,11 @@ import { PDFLoader } from "@langchain/community/document_loaders/fs/pdf";
 import type { RecursiveCharacterTextSplitterParams } from "langchain/text_splitter";
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
 
-type IndexUpsertPayload = { input: string | number[]; id?: string; metadata?: string };
+type IndexUpsertPayload = { input: number[]; id?: string | number; metadata?: string };
 type FilePath = string;
 
 export type AddContextPayload =
-  | { dataType: "text"; data: string }
+  | { dataType: "text"; data: string; id?: string | number }
   | { dataType: "embedding"; data: IndexUpsertPayload[] }
   | {
       dataType: "pdf";
@@ -88,23 +88,16 @@ export class RetrievalService {
       case "text": {
         return this.index.upsert({
           data: input.data,
-          id: nanoid(),
+          id: input.id ?? nanoid(),
           metadata: { [metadataKey]: input.data },
         });
       }
       case "embedding": {
         const items = input.data.map((context) => {
-          const isText = typeof context.input === "string";
-          const metadata = context.metadata
-            ? { [metadataKey]: context.metadata }
-            : isText
-              ? { [metadataKey]: context.input }
-              : {};
-
           return {
-            [isText ? "data" : "vector"]: context.input,
+            vector: context.input,
             id: context.id ?? nanoid(),
-            metadata,
+            metadata: { [metadataKey]: context.metadata },
           };
         });