openai_client.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""OpenAI client for GPU tools — local key or Aquin web proxy (same as aquin chat)."""
from __future__ import annotations
 
import json
import os
from pathlib import Path
from types import SimpleNamespace
from typing import Any
 
import httpx
 
_CONFIG_PATH = Path.home() / ".aquin" / "config.json"
_ENGINE_STATE = Path.home() / ".aquin" / "engine.json"
 
 
def _read_json(path: Path) -> dict[str, str]:
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}
 
 
def _saved_api_key() -> str:
    from aquin.auth_config import resolve_api_key
    return resolve_api_key(allow_missing=True)
 
 
class _ProxyChatCompletions:
    def __init__(self, proxy: "AquinOpenAIProxy") -> None:
        self._proxy = proxy
 
    def create(self, **kwargs: Any) -> Any:
        payload = {k: v for k, v in kwargs.items() if v is not None}
        data = self._proxy._chat(payload)
        choice = (data.get("choices") or [{}])[0]
        message = choice.get("message") or {}
        return SimpleNamespace(
            choices=[SimpleNamespace(message=SimpleNamespace(content=message.get("content", "")))]
        )
 
 
class _ProxyEmbeddings:
    def __init__(self, proxy: "AquinOpenAIProxy") -> None:
        self._proxy = proxy
 
    def create(self, **kwargs: Any) -> Any:
        payload = {k: v for k, v in kwargs.items() if v is not None}
        data = self._proxy._embeddings(payload)
        rows = data.get("data") or []
        return SimpleNamespace(
            data=[SimpleNamespace(embedding=row.get("embedding", [])) for row in rows]
        )
 
 
class AquinOpenAIProxy:
    """Uses OPENAI_API_KEY from the Next.js app (Bearer AQUIN_API_KEY), like aquin chat."""
 
    def __init__(self, api_key: str, base_url: str, session_id: str = "") -> None:
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._session_id = session_id.strip()
        self.chat = SimpleNamespace(completions=_ProxyChatCompletions(self))
        self.embeddings = _ProxyEmbeddings(self)
 
    def _headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {self._api_key}"}
 
    def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        url = f"{self._base_url}{path}"
        resp = httpx.post(
            url,
            json=payload,
            headers=self._headers(),
            timeout=120.0,
        )
        if resp.status_code >= 400:
            raise RuntimeError(
                f"OpenAI proxy {path} failed ({resp.status_code}) at {url}: {resp.text[:300]}"
            )
        return resp.json()
 
    def _chat(self, payload: dict[str, Any]) -> dict[str, Any]:
        body = {**payload, "stream": False}
        if self._session_id:
            return self._post(f"/api/sync/sessions/{self._session_id}/agent", body)
        return self._post("/api/sync/openai/chat", body)
 
    def _embeddings(self, payload: dict[str, Any]) -> dict[str, Any]:
        if self._session_id:
            return self._post(f"/api/sync/sessions/{self._session_id}/embeddings", payload)
        return self._post("/api/sync/openai/embeddings", payload)
 
 
def get_openai_client(ctx: dict[str, Any] | None = None) -> Any | None:
    """
    Prefer local OPENAI_API_KEY; otherwise proxy via Aquin web (AQUIN_API_KEY + base URL).
    Chat uses /api/sync/sessions/{id}/agent — same path as aquin chat.
    """
    local_key = os.environ.get("OPENAI_API_KEY", "").strip()
    if local_key:
        try:
            from openai import OpenAI
            return OpenAI(api_key=local_key)
        except Exception:
            return None
 
    cfg = _read_json(_ENGINE_STATE)
    c = ctx or {}
    api_key = (
        c.get("api_key")
        or _saved_api_key()
        or os.environ.get("AQUIN_API_KEY", "").strip()
        or cfg.get("api_key", "").strip()
    )
    base_url = (
        c.get("base_url")
        or os.environ.get("AQUIN_BASE_URL", "").strip()
        or cfg.get("base_url", "").strip()
        or "https://api.aquin.app"
    )
    session_id = c.get("session_id") or ""
    if api_key:
        return AquinOpenAIProxy(api_key, base_url, session_id=session_id)
    return None