add full input chains validation

Author: harry_squater

langchain/src/chains/simple_sequential_chain.ts

@@ -64,6 +64,10 @@ export class SimpleSequentialChain
     return [this.inputKey];
   }
 
+  get outputKeys(): string[] {
+    return [this.outputKey];
+  }
+
   constructor(fields: SimpleSequentialChainInput) {
     super(fields.memory, fields.verbose, fields.callbackManager);
     this.chains = fields.chains;
@@ -80,6 +84,13 @@ export class SimpleSequentialChain
           } for ${chain._chainType()}.`
         );
       }
+      if (chain.outputKeys.length !== 1) {
+        throw new Error(
+          `Chains used in SimpleSequentialChain should all have one output, got ${
+            chain.outputKeys.length
+          } for ${chain._chainType()}.`
+        );
+      }
     }
   }
 

langchain/src/chains/tests/simple_sequential_chain.test.ts

@@ -4,6 +4,11 @@ import { LLMResult } from "../../schema/index.js";
 import { LLMChain } from "../llm_chain.js";
 import { PromptTemplate } from "../../prompts/index.js";
 import { SimpleSequentialChain } from "../simple_sequential_chain.js";
+import { AnalyzeDocumentChain } from "../analyze_documents_chain.js";
+import { ConversationalRetrievalQAChain } from "../conversational_retrieval_chain.js";
+import { VectorStoreRetriever } from "../../vectorstores/base.js";
+import { FakeEmbeddings } from "../../embeddings/fake.js";
+import { MemoryVectorStore } from "../../vectorstores/memory.js";
 
 class FakeLLM1 extends BaseLLM {
   nrMapCalls = 0;
@@ -64,3 +69,50 @@ test("Test SimpleSequentialChain", async () => {
   const response = await combinedChain.run("initial question");
   expect(response).toEqual("final answer");
 });
+
+test("Test SimpleSequentialChain input chains' single input validation", async () => {
+  const model1 = new FakeLLM1({});
+  const model2 = new FakeLLM2({});
+  const template = "Some arbitrary template with fake {input1} and {input2}.";
+  const prompt = new PromptTemplate({
+    template,
+    inputVariables: ["input1", "input2"],
+  });
+  const chain1 = new LLMChain({ llm: model1, prompt });
+  const chain2 = new LLMChain({ llm: model2, prompt });
+  expect(() => {
+    /* eslint-disable no-new */
+    new SimpleSequentialChain({ chains: [chain1, chain2] });
+  }).toThrowErrorMatchingInlineSnapshot(
+    `"Chains used in SimpleSequentialChain should all have one input, got 2 for llm_chain."`
+  );
+});
+
+test("Test SimpleSequentialChain input chains' single ouput validation", async () => {
+  const model1 = new FakeLLM1({});
+  const fakeEmbeddings = new FakeEmbeddings();
+  const anyStore = new MemoryVectorStore(fakeEmbeddings);
+  const retriever = new VectorStoreRetriever({
+    vectorStore: anyStore,
+  });
+  const template = "Some arbitrary template with fake {input}.";
+  const prompt = new PromptTemplate({ template, inputVariables: ["input"] });
+  const chain1 = new LLMChain({ llm: model1, prompt });
+  const chain2 = new ConversationalRetrievalQAChain({
+    retriever,
+    combineDocumentsChain: chain1,
+    questionGeneratorChain: chain1,
+    returnSourceDocuments: true,
+  });
+  // Chain below is is not meant to work in a real-life scenario.
+  // It's only combined this way to get one input/multiple outputs chain.
+  const multipleOutputChain = new AnalyzeDocumentChain({
+    combineDocumentsChain: chain2,
+  });
+  expect(() => {
+    /* eslint-disable no-new */
+    new SimpleSequentialChain({ chains: [chain1, multipleOutputChain] });
+  }).toThrowErrorMatchingInlineSnapshot(
+    `"Chains used in SimpleSequentialChain should all have one output, got 2 for analyze_document_chain."`
+  );
+});