from typing import Any


class _MockMessage:
  content: str


class _MockChoice:
  message: _MockMessage

  def __init__(self):
    self.message = _MockMessage()
    self.message.content = 'mock response'


class _MockResponse:
  choices: list[_MockChoice]

  def __init__(self):
    self.choices = [_MockChoice()]


class MockChat:
  """Mock chat completion handler."""

  def complete(
      self, model: str, messages: list[Any], max_tokens: int | None = None,
      temperature: float | None = None, stop: list[str] | None = None
  ) -> _MockResponse:
    return _MockResponse()


class MistralMock:
  """Mock Mistral API client for testing."""

  chat: MockChat

  def __init__(self):
    self.chat = MockChat()
