fix

pingcap · Jun 3, 2024 · 2213f7f · 2213f7f
1 parent 241927d
commit 2213f7f
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 31 deletions.
diff --git a/src/app/api/v1/chats/route.ts b/src/app/api/v1/chats/route.ts
@@ -1,5 +1,5 @@
 import {type Chat, createChat, getChatByUrlKey, listChats} from '@/core/repositories/chat';
-import {getChatEngineByNameOrDefault} from '@/core/repositories/chat_engine';
+import {getChatEngineByIdOrName} from '@/core/repositories/chat_engine';
 import {getIndexByNameOrThrow} from '@/core/repositories/index_';
 import {LlamaindexChatService} from '@/core/services/llamaindex/chating';
 import {toPageRequest} from '@/lib/database';
@@ -21,6 +21,9 @@ const ChatRequest = z.object({
   // TODO: using `title` instead of.
   name: z.string().optional(),
   index: z.string().optional(),
+  // TODO: remove it
+  // @Deprecated
+  engine: z.number().optional(),
   chat_engine: z.string().optional(),
   regenerate: z.boolean().optional(),
   messageId: z.coerce.number().int().optional(),
@@ -42,7 +45,7 @@ export const POST = defineHandler({
     messages,
   } = body;
 
-  const engine = await getChatEngineByNameOrDefault(body.chat_engine);
+  const engine = await getChatEngineByIdOrName(body.chat_engine ?? body.engine)
 
   // TODO: need refactor, it is too complex now
   // For chat page, create a chat and return the session ID (url_key) first.
@@ -75,6 +78,7 @@ export const POST = defineHandler({
   if (!sessionId) {
     chat = await createChat({
       engine: engine.engine,
+      engine_id: engine.id,
       engine_name: engine.name,
       engine_options: JSON.stringify(engine.engine_options),
       created_at: new Date(),

diff --git a/src/app/api/v1/indexes/[name]/retrieve/route.ts b/src/app/api/v1/indexes/[name]/retrieve/route.ts
@@ -1,8 +1,5 @@
 import {
-  ChatEngine,
-  getChatEngineByIdOrDefault,
-  getChatEngineByNameOrDefault,
-  getDefaultChatEngine
+  getChatEngineByIdOrName,
 } from "@/core/repositories/chat_engine";
 import {getIndexByName} from '@/core/repositories/index_';
 import {LlamaindexRetrieveService} from '@/core/services/llamaindex/retrieving';
@@ -31,17 +28,7 @@ export const POST = defineHandler({
     notFound();
   }
 
-  let engine: ChatEngine;
-  if (body.chat_engine) {
-    engine = await getChatEngineByNameOrDefault(body.chat_engine);
-  } else if (body.engine) {
-    // TODO: remove it after migration is finished.
-    engine = await getChatEngineByIdOrDefault(body.engine)
-  } else {
-    engine = await getDefaultChatEngine();
-  }
-  const { engine_options } = engine;
-
+  const { engine_options } = await getChatEngineByIdOrName(body.chat_engine ?? body.engine)
   const {
     llm: llmConfig = {
       provider: LLMProvider.OPENAI,

diff --git a/src/core/repositories/chat_engine.ts b/src/core/repositories/chat_engine.ts
@@ -1,3 +1,4 @@
+import {typeOf} from "react-is";
 import {ChatEngineProvider} from "../schema/chat_engines";
 import {
   CondenseQuestionChatEngineOptions,
@@ -52,23 +53,17 @@ export async function getDefaultChatEngine () {
     .executeTakeFirstOrThrow();
 }
 
-export async function getChatEngineByIdOrDefault (engineId?: number): Promise<ChatEngine> {
-  if (engineId) {
-    const engine = await getChatEngineById(engineId);
+export async function getChatEngineByIdOrName (engineIdOrName?: number | string): Promise<ChatEngine> {
+  if (typeof engineIdOrName === 'number') {
+    const engine = await getChatEngineById(engineIdOrName);
     if (!engine) {
-      throw CHAT_ENGINE_NOT_FOUND_ERROR.format(engineId);
+      throw CHAT_ENGINE_NOT_FOUND_ERROR.format(engineIdOrName);
     }
     return engine;
-  } else {
-    return await getDefaultChatEngine();
-  }
-}
-
-export async function getChatEngineByNameOrDefault (engineName?: string): Promise<ChatEngine> {
-  if (engineName) {
-    const engine = await getChatEngineByName(engineName);
+  } else if (typeof engineIdOrName === 'string') {
+    const engine = await getChatEngineByName(engineIdOrName);
     if (!engine) {
-      throw CHAT_ENGINE_NOT_FOUND_ERROR.format(engineName);
+      throw CHAT_ENGINE_NOT_FOUND_ERROR.format(engineIdOrName);
     }
     return engine;
   } else {

diff --git a/src/core/services/llamaindex/chating.ts b/src/core/services/llamaindex/chating.ts
@@ -105,7 +105,7 @@ export class LlamaindexChatService extends AppChatService {
       prompts
     } = engineOptions;
 
-    const chatEngineName = await getChatEngineNameByID(chat.engine_id);
+    const chatEngineName = chat.engine_name || await getChatEngineNameByID(chat.engine_id);
 
     // Init tracing.
     const trace = this.langfuse?.trace({