feat: support langchain serialization

Signed-off-by: Tomas Dvorak <[email protected]>
IBM · May 28, 2024 · b7ed959 · b7ed959
1 parent b89924a
commit b7ed959
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 3 deletions.
diff --git a/src/langchain/llm-chat.ts b/src/langchain/llm-chat.ts
@@ -11,6 +11,8 @@ import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
 import { ChatGenerationChunk, ChatResult } from '@langchain/core/outputs';
 import { BaseLanguageModelCallOptions as BaseChatModelCallOptions } from '@langchain/core/language_models/base';
 import merge from 'lodash/merge.js';
+import { load } from '@langchain/core/load';
+import type { Serialized } from '@langchain/core/load/serializable';
 
 import { Client, Configuration } from '../client.js';
 import { TextChatCreateInput, TextChatCreateStreamInput } from '../schema.js';
@@ -218,11 +220,40 @@ export class GenAIChatModel extends BaseChatModel<GenAIChatModelOptions> {
     });
   }
 
-  _llmType(): string {
-    return 'GenAIChat';
+  lc_serializable = true;
+  lc_namespace = ['@ibm-generative-ai/node-sdk', 'langchain', 'llm-chat'];
+
+  get lc_id(): string[] {
+    return [...this.lc_namespace, 'GenAIChatModel'];
+  }
+
+  lc_kwargs = {
+    modelId: undefined,
+    promptId: undefined,
+    conversationId: undefined,
+    parameters: undefined,
+    moderations: undefined,
+    useConversationParameters: undefined,
+    parentId: undefined,
+    trimMethod: undefined,
+  };
+
+  static async fromJSON(value: string | Serialized) {
+    const input = typeof value === 'string' ? value : JSON.stringify(value);
+    return await load(input, {
+      optionalImportsMap: {
+        '@ibm-generative-ai/node-sdk/langchain/llm-chat': {
+          GenAIModel: GenAIChatModel,
+        },
+      },
+    });
   }
 
   _modelType(): string {
     return this.modelId;
   }
+
+  _llmType(): string {
+    return 'GenAIChatModel';
+  }
 }
diff --git a/src/langchain/llm.ts b/src/langchain/llm.ts
@@ -7,6 +7,8 @@ import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
 import type { LLMResult } from '@langchain/core/outputs';
 import { GenerationChunk } from '@langchain/core/outputs';
 import merge from 'lodash/merge.js';
+import { load } from '@langchain/core/load';
+import type { Serialized } from '@langchain/core/load/serializable';
 
 import { Client, Configuration } from '../client.js';
 import { concatUnique, isNullish } from '../helpers/common.js';
@@ -183,11 +185,40 @@ export class GenAIModel extends BaseLLM<GenAIModelOptions> {
     return result.results.at(0)?.token_count ?? 0;
   }
 
+  static async fromJSON(value: string | Serialized) {
+    const input = typeof value === 'string' ? value : JSON.stringify(value);
+    return await load(input, {
+      optionalImportsMap: {
+        '@ibm-generative-ai/node-sdk/langchain/llm': {
+          GenAIModel: GenAIModel,
+        },
+      },
+    });
+  }
+
   _modelType(): string {
     return this.modelId;
   }
 
   _llmType(): string {
-    return 'GenAI';
+    return 'GenAIModel';
   }
+
+  lc_serializable = true;
+  lc_namespace = ['@ibm-generative-ai/node-sdk', 'langchain', 'llm'];
+
+  get lc_id(): string[] {
+    return [...this.lc_namespace, 'GenAIModel'];
+  }
+
+  lc_kwargs = {
+    modelId: undefined,
+    promptId: undefined,
+    conversationId: undefined,
+    parameters: undefined,
+    moderations: undefined,
+    useConversationParameters: undefined,
+    parentId: undefined,
+    trimMethod: undefined,
+  };
 }
diff --git a/tests/e2e/langchain/llm-chat.test.ts b/tests/e2e/langchain/llm-chat.test.ts
@@ -122,4 +122,11 @@ describe('LangChain Chat', () => {
       expect(tokens).toStrictEqual(contents);
     });
   });
+
+  it('Serializes', async () => {
+    const model = makeModel();
+    const serialized = model.toJSON();
+    const deserialized = await GenAIChatModel.fromJSON(serialized);
+    expect(deserialized).toBeInstanceOf(GenAIChatModel);
+  });
 });
diff --git a/tests/e2e/langchain/llm.test.ts b/tests/e2e/langchain/llm.test.ts
@@ -34,6 +34,13 @@ describe('Langchain', () => {
     });
   });
 
+  it('Serializes', async () => {
+    const client = makeClient('google/flan-ul2');
+    const serialized = client.toJSON();
+    const deserialized = await GenAIModel.fromJSON(serialized);
+    expect(deserialized).toBeInstanceOf(GenAIModel);
+  });
+
   describe('generate', () => {
     // TODO: enable once we will set default model for the test account
     test.skip('should handle empty modelId', async () => {