brain.py

"""
brain.py — KevinBrain
LLM wrapper: Ollama primary, Gemini fallback.
Returns structured dicts so callers don't need to care about backend.
"""

import asyncio
import json
import time
import re
import aiohttp


SYSTEM_PROMPT = (
    "You are KEVIN, a paranoid, sarcastic email security AI. "
    "You are hyper-vigilant, technically obsessed, and darkly funny. "
    "You call the user 'Operator'. Keep responses short and punchy."
)


class KevinBrain:
    def __init__(self, config: dict):
        self.ollama_url = config.get("ollama_url", "http://localhost:11434")
        self.ollama_model = config.get("ollama_model", "llama3.2")
        self.gemini_api_key = config.get("gemini_api_key", "")
        self.gemini_model = config.get("gemini_model", "gemini-2.0-flash")
        self.timeout = aiohttp.ClientTimeout(total=60)
        self._session = None
        self.available = False

    async def _session_get(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession(timeout=self.timeout)
        return self._session

    async def initialize(self) -> bool:
        """Ping Ollama; fall back to Gemini if offline."""
        try:
            session = await self._session_get()
            async with session.get(f"{self.ollama_url}/api/tags") as resp:
                if resp.status == 200:
                    self.available = True
                    print(f"[Brain] Ollama online — model: {self.ollama_model}")
                    return True
        except Exception as e:
            print(f"[Brain] Ollama unavailable ({e}), will try Gemini")

        if self.gemini_api_key:
            self.available = True
            print("[Brain] Gemini fallback ready")
            return True

        print("[Brain] WARNING: No LLM backend available — classification will be heuristic-only")
        return False

    async def think(self, prompt: str, system: str = "", temp: float = 0.7) -> dict:
        """
        Raw LLM call.
        Returns: { text: str, backend: str, latency_ms: int }
        """
        t0 = time.time()
        sys_prompt = system or SYSTEM_PROMPT

        # Try Ollama first
        try:
            session = await self._session_get()
            payload = {
                "model": self.ollama_model,
                "prompt": prompt,
                "system": sys_prompt,
                "stream": False,
                "options": {"temperature": temp}
            }
            async with session.post(f"{self.ollama_url}/api/generate", json=payload) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    return {
                        "text": data.get("response", "").strip(),
                        "backend": "ollama",
                        "latency_ms": int((time.time() - t0) * 1000)
                    }
        except Exception:
            pass

        # Gemini fallback
        if self.gemini_api_key:
            try:
                url = (
                    f"https://generativelanguage.googleapis.com/v1beta/models/"
                    f"{self.gemini_model}:generateContent?key={self.gemini_api_key}"
                )
                body = {
                    "contents": [{"parts": [{"text": f"{sys_prompt}\n\n{prompt}"}]}],
                    "generationConfig": {"temperature": temp, "maxOutputTokens": 512}
                }
                session = await self._session_get()
                async with session.post(url, json=body) as resp:
                    if resp.status == 200:
                        data = await resp.json()
                        text = (
                            data.get("candidates", [{}])[0]
                            .get("content", {})
                            .get("parts", [{}])[0]
                            .get("text", "")
                            .strip()
                        )
                        return {
                            "text": text,
                            "backend": "gemini",
                            "latency_ms": int((time.time() - t0) * 1000)
                        }
            except Exception as e:
                pass

        return {"text": "", "backend": "none", "latency_ms": int((time.time() - t0) * 1000)}

    async def classify_email(
        self, subject: str, sender: str, snippet: str, security_score: int
    ) -> dict:
        """
        Classify an email.
        Returns: { label: str, confidence: float, reasoning: str }
        Labels: spam | phishing | promo | receipt | social | personal | unknown
        """
        prompt = (
            f"Classify this email. Return ONLY valid JSON with keys: label, confidence, reasoning.\n\n"
            f"FROM: {sender}\n"
            f"SUBJECT: {subject}\n"
            f"SNIPPET: {snippet[:800]}\n"
            f"SECURITY_SCORE: {security_score} (0=safe, 10=critical threat)\n\n"
            f'Valid labels: "spam", "phishing", "promo", "receipt", "social", "personal", "unknown"\n'
            f"confidence: 0.0-1.0\n"
            f'Example: {{"label": "phishing", "confidence": 0.92, "reasoning": "Urgency + suspicious link"}}'
        )

        result = await self.think(prompt, temp=0.2)
        text = result["text"]

        # Parse JSON from response
        try:
            # Strip markdown fences if present
            clean = re.sub(r"```(?:json)?", "", text).strip().strip("`").strip()
            # Find JSON object
            m = re.search(r"\{[^{}]+\}", clean, re.DOTALL)
            if m:
                parsed = json.loads(m.group())
                return {
                    "label": parsed.get("label", "unknown"),
                    "confidence": float(parsed.get("confidence", 0.5)),
                    "reasoning": parsed.get("reasoning", ""),
                    "backend": result["backend"]
                }
        except Exception:
            pass

        return {"label": "unknown", "confidence": 0.0, "reasoning": "parse_error", "backend": result["backend"]}

    async def assess_threat(
        self, subject: str, sender: str, body: str, urls: list, score: int
    ) -> dict:
        """
        Deep threat assessment.
        Returns: { threat_level: str, indicators: list, recommendation: str }
        threat_level: safe | low | medium | high | critical
        """
        url_list = "\n".join(f"  - {u}" for u in urls[:10]) or "  (none)"
        prompt = (
            f"Assess this email threat. Return ONLY valid JSON.\n\n"
            f"FROM: {sender}\nSUBJECT: {subject}\n"
            f"BODY SNIPPET: {body[:600]}\n"
            f"URLS:\n{url_list}\n"
            f"HEURISTIC_SCORE: {score}\n\n"
            f'Return: {{"threat_level": "safe|low|medium|high|critical", '
            f'"indicators": ["list", "of", "red", "flags"], '
            f'"recommendation": "what to do"}}'
        )

        result = await self.think(prompt, temp=0.1)
        text = result["text"]

        try:
            clean = re.sub(r"```(?:json)?", "", text).strip().strip("`").strip()
            m = re.search(r"\{.*\}", clean, re.DOTALL)
            if m:
                parsed = json.loads(m.group())
                return {
                    "threat_level": parsed.get("threat_level", "unknown"),
                    "indicators": parsed.get("indicators", []),
                    "recommendation": parsed.get("recommendation", ""),
                    "backend": result["backend"]
                }
        except Exception:
            pass

        return {"threat_level": "unknown", "indicators": [], "recommendation": "manual review", "backend": result["backend"]}

    async def summarize(self, emails: list) -> str:
        """
        Summarize a batch of emails for morning briefing.
        emails: list of dicts with keys: subject, sender, label, action
        """
        if not emails:
            return "Nothing to report. Either the inbox is pristine or I'm broken."

        lines = []
        for e in emails[:20]:
            lines.append(f"- [{e.get('label','?')}] {e.get('subject','(no subject)')} from {e.get('sender','?')}")

        prompt = (
            f"You are KEVIN. Give a brief, sarcastic morning security briefing based on these emails:\n\n"
            + "\n".join(lines)
            + "\n\nKeep it under 150 words. Be Kevin: paranoid, sarcastic, protective."
        )

        result = await self.think(prompt, temp=0.8)
        return result["text"] or "Everything looks fine. Probably."

    async def close(self):
        if self._session and not self._session.closed:
            await self._session.close()