Refactor RunnableExecutor.py to improve output handling

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-27 17:52:34 -03:00
commit 77d1f7e731

View file

@ -5,15 +5,22 @@ from langflow.interface.custom.custom_component import CustomComponent
class RunnableExecComponent(CustomComponent):
documentation: str = "http://docs.langflow.org/components/custom"
description = "Execute a runnable. It will try to guess the input and output keys."
display_name = "Runnable Executor"
beta: bool = True
field_order = [
"input_key",
"output_key",
"input_value",
"runnable",
]
def build_config(self):
return {
"input_key": {
"display_name": "Input Key",
"info": "The key to use for the input.",
"advanced": True,
},
"input_value": {
"display_name": "Inputs",
@ -27,9 +34,54 @@ class RunnableExecComponent(CustomComponent):
"output_key": {
"display_name": "Output Key",
"info": "The key to use for the output.",
"advanced": True,
},
}
def get_output(self, result, input_key, output_key):
"""
Retrieves the output value from the given result dictionary based on the specified input and output keys.
Args:
result (dict): The result dictionary containing the output value.
input_key (str): The key used to retrieve the input value from the result dictionary.
output_key (str): The key used to retrieve the output value from the result dictionary.
Returns:
tuple: A tuple containing the output value and the status message.
"""
possible_output_keys = ["answer", "response", "output", "result", "text"]
status = ""
result_value = None
if output_key in result:
result_value = result.get(output_key)
elif len(result) == 2 and input_key in result:
# get the other key from the result dict
other_key = [k for k in result if k != input_key][0]
if other_key == output_key:
result_value = result.get(output_key)
else:
status += f"Warning: The output key is not '{output_key}'. The output key is '{other_key}'."
result_value = result.get(other_key)
elif len(result) == 1:
result_value = list(result.values())[0]
elif any(k in result for k in possible_output_keys):
for key in possible_output_keys:
if key in result:
result_value = result.get(key)
status += f"Output key: '{key}'."
break
if result_value is None:
result_value = result
status += f"Warning: The output key is not '{output_key}'."
else:
result_value = result
status += f"Warning: The output key is not '{output_key}'."
return result_value, status
def build(
self,
input_value: Text,
@ -38,6 +90,7 @@ class RunnableExecComponent(CustomComponent):
output_key: str = "output",
) -> Text:
result = runnable.invoke({input_key: input_value})
result = result.get(output_key)
self.status = result
return result
result_value, status = self.get_output(result, input_key, output_key)
status += f"\nOutput: {result_value}\n\nRaw Output: {result}"
self.status = status
return result_value