Model/infrastructure/chatgpt/chatgpt.py

from __future__ import annotations

import os
from typing import Optional

from openai import OpenAI
from openai.types.chat import ChatCompletionMessageParam

from infrastructure.chatgpt.exceptions import ChatGPTClientError


class ChatGPT:
    """Thin wrapper over the OpenAI Chat Completions API.

    Sends a single prompt and returns the assistant's reply as plain text.
    """

    DEFAULT_MODEL = "gpt-4o-mini"

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: Optional[str] = None,
    ) -> None:
        key = api_key or os.environ.get("OPENAI_API_KEY")
        if not key:
            raise ChatGPTClientError(
                "No OpenAI API key provided. "
                "Pass api_key or set the OPENAI_API_KEY environment variable."
            )
        self._client = OpenAI(api_key=key)
        self._model = model or self.DEFAULT_MODEL

    def generate(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
    ) -> str:
        """Send a prompt to the model and return its reply text.

        Args:
            prompt: The user message to send.
            system_prompt: Optional instruction that sets the model's behaviour.

        Raises:
            OpenAiClientError: If the model returns an empty response.
        """
        messages: list[ChatCompletionMessageParam] = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        response = self._client.chat.completions.create(
            model=self._model,
            messages=messages,
        )
        content = response.choices[0].message.content
        if content is None:
            raise ChatGPTClientError("ChatGPT returned an empty response.")
        return content