"""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": ""}, ]}] 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": ""}, ]}] 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}"