Add support for fork chains

langchain/chains/base.py

@@ -63,11 +63,9 @@ class Chain(BaseModel, ABC):
             raise ValueError(f"Missing some input keys: {missing_keys}")
 
     def _validate_outputs(self, outputs: Dict[str, str]) -> None:
-        if set(outputs) != set(self.output_keys):
-            raise ValueError(
-                f"Did not get output keys that were expected. "
-                f"Got: {set(outputs)}. Expected: {set(self.output_keys)}."
-            )
+        missing_keys = set(self.output_keys).difference(outputs)
+        if missing_keys:
+            raise ValueError(f"Missing some output keys: {missing_keys}")
 
     @abstractmethod
     def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:

langchain/chains/fork_with_follow_ups.py

@@ -0,0 +1,55 @@
+"""Conditionally executes a follow up chain based on the output of a decision chain."""
+from typing import Dict, List
+
+from pydantic import BaseModel, Extra, validator
+
+from langchain.chains import LLMChain
+from langchain.chains.base import Chain
+
+
+class ForkChain(Chain, BaseModel):
+    """Conditionally executes a follow up chain based on the output of a decision chain."""
+
+    decision_chain: LLMChain
+    follow_up_chains: Dict[str, Chain]
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
+    @validator("follow_up_chains")
+    def default_in_follow_up_chains(cls, v):
+        if "default" not in v:
+            raise ValueError(
+                "`follow_up_chains` must contain a 'default' option. "
+                "This is the chain that is called when the output of the "
+                "decision chain doesn't match any key in `follow_up_chains`."
+            )
+        return v
+
+    @property
+    def input_keys(self) -> List[str]:
+        """Will be whatever keys the prompt expects.
+
+        :meta private:
+        """
+        return self.decision_chain.input_keys
+
+    @property
+    def output_keys(self) -> List[str]:
+        """Will always return text key.
+
+        :meta private:
+        """
+        return []
+
+    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+        decision_chain_output = self.decision_chain.run(inputs)
+        classification_output = decision_chain_output
+
+        try:
+            return self.follow_up_chain[classification_output](inputs)
+        except KeyError:
+            return self.follow_up_chains['default'](inputs)