| from .base_agent import BaseAgent | |
| from prompt.template import MODEL_SELECTION_PROMPT | |
| from prompt.constants import modeling_methods | |
| class ModelSelection(BaseAgent): | |
| def __init__(self, llm): | |
| super().__init__(llm) | |
| def select_models(self, modeling_problem: str): | |
| prompt = MODEL_SELECTION_PROMPT.format(modeling_methods=modeling_methods, modeling_problem=modeling_problem) | |
| answer = self.llm.generate(prompt) | |
| selected_models = answer.split('Conclusion:')[-1] | |
| return selected_models | |