agents/nvidia_llm.py

import os
import logging
import threading
import queue as _queue_module
import time
import requests
import keyring
import base64
from flask import Flask, request, jsonify

# ── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(levelname)s %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("nvidia_llm")
logging.getLogger("werkzeug").setLevel(logging.ERROR)
logging.getLogger("httpx").setLevel(logging.WARNING)

# ── Config ────────────────────────────────────────────────────────────────────
HOST         = os.getenv("NVIDIA_HOST",  "127.0.0.1")
PORT         = int(os.getenv("NVIDIA_PORT", "8002"))
try:
    NVIDIA_API_KEY = keyring.get_password("system", "NVIDIA_API_KEY") or os.getenv("NVIDIA_API_KEY")
except Exception as e:
    NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")

_QUEUE_MAX_SIZE    = int(os.getenv("NVIDIA_QUEUE_MAX", "8"))
_REQUEST_TIMEOUT_S = int(os.getenv("NVIDIA_LLM_TIMEOUT", "600"))

# DEFAULT_MODEL = "google/diffusiongemma-26b-a4b-it"
DEFAULT_MODEL = "google/diffusiongemma-26b-a4b-it"

# ── Flask app ─────────────────────────────────────────────────────────────────
app = Flask(__name__)

@app.after_request
def after_request(response):
    response.headers.add('Access-Control-Allow-Origin', '*')
    response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
    response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
    return response

# ── Inference Queue (multi-user serialization) ─────────────────────────────────
_inference_queue: _queue_module.Queue = _queue_module.Queue(maxsize=_QUEUE_MAX_SIZE)

if not NVIDIA_API_KEY:
    log.warning("NVIDIA_API_KEY is not set. API calls might fail if the token is required.")

def read_b64(path):
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode()

def _run_inference(data: dict) -> dict:
    raw_prompt = data.get("prompt", "")
    if not raw_prompt:
        return {"error": "Field 'prompt' is required."}

    prompts = raw_prompt if isinstance(raw_prompt, list) and (len(raw_prompt) == 0 or not isinstance(raw_prompt[0], dict)) else [raw_prompt]
    
    llm_mode = data.get("llm_mode", "expert")
    if llm_mode == "assistant":
        model_name = "minimaxai/minimax-m3"
        max_tokens = int(data.get("max_tokens", 8192))
        chat_template_kwargs = None
        temperature = float(data.get("temperature", 1.0))
        top_p = float(data.get("top_p", 0.95))
        reasoning_effort = None
    else:
        model_name = data.get("model", DEFAULT_MODEL)
        max_tokens = int(data.get("max_tokens", 4096))
        chat_template_kwargs = {"enable_thinking": True}
        temperature = float(data.get("temperature", 1.00))
        top_p = float(data.get("top_p", 0.95))
        reasoning_effort = None

    invoke_url = "https://integrate.api.nvidia.com/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {NVIDIA_API_KEY}",
        "Accept": "application/json"
    }

    choices = []
    total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
    
    for i, prompt in enumerate(prompts):
        if isinstance(prompt, list):
            messages = prompt
        else:
            messages = [
                {"role": "user", "content": prompt}
            ]

        try:
            print(f"\n[CONSOLE STREAM] Generating via NVIDIA for: {model_name}")
            print("-" * 30)
            
            payload = {
                "model": model_name,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": temperature,
                "top_p": top_p,
                "stream": False,
            }
            if chat_template_kwargs:
                payload["chat_template_kwargs"] = chat_template_kwargs
            if reasoning_effort:
                payload["reasoning_effort"] = reasoning_effort

            response = requests.post(invoke_url, headers=headers, json=payload)
            response.raise_for_status()
            resp_data = response.json()

            usage = resp_data.get("usage", {})
            total_usage["prompt_tokens"] += usage.get("prompt_tokens", 0)
            total_usage["completion_tokens"] += usage.get("completion_tokens", 0)
            total_usage["total_tokens"] += usage.get("total_tokens", 0)

            full_output = ""
            if "choices" in resp_data and len(resp_data["choices"]) > 0:
                message = resp_data["choices"][0].get("message", {})
                if message.get("content"):
                    full_output = message["content"]
                    print(full_output)

            print("\n" + "-" * 30)
            
            choices.append({
                "index": i,
                "text": full_output.strip(),
                "thinking": ""
            })

        except Exception as e:
            log.error(f"Error calling NVIDIA API: {e}")
            choices.append({
                "index": i,
                "text": f"Error: {str(e)}",
                "thinking": ""
            })

    return {
        "model": model_name,
        "choices": choices,
        "usage": total_usage,
        "device": "cloud_nvidia"
    }

def _inference_worker() -> None:
    log.info("Inference worker thread started (pid=%d)", os.getpid())
    while True:
        try:
            item = _inference_queue.get(timeout=1.0)
        except _queue_module.Empty:
            continue

        req_data, result_holder, done_event = item
        try:
            result_holder[0] = _run_inference(req_data)
        except Exception as exc:
            log.error("Inference worker error: %s", exc)
            result_holder[0] = {"error": f"Inference failed: {exc}"}
        finally:
            done_event.set()
            _inference_queue.task_done()

_worker_thread = threading.Thread(target=_inference_worker, name="inference-worker", daemon=True)
_worker_thread.start()

# ── Routes ────────────────────────────────────────────────────────────────────

@app.route("/health", methods=["GET"])
def health():
    return jsonify({
        "status": "ok",
        "queue_depth": _inference_queue.qsize(),
        "queue_max": _QUEUE_MAX_SIZE
    })

@app.route("/v1/completions", methods=["POST", "OPTIONS"])
def completions():
    if request.method == "OPTIONS":
        return jsonify({}), 200

    data: dict = request.get_json(force=True) or {}

    current_depth = _inference_queue.qsize()
    if current_depth >= _QUEUE_MAX_SIZE:
        return jsonify({
            "error": "Server busy — all inference slots are occupied. Please try again shortly.",
            "retry_after": 5
        }), 503

    result_holder: list = [None]
    done_event = threading.Event()

    try:
        _inference_queue.put_nowait((data, result_holder, done_event))
    except _queue_module.Full:
        return jsonify({
            "error": "Server busy — inference queue full. Please try again shortly.",
            "retry_after": 5,
        }), 503

    completed = done_event.wait(timeout=_REQUEST_TIMEOUT_S)

    if not completed:
        return jsonify({
            "error": f"Request timed out after {_REQUEST_TIMEOUT_S}s. ",
            "retry_after": 10,
        }), 503

    result = result_holder[0]
    if result is None:
        return jsonify({"error": "Internal error: inference worker returned no result."}), 500

    if "error" in result:
        return jsonify(result), 500

    return jsonify(result)

if __name__ == "__main__":
    import signal, sys

    def sigint_handler(sig, frame):
        sys.exit(0)
    signal.signal(signal.SIGINT, sigint_handler)

    log.info(f"Starting NVIDIA LLM agent on http://{HOST}:{PORT}")
    app.run(host=HOST, port=PORT, debug=False, threaded=True)