request.py

import os
import json
from typing import Any, Dict, Iterable, List, Optional, Union

import requests


class Client:
    """
    Minimal requests-based client for Novita AI's OpenAI-compatible LLM API.

    Docs:
      - LLM API overview: https://novita.ai/docs/guides/llm-api
      - Base URL (OpenAI-compatible): https://api.novita.ai/v3/openai
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: str = "https://api.novita.ai/v3/openai",
        timeout: float = 300.0,
    ) -> None:
        """
        Args:
            api_key: Your Novita AI API key. If None, reads API_KEY from env.
            base_url: Base URL for the OpenAI-compatible API.
            timeout: Request timeout in seconds.
        """
        self.api_key = api_key or os.getenv("API_KEY")
        if not self.api_key:
            raise ValueError("Novita API key not provided and API_KEY not set.")

        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    # ------------------------------------------------------------------ #
    # Internal helpers
    # ------------------------------------------------------------------ #

    @property
    def _headers(self) -> Dict[str, str]:
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

    def _post(
        self,
        path: str,
        json_body: Dict[str, Any],
        stream: bool = False,
    ) -> requests.Response:
        url = f"{self.base_url}{path}"
        resp = requests.post(
            url,
            headers=self._headers,
            json=json_body,
            timeout=self.timeout,
            stream=stream,
        )
        # Let the caller decide how to handle stream errors,
        # but for non-stream we raise here:
        if not stream:
            resp.raise_for_status()
        return resp

    def _get(self, path: str) -> Dict[str, Any]:
        url = f"{self.base_url}{path}"
        resp = requests.get(url, headers=self._headers, timeout=self.timeout)
        resp.raise_for_status()
        return resp.json()

    # ------------------------------------------------------------------ #
    # Public API methods
    # ------------------------------------------------------------------ #

    def chat_completion(
        self,
        model: str,
        messages: List[Dict[str, str]],
        *,
        stream: bool = False,
        response_format: Optional[Dict[str, Any]] = None,
        max_tokens: Optional[int] = None,
        **params: Any,
    ):
        body: Dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": stream,
            **params,
        }

        if max_tokens is not None:
            body["max_tokens"] = max_tokens

        if response_format is not None:
            body["response_format"] = response_format

        resp = self._post("/chat/completions", body, stream=stream)

        if not stream:
            return resp.json()

        def _iter_stream():
            try:
                for line in resp.iter_lines(decode_unicode=True):
                    if not line:
                        continue
                    if isinstance(line, bytes):
                        line = line.decode("utf-8", errors="ignore")
                    if line.startswith("data:"):
                        data = line[len("data:"):].strip()
                        if data == "[DONE]":
                            break
                        import json
                        try:
                            yield json.loads(data)
                        except json.JSONDecodeError:
                            continue
            finally:
                resp.close()

        return _iter_stream()


    def completion(
        self,
        model: str,
        prompt: Union[str, List[str]],
        stream: bool = False,
        **params: Any,
    ) -> Union[Dict[str, Any], Iterable[Dict[str, Any]]]:
        """
        Create a text completion (non-chat).

        Args:
            model: Model name.
            prompt: Single prompt string or list of prompts.
            stream: If True, return generator for streamed chunks.
            **params: OpenAI-compatible params (max_tokens, temperature, etc.).

        Returns:
            - If stream=False: dict
            - If stream=True: iterator of dicts
        """
        body = {
            "model": model,
            "prompt": prompt,
            "stream": stream,
            **params,
        }

        resp = self._post("/completions", body, stream=stream)

        if not stream:
            return resp.json()

        def _iter_stream() -> Iterable[Dict[str, Any]]:
            try:
                for line in resp.iter_lines(decode_unicode=True):
                    if not line:
                        continue
                    if isinstance(line, bytes):
                        line = line.decode("utf-8", errors="ignore")
                    if line.startswith("data:"):
                        data = line[len("data:") :].strip()
                        if data == "[DONE]":
                            break
                        try:
                            yield json.loads(data)
                        except json.JSONDecodeError:
                            continue
            finally:
                resp.close()

        return _iter_stream()

    def list_models(self) -> Dict[str, Any]:
        """
        List available models.
        """
        return self._get("/models")

    def retrieve_model(self, model: str) -> Dict[str, Any]:
        """
        Retrieve details for a specific model.

        Args:
            model: Model id/name.
        """
        return self._get(f"/models/{model}")