main.py

"""
main.py
Calls the NVIDIA Integrate chat completions API. File-to-base64 encoding
(e.g. for image inputs) is delegated to the compiled C++ extension `fastb64`
for speed -- build it first with `python setup.py build_ext --inplace`.

The API key is read from the NVIDIA_API_KEY environment variable.
Never hardcode API keys in source files.
"""

import os
import json
import requests

import fastb64  # compiled C++ extension (base64_module.cpp)

INVOKE_URL = "https://integrate.api.nvidia.com/v1/chat/completions"
MODEL = "qwen/qwen3.5-397b-a17b"
STREAM = False

API_KEY = os.environ.get("NVIDIA_API_KEY")
if not API_KEY:
    raise RuntimeError(
        "Set the NVIDIA_API_KEY environment variable before running this script, "
        "e.g.: export NVIDIA_API_KEY='nvapi-...'"
    )


def encode_file(path: str) -> str:
    """Base64-encode a local file using the fast C++ extension."""
    return fastb64.read_file_b64(path)


def build_payload(prompt: str, image_path: str | None = None) -> dict:
    content = prompt
    if image_path:
        b64 = encode_file(image_path)
        content = f'{prompt} <img src="data:image/png;base64,{b64}" />'

    return {
        "model": MODEL,
        "messages": [{"role": "user", "content": content}],
        "max_tokens": 16384,
        "temperature": 0.60,
        "top_p": 0.95,
        "top_k": 20,
        "presence_penalty": 0,
        "repetition_penalty": 1,
        "stream": STREAM,
    }


def call_api(prompt: str, image_path: str | None = None):
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Accept": "text/event-stream" if STREAM else "application/json",
    }
    payload = build_payload(prompt, image_path)

    response = requests.post(INVOKE_URL, headers=headers, json=payload, stream=STREAM)
    response.raise_for_status()

    if STREAM:
        for line in response.iter_lines():
            if line:
                print(line.decode("utf-8"))
        return None

    data = response.json()
    print(json.dumps(data, indent=2))
    return data


if __name__ == "__main__":
    # Example: call_api("Describe this image", image_path="photo.png")
    call_api("Hello! Just testing the API connection.")