utils/api.py

"""Anthropic API wrapper with streaming support."""

import anthropic


MODELS = {
    "Claude Opus 4.6 (Best quality)": "claude-opus-4-6",
    "Claude Sonnet 4.6 (Balanced)": "claude-sonnet-4-6",
    "Claude Haiku 4.5 (Fastest / cheapest)": "claude-haiku-4-5-20251001",
}


def validate_key(api_key: str) -> tuple[bool, str]:
    """Test the API key with a minimal call. Returns (valid, message)."""
    try:
        client = anthropic.Anthropic(api_key=api_key)
        client.messages.create(
            model="claude-haiku-4-5-20251001",
            max_tokens=10,
            messages=[{"role": "user", "content": "hi"}],
        )
        return True, "✅ API key verified"
    except anthropic.AuthenticationError:
        return False, "❌ Invalid API key — check and try again"
    except Exception as e:
        return False, f"❌ Error: {str(e)}"


def stream_to_placeholder(placeholder, api_key: str, model: str,
                           system_prompt: str, user_message: str,
                           max_tokens: int = 4000) -> str:
    """
    Stream a Claude response into a Streamlit placeholder.
    Returns the full completed text.
    """
    client = anthropic.Anthropic(api_key=api_key)
    full_text = ""

    with client.messages.stream(
        model=model,
        max_tokens=max_tokens,
        system=system_prompt,
        messages=[{"role": "user", "content": user_message}],
    ) as stream:
        for chunk in stream.text_stream:
            full_text += chunk
            placeholder.markdown(full_text + "▌")

    placeholder.markdown(full_text)
    return full_text


def count_tokens(api_key: str, model: str, system_prompt: str,
                 user_message: str) -> int:
    """Estimate token count before sending (uses beta count_tokens endpoint)."""
    try:
        client = anthropic.Anthropic(api_key=api_key)
        response = client.messages.count_tokens(
            model=model,
            system=system_prompt,
            messages=[{"role": "user", "content": user_message}],
        )
        return response.input_tokens
    except Exception:
        return 0