feat(core): add free function support and active state management to …

…ChatLunaSaveableVectorStore
ChatLunaLab · Dec 29, 2024 · 8cff498 · 8cff498
1 parent bb7feb8
commit 8cff498
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 2 deletions.
diff --git a/packages/core/src/llm-core/model/base.ts b/packages/core/src/llm-core/model/base.ts
@@ -21,6 +21,10 @@ export class ChatLunaSaveableVectorStore<T extends VectorStore = VectorStore>
         ...args: Parameters<T['similaritySearchVectorWithScore']>
     ) => Promise<[Document, number][]>
 
+    freeFunction?: () => Promise<void>
+
+    private _isActive = true
+
     constructor(
         private _store: T,
         input: ChatLunaSaveableVectorStoreInput<T>
@@ -31,13 +35,16 @@ export class ChatLunaSaveableVectorStore<T extends VectorStore = VectorStore>
         this.addDocumentsFunction = input.addDocumentsFunction
         this.similaritySearchVectorWithScoreFunction =
             input.similaritySearchVectorWithScoreFunction
+        this.freeFunction = input.freeFunction
     }
 
     addVectors(...args: Parameters<typeof this._store.addVectors>) {
+        this._checkActive()
         return this._store.addVectors(...args)
     }
 
     addDocuments(...args: Parameters<T['addDocuments']>) {
+        this._checkActive()
         if (this.addDocumentsFunction) {
             return this.addDocumentsFunction(this._store, ...args)
         }
@@ -47,6 +54,7 @@ export class ChatLunaSaveableVectorStore<T extends VectorStore = VectorStore>
     similaritySearchVectorWithScore(
         ...args: Parameters<T['similaritySearchVectorWithScore']>
     ) {
+        this._checkActive()
         if (this.similaritySearchVectorWithScoreFunction) {
             return this.similaritySearchVectorWithScoreFunction(
                 this._store,
@@ -61,6 +69,8 @@ export class ChatLunaSaveableVectorStore<T extends VectorStore = VectorStore>
     }
 
     async editDocument(oldDocumentId: string, newDocument: Document) {
+        this._checkActive()
+
         // delete
         await this.delete({ ids: [oldDocumentId] })
 
@@ -71,11 +81,13 @@ export class ChatLunaSaveableVectorStore<T extends VectorStore = VectorStore>
     }
 
     save() {
+        this._checkActive()
         return this?.saveableFunction(this._store)
     }
 
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     delete(input: ChatLunaSaveableVectorDelete) {
+        this._checkActive()
         return (
             this?.deletableFunction?.(this._store, input) ??
             this._store.delete(input)
@@ -85,6 +97,20 @@ export class ChatLunaSaveableVectorStore<T extends VectorStore = VectorStore>
     _vectorstoreType(): string {
         return this._store?._vectorstoreType() ?? '?'
     }
+
+    private _checkActive() {
+        if (!this._isActive) {
+            throw new Error('VectorStore is not active')
+        }
+    }
+
+    async free() {
+        if (this.freeFunction) {
+            await this.freeFunction()
+        }
+        this._store = undefined
+        this._isActive = false
+    }
 }
 
 export interface ChatLunaSaveableVectorStoreInput<T extends VectorStore> {
@@ -101,6 +127,7 @@ export interface ChatLunaSaveableVectorStoreInput<T extends VectorStore> {
         store: T,
         ...args: Parameters<T['similaritySearchVectorWithScore']>
     ) => Promise<[Document, number][]>
+    freeFunction?: () => Promise<void>
 }
 
 export interface ChatLunaSaveableVectorDelete

diff --git a/packages/core/src/llm-core/platform/service.ts b/packages/core/src/llm-core/platform/service.ts
@@ -41,7 +41,10 @@ export class PlatformService {
         string,
         ChatLunaSaveableVectorStore
     >({
-        max: 20
+        max: 20,
+        dispose: (value, key, reason) => {
+            value.free()
+        }
     })
 
     constructor(private ctx: Context) {

diff --git a/packages/core/src/llm-core/prompt/type.ts b/packages/core/src/llm-core/prompt/type.ts
@@ -37,6 +37,7 @@ export interface RawPreset {
               recursiveScan?: boolean
               maxRecursionDepth?: number
               matchWholeWord?: boolean
+              constant?: boolean
               caseSensitive?: boolean
               enabled?: boolean
               order?: number

diff --git a/packages/vector-store-service/src/vectorstore/faiss.ts b/packages/vector-store-service/src/vectorstore/faiss.ts
@@ -122,6 +122,9 @@ export async function apply(
                     await store.addDocuments(documents, {
                         ids
                     })
+                },
+                async freeFunction() {
+                    faissStore = undefined
                 }
             }
         )

diff --git a/packages/vector-store-service/src/vectorstore/milvus.ts b/packages/vector-store-service/src/vectorstore/milvus.ts
@@ -26,7 +26,7 @@ export async function apply(
 
         const key = params.key ?? 'chatluna'
 
-        const vectorStore = new MilvusClass(embeddings, {
+        let vectorStore = new MilvusClass(embeddings, {
             collectionName: 'chatluna_collection',
             partitionName: key,
             url: config.milvusUrl,
@@ -266,6 +266,9 @@ export async function apply(
                     })
                     // console.log("Search result: " + JSON.stringify(results, null, 2));
                     return results
+                },
+                async freeFunction() {
+                    vectorStore = undefined
                 }
             }
         )

diff --git a/packages/vector-store-service/src/vectorstore/voy.ts b/packages/vector-store-service/src/vectorstore/voy.ts
@@ -130,6 +130,10 @@ export async function apply(
                     await store.addDocuments(documents, {
                         ids
                     })
+                },
+                async freeFunction() {
+                    voyStore.client.free()
+                    voyStore = undefined
                 }
             }
         )
-Original file line number
+Diff line change
@@ Expand Up / @@ -122,6 +122,9 @@ export async function apply( @@
                         await store.addDocuments(documents, {
                             ids
                         })
+                    },
+                    async freeFunction() {
+                        faissStore = undefined
                     }
                 }
             )
@@ Expand Down @@