query_llm/rits_query.py

#!/usr/bin/env python3
"""Minimal RITS chat-completions client."""

from __future__ import annotations

import argparse
import json
import os
import sys
import urllib.error
import urllib.request
from pathlib import Path
from urllib.parse import urlsplit, urlunsplit


def load_dotenv(path: Path) -> None:
    if not path.exists():
        return

    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue

        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        os.environ.setdefault(key, value)


def model_to_route(model: str) -> str:
    model_name = model.strip("/").rsplit("/", 1)[-1]
    return model_name.replace(".", "-")


def base_url_for_model(base_url: str, model: str) -> str:
    """Build the RITS base URL from .env and the requested model."""
    normalized = base_url.rstrip("/")
    model_route = model_to_route(model)

    split = urlsplit(normalized)
    parts = [part for part in split.path.split("/") if part]
    if parts and parts[-1] == "v1":
        parts = model_route.split("/") + ["v1"]
    else:
        parts = parts + model_route.split("/") + ["v1"]

    return urlunsplit((split.scheme, split.netloc, "/" + "/".join(parts), "", ""))


def chat_completion(
    *,
    base_url: str,
    api_key: str,
    model: str,
    prompt: str,
    temperature: float,
    max_tokens: int,
) -> str:
    url = f"{base_url.rstrip('/')}/chat/completions"
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {api_key}",
            "RITS_API_KEY": api_key,
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
        method="POST",
    )

    try:
        with urllib.request.urlopen(request, timeout=60) as response:
            data = json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"RITS returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"Could not connect to RITS: {exc.reason}") from exc

    try:
        return data["choices"][0]["message"]["content"]
    except (KeyError, IndexError, TypeError) as exc:
        raise RuntimeError(f"Unexpected RITS response: {json.dumps(data)}") from exc


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Call a RITS-hosted chat model.")
    parser.add_argument("prompt", nargs="?", help="Prompt to send to the model.")
    parser.add_argument(
        "-m",
        "--model",
        help="Model name to request. Defaults to LLM_MODEL from .env.",
    )
    parser.add_argument(
        "--base-url",
        help="RITS OpenAI-compatible base URL template. Defaults to LLM_BASE_URL from .env.",
    )
    parser.add_argument(
        "--api-key",
        help="RITS API key. Defaults to LLM_API_KEY from .env.",
    )
    parser.add_argument(
        "--input-file",
        help="Path to a text file containing the prompt to send to the model.",
    )
    parser.add_argument(
        "--gpus",
        type=int,
        default=1,
        help="Accepted for job compatibility. Defaults to 1 and is not sent to RITS.",
    )
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--max-tokens", type=int, default=512)
    return parser.parse_args()


def prompt_from_args(args: argparse.Namespace) -> str:
    if args.input_file:
        try:
            return Path(args.input_file).read_text(encoding="utf-8")
        except OSError as exc:
            raise RuntimeError(f"Could not read input file {args.input_file}: {exc}") from exc

    if args.prompt is None:
        raise RuntimeError("Missing prompt. Pass a prompt argument or use --input-file.")

    return args.prompt


def main() -> int:
    load_dotenv(Path(".env"))
    args = parse_args()
    model = args.model or os.getenv("LLM_MODEL")
    base_url_template = args.base_url or os.getenv("LLM_BASE_URL")
    api_key = args.api_key or os.getenv("LLM_API_KEY")

    missing = [
        name
        for name, value in {
            "LLM_API_KEY": api_key,
            "LLM_BASE_URL": base_url_template,
            "LLM_MODEL or --model": model,
        }.items()
        if not value
    ]
    if missing:
        print(f"Missing required configuration: {', '.join(missing)}", file=sys.stderr)
        return 2

    base_url = base_url_for_model(base_url_template, model)
    try:
        prompt = prompt_from_args(args)
        result = chat_completion(
            base_url=base_url,
            api_key=api_key,
            model=model,
            prompt=prompt,
            temperature=args.temperature,
            max_tokens=args.max_tokens,
        )
    except RuntimeError as exc:
        print(exc, file=sys.stderr)
        return 1

    print(result)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())