llm-connect/llm_connect/openai.py

"""
OpenAI (ChatGPT) adapter — calls the OpenAI chat completions API.
"""

import asyncio
import time
from typing import Optional, Dict, Any

from llm_connect.adapter import LLMAdapter
from llm_connect.models import RunConfig, LLMResponse
from llm_connect.config import resolve_api_key, find_project_root
from llm_connect._http import post_json
from llm_connect.exceptions import (
    LLMConfigurationError,
    LLMAPIError,
    LLMRateLimitError,
)

_DEFAULT_MODEL = "gpt-4.1-mini"
_API_BASE = "https://api.openai.com/v1"


class OpenAIAdapter(LLMAdapter):
    """LLM adapter that calls the OpenAI chat completions endpoint."""

    def __init__(
        self,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        system_prompt: Optional[str] = None,
        max_retries: int = 3,
        **_kwargs: Any,
    ):
        self._model = model or _DEFAULT_MODEL
        self._system_prompt = system_prompt
        self._max_retries = max_retries

        root = find_project_root()
        key_file_paths = [root / "apikey-chatgpt.txt"] if root else []
        self._api_key = resolve_api_key(
            explicit=api_key,
            env_var="OPENAI_API_KEY",
            key_file_paths=key_file_paths,
        )
        if not self._api_key:
            raise LLMConfigurationError(
                "No OpenAI API key found. Set OPENAI_API_KEY or create "
                "apikey-chatgpt.txt in the project root.",
                context={"provider": "openai"},
            )

    # ── LLMAdapter interface ────────────────────────────────────────

    def execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse:
        self._preflight_budget(config)
        model = self._model

        messages: list[Dict[str, str]] = []
        if self._system_prompt:
            messages.append({"role": "system", "content": self._system_prompt})
        messages.append({"role": "user", "content": prompt})

        payload: Dict[str, Any] = {
            "model": model,
            "messages": messages,
            "temperature": config.temperature,
            "max_tokens": config.max_tokens,
        }

        headers = {
            "Authorization": f"Bearer {self._api_key}",
        }
        url = f"{_API_BASE}/chat/completions"

        start = time.time()
        data = self._post_with_retries(url, payload, headers, config.timeout_seconds)
        latency = time.time() - start

        # Parse response (OpenAI chat completions format)
        choice = data.get("choices", [{}])[0]
        content = choice.get("message", {}).get("content", "")
        finish_reason = choice.get("finish_reason", "stop")
        usage = data.get("usage", {})

        response = LLMResponse(
            content=content,
            model=data.get("model", model),
            usage={
                "prompt_tokens": usage.get("prompt_tokens", 0),
                "completion_tokens": usage.get("completion_tokens", 0),
                "total_tokens": usage.get("total_tokens", 0),
            },
            finish_reason=finish_reason,
            metadata={
                "provider": "openai",
                "latency_seconds": round(latency, 3),
                "response_id": data.get("id", ""),
            },
        )
        self._consume_budget(config, response)
        return response

    async def async_execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse:
        """Async wrapper — runs execute_prompt in a thread executor."""
        return await asyncio.to_thread(self.execute_prompt, prompt, config)

    def validate_config(self, config: RunConfig) -> bool:
        if not self._api_key:
            return False
        if not (0.0 <= config.temperature <= 2.0):
            return False
        return True

    # ── Internals ───────────────────────────────────────────────────

    def _post_with_retries(
        self,
        url: str,
        payload: Dict[str, Any],
        headers: Dict[str, str],
        timeout: int,
    ) -> Dict[str, Any]:
        last_exc: Optional[Exception] = None
        for attempt in range(self._max_retries + 1):
            try:
                return post_json(url, payload, headers, timeout=timeout)
            except LLMRateLimitError as exc:
                last_exc = exc
                if attempt < self._max_retries:
                    time.sleep(2 ** attempt)
            except LLMAPIError as exc:
                if exc.status_code >= 500 and attempt < self._max_retries:
                    last_exc = exc
                    time.sleep(2 ** attempt)
                else:
                    raise
        raise last_exc  # type: ignore[misc]