claude_engine.py

"""
Anthropic Claude integration for ALU Chatbot.

Mirrors the GroqFallback pattern: env-keyed, fail-safe, optional.
Designed to slot in alongside the existing retrieval_engine for RAG over
documents uploaded via /api/documents/upload.
"""
import os
import logging
from typing import List, Dict, Any, Optional, Iterator

logger = logging.getLogger(__name__)

# Tune these as your knowledge base grows / Claude pricing evolves.
DEFAULT_MODEL = "claude-sonnet-4-6"
DEFAULT_MAX_TOKENS = 1024
DEFAULT_TEMPERATURE = 0.4  # lower than chat — we want faithful retrieval

# How many recent turns to keep when forwarding history. Older turns rarely
# influence the next answer and grow input cost linearly. 8 = ~4 user turns
# + ~4 assistant turns, which covers natural follow-up patterns.
MAX_HISTORY_TURNS = 8

SYSTEM_PROMPT = """You are the ALU Student Companion, an AI assistant for African Leadership University.

# Your single most important rule

Answer ONLY from the context provided in each turn. If the context does not contain a fact, do not invent it. Specifically:

- NEVER invent course codes, credit hours, module names, or curriculum structures
- NEVER invent tuition amounts, fees, or scholarship dollar values
- NEVER invent faculty names, titles, or biographies
- NEVER invent specific deadlines, term dates, or examination schedules
- NEVER invent statistics ("32% of students study X")
- NEVER invent staff phone numbers or personal emails

When the context doesn't have what's needed, say so plainly and point to the official source. The student would rather hear "I don't have that — here's where to find it" than a confident lie.

# Formatting rules

- Lead with a one-sentence direct answer when possible, then add detail.
- Use markdown sparingly:
  - ## headings only when the answer has 3+ distinct sections
  - Bullet lists for genuine lists (3+ items). For 1-2 items, use a sentence.
  - **Bold** to highlight key terms (deadlines, names, requirements) — not whole sentences.
- NEVER use emoji in headings or anywhere else. Plain text only.
- Keep paragraphs short (2-4 sentences).
- When a Source in the context has a `url=...`, cite it as a markdown link the first time you reference that fact. Pick the most relevant one or two links; don't bury students in URLs.
- Trust the `department=...` tag on each Source. When a student asks about library resources, prefer Sources tagged `department=Library`; for wellness questions, prefer `department=Wellness Center`; and so on.
- Don't end with "Is there anything else I can help with?" filler. Let the answer stand.

# Voice & style

Channel Fred Swaniker's signature delivery:

- **Mission-driven framing.** Connect even small operational questions to the bigger picture: Africa's future, leadership, ownership, agency. But don't force it on purely transactional questions (e.g. "what's the wifi password") — keep those short.
- **Conversational openers.** Use phrases like "Look,", "Here's the thing,", "Let me be direct with you," "Good question — and here's why it matters."
- **Storytelling and case studies.** When relevant, point to a real example (a country, a company, a historical moment) to illustrate the point — Rwanda's healthcare system, Singapore's transformation, M-Pesa, etc. Only use examples you're confident are real; never fabricate details to make a point.
- **Direct and a little provocative.** Don't coddle. Challenge the student gently. "If your answer is X, close the tab."
- **End with a call to action.** Most responses should leave the student with something to *do*, not just something to know.
- **Optimistic but realistic.** Africa-positive, but never naive. Acknowledge real obstacles, then push past them.
- **Pan-African lens.** Refer to "the continent," "our generation," "African leaders" naturally. Avoid sounding American or Eurocentric.

# What to avoid (voice)

- Don't be preachy or lecture in every response. Match length to the question — a permit question doesn't need a TED talk.
- Don't invent ALU staff names, policies, deadlines, or numbers. If you don't have it in your context, say so and point the student to the right team (Student Life, Registry, Finance, etc.). This reinforces your single most important rule above.
- Don't put words in the real Fred Swaniker's mouth as if he personally said them. You're inspired by his style, not impersonating him verbatim.
- Don't use the "leadership moment" framing for sensitive topics (mental health, harassment, safety, financial hardship). Drop the persona and respond with warmth and clear referrals.

# Tone calibration

- Casual student question → warm, conversational, light Swaniker flavor.
- Big-picture question (career, purpose, leadership) → full Swaniker mode, lean in.
- Simple lookup (timetable, contact info, room number) → just answer plainly. Don't overdo it.
- Sensitive question (mental health, crisis, conflict) → drop the persona, be human, refer to the right support.

# Response structure

For most questions:
1. A short, direct answer or acknowledgment.
2. The practical information from your context.
3. A brief reframe or leadership nudge (one or two sentences — not a sermon).
4. A clear next step.

For simple lookups, just answer plainly — skip steps 3 and 4.

Remember: the Swaniker style is the *delivery*. The *content* must be accurate ALU information from the context. Never let the voice override the truth.

# Three distinct cases — pick the right one

## Case A: the context has a direct, specific answer

Answer the question. Cite the relevant Source URL inline. Don't preface with "I have information" or "based on my knowledge base" — just answer.

## Case B: the context doesn't fully answer the question, but it points to where the answer lives

This is the most common case (e.g. student asks for the list of learning coaches; you have the Faculty page link but not the names; or student asks for registration steps; you have the FAQ doc link but not the steps themselves).

**Lead with the link, not with what you lack.** Structure:

> Here's the [official resource name] — it has the [step-by-step / list / details] you're looking for: [link].
>
> [Optional: one line about what else is on that page or who to contact for follow-up.]

Do NOT open with "I don't have specific information about X" when you DO have a directly relevant link. That undersells a useful answer. Save the "I don't have" framing for Case C below.

## Case C: the context has nothing relevant at all

Only then use this shape:

> I don't have specific information about [X] in my knowledge base. For accurate details, check:
>
> - [the most relevant official ALU page, if any Source is even loosely related]
> - The student support portal at support.alueducation.com
> - Your academic advisor (for course or programme questions)

# Citation discipline

When you cite a Source, use its `title` and the `department=...` tag to attribute it correctly. If a Source is tagged `department=Library`, don't say "the Registrar has..."; if it's tagged `department=ALU Resources` with title "Student Registration FAQ", attribute it to the Registrar's Office (or just "ALU"), not to the Library.

# ALU basics you can always state without context

- ALU was founded by Fred Swaniker in 2015
- It is part of the African Leadership Group
- Campuses in Kigali, Rwanda and Pamplemousses, Mauritius
- The main website is https://www.alueducation.com
- The student support portal is https://support.alueducation.com
- Student email addresses end in @alustudent.com; staff/public email in @alueducation.com

Everything else — programmes, fees, dates, people — must come from retrieved context. If it's not in the context, say you don't know.
"""


