🔨 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:
"""Handles specific field values for certain fields."""
if key == "headers":
field.value = """{'Authorization':
'Bearer <token>'}"""
if name == "OpenAI" and key == "model_name":
field.options = constants.OPENAI_MODELS
field.is_list = True
elif name == "ChatOpenAI" and key == "model_name":
field.options = constants.CHAT_OPENAI_MODELS
field.is_list = True
elif (name == "Anthropic" or name == "ChatAnthropic") and key == "model_name":
field.options = constants.ANTHROPIC_MODELS
field.value = """{'Authorization': 'Bearer <token>'}"""
FrontendNode._handle_model_specific_field_values(field, key, name)
FrontendNode._handle_api_key_specific_field_values(field, key, name)
@staticmethod
def _handle_model_specific_field_values(
field: TemplateField, key: str, name: Optional[str] = None
) -> None:
"""Handles specific field values related to 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
@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):
field.display_name = "OpenAI API Key"
field.required = False