Improve RAG with configurable seasrch prompt params via the API

verida · Nov 6, 2024 · d3c4bc8 · d3c4bc8
1 parent d5d3d10
commit d3c4bc8
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 18 deletions.
diff --git a/src/api/rest/v1/llm/controller.ts b/src/api/rest/v1/llm/controller.ts
@@ -4,6 +4,7 @@ import { PromptSearchService } from '../../../../services/assistants/search'
 import { Utils } from "../../../../utils";
 import { HotLoadProgress } from "../../../../services/data";
 import { DataService } from "../../../../services/data";
+import { PromptSearchServiceConfig } from "../../../../services/assistants/interfaces";
 const _ = require('lodash')
 
 export interface LLMConfig {
@@ -79,6 +80,7 @@ export class LLMController {
             const { context, account } = await Utils.getNetworkConnectionFromRequest(req)
             const did = await account.did()
             const prompt = req.body.prompt
+            const promptConfig: PromptSearchServiceConfig = req.body.promptConfig
 
             const {
                 customEndpoint,
@@ -89,7 +91,7 @@ export class LLMController {
             const llm = getLLM(llmProvider, llmModel, customEndpoint)
 
             const promptService = new PromptSearchService(did, context)
-            const promptResult = await promptService.prompt(prompt, llm)
+            const promptResult = await promptService.prompt(prompt, llm, promptConfig)
 
             return res.json(promptResult)
         } catch (error) {

diff --git a/src/services/assistants/interfaces.ts b/src/services/assistants/interfaces.ts
@@ -20,10 +20,7 @@ export interface PromptSearchServiceDataTypes {
 }
 
 export interface PromptSearchServiceConfig {
-    searchType?: PromptSearchType
     maxContextLength?: number
     dataTypes?: PromptSearchServiceDataTypes
-    promptSearchConfig?: PromptSearchLLMResponseOptional
-}
-
-export interface PromptSearchLLMResponseOptional extends Partial<PromptSearchLLMResponse> {}
+    promptSearchConfig?: PromptSearchLLMResponse
+}
diff --git a/src/services/assistants/search.ts b/src/services/assistants/search.ts
@@ -37,16 +37,22 @@ const DEFAULT_PROMPT_SEARCH_SERVICE_CONFIG: PromptSearchServiceConfig = {
 
 export class PromptSearchService extends VeridaService {
 
-    public async prompt(prompt: string, llm: LLM, config: PromptSearchServiceConfig = {}): Promise<{
+    public async prompt(prompt: string, llm: LLM, config?: PromptSearchServiceConfig): Promise<{
         result: string,
         duration: number,
         process: PromptSearchLLMResponse
     }> {
+        const start = Date.now()
+
         config = _.merge({}, DEFAULT_PROMPT_SEARCH_SERVICE_CONFIG, config)
 
-        const start = Date.now()
-        const promptSearch = new PromptSearch(llm)
-        const promptSearchResult = await promptSearch.search(prompt, undefined, config.promptSearchConfig)
+        let promptSearchResult
+        if (config?.promptSearchConfig) {
+            promptSearchResult = config.promptSearchConfig
+        } else {
+            const promptSearch = new PromptSearch(llm)
+            promptSearchResult = await promptSearch.search(prompt)
+        }
 
         console.log(promptSearchResult)
 
@@ -86,8 +92,10 @@ export class PromptSearchService extends VeridaService {
             console.log(`Searching by timeframe: ${maxDatetime} ${sort}`)
             if (promptSearchResult.databases.indexOf(SearchType.EMAILS) !== -1) {
                 emails = await searchService.schemaByDateRange<SchemaEmail>(SearchType.EMAILS, maxDatetime, sort, config.dataTypes.emails.limit*3)
-                const emailShortlist = new EmailShortlist(llm)
-                emails = await emailShortlist.shortlist(prompt, emails, config.dataTypes.emails.limit)
+                if (emails.length > config.dataTypes.emails.limit) {
+                    const emailShortlist = new EmailShortlist(llm)
+                    emails = await emailShortlist.shortlist(prompt, emails, config.dataTypes.emails.limit)
+                }
             }
             if (promptSearchResult.databases.indexOf(SearchType.FILES) !== -1) {
                 files = await searchService.schemaByDateRange<SchemaFile>(SearchType.FILES, maxDatetime, sort, config.dataTypes.files.limit)
@@ -170,9 +178,9 @@ export class PromptSearchService extends VeridaService {
         // console.log(finalPrompt)
 
         const finalResponse = await llm.prompt(finalPrompt, undefined, false)
-        const duration = Date.now() - start
+        const duration = ((Date.now() - start) / 1000.0)
 
-        console.log(contextString)
+        // console.log(contextString)
 
         return {
             result: finalResponse.choices[0].message.content!,

diff --git a/src/services/tools/promptSearch.ts b/src/services/tools/promptSearch.ts
@@ -1,6 +1,5 @@
 const _ = require('lodash')
 import { KeywordSearchTimeframe } from "../../helpers/interfaces";
-import { PromptSearchLLMResponseOptional } from "../assistants/interfaces";
 import { LLM } from "../llm"
 import { SearchType } from "../search";
 
@@ -56,12 +55,11 @@ export class PromptSearch {
         this.llm = llm
     }
 
-    public async search(userPrompt: string, retries = 3, defaultResponse?: PromptSearchLLMResponseOptional): Promise<PromptSearchLLMResponse> {
+    public async search(userPrompt: string, retries = 3): Promise<PromptSearchLLMResponse> {
         const response = await this.llm.prompt(userPrompt, systemPrompt)
 
         try {
-          const searchResponse = <PromptSearchLLMResponse> JSON.parse(response.choices[0].message.content!)
-          return _.merge({}, searchResponse, defaultResponse ? defaultResponse : {})
+          return <PromptSearchLLMResponse> JSON.parse(response.choices[0].message.content!)
         } catch (err: any) {
           if (retries === 0) {
             throw new Error(`No user data query available`)