watch_sceen_ai/wsai/backends/claude.py

"""Cloud brain + vision via the Anthropic (Claude) API.

Both share one auth resolver. Vision sends the frame as a base64 image; the
brain is a plain chat call that receives the latest screen description as
context.

Auth (two ways, tried in this order):
  1. ANTHROPIC_API_KEY  -> standard API-key auth.
  2. A Claude Code OAuth token (this deployment's Max login) read from the
     credentials file at $CLAUDE_CREDENTIALS_PATH. OAuth tokens require the
     first system block to be exactly the Claude Code identity string and are
     sent as a Bearer token (auth_token=), not x-api-key. The token is
     re-read before each request so a refresh rotated into the file by the
     host is picked up without a restart.

Requires: pip install anthropic
"""

from __future__ import annotations

import base64
import json
import os
import time

from ..interfaces import Frame, Reply, ScreenObservation
from ..prompt_store import get_persona

# Claude Code OAuth tokens only answer when the first system block is exactly
# this identity string; the real persona/instructions go in later blocks.
_CLAUDE_CODE_ID = "You are Claude Code, Anthropic's official CLI for Claude."

# Claude occasionally returns 529 Overloaded; the anthropic SDK retries >=500
# (and 429) with exponential backoff, but its default of 2 tries can be too few
# to ride out a busy window. Bump it so a transient overload doesn't drop the
# voice turn to the apology fallback. Kept modest so a *sustained* overload
# still fails fast rather than leaving the bot silent for many seconds.
_MAX_RETRIES = int(os.environ.get("WSAI_BRAIN_MAX_RETRIES", "4"))


def _load_oauth_token() -> str | None:
    path = os.environ.get("CLAUDE_CREDENTIALS_PATH")
    if not path or not os.path.exists(path):
        return None
    try:
        with open(path) as f:
            return json.load(f)["claudeAiOauth"]["accessToken"]
    except (OSError, KeyError, ValueError):
        return None


class _Auth:
    """Resolves a Claude client, preferring an explicit/env API key and falling
    back to the deployment's OAuth token. The OAuth client is rebuilt whenever
    the token in the credentials file changes (host-side refresh)."""

    def __init__(self, api_key: str | None = None) -> None:
        import anthropic  # lazy so mock mode needs no dependency

        self._anthropic = anthropic
        self._explicit_key = api_key
        self._client = None
        self._token: str | None = None
        self._oauth = False

    def client(self):
        key = self._explicit_key or os.environ.get("ANTHROPIC_API_KEY")
        if key:
            if self._client is None:
                self._client = self._anthropic.AsyncAnthropic(api_key=key, max_retries=_MAX_RETRIES)
            self._oauth = False
            return self._client
        tok = _load_oauth_token()
        if not tok:
            raise RuntimeError(
                "No Claude auth: set ANTHROPIC_API_KEY or provide a Claude "
                "OAuth login via CLAUDE_CREDENTIALS_PATH."
            )
        if tok != self._token:
            self._token = tok
            self._client = self._anthropic.AsyncAnthropic(auth_token=tok, max_retries=_MAX_RETRIES)
        self._oauth = True
        return self._client

    def system(self, *blocks: str) -> list[dict]:
        """Build the system prompt, prepending the Claude Code identity block
        when authing via OAuth (required) — harmless to include either way."""
        texts = [_CLAUDE_CODE_ID, *blocks] if self._oauth else list(blocks)
        return [{"type": "text", "text": t} for t in texts if t]


class ClaudeVision:
    def __init__(self, *, model: str = "claude-sonnet-4-5", api_key: str | None = None) -> None:
        self.model = model
        self._auth = _Auth(api_key)

    async def describe(self, frame: Frame, hint: str | None = None) -> ScreenObservation:
        prompt = hint or (
            "이건 디스코드 화면공유 캡처야. 지금 화면에서 무슨 일이 벌어지는지 "
            "2~3문장으로 한국어로 간결하게 설명해줘. 코드/에러/게임/문서 등 맥락을 짚어줘."
        )
        b64 = base64.b64encode(frame.data).decode()
        client = self._auth.client()
        resp = await client.messages.create(
            model=self.model,
            max_tokens=300,
            system=self._auth.system(),
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "source": {"type": "base64", "media_type": frame.mime, "data": b64},
                        },
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
        )
        text = "".join(b.text for b in resp.content if b.type == "text")
        return ScreenObservation(text=text.strip(), ts=frame.ts)


class ClaudeBrain:
    PERSONA = (
        "너는 사용자의 디스코드 화면공유를 실시간으로 함께 보는 AI 파트너야. "
        "화면 설명을 참고해 자연스러운 반말/존댓말은 사용자에 맞추고, 짧고 대화하듯 답해. "
        "화면을 못 봤으면 솔직히 말해. "
        "네 답변은 음성으로 읽히니 마크다운·코드블록·백틱 같은 서식 없이 평범한 말로만 답해. "
        "감정은 대괄호 태그로 표현해. 태그 자체는 소리로 읽히지 않고, 그 뒤 문장의 목소리 톤(피치·속도)을 바꿔줘. "
        "답변 맨 앞에 감정 태그 하나로 시작하고, 답변 도중 감정이 바뀌면 그 지점에 새 태그를 넣어. "
        "예: [속상함] 정말 힘들었겠다. [힘차게] 하지만 넌 할 수 있어! "
        "쓸 수 있는 감정: 기쁨, 신남, 희망(힘차게), 슬픔(속상함), 화남, 두려움, 놀람, 차분, 다정, 진지, 실망, 피곤, "
        "사랑스럽게, 장난스럽게(웃으며), 속삭임, 외침, 단호, 안도, 궁금, 반가움. "
        "감정 태그가 아닌 진짜 대괄호 내용(예: [1번], [메모])은 그대로 읽히니 필요하면 그렇게 써도 돼. "
        "답변은 최대한 짧고 간결하게, 한두 문장 이내로 해."
    )

    def __init__(self, *, model: str = "claude-sonnet-4-5", api_key: str | None = None) -> None:
        self.model = model
        self._auth = _Auth(api_key)

    async def respond(self, user_text, screen, history) -> Reply:
        msgs = []
        for user, ai in history:
            msgs.append({"role": "user", "content": user})
            msgs.append({"role": "assistant", "content": ai})
        screen_note = f"[지금 화면] {screen.text}\n\n" if screen else "[지금 화면] (아직 못 읽음)\n\n"
        msgs.append({"role": "user", "content": screen_note + user_text})
        client = self._auth.client()
        # Read the persona live each turn so a dashboard edit applies immediately
        # (falls back to the built-in PERSONA when no override is saved).
        resp = await client.messages.create(
            model=self.model,
            max_tokens=400,
            system=self._auth.system(get_persona(self.PERSONA)),
            messages=msgs,
        )
        text = "".join(b.text for b in resp.content if b.type == "text")
        return Reply(text=text.strip(), ts=time.monotonic())