made output conform

This commit is contained in:
Shreya Rajpal 2022-12-23 21:05:59 +05:30
parent 3fcbae90ed
commit bb5427d2cd

View File

@ -47,9 +47,7 @@ class ForkChain(Chain, BaseModel):
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)
return self.follow_up_chains[decision_chain_output.strip()](inputs)
except KeyError:
return self.follow_up_chains['default'](inputs)