feat: switch generate provider based on attachment

toeverything · Apr 9, 2024 · 76f2196 · 76f2196
1 parent 31a723d
commit 76f2196
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git a/packages/backend/server/src/plugins/copilot/controller.ts b/packages/backend/server/src/plugins/copilot/controller.ts
@@ -24,7 +24,7 @@ import { Public } from '../../core/auth';
 import { CurrentUser } from '../../core/auth/current-user';
 import { ProviderService } from './providers';
 import { ChatSession, ChatSessionService } from './session';
-import { CopilotCapability, CopilotProviderType } from './types';
+import { CopilotCapability } from './types';
 
 export interface ChatEvent {
   type: 'attachment' | 'message';
@@ -39,6 +39,21 @@ export class CopilotController {
     private readonly provider: ProviderService
   ) {}
 
+  private async hasAttachment(sessionId: string, messageId?: string) {
+    const session = await this.chatSession.get(sessionId);
+    if (!session) {
+      throw new BadRequestException('Session not found');
+    }
+
+    if (messageId) {
+      const message = await session.getMessageById(messageId);
+      if (Array.isArray(message.attachments) && message.attachments.length) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   private async appendSessionMessage(
     sessionId: string,
     message?: string,
@@ -180,8 +195,9 @@ export class CopilotController {
     @Query() params: Record<string, string>
   ): Promise<Observable<ChatEvent>> {
     const provider = this.provider.getProviderByCapability(
-      CopilotCapability.ImageToImage,
-      CopilotProviderType.FAL
+      (await this.hasAttachment(sessionId, messageId))
+        ? CopilotCapability.ImageToImage
+        : CopilotCapability.TextToImage
     );
     if (!provider) {
       throw new InternalServerErrorException('No provider available');

diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts
@@ -42,6 +42,14 @@ export class ChatSession implements AsyncDisposable {
     this.state.messages.push(message);
   }
 
+  async getMessageById(messageId: string) {
+    const message = await this.messageCache.get(messageId);
+    if (!message || message.sessionId !== this.state.sessionId) {
+      throw new Error(`Message not found: ${messageId}`);
+    }
+    return message;
+  }
+
   async pushByMessageId(messageId: string) {
     const message = await this.messageCache.get(messageId);
     if (!message || message.sessionId !== this.state.sessionId) {