scan4llm/src/postprocess/qwen_client.py

from __future__ import annotations

import base64
import logging
import os
import time as _time

import cv2
import requests

from src.env import load_env

logger = logging.getLogger(__name__)

API_URL: str = "https://opencode.ai/zen/go/v1/messages"
DEFAULT_MODEL: str = "qwen3.8-max"
IMAGE_QUALITY: int = 90

_ENV_LOADED: bool = False


def _ensure_env() -> None:
    global _ENV_LOADED
    if _ENV_LOADED:
        return
    load_env()
    _ENV_LOADED = True


def _api_key() -> str:
    _ensure_env()
    return os.environ.get("OPENCODE_API_KEY", "")


def image_to_base64(image) -> str:
    _, buf = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, IMAGE_QUALITY])
    return base64.b64encode(buf).decode("utf-8")


def call_qwen_vlm(
    image_b64: str,
    prompt: str,
    model: str = DEFAULT_MODEL,
    max_tokens: int = 300,
) -> str | None:
    """Отправляет изображение + prompt в Qwen3.8 Max API и возвращает текст ответа."""
    if not _api_key():
        logger.warning("OPENCODE_API_KEY не задан")
        return None

    messages = [{
        "role": "user",
        "content": [
            {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": image_b64}},
            {"type": "text", "text": prompt},
        ],
    }]

    for attempt in (1, 2):
        try:
            r = requests.post(
                API_URL,
                headers={"x-api-key": _api_key(), "anthropic-version": "2023-06-01", "Content-Type": "application/json"},
                json={"model": model, "max_tokens": max_tokens, "messages": messages},
                timeout=120,
            )
            r.raise_for_status()
            data = r.json()
            for item in data.get("content", []):
                if item.get("type") == "text":
                    return item.get("text", "").strip()
            return None
        except requests.HTTPError as e:
            if e.response is not None and e.response.status_code == 500 and attempt == 1:
                logger.info("Qwen API: HTTP 500, ожидание 5с и повтор...")
                _time.sleep(5)
                continue
            logger.warning("Qwen API: ошибка — %s", e)
            return None
        except (requests.RequestException, KeyError, IndexError) as e:
            logger.warning("Qwen API: ошибка — %s", e)
            return None
    return None