src/infer_api.py

"""Remote inference via an OpenAI-compatible API endpoint.

Used when MODEL_API_URL and RITS_API_KEY environment variables are set.
Provides drop-in replacements for the local inference functions in
infer_vision_qa and infer_chart2csv.
"""

from __future__ import annotations

import base64
import io

import json
import logging

import os
import traceback
from collections.abc import Generator
from typing import Any

import requests
from PIL import Image

from model_loader import get_model_name

logger = logging.getLogger(__name__)


def _pil_to_data_uri(image: Image.Image) -> str:
    """Convert a PIL Image to a base64 data URI."""
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{b64}"


def _build_api_url() -> str:
    """Build the chat completions endpoint URL from MODEL_API_URL."""
    url = os.environ["MODEL_API_URL"].rstrip("/")
    if url.endswith("/chat/completions"):
        return url
    if url.endswith("/v1"):
        return f"{url}/chat/completions"
    return f"{url}/v1/chat/completions"


def _translate_messages(conversation: list[dict], image: Image.Image) -> list[dict]:
    """Convert local-format messages to OpenAI vision API format.

    Replaces ``{"type": "image"}`` entries with the ``image_url`` structure
    expected by the OpenAI vision API.
    """
    translated = []
    for msg in conversation:
        new_content = []
        for item in msg["content"]:
            if item.get("type") == "image":
                new_content.append({
                    "type": "image_url",
                    "image_url": {"url": _pil_to_data_uri(image)},
                })
            else:
                new_content.append(item)
        translated.append({"role": msg["role"], "content": new_content})
    return translated


def _call_chat_api(messages: list[dict]) -> str:
    """Send a chat completion request to the remote API and return the response text."""
    url = _build_api_url()
    headers = {
        "Content-Type": "application/json",
        "RITS_API_KEY": os.environ["RITS_API_KEY"],
    }
    payload = {
        "model": get_model_name(),
        "messages": messages,
        "max_tokens": 4096,
    }
    resp = requests.post(url, json=payload, headers=headers, timeout=120)
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"]


def _call_chat_api_stream(messages: list[dict]) -> Generator[str, None, None]:
    """Send a streaming chat completion request and yield token chunks."""
    url = _build_api_url()
    model = get_model_name()
    headers = {
        "Content-Type": "application/json",
        "RITS_API_KEY": os.environ["RITS_API_KEY"],
    }
    payload = {
        "model": model,
        "messages": messages,
        "max_tokens": 4096,
        "stream": True,
    }

    logger.debug("Stream request: POST %s model=%s messages=%d", url, model, len(messages))
    for i, msg in enumerate(messages):
        content_types = [c.get("type") for c in msg.get("content", [])]
        logger.debug("  message[%d] role=%s content_types=%s", i, msg.get("role"), content_types)

    try:
        resp = requests.post(url, json=payload, headers=headers, timeout=120, stream=True)
        logger.debug("Stream response: status=%d headers=%s", resp.status_code, dict(resp.headers))
        resp.raise_for_status()
    except requests.RequestException:
        logger.exception("Stream request failed: POST %s", url)
        raise

    for line in resp.iter_lines(decode_unicode=True):
        if not line or not line.startswith("data:"):
            continue
        data = line[len("data:"):].strip()
        if data == "[DONE]":
            logger.debug("Stream complete: received [DONE]")
            break
        try:
            chunk = json.loads(data)
            if "error" in chunk:
                err_msg = chunk["error"].get("message", "Unknown server error")
                err_code = chunk["error"].get("code", "unknown")
                logger.error("Server error in stream: code=%s message=%s full_error=%s", err_code, err_msg, chunk["error"])
                yield f"Error: Remote API returned {err_code} — {err_msg}"
                return
            content = chunk["choices"][0]["delta"].get("content", "")
            if content:
                yield content
        except (json.JSONDecodeError, KeyError, IndexError):
            logger.warning("Failed to parse stream chunk: %s", data[:200])
            continue



def answer_question_api(
    image: Image.Image,
    question: str,
    conversation_history: list[dict[str, Any]],
    current_image_path: str | None,
) -> tuple[str, list[dict[str, Any]], str | None]:
    """Answer a question about an image via the remote API.

    Drop-in replacement for ``infer_vision_qa.answer_question``.
    """
    try:
        image = image.convert("RGB")

        if not conversation_history:
            new_user_turn: dict[str, Any] = {"role": "user", "content": [
                {"type": "image"},
                {"type": "text", "text": question},
            ]}
        else:
            new_user_turn = {"role": "user", "content": [
                {"type": "text", "text": question},
            ]}

        conversation = [*conversation_history, new_user_turn]
        api_messages = _translate_messages(conversation, image)
        answer = _call_chat_api(api_messages)

        updated_history = [
            *conversation_history,
            new_user_turn,
            {"role": "assistant", "content": [{"type": "text", "text": answer}]},
        ]
        return answer, updated_history, None

    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        return f"Error: {e!s}", conversation_history, None


def answer_question_stream_api(
    image: Image.Image,
    question: str,
    conversation_history: list[dict[str, Any]],
    current_image_path: str | None,
) -> Generator[str, None, None]:
    """Stream an answer token-by-token via the remote API.
    Drop-in replacement for ``infer_vision_qa.answer_question_stream``.
    """
    try:
        image = image.convert("RGB")

        if not conversation_history:
            new_user_turn: dict[str, Any] = {"role": "user", "content": [
                {"type": "image"},
                {"type": "text", "text": question},
            ]}
        else:
            new_user_turn = {"role": "user", "content": [
                {"type": "text", "text": question},
            ]}

        conversation = [*conversation_history, new_user_turn]
        api_messages = _translate_messages(conversation, image)
        yield from _call_chat_api_stream(api_messages)

    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        yield f"Error: {e!s}"


def extract_csv_stream_api(image: Image.Image) -> Generator[str, None, None]:
    """Stream CSV extraction token-by-token via the remote API.
    Drop-in replacement for ``infer_chart2csv.extract_csv_stream``.
    """
    try:
        image = image.convert("RGB")
        conversation = [{"role": "user", "content": [
            {"type": "image"},
            {"type": "text", "text": "<chart2csv>"},
        ]}]
        api_messages = _translate_messages(conversation, image)
        yield from _call_chat_api_stream(api_messages)

    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        yield f"Error: {e!s}"

def extract_csv_api(image: Image.Image) -> str:
    """Extract CSV data from a chart image via the remote API.

    Drop-in replacement for ``infer_chart2csv.extract_csv``.
    """
    try:
        image = image.convert("RGB")
        conversation = [{"role": "user", "content": [
            {"type": "image"},
            {"type": "text", "text": "<chart2csv>"},
        ]}]
        api_messages = _translate_messages(conversation, image)
        return _call_chat_api(api_messages)

    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        return f"Error: {e!s}"