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