feat: save document source on DB (#186)

* feat: save document source on DB * feat: update chat schema * feat: fix PR comments
xgeekshq · Sep 21, 2023 · dafccaf · dafccaf
1 parent ce3c3c2
commit dafccaf
Show file tree

Hide file tree

Showing 11 changed files with 67 additions and 11 deletions.
diff --git a/packages/api/src/ai/facades/ai.service.ts b/packages/api/src/ai/facades/ai.service.ts
@@ -5,6 +5,7 @@ import { SimpleConversationChainService } from '@/ai/services/simple-conversatio
 import { VectorDbService } from '@/ai/services/vector-db.service';
 import { AppConfigService } from '@/app-config/app-config.service';
 import { RedisChatMemoryNotFoundException } from '@/chats/exceptions/redis-chat-memory-not-found.exception';
+import { AiResponse } from '@/common/types/ai-response';
 import { ChatDocument } from '@/common/types/chat';
 import { Injectable, Logger } from '@nestjs/common';
 import {
@@ -42,7 +43,7 @@ Helpful answer:`
     chatLlmId: string,
     shouldSummarize: boolean,
     documents: ChatDocument[]
-  ): Promise<string> {
+  ): Promise<AiResponse | ChainValues> {
     let summary: ChainValues;
     let aiExecutor: AIExecutor;
 
@@ -69,8 +70,7 @@ Helpful answer:`
     }
 
     try {
-      const aiResponse = await aiExecutor.call({ input });
-      return aiResponse.output;
+      return aiExecutor.call({ input });
     } catch (e) {
       this.logger.error(
         `Failed to ask AI in room ${roomId} for a response for question: `,

diff --git a/packages/api/src/ai/services/document-conversation-chain.service.ts b/packages/api/src/ai/services/document-conversation-chain.service.ts
@@ -58,7 +58,11 @@ export class DocumentConversationChain {
           res.text,
           this.args.summary
         );
-        return { output: res.text, source: res.sourceDocuments };
+        return {
+          output: res.text,
+          source: res.sourceDocuments,
+          document: document.meta.filename,
+        };
       }
     }
     const simpleConversationChain =

diff --git a/packages/api/src/chats/chat-socket.gateway.ts b/packages/api/src/chats/chat-socket.gateway.ts
@@ -89,12 +89,11 @@ export class ChatSocketGateway {
       await this.shouldSummarizeChat(data.roomId),
       await this.fetchDocumentsForContext(data.roomId)
     );
-
     this.server.to(data.roomId).emit(
       'message',
       createSocketMessageResponseFactory({
         id: data.id,
-        response: aiResponse,
+        response: aiResponse.output,
         isAi: true,
         createdAt: new Date().toISOString(),
       })
@@ -112,12 +111,22 @@ export class ChatSocketGateway {
         }),
         AddMessageToChatRequestSchema.parse({
           sender: { isAi: true },
-          content: aiResponse,
+          content: aiResponse.output,
           meta: {
-            tokens: encode(aiResponse).length,
+            tokens: encode(aiResponse.output).length,
             ai: {
               llmModel: this.appConfigService.getAiAppConfig().defaultAiModel,
             },
+            ...(aiResponse.document
+              ? {
+                  source: {
+                    filename: aiResponse.document,
+                    snippets: aiResponse.source.map(
+                      (source) => source.pageContent
+                    ),
+                  },
+                }
+              : {}),
           },
         }),
         data.userId

diff --git a/packages/api/src/chats/job-consumers/add-message-pair-to-history.job-consumer.ts b/packages/api/src/chats/job-consumers/add-message-pair-to-history.job-consumer.ts
@@ -35,7 +35,6 @@ export class AddMessagePairToHistoryJobConsumer {
     });
 
     answer.meta.replyTo = savedQuestion.id;
-
     await this.addMessageToChatUsecase.execute(roomId, answer);
 
     this.logger.debug(`Answer saved in chat history: `, {

diff --git a/packages/api/src/chats/schemas/chat.mongoose.schema.ts b/packages/api/src/chats/schemas/chat.mongoose.schema.ts
@@ -20,6 +20,20 @@ const MessageMetaSchema = new mongoose.Schema(
       _id: false,
       required: false,
     },
+    source: {
+      type: {
+        filename: {
+          type: String,
+          required: true,
+        },
+        snippets: {
+          type: [String],
+          required: true,
+        },
+      },
+      _id: false,
+      required: false,
+    },
   },
   { _id: false, required: true }
 );

diff --git a/packages/api/src/common/factories/create-chat-message.factory.ts b/packages/api/src/common/factories/create-chat-message.factory.ts
@@ -6,7 +6,6 @@ export function createChatMessageFactory(
   userId?: string
 ): Omit<ChatMessage, 'id'> {
   const { sender, content, meta } = addMessageToChatRequestDto;
-
   return {
     sender: {
       ...(userId ? { userId } : {}),
@@ -23,6 +22,14 @@ export function createChatMessageFactory(
             },
           }
         : {}),
+      ...(!!meta.source
+        ? {
+            source: {
+              filename: meta.source.filename,
+              snippets: meta.source.snippets,
+            },
+          }
+        : {}),
     },
     createdAt: new Date().toISOString(),
   };

diff --git a/packages/api/src/common/types/ai-response.ts b/packages/api/src/common/types/ai-response.ts
@@ -0,0 +1,5 @@
+export interface AiResponse {
+  output: string;
+  source?: string[];
+  document?: string;
+}
diff --git a/packages/api/src/common/types/chat.ts b/packages/api/src/common/types/chat.ts
@@ -4,10 +4,16 @@ export interface AiMeta {
   llmModel: string;
 }
 
+export interface SourceMeta {
+  filename: string;
+  snippets: string[];
+}
+
 export interface ChatMessageMeta {
   tokens: number;
   replyTo?: string;
   ai?: AiMeta;
+  source?: SourceMeta;
 }
 
 export interface Sender {

diff --git a/packages/contract/chats/add-message-to-chat.request.dto.ts b/packages/contract/chats/add-message-to-chat.request.dto.ts
@@ -13,5 +13,11 @@ export const AddMessageToChatRequestSchema = z.object({
         llmModel: z.string(),
       })
       .optional(),
+    source: z
+      .object({
+        filename: z.string(),
+        snippets: z.array(z.string()),
+      })
+      .optional(),
   }),
 });
diff --git a/packages/contract/chats/chat-message.response.dto.ts b/packages/contract/chats/chat-message.response.dto.ts
@@ -15,6 +15,12 @@ export const ChatMessageResponseSchema = z.object({
         llmModel: z.string(),
       })
       .optional(),
+    source: z
+      .object({
+        filename: z.string(),
+        snippets: z.array(z.string()),
+      })
+      .optional(),
   }),
   createdAt: z.string().datetime(),
 });
diff --git a/packages/web-ui/components/error-layout.tsx b/packages/web-ui/components/error-layout.tsx
@@ -18,7 +18,7 @@ export default function ErrorLayout({ minHeight }: ErrorLayoutProps) {
               An unexpected error has occurred
             </h1>
             <p className="mt-4 text-gray-500">
-              Unfortunately something goes wrong. Please attempt to refresh the
+              Unfortunately something went wrong. Please attempt to refresh the
               page, if the problem persists, kindly consider creating a GitHub
               issue for further assistance.
             </p>