config/clients/nebius_cli.py

import base64
import os
from io import BytesIO
from typing import Dict, List, Optional

from openai import APIConnectionError, AuthenticationError, OpenAI, RateLimitError
from PIL import Image


# Configuration
API_KEY_ENV_VAR = "NEBIUS_API_KEY"
BASE_URL = os.getenv("NEBIUS_BASE_URL", "https://api.studio.nebius.ai/v1/")
DEFAULT_TEXT_MODEL = "meta-llama/Meta-Llama-3.1-70B-Instruct"
DEFAULT_VISION_MODEL = "meta-llama/Llama-3.2-11B-Vision-Instruct"


def _get_api_key(provided_key: Optional[str] = None) -> Optional[str]:
    """
    Returns the provided key if valid, otherwise falls back to the environment variable.
    """
    if provided_key and provided_key.strip():
        return provided_key.strip()
    return os.getenv(API_KEY_ENV_VAR)


def call_api(messages: List[Dict], model_name: Optional[str] = None, image: Optional[Image.Image] = None, api_key: Optional[str] = None, **kwargs):
    """
    Calls the Nebius AI Studio API (OpenAI-compatible).
    """
    # 1. Resolve API Key
    final_api_key = _get_api_key(api_key)

    if not final_api_key:
        return "Error: Authentication required. Please provide a Nebius API Key."

    # 2. Initialize Client
    try:
        client = OpenAI(base_url=BASE_URL, api_key=final_api_key)
    except Exception as e:
        return f"Error: Failed to initialize Nebius client. {e}"

    # 3. Prepare Messages
    final_messages = []

    if image:
        print("Making a VISION call to Nebius.")
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

        user_text = "Analyze this image."
        for msg in messages:
            if msg["role"] == "user":
                user_text = msg["content"]
                break

        content = [{"type": "text", "text": user_text}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}}]

        for msg in messages:
            if msg["role"] == "system":
                final_messages.append(msg)

        final_messages.append({"role": "user", "content": content})
        final_model = model_name or DEFAULT_VISION_MODEL
    else:
        print("Making a TEXT call to Nebius.")
        final_messages = messages
        final_model = model_name or DEFAULT_TEXT_MODEL

    # 4. Call API
    try:
        print(f"Calling Nebius API with model: {final_model}")
        completion = client.chat.completions.create(model=final_model, messages=final_messages, **kwargs)
        return completion.choices[0].message.content
    except AuthenticationError as e:
        raise ValueError(f"Authentication Failed: {e.body.get('message', str(e)) if e.body else str(e)}")
    except RateLimitError:
        raise RuntimeError("Rate Limit Exceeded (429). Please try again later.")
    except APIConnectionError:
        raise ConnectionError("Failed to connect to API. Check your internet.")
    except Exception as e:
        raise RuntimeError(f"API Error: {str(e)}")