from yera.models.llm_interfaces.base import BaseLLMInterface


class MockLLMInterface(BaseLLMInterface):
    def __init__(self, config, overrides, creds_map):
        args = {**config.inference.model_dump(), **overrides}

        self.chat_response = args["chat_response"]
        self.struct_response = args["struct_response"]
        self.started = False
        self.stopped = False
        self.chat_calls = []
        self.struct_calls = []
        self.creds_map = creds_map
        super().__init__(**args)

    def chat(self, messages, **kwargs):
        self.chat_calls.append((messages.copy(), kwargs))
        yield from self.chat_response

    def make_struct(self, messages, cls, **kwargs):
        self.struct_calls.append((messages.copy(), cls, kwargs))
        yield from self.struct_response

    def start(self):
        self.started = True

    def stop(self):
        self.stopped = True
