cowork-local/presentation/folder/image_model_picker.py

"""Chọn model sinh ảnh cho AI sửa file — R08-T12.

Dò khắp mọi provider đã cấu hình xem cái nào sinh được ảnh, rồi gợi ý khi
câu người dùng gõ nghe như đang muốn tạo ảnh. Có thể gợi ý model của
provider KHÁC provider đang chọn — nên nó tách riêng: đây là chỗ duy nhất
trong màn Thư mục biết tới nhiều provider cùng lúc.
"""
from __future__ import annotations

from .file_helpers import (
    DOC_SUFFIXES, _HTML_SUFFIXES, _IMAGE_SUFFIXES, _PPTX_SUFFIXES, _is_probably_text, _pptx_available,
)
import os
from pathlib import Path
from PySide6.QtCore import Qt, Signal
from PySide6.QtWidgets import QComboBox, QHBoxLayout, QLabel, QLineEdit, QPushButton, QVBoxLayout, QWidget
from ...core.worker import AgentWorker
from ...i18n import tr
from ...theme import current_palette
from ...ui.chat_view import ChatView
from ...ui.libreoffice_view import DOC_SUFFIXES


class ImageModelPickerMixin:
    """Trộn vào FolderTab."""

    def _scan_all_image_models(self, then_suggest: bool = False) -> None:
        """Background: find image-capable models across EVERY configured provider
        (not just the active one), so we can suggest one when an edit involves
        images even if the active provider has none. Caches
        ``self._all_image_models = [(provider_key, model)]``."""
        if self._img_scan_worker is not None:
            if then_suggest:
                self._pending_img_suggest = True
            return
        providers = dict(self.ctx.config.data.get("providers", {}))
        # Only providers that actually have an endpoint/key configured.
        candidates = [k for k, c in providers.items()
                      if (c.get("base_url") or c.get("api_key"))]

        def job(worker):
            from ...core import image_gen
            found = []
            for key in candidates:
                try:
                    prov = self.ctx.build_provider_for(key)
                    models = list(getattr(prov, "list_models", lambda: [])() or [])
                except Exception:  # noqa: BLE001 - a broken provider must not block the scan
                    models = []
                for m in models:
                    if image_gen.looks_like_image_model(m):
                        found.append((key, m))
            return {"found": found}

        def done(res):
            self._img_scan_worker = None
            self._all_image_models = list(res.get("found", []))
            if getattr(self, "_pending_img_suggest", False):
                self._pending_img_suggest = False
                self._suggest_cross_provider_image()

        w = AgentWorker(job)
        w.finished_ok.connect(done)
        w.failed.connect(lambda _e: setattr(self, "_img_scan_worker", None))
        self._img_scan_worker = w
        if then_suggest:
            self._pending_img_suggest = True
        w.start()

    def _maybe_suggest_image_model(self, instruction: str) -> None:
        """If the request looks image-related, suggest a suitable image model
        BEFORE running — searching the active provider first, then ALL providers.
        The suggested model is what image generation will auto-use."""
        from ...core import image_gen
        low = (instruction or "").lower()
        if not any(w in low for w in self._IMAGE_WORDS):
            return
        picked = self.ai_model_combo.currentData()
        if picked and image_gen.looks_like_image_model(picked):
            return
        local = image_gen.suggest_image_model(self._ai_models)
        if local:
            self.ai_chat.add_status(tr("folder.ai_image_suggest", model=local))
            return
        # None on the active provider → look across ALL providers (cached, or scan
        # now and suggest when the scan returns).
        if self._all_image_models:
            self._suggest_cross_provider_image()
        elif self._img_scan_worker is not None:
            self._pending_img_suggest = True          # a scan is already running
        else:
            self._scan_all_image_models(then_suggest=True)

    def _suggest_cross_provider_image(self) -> None:
        """Post a suggestion listing image models found on OTHER providers. When
        none exist anywhere, fall back to telling the user their PICKED model
        will be used for image generation (or that there's nothing to use)."""
        from ...config import PROVIDER_LABELS
        if not self._all_image_models:
            picked = self.ai_model_combo.currentData()
            if picked:
                self.ai_chat.add_status(tr("folder.ai_image_use_selected", model=picked))
            else:
                self.ai_chat.add_status(tr("folder.ai_image_none"))
            return
        seen, lines = set(), []
        for key, model in self._all_image_models:
            tag = (key, model)
            if tag in seen:
                continue
            seen.add(tag)
            lines.append(f"• {model}  ({PROVIDER_LABELS.get(key, key)})")
            if len(lines) >= 5:
                break
        self.ai_chat.add_status(tr("folder.ai_image_suggest_all") + "\n" + "\n".join(lines))