| """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: |
| 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: |
| 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: |
| 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: |
| traceback.print_exc() |
| return f"Error: {e!s}" |
|
|
|
|
