🔨 refactor(base.py): extract methods to handle model-specific and API key-specific field values

This commit refactors the FrontendNode class by extracting two methods to handle specific field values related to models and API keys. The _handle_model_specific_field_values method handles the options and is_list properties for fields related to models, while the _handle_api_key_specific_field_values method handles the display_name and required properties for fields related to API keys. This improves the readability and maintainability of the code.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-05 13:48:17 -03:00
commit 5a103f0f6f

View file

@ -117,17 +117,30 @@ class FrontendNode(BaseModel):
) -> None: ) -> None:
"""Handles specific field values for certain fields.""" """Handles specific field values for certain fields."""
if key == "headers": if key == "headers":
field.value = """{'Authorization': field.value = """{'Authorization': 'Bearer <token>'}"""
'Bearer <token>'}""" FrontendNode._handle_model_specific_field_values(field, key, name)
if name == "OpenAI" and key == "model_name": FrontendNode._handle_api_key_specific_field_values(field, key, name)
field.options = constants.OPENAI_MODELS
field.is_list = True @staticmethod
elif name == "ChatOpenAI" and key == "model_name": def _handle_model_specific_field_values(
field.options = constants.CHAT_OPENAI_MODELS field: TemplateField, key: str, name: Optional[str] = None
field.is_list = True ) -> None:
elif (name == "Anthropic" or name == "ChatAnthropic") and key == "model_name": """Handles specific field values related to models."""
field.options = constants.ANTHROPIC_MODELS model_dict = {
"OpenAI": constants.OPENAI_MODELS,
"ChatOpenAI": constants.CHAT_OPENAI_MODELS,
"Anthropic": constants.ANTHROPIC_MODELS,
"ChatAnthropic": constants.ANTHROPIC_MODELS,
}
if name in model_dict and key == "model_name":
field.options = model_dict[name]
field.is_list = True field.is_list = True
@staticmethod
def _handle_api_key_specific_field_values(
field: TemplateField, key: str, name: Optional[str] = None
) -> None:
"""Handles specific field values related to API keys."""
if "api_key" in key and "OpenAI" in str(name): if "api_key" in key and "OpenAI" in str(name):
field.display_name = "OpenAI API Key" field.display_name = "OpenAI API Key"
field.required = False field.required = False