make format async (langchain-ai#161)

* make format async * make integration test
NJordan72 · Feb 28, 2023 · 4396c15 · 4396c15
1 parent bff3338
commit 4396c15
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 11 deletions.
diff --git a/langchain/src/chains/combine_docs_chain.ts b/langchain/src/chains/combine_docs_chain.ts
@@ -171,11 +171,14 @@ export class MapReduceDocumentsChain
         [this.documentVariableName]: d.pageContent,
         ...rest,
       }));
-      const length = inputs
-        .map((i) =>
-          this.llmChain.llm.getNumTokens(this.llmChain.prompt.format(i))
-        )
-        .reduce((a, b) => a + b, 0);
+      const promises = inputs.map(async (i) => {
+        const prompt = await this.llmChain.prompt.format(i);
+        return this.llmChain.llm.getNumTokens(prompt);
+      });
+
+      const length = await Promise.all(promises).then((results) =>
+        results.reduce((a, b) => a + b, 0)
+      );
 
       if (length < this.maxTokens) {
         break;

diff --git a/langchain/src/chains/llm_chain.ts b/langchain/src/chains/llm_chain.ts
@@ -69,7 +69,7 @@ export class LLMChain extends BaseChain implements LLMChainInput {
     if ("stop" in values && Array.isArray(values.stop)) {
       stop = values.stop;
     }
-    const formattedString = this.prompt.format(values);
+    const formattedString = await this.prompt.format(values);
     const llmResult = await this.llm.call(formattedString, stop);
     const result = { [this.outputKey]: llmResult };
     return result;

diff --git a/langchain/src/prompts/base.ts b/langchain/src/prompts/base.ts
@@ -56,7 +56,7 @@ export abstract class BasePromptTemplate implements BasePromptTemplateInput {
    * prompt.format({ foo: "bar" });
    * ```
    */
-  abstract format(values: InputValues): string;
+  abstract format(values: InputValues): Promise<string>;
 
   /**
    * Return the string type key uniquely identifying this class of prompt template.

diff --git a/langchain/src/prompts/few_shot.ts b/langchain/src/prompts/few_shot.ts
@@ -152,7 +152,7 @@ export class FewShotPromptTemplate
     );
   }
 
-  format(values: InputValues): string {
+  async format(values: InputValues): Promise<string> {
     const examples = this.getExamples(values);
 
     const exampleStrings = examples.map((example) =>

diff --git a/langchain/src/prompts/prompt.ts b/langchain/src/prompts/prompt.ts
@@ -88,7 +88,7 @@ export class PromptTemplate
     return "prompt";
   }
 
-  format(values: InputValues): string {
+  async format(values: InputValues): Promise<string> {
     return renderTemplate(this.template, this.templateFormat, values);
   }
 

diff --git a/langchain/src/prompts/tests/load.test.ts b/langchain/src/prompts/tests/load.test.ts
@@ -12,13 +12,13 @@ test("Load Hello World Prompt", async () => {
   const helloWorld = path.join(PROMPTS_DIR, "hello_world.yaml");
   const prompt = await loadPrompt(helloWorld);
   expect(prompt._getPromptType()).toBe("prompt");
-  expect(prompt.format({})).toBe("Say hello world.");
+  expect(await prompt.format({})).toBe("Say hello world.");
 });
 
 test("Load hub prompt", async () => {
   const prompt = await loadPrompt(
     "lc@abb92d8://prompts/hello-world/prompt.yaml"
   );
   expect(prompt._getPromptType()).toBe("prompt");
-  expect(prompt.format({})).toBe("Say hello world.");
+  expect(await prompt.format({})).toBe("Say hello world.");
 });