agents/cognitive_annotator.py

from __future__ import annotations

import asyncio
import base64
import imghdr
import json
import logging
import os
from pathlib import Path
from typing import Any

import httpx
from openai import AsyncOpenAI
from tqdm.asyncio import tqdm

from base_agent import BaseAgent
from schemas.annotation import AnnotationResult

logger = logging.getLogger(__name__)


class CognitiveAnnotatorAgent(BaseAgent):
    """
    基于 Qwen 兼容接口的高并发认知标注智能体。

    当前版本面向“单图多目标”西红柿认知标注任务：
    1. 模型返回图中全部尽可能清晰可判断的西红柿目标。
    2. 模型侧输出的 bbox 视为 0 到 1000 的相对坐标。
    3. 审核阶段负责将相对坐标换算为图像像素坐标。
    4. 每个目标的成熟度、遮挡与推理说明都必须与对应 bbox 一一对齐。
    """

    SYSTEM_PROMPT = "你是一个严格遵守要求并只返回合法 JSON 的助手。"

    USER_PROMPT = """请对这张图像中的西红柿目标进行尽可能完整的结构化标注。

你必须识别图中所有你能清晰判断的西红柿目标，而不是只选一个主目标。
如果图中存在多个西红柿，请输出多个目标。
如果某些目标严重遮挡、极小、极模糊、无法可靠判断，可以跳过，但应优先保证“尽可能全”。

请先在心中完成“逐个目标定位”，再分别输出每个目标的 bbox 与属性；不要先写属性再随意补 bbox。

请只返回合法 JSON，不要输出 Markdown，不要输出代码块，不要输出额外解释文字。

输出格式必须为一个 JSON 对象，且只能包含一个 key：
- annotations: 一个数组，数组中每个元素对应一个西红柿目标

annotations 中每个目标必须包含以下字段：
1. bbox: [x_min, y_min, x_max, y_max]
2. maturity_level: 只能是“未成熟”/“半成熟”/“完熟”
3. maturity_ratio: 0.0 到 1.0 之间的浮点数
4. occlusion_degree: 只能是“无”/“轻度”/“重度”
5. reasoning: 字符串，说明你对该目标判断的依据

严格要求：
1. 每个 bbox 必须和对应的那个西红柿目标一一对应，不能张冠李戴。
2. reasoning、maturity_level、maturity_ratio、occlusion_degree 都必须只针对该 bbox 对应的目标本身，不要混入其他果实的信息。
3. maturity_ratio 必须和 maturity_level 保持一致：
   - 未成熟：通常小于 0.3
   - 半成熟：通常在 0.3 到 0.7 之间
   - 完熟：通常大于 0.7
4. 如果一个目标无法可靠判断，可以不输出该目标，也不要编造。
5. bbox 默认按 0 到 1000 的归一化相对坐标输出。
6. 输出必须是可被 json.loads 直接解析的合法 JSON。"""

    def __init__(
        self,
        agent_name: str = "cognitive_annotator",
        config: dict[str, Any] | None = None,
    ) -> None:
        """初始化认知标注智能体。"""
        super().__init__(agent_name=agent_name, config=config)

        self.api_key = self.config.get("api_key") or os.getenv("QWEN_API_KEY") or os.getenv("OPENAI_API_KEY")
        self.base_url = self.config.get("base_url") or os.getenv("QWEN_BASE_URL")
        self.model = self.config.get("model", "qwen-vl-max")
        self.max_concurrency = int(self.config.get("max_concurrency", 10))
        self.max_retries = int(self.config.get("max_retries", 5))
        self.request_timeout = float(self.config.get("request_timeout", 120))
        self.output_path = Path(self.config.get("output_path", "raw_annotations.jsonl"))

        if not self.api_key:
            raise ValueError("缺少 API Key，请在 config['api_key'] 或环境变量中提供。")
        if not self.base_url:
            raise ValueError("缺少 base_url，请在 config['base_url'] 或环境变量中提供。")

        self.client = AsyncOpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
            http_client=httpx.AsyncClient(),
        )
        self._semaphore = asyncio.Semaphore(self.max_concurrency)
        self._file_lock = asyncio.Lock()

    def run(self, image_path: str) -> dict[str, Any]:
        """同步处理单张图像。"""
        try:
            asyncio.get_running_loop()
        except RuntimeError:
            return asyncio.run(self.annotate_single(image_path))

        raise RuntimeError("当前已处于事件循环中，请使用 await annotate_single(image_path)。")

    async def annotate_single(self, image_path: str) -> dict[str, Any]:
        """异步处理单张图像并返回多目标结构化结果。"""
        image_file = Path(image_path)
        annotations = await self._annotate_with_retry(image_file)
        record = self._build_output_record(image_file=image_file, annotations=annotations)
        await self._append_record(record)
        return record

    async def annotate_images_async(self, image_paths: list[str]) -> dict[str, int]:
        """并发处理多张图像并实时保存结果。"""
        processed_paths = self._load_processed_paths()
        pending_paths = [Path(path) for path in image_paths if str(Path(path).resolve()) not in processed_paths]

        success_count = 0
        failure_count = 0
        skipped_count = len(image_paths) - len(pending_paths)

        if not pending_paths:
            return {
                "total": len(image_paths),
                "skipped": skipped_count,
                "success": success_count,
                "failure": failure_count,
            }

        progress_bar = tqdm(total=len(pending_paths), desc="Qwen 标注进度", unit="img")
        tasks = [asyncio.create_task(self._process_image(image_file)) for image_file in pending_paths]

        for future in asyncio.as_completed(tasks):
            try:
                is_success = await future
                if is_success:
                    success_count += 1
                else:
                    failure_count += 1
            except Exception:
                failure_count += 1

            progress_bar.update(1)
            progress_bar.set_postfix(
                success=success_count,
                failed=failure_count,
                skipped=skipped_count,
            )

        progress_bar.close()

        return {
            "total": len(image_paths),
            "skipped": skipped_count,
            "success": success_count,
            "failure": failure_count,
        }

    def annotate_images(self, image_paths: list[str]) -> dict[str, int]:
        """同步入口：批量处理图像。"""
        try:
            asyncio.get_running_loop()
        except RuntimeError:
            return asyncio.run(self.annotate_images_async(image_paths))

        raise RuntimeError("当前已处于事件循环中，请使用 await annotate_images_async(image_paths)。")

    async def aclose(self) -> None:
        """关闭底层异步客户端，释放网络连接资源。"""
        await self.client.close()

    async def _process_image(self, image_file: Path) -> bool:
        """处理单张图片并捕获异常。"""
        try:
            annotations = await self._annotate_with_retry(image_file)
            await self._append_record(self._build_output_record(image_file=image_file, annotations=annotations))
            return True
        except Exception:
            logger.exception("处理图片失败: %s", image_file)
            return False

    async def _annotate_with_retry(self, image_file: Path) -> list[AnnotationResult]:
        """执行带指数退避的单图推理请求。"""
        last_error: Exception | None = None

        for attempt in range(1, self.max_retries + 1):
            try:
                async with self._semaphore:
                    return await self._call_qwen_api(image_file)
            except Exception as exc:
                last_error = exc
                logger.warning(
                    "第 %s/%s 次调用失败，准备重试: %s, error=%s",
                    attempt,
                    self.max_retries,
                    image_file,
                    exc,
                )
                if attempt >= self.max_retries:
                    break

                backoff_seconds = min(2 ** (attempt - 1), 16)
                jitter = 0.1 * attempt
                await asyncio.sleep(backoff_seconds + jitter)

        raise RuntimeError(f"图像处理失败: {image_file}") from last_error

    async def _call_qwen_api(self, image_file: Path) -> list[AnnotationResult]:
        """调用 Qwen 兼容多模态接口并解析结果。"""
        image_data_url = self._build_image_data_url(image_file)

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": self.SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": self.USER_PROMPT},
                        {"type": "image_url", "image_url": {"url": image_data_url}},
                    ],
                },
            ],
            temperature=0.1,
            response_format={"type": "json_object"},
            timeout=self.request_timeout,
        )

        content = response.choices[0].message.content
        if not content:
            raise ValueError(f"模型未返回内容: {image_file}")

        if isinstance(content, list):
            text_chunks = [item.get("text", "") for item in content if isinstance(item, dict)]
            content = "".join(text_chunks).strip()

        payload = json.loads(content)
        return self._validate_annotation_batch(payload)

    def _build_output_record(self, image_file: Path, annotations: list[AnnotationResult]) -> dict[str, Any]:
        """构造写入 JSONL 的最终记录。"""
        return {
            "image_path": str(image_file.resolve()),
            "annotations": annotations,
        }

    async def _append_record(self, record: dict[str, Any]) -> None:
        """以追加模式实时写入 JSONL 文件。"""
        self.output_path.parent.mkdir(parents=True, exist_ok=True)
        line = json.dumps(record, ensure_ascii=False)

        async with self._file_lock:
            with self.output_path.open("a", encoding="utf-8") as file:
                file.write(line)
                file.write("\n")

    def _load_processed_paths(self) -> set[str]:
        """读取已存在 JSONL，构建断点续传索引。"""
        if not self.output_path.exists():
            return set()

        processed_paths: set[str] = set()
        with self.output_path.open("r", encoding="utf-8") as file:
            for line in file:
                line = line.strip()
                if not line:
                    continue
                try:
                    payload = json.loads(line)
                except json.JSONDecodeError:
                    continue
                image_path = payload.get("image_path")
                if isinstance(image_path, str):
                    processed_paths.add(str(Path(image_path).resolve()))
        return processed_paths

    def _build_image_data_url(self, image_file: Path) -> str:
        """将本地图像编码为 data URL。"""
        with image_file.open("rb") as file:
            image_bytes = file.read()

        image_format = imghdr.what(None, h=image_bytes)
        mime_type = self._guess_mime_type(image_file=image_file, image_format=image_format)
        encoded = base64.b64encode(image_bytes).decode("utf-8")
        return f"data:{mime_type};base64,{encoded}"

    def _guess_mime_type(self, image_file: Path, image_format: str | None) -> str:
        """推断图像 MIME 类型。"""
        format_name = (image_format or image_file.suffix.lstrip(".")).lower()
        if format_name in {"jpg", "jpeg"}:
            return "image/jpeg"
        if format_name == "png":
            return "image/png"
        if format_name == "webp":
            return "image/webp"
        if format_name == "bmp":
            return "image/bmp"
        raise ValueError(f"不支持的图像格式: {image_file}")

    def _validate_annotation_batch(self, payload: dict[str, Any]) -> list[AnnotationResult]:
        """校验多目标结构化结果。"""
        annotations = payload.get("annotations")
        if not isinstance(annotations, list):
            raise ValueError("模型返回必须包含 annotations 数组。")
        if not annotations:
            raise ValueError("模型未输出任何西红柿目标。")

        validated: list[AnnotationResult] = []
        for item in annotations:
            if not isinstance(item, dict):
                raise ValueError("annotations 中每个元素必须为对象。")
            validated.append(self._validate_single_annotation(item))
        return validated

    def _validate_single_annotation(self, payload: dict[str, Any]) -> AnnotationResult:
        """校验单个目标标注结果。"""
        required_keys = {"bbox", "maturity_level", "maturity_ratio", "occlusion_degree", "reasoning"}
        missing_keys = required_keys.difference(payload.keys())
        if missing_keys:
            raise ValueError(f"模型返回缺少字段: {sorted(missing_keys)}")

        maturity_level = str(payload["maturity_level"])
        if maturity_level not in {"未成熟", "半成熟", "完熟"}:
            raise ValueError("maturity_level 取值非法。")

        maturity_ratio = float(payload["maturity_ratio"])
        if not 0.0 <= maturity_ratio <= 1.0:
            raise ValueError("maturity_ratio 必须位于 0.0 到 1.0 之间。")

        occlusion_degree = str(payload["occlusion_degree"])
        if occlusion_degree not in {"无", "轻度", "重度"}:
            raise ValueError("occlusion_degree 取值非法。")

        reasoning = str(payload["reasoning"]).strip()
        if not reasoning:
            raise ValueError("reasoning 不能为空。")

        return {
            "bbox": self._validate_bbox_1000(payload["bbox"]),
            "maturity_level": maturity_level,
            "maturity_ratio": maturity_ratio,
            "occlusion_degree": occlusion_degree,
            "reasoning": reasoning,
        }

    def _validate_bbox_1000(self, bbox: Any) -> list[float]:
        """校验 0-1000 归一化 bbox 合法性。"""
        if not isinstance(bbox, list) or len(bbox) != 4:
            raise ValueError("bbox 必须是长度为 4 的列表。")
        if not all(isinstance(value, (int, float)) for value in bbox):
            raise ValueError("bbox 的 4 个坐标必须全部为数值类型。")

        x1, y1, x2, y2 = [float(value) for value in bbox]
        if max(x1, y1, x2, y2) <= 1.0:
            x1, y1, x2, y2 = [value * 1000.0 for value in (x1, y1, x2, y2)]
        if min(x1, y1, x2, y2) < 0.0 or max(x1, x2) > 1000.0 or max(y1, y2) > 1000.0:
            raise ValueError("bbox 超出 0 到 1000 的归一化坐标范围。")
        if x2 <= x1 or y2 <= y1:
            raise ValueError("bbox 坐标非法，必须满足 x2 > x1 且 y2 > y1。")
        return [x1, y1, x2, y2]


CognitiveAnnotator = CognitiveAnnotatorAgent