feat(langchain): add promptId parameter (#27)

Signed-off-by: Tomáš Dvořák <toomas2d@gmail.com>
IBM · Jun 28, 2023 · 53b9853 · 53b9853
1 parent 5bef63d
commit 53b9853
Showing 1 changed file with 17 additions and 5 deletions.
diff --git a/src/langchain/llm.ts b/src/langchain/llm.ts
@@ -6,24 +6,29 @@ import type { LLMResult, Generation } from 'langchain/schema';
 import type { GenerateOutput } from '../client-types.js';
 import { GenerateInput } from '../client-types.js';
 
-export interface GenAIModelOptions {
-  modelId?: string;
+interface BaseGenAIModelOptions {
   stream?: boolean;
   parameters?: Record<string, any>;
   timeout?: number;
   configuration?: Configuration;
 }
 
+export type GenAIModelOptions =
+  | (BaseGenAIModelOptions & { modelId?: string; promptId?: never })
+  | (BaseGenAIModelOptions & { modelId?: never; promptId: string });
+
 export class GenAIModel extends BaseLLM {
   #client: Client;
 
   protected modelId?: string;
+  protected promptId?: string;
   protected stream: boolean;
   protected timeout: number | undefined;
   protected parameters: Record<string, any>;
 
   constructor({
     modelId,
+    promptId,
     stream = false,
     parameters,
     timeout,
@@ -33,6 +38,7 @@ export class GenAIModel extends BaseLLM {
     super(baseParams ?? {});
 
     this.modelId = modelId;
+    this.promptId = promptId;
     this.timeout = timeout;
     this.parameters = parameters || {};
     this.stream = Boolean(stream);
@@ -46,9 +52,15 @@ export class GenAIModel extends BaseLLM {
     const stopSequences = concatUnique(this.parameters.stop, options.stop);
 
     return prompts.map((input) => ({
-      ...(!isNullish(this.modelId) && {
-        model_id: this.modelId,
-      }),
+      ...(!isNullish(this.promptId)
+        ? {
+            prompt_id: this.promptId,
+          }
+        : !isNullish(this.modelId)
+        ? {
+            model_id: this.modelId,
+          }
+        : {}),
       input,
       parameters: {
         ...this.parameters,