integrations/quantaalpha/client.py

"""OpenAI-compatible LLM client (QuantaAlpha style)."""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import yaml

from config.settings import PROJECT_ROOT


@dataclass
class LLMConfig:
    api_key: str
    base_url: str
    chat_model: str
    reasoning_model: str | None = None
    temperature: float = 0.3
    max_tokens: int = 4000
    max_retry: int = 3


def load_llm_config(config_path: str | None = None) -> LLMConfig:
    path = PROJECT_ROOT / "config" / "quantaalpha.yaml"
    if config_path:
        path = Path(config_path) if Path(config_path).is_absolute() else PROJECT_ROOT / config_path

    raw: dict[str, Any] = {}
    if path.exists():
        with open(path, encoding="utf-8") as f:
            raw = yaml.safe_load(f) or {}

    llm = raw.get("llm", {})
    return LLMConfig(
        api_key=os.environ.get("OPENAI_API_KEY", llm.get("api_key", "")),
        base_url=os.environ.get("OPENAI_BASE_URL", llm.get("base_url", "https://api.openai.com/v1")),
        chat_model=os.environ.get("CHAT_MODEL", llm.get("chat_model", "gpt-4o-mini")),
        reasoning_model=os.environ.get("REASONING_MODEL", llm.get("reasoning_model")),
        temperature=float(os.environ.get("CHAT_TEMPERATURE", llm.get("temperature", 0.3))),
        max_tokens=int(os.environ.get("CHAT_MAX_TOKENS", llm.get("max_tokens", 4000))),
        max_retry=int(os.environ.get("MAX_RETRY", llm.get("max_retry", 3))),
    )


class QuantaAlphaLLMClient:
    """Thin wrapper over OpenAI-compatible chat completions API."""

    def __init__(self, config: LLMConfig | None = None):
        self.config = config or load_llm_config()
        if not self.config.api_key:
            raise ValueError(
                "OPENAI_API_KEY not set. Configure config/quantaalpha.yaml or export OPENAI_API_KEY."
            )

        try:
            from openai import OpenAI
        except ImportError as exc:
            raise ImportError("Install openai: pip install openai") from exc

        self._client = OpenAI(api_key=self.config.api_key, base_url=self.config.base_url)

    def chat(
        self,
        messages: list[dict[str, str]],
        model: str | None = None,
        temperature: float | None = None,
        response_json: bool = False,
    ) -> str:
        model = model or self.config.chat_model
        temperature = self.config.temperature if temperature is None else temperature

        kwargs: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": self.config.max_tokens,
        }
        if response_json:
            kwargs["response_format"] = {"type": "json_object"}

        last_err = None
        for _ in range(self.config.max_retry):
            try:
                resp = self._client.chat.completions.create(**kwargs)
                return resp.choices[0].message.content or ""
            except Exception as exc:
                last_err = exc
        raise RuntimeError(f"LLM request failed after retries: {last_err}")

    def chat_json(self, messages: list[dict[str, str]], **kwargs) -> dict[str, Any]:
        text = self.chat(messages, response_json=True, **kwargs)
        return json.loads(text)