from typing import Any, List, Optional


class MockAI:
    def __init__(self, response: List):
        self.responses = iter(response)

    def start(self, system: str, user: Any, *, step_name: str) -> List[str]:
        return [next(self.responses)]

    def next(
        self, messages: List[str], prompt: Optional[str] = None, *, step_name: str
    ) -> List[str]:
        return [next(self.responses)]