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


def _api_key() -> str:
    load_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,
) -> tuple[str | None, dict]:
    """Отправляет изображение + prompt в Qwen3.8 Max API.
    Возвращает (текст, usage_dict)."""
    empty_usage: dict = {}
    if not _api_key():
        logger.warning("OPENCODE_API_KEY не задан")
        return None, empty_usage

    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, 3):
        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=300,
            )
            r.raise_for_status()
            data = r.json()
            text = ""
            for item in data.get("content", []):
                if item.get("type") == "text":
                    text = item.get("text", "").strip()
            stop = data.get("stop_reason", "")
            usage = data.get("usage", {})
            if text and stop == "max_tokens":
                logger.warning("Qwen API: ответ обрезан (max_tokens=%d), stop_reason=%s", max_tokens, stop)
            return text or None, usage
        except requests.HTTPError as e:
            status = e.response.status_code if e.response is not None else 0
            if status in (500, 503) and attempt < 3:
                delay = 5 * attempt
                logger.info("Qwen API: HTTP %d, ожидание %dс и повтор...", status, delay)
                _time.sleep(delay)
                continue
            logger.warning("Qwen API: ошибка — %s", e)
            return None, {}
        except (requests.Timeout, requests.ConnectionError) as e:
            if attempt < 3:
                delay = 5 * attempt
                logger.info("Qwen API: таймаут/соединение, ожидание %dс и повтор...", delay)
                _time.sleep(delay)
                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, {}