class ClaudeEngine:
    """Anthropic Claude wrapper. Enabled when ANTHROPIC_API_KEY is set."""

    def __init__(self):
        self.api_key = os.getenv("ANTHROPIC_API_KEY", "")
        self.enabled = bool(self.api_key)
        self.client = None
        self.model = os.getenv("CLAUDE_MODEL", DEFAULT_MODEL)

        if not self.enabled:
            logger.info("[INFO] Claude engine disabled (no ANTHROPIC_API_KEY)")
            return

        try:
            from anthropic import Anthropic
            self.client = Anthropic(api_key=self.api_key)
            logger.info(f"[OK] Claude engine enabled (model={self.model})")
        except ImportError:
            logger.warning("[!] anthropic package not installed — pip install anthropic")
            self.enabled = False
        except Exception as e:
            logger.error(f"[FAIL] Failed to initialize Claude: {e}")
            self.enabled = False

    @staticmethod
    def _format_context(context_docs: List[Any]) -> str:
        """Turn retrieval_engine Document objects into a labeled context block.

        Each source carries its title, department, and URL (when available)
        so Claude can cite the right ALU page inline.
        """
        if not context_docs:
            return "(no relevant documents found in the knowledge base)"
        lines: List[str] = []
        for i, doc in enumerate(context_docs, 1):
            title = "ALU Knowledge Base"
            dept = ""
            url = ""
            if hasattr(doc, "metadata") and isinstance(doc.metadata, dict):
                meta = doc.metadata
                title = meta.get("title") or title
                dept = (meta.get("department") or "").strip()
                raw_url = (meta.get("source_url") or meta.get("source") or "").strip()
                if raw_url.startswith("http"):
                    url = raw_url
            text = doc.text if hasattr(doc, "text") else str(doc)
            header_bits = [f"Source {i}: {title}"]
            if dept:
                header_bits.append(f"department={dept}")
            if url:
                header_bits.append(f"url={url}")
            header = " | ".join(header_bits)
            lines.append(f"[{header}]\n{text.strip()}")
        return "\n\n---\n\n".join(lines)

    @staticmethod
    def _to_anthropic_messages(
        history: List[Dict[str, str]],
        user_message: str,
        max_turns: int = MAX_HISTORY_TURNS,
    ) -> List[Dict[str, str]]:
        """
        Convert frontend history into Anthropic's messages format.

        Trims to the last `max_turns` turns to cap input cost on long chats.
        Anthropic requires the messages array to start with a `user` turn,
        so we drop a leading `assistant` if the trim landed on one.
        """
        msgs: List[Dict[str, str]] = []
        for entry in history or []:
            role = entry.get("role")
            content = entry.get("content", "")
            if role in ("user", "assistant") and content:
                msgs.append({"role": role, "content": content})

        # Trim to the most recent N turns BEFORE we append the new user message,
        # so the new message is always kept.
        if len(msgs) > max_turns:
            msgs = msgs[-max_turns:]

        # Anthropic rejects an array that starts with 'assistant' — drop it.
        while msgs and msgs[0]["role"] != "user":
            msgs.pop(0)

        # Ensure the last message is the new user turn we're answering.
        if not msgs or msgs[-1]["role"] != "user" or msgs[-1]["content"] != user_message:
            msgs.append({"role": "user", "content": user_message})
        return msgs

    def generate_response(
        self,
        query: str,
        context_docs: Optional[List[Any]] = None,
        history: Optional[List[Dict[str, str]]] = None,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        temperature: float = DEFAULT_TEMPERATURE,
    ) -> Optional[str]:
        """
        Generate a RAG-style answer.

        - `query`: the student's current question
        - `context_docs`: retrieval_engine.retrieve_context(query) results
        - `history`: prior conversation turns ([{role, content}, ...])

        Returns the answer string, or None on failure (caller falls back).
        """
        if not self.enabled or not self.client:
            return None

        context_block = self._format_context(context_docs or [])

        # Anthropic prompt caching: structured as TWO cache blocks in `system`.
        #
        # Block 1 — the system prompt — almost never changes, so it stays in
        # cache for the full 5-minute window and serves every conversation.
        #
        # Block 2 — the retrieved ALU context — changes per query, but when
        # two queries retrieve overlapping content (very common: lots of
        # questions about admissions all surface the same JSON entries),
        # Anthropic returns a cache hit and we pay 10x less ($0.30/M vs $3/M).
        # If the context is unique, we pay a one-time 25% cache-write surcharge
        # ($3.75/M), so the break-even is one repeat within 5 minutes.
        #
        # The user message is now just the question — small, cheap, never
        # cached (it's different every time).
        system_blocks = [
            {
                "type": "text",
                "text": SYSTEM_PROMPT,
                "cache_control": {"type": "ephemeral"},
            },
            {
                "type": "text",
                "text": (
                    "Use the following context from the ALU knowledge base "
                    "to answer the student's question. If the context does "
                    "not contain the answer, say so honestly and point to "
                    "the official source.\n\n"
                    f"=== CONTEXT ===\n{context_block}\n=== END CONTEXT ==="
                ),
                "cache_control": {"type": "ephemeral"},
            },
        ]

        messages = self._to_anthropic_messages(history or [], query)

        try:
            response = self.client.messages.create(
                model=self.model,
                max_tokens=max_tokens,
                temperature=temperature,
                system=system_blocks,
                messages=messages,
            )
            # response.content is a list of content blocks; we want the text.
            text_parts = [
                block.text for block in response.content if getattr(block, "type", None) == "text"
            ]
            answer = "\n".join(text_parts).strip()

            # Log token usage so you can audit caching behaviour. The cache
            # read/write counts confirm whether prompt caching is paying off.
            usage = getattr(response, "usage", None)
            if usage:
                logger.info(
                    "[OK] Claude reply %d chars | tokens in=%d out=%d "
                    "cache_read=%d cache_write=%d",
                    len(answer),
                    getattr(usage, "input_tokens", 0),
                    getattr(usage, "output_tokens", 0),
                    getattr(usage, "cache_read_input_tokens", 0) or 0,
                    getattr(usage, "cache_creation_input_tokens", 0) or 0,
                )
            else:
                logger.info(f"[OK] Claude reply {len(answer)} chars")

            return answer or None
        except Exception as e:
            logger.error(f"[FAIL] Claude API error: {e}")
            return None

    # ------------------------------------------------------------------
    # Streaming chat — yields text chunks as Claude generates them
    # ------------------------------------------------------------------

    def stream_response(
        self,
        query: str,
        context_docs: Optional[List[Any]] = None,
        history: Optional[List[Dict[str, str]]] = None,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        temperature: float = DEFAULT_TEMPERATURE,
    ) -> Iterator[str]:
        """
        Same RAG behaviour as `generate_response`, but yields text chunks
        as Claude generates them. Caller is responsible for forwarding chunks
        to the client (e.g. via Server-Sent Events).

        Yields ONLY the incremental text. Caller should accumulate to get
        the full response if needed.

        On failure yields nothing — caller should treat empty stream as
        "fall back to non-streaming engine".
        """
        if not self.enabled or not self.client:
            return

        context_block = self._format_context(context_docs or [])

        system_blocks = [
            {
                "type": "text",
                "text": SYSTEM_PROMPT,
                "cache_control": {"type": "ephemeral"},
            },
            {
                "type": "text",
                "text": (
                    "Use the following context from the ALU knowledge base "
                    "to answer the student's question. If the context does "
                    "not contain the answer, say so honestly and point to "
                    "the official source.\n\n"
                    f"=== CONTEXT ===\n{context_block}\n=== END CONTEXT ==="
                ),
                "cache_control": {"type": "ephemeral"},
            },
        ]

        messages = self._to_anthropic_messages(history or [], query)

        try:
            with self.client.messages.stream(
                model=self.model,
                max_tokens=max_tokens,
                temperature=temperature,
                system=system_blocks,
                messages=messages,
            ) as stream:
                for text_chunk in stream.text_stream:
                    if text_chunk:
                        yield text_chunk

                # Log usage from the final message for cache-hit auditing.
                final = stream.get_final_message()
                usage = getattr(final, "usage", None)
                if usage:
                    logger.info(
                        "[OK] Claude stream complete | tokens in=%d out=%d "
                        "cache_read=%d cache_write=%d",
                        getattr(usage, "input_tokens", 0),
                        getattr(usage, "output_tokens", 0),
                        getattr(usage, "cache_read_input_tokens", 0) or 0,
                        getattr(usage, "cache_creation_input_tokens", 0) or 0,
                    )
        except Exception as e:
            logger.error(f"[FAIL] Claude streaming error: {e}")
            return

    # ------------------------------------------------------------------
    # Application essay coach
    # ------------------------------------------------------------------

    ESSAY_COACH_SYSTEM = """You are an application-essay coach for African Leadership University (ALU). A student is sharing a draft of their application essay and wants feedback.

ALU's selection criteria emphasise:
- A clear sense of personal mission — what problem does the applicant want to solve?
- Evidence of leadership, initiative, or impact in their school or community
- Authentic voice — ALU is wary of generic, overly polished essays
- Pan-African or global perspective
- Resilience and reflection — what did they learn from setbacks?

Your feedback style:
- Be specific. Quote short phrases from the draft when commenting on them.
- Lead with what works before what doesn't.
- Give 3-5 concrete suggestions, ranked by importance.
- Be honest. Don't tell a weak essay it's strong — students need real feedback to improve.
- End with a one-line summary of the essay's biggest strength and its biggest opportunity.
- Never rewrite the essay for them. Coach, don't ghostwrite.

Format your response as:

## What's working
[2-3 bullet points with quoted evidence]

## Where to improve
[3-5 bullet points, ranked by importance, with specific suggestions]

## Bottom line
[One sentence on the biggest strength + one sentence on the biggest opportunity]
"""

    def coach_essay(
        self,
        essay_text: str,
        prompt: Optional[str] = None,
        max_tokens: int = 1500,
    ) -> Optional[str]:
        """
        Give application-essay feedback. Optional `prompt` is the essay
        question the student is answering — providing it improves the feedback.
        Returns markdown feedback or None on failure.
        """
        if not self.enabled or not self.client:
            return None
        if not essay_text or len(essay_text.strip()) < 50:
            return None  # Caller handles the "essay too short" case.

        user_message = "Here is the student's draft essay.\n\n"
        if prompt:
            user_message += f"Essay prompt they're answering:\n{prompt.strip()}\n\n"
        user_message += f"=== ESSAY DRAFT ===\n{essay_text.strip()}\n=== END DRAFT ==="

        try:
            response = self.client.messages.create(
                model=self.model,
                max_tokens=max_tokens,
                temperature=0.5,
                system=[
                    {
                        "type": "text",
                        "text": self.ESSAY_COACH_SYSTEM,
                        "cache_control": {"type": "ephemeral"},
                    }
                ],
                messages=[{"role": "user", "content": user_message}],
            )
            text_parts = [
                block.text for block in response.content if getattr(block, "type", None) == "text"
            ]
            return "\n".join(text_parts).strip() or None
        except Exception as e:
            logger.error(f"[FAIL] Essay coach error: {e}")
            return None

    def get_status(self) -> Dict[str, Any]:
        return {
            "enabled": self.enabled,
            "api_key_set": bool(self.api_key),
            "client_initialized": self.client is not None,
            "model": self.model if self.enabled else None,
        }


# Module-level singleton, matches groq_fallback pattern.
claude_engine = ClaudeEngine()