Artemis-Eapil/artemis/core/diagnostics/probes/credentials_probe.py

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""LLM & Multimodal Vision Credentials Readiness Probe."""

from typing import Any
from artemis.config import settings
from artemis.core.diagnostics.probes.base import BaseProbe
from artemis.core.diagnostics.schema import (
    ProbeAction,
    ProbeCategory,
    ProbeResult,
    ProbeStatus,
)


class LLMCredentialsProbe(BaseProbe):
    """Readiness depends on an OpenAI credential, never merely a base URL."""

    @property
    def probe_id(self) -> str:
        return "openai_api_key"

    @property
    def category(self) -> ProbeCategory:
        return ProbeCategory.CREDENTIALS

    @property
    def is_blocker(self) -> bool:
        return True

    async def probe(self) -> ProbeResult:
        key = settings.get_api_key("openai")
        configured = bool(key)
        providers = (
            [
                {
                    "provider": "openai",
                    "label": "OpenAI Responses",
                    "masked": "****" + key.get_secret_value()[-4:],
                }
            ]
            if key
            else []
        )
        return ProbeResult(
            id=self.probe_id,
            category=self.category,
            title="OpenAI API 密钥",
            status=ProbeStatus.PASS if configured else ProbeStatus.FAIL,
            is_blocker=True,
            summary="已配置 OpenAI Responses" if configured else "缺少 API 密钥",
            description="模型统一通过 OpenAI Responses 接口调用。"
            if configured
            else "请在环境变量或 .env 中设置 OPENAI_API_KEY。",
            metadata={
                "configured_count": len(providers),
                "providers": providers,
                "has_ocr_key": bool(settings.get_api_key("ocr")),
            },
            actions=[]
            if configured
            else [ProbeAction(action_type="command", label="初始化配置", payload="artemis init")],
        )


class VisionOCRProbe(BaseProbe):
    """Probe verifying optional Google Cloud Vision OCR credentials."""

    @property
    def probe_id(self) -> str:
        return "vision_ocr_key"

    @property
    def category(self) -> ProbeCategory:
        return ProbeCategory.CREDENTIALS

    @property
    def is_blocker(self) -> bool:
        return False

    def _mask_key(self, key_str: str) -> str:
        """Helper to safely mask an API credential for display."""
        if len(key_str) > 10:
            return f"{key_str[:6]}...{key_str[-4:]}"
        return "***"

    async def probe(self) -> ProbeResult:
        from artemis.utils.ocr_api import is_ocr_configured

        ocr_key = settings.get_api_key("ocr")
        is_configured = is_ocr_configured() and ocr_key is not None

        if is_configured and ocr_key:
            val = ocr_key.get_secret_value()
            masked = self._mask_key(val)
            return ProbeResult(
                id=self.probe_id,
                category=self.category,
                title="Vision OCR API Key (Optional)",
                status=ProbeStatus.PASS,
                is_blocker=False,
                summary="Active & Configured",
                description=f"Google Cloud Vision OCR ({masked}) is active for image text recognition.",
                metadata={"configured": True, "masked_key": masked},
                actions=[
                    ProbeAction(
                        action_type="hint",
                        label="OCR Enabled",
                        payload="OCR is active and will fuse text with UI XML hierarchy.",
                    )
                ],
            )

        return ProbeResult(
            id=self.probe_id,
            category=self.category,
            title="Vision OCR API Key (Optional)",
            status=ProbeStatus.PASS,
            is_blocker=False,
            summary="Not Configured (Optional)",
            description="OCR_API_KEY is not set. Perception uses the UI XML hierarchy without OCR.",
            metadata={"configured": False},
            actions=[
                ProbeAction(
                    action_type="hint",
                    label="Standard XML Perception",
                    payload="Artemis uses pure UI layout parsing and Set-of-Marks visual grounding.",
                )
            ],
        )