from agents import Agent

from diona.project.agent.model import BaseModelProvider


class TestAgent:
    @staticmethod
    def create_agent(model_provider: BaseModelProvider = None) -> Agent:
        model_provider = model_provider or BaseModelProvider.get_default()

        try:
            model = model_provider.get_model(None)
        except Exception as e:
            raise ValueError(f"Failed to get model: {e}. Check API key and base URL configuration.")

        instructions = """
        You are a test agent. Simply repeat the user's input with a prefix 'Test Agent: '.
        
        If the user provides no input, respond with 'Test Agent: Hello! This is a test agent.'
        """

        agent = Agent(
            name="TestAgent",
            instructions=instructions,
            model=model,
            handoff_description="A test agent that echoes your input"
        )

        return agent


def get_test_agent(model_provider: BaseModelProvider = None) -> Agent:
    return TestAgent.create_agent(model_provider)
