Merge pull request #1559 from benjidotsh/ai/chat-completion-models

feat(ai): add support for chat completion models

.changeset/chatty-pumpkins-shave.md

@@ -0,0 +1,5 @@
+---
+'@roadiehq/rag-ai-backend': patch
+---
+
+Added support for chat completion models

plugins/backend/rag-ai-backend/src/service/LlmService.ts

@@ -14,13 +14,14 @@
  * limitations under the License.
  */
 import { BaseLLM } from '@langchain/core/language_models/llms';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { EmbeddingDoc } from '@roadiehq/rag-ai-node';
 import { Logger } from 'winston';
 import { createPromptTemplates } from './prompts';
 
 export class LlmService {
   private readonly logger: Logger;
-  private readonly model: BaseLLM;
+  private readonly model: BaseLLM | BaseChatModel;
   private readonly prompts: {
     prefixPrompt: (embedding: string) => string;
     suffixPrompt: (input: string) => string;
@@ -32,7 +33,7 @@ export class LlmService {
     configuredPrompts,
   }: {
     logger: Logger;
-    model: BaseLLM;
+    model: BaseLLM | BaseChatModel;
     configuredPrompts?: {
       prefix?: string;
       suffix?: string;
@@ -51,6 +52,7 @@ export class LlmService {
     const prompt = `Human:\n${this.prompts.prefixPrompt(
       promptEmbeddings,
     )}\n ---\n${this.prompts.suffixPrompt(query)}\nAssistant:`;
-    return await this.model.invoke(prompt);
+    const response = await this.model.invoke(prompt);
+    return typeof response === 'string' ? response : response.content;
   }
 }

plugins/backend/rag-ai-backend/src/service/router.ts

@@ -19,6 +19,7 @@ import Router from 'express-promise-router';
 import { Logger } from 'winston';
 import { AugmentationIndexer, RetrievalPipeline } from '@roadiehq/rag-ai-node';
 import { BaseLLM } from '@langchain/core/language_models/llms';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { LlmService } from './LlmService';
 import { RagAiController } from './RagAiController';
 import { isEmpty } from 'lodash';
@@ -36,7 +37,7 @@ export interface RouterOptions {
   logger: Logger;
   augmentationIndexer: AugmentationIndexer;
   retrievalPipeline: RetrievalPipeline;
-  model: BaseLLM;
+  model: BaseLLM | BaseChatModel;
   config: Config;
 }
 

plugins/backend/rag-ai-backend/src/service/types.ts

@@ -17,13 +17,14 @@ import { TokenManager } from '@backstage/backend-common';
 import { AugmentationIndexer, RetrievalPipeline } from '@roadiehq/rag-ai-node';
 import { Logger } from 'winston';
 import { BaseLLM } from '@langchain/core/language_models/llms';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { Config } from '@backstage/config';
 
 export interface RagAiConfig {
   logger: Logger;
   tokenManager: TokenManager;
   augmentationIndexer: AugmentationIndexer;
   retrievalPipeline: RetrievalPipeline;
-  model: BaseLLM;
+  model: BaseLLM | BaseChatModel;
   config: Config;
 }