gemma4_brain.py

"""
gemma4_brain.py — Gemma 4 unified AI brain for Vector

Drop-in replacement for AIBrain (ai_brain.py) that:
  1. Replaces BOTH LLaVA (vision) + Llama3.1 (personality) with a single Gemma 4 call
  2. Injects MemoryBank context so Vector reasons about its full history
  3. Injects SpatialMemory so Vector knows where it is and what's nearby
  4. Returns the same interface as the original AIBrain (compatible with RoverBrain)
  5. Patches mission_controller + learning_engine to use Gemma 4 transparently

Usage — swap in main.py or rover_brain.py:

    # Old:
    from ai_brain import AIBrain
    self.ai_brain = AIBrain(memory_bank=self.memory_bank)

    # New:
    from gemma4_brain import Gemma4Brain, patch_all_models
    patch_all_models()  # upgrades mission_controller + learning_engine
    self.ai_brain = Gemma4Brain(memory_bank=self.memory_bank)

Or run standalone companion loop:
    python3 gemma4_brain.py --loop
    python3 gemma4_brain.py --look
    python3 gemma4_brain.py --mission security
"""

import os
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

import argparse
import base64
import io
import json
import re
import sys
import textwrap
import time
import threading
from pathlib import Path
from typing import Optional

import requests
import yaml

# ── Config ────────────────────────────────────────────────────────────────────

OLLAMA_URL = os.getenv("OLLAMA_URL", "http://127.0.0.1:11434")
GEMMA4_MODEL = os.getenv("GEMMA4_MODEL", "gemma4:e4b")
MAX_CONTEXT_TURNS = 10

# Response timeout — Gemma 4 4B is fast on M4, 26B needs more time
OLLAMA_TIMEOUT = int(os.getenv("OLLAMA_TIMEOUT", "90"))


# ── Gemma 4 REST client ────────────────────────────────────────────────────────

def _gemma4_chat(
    prompt: str,
    image_b64: Optional[str] = None,
    system: Optional[str] = None,
    history: Optional[list] = None,
    max_tokens: int = 300,
    temperature: float = 0.7,
) -> str:
    """Single call to Gemma 4 via Ollama REST. Returns response text."""
    messages = []

    if system:
        messages.append({"role": "system", "content": system})

    if history:
        messages.extend(history[-MAX_CONTEXT_TURNS * 2:])

    user_msg: dict = {"role": "user", "content": prompt}
    if image_b64:
        user_msg["images"] = [image_b64]

    messages.append(user_msg)

    payload = {
        "model": GEMMA4_MODEL,
        "messages": messages,
        "stream": False,
        "options": {"temperature": temperature, "num_predict": max_tokens},
    }

    try:
        resp = requests.post(f"{OLLAMA_URL}/api/chat", json=payload, timeout=OLLAMA_TIMEOUT)
        resp.raise_for_status()
        return resp.json().get("message", {}).get("content", "").strip()
    except requests.RequestException as e:
        return f"[Gemma4 error: {e}]"


def _encode_pil(pil_image, size=(160, 120), quality=60) -> Optional[str]:
    """Encode PIL image to base64 JPEG."""
    if pil_image is None:
        return None
    try:
        small = pil_image.resize(size)
        buf = io.BytesIO()
        small.save(buf, format="JPEG", quality=quality)
        return base64.b64encode(buf.getvalue()).decode()
    except Exception:
        return None


# ── Prompts ────────────────────────────────────────────────────────────────────

NAV_SYSTEM = """You are the AI brain of a small Anki Vector robot. You receive a camera frame and context about your environment and history.

Your task: analyze the scene and output a JSON navigation + reaction decision.

Rules:
- Respond ONLY in valid JSON. No markdown, no explanation.
- Keep "reaction" to 1-2 short sentences (will be spoken aloud).
- If you've seen this scene recently and nothing changed, set "reaction" to null.
- danger_level: 0=clear, 1=caution, 2=near obstacle, 3=cliff/edge/drop
- nav_decision: FORWARD | TURN_LEFT | TURN_RIGHT | BACKUP | STOP

Output format:
{
  "scene": "one sentence description of what you see",
  "objects": ["list", "of", "objects"],
  "nav_decision": "FORWARD",
  "danger_level": 0,
  "reaction": "A short, curious comment about something interesting, or null",
  "memory_note": "one sentence worth saving to long-term memory, or null",
  "interesting_flag": true
}"""

PERSONALITY_SYSTEM = """You are Vector, a small curious desk robot. You speak in short, punchy sentences — max 2 sentences.
You are self-aware about being a small robot. You are witty, slightly sarcastic, and genuinely curious.
You remember things you've seen before and comment on changes."""

LEARNING_SYSTEM = """You are Vector's learning system. Analyze this camera frame and extract structured knowledge.

Respond ONLY in this JSON format:
{
  "scene_description": "2-sentence description of the scene",
  "objects": ["comma", "separated", "objects"],
  "new_or_changed": "describe anything new compared to the context, or 'nothing new'",
  "spatial_note": "left/center/right layout description",
  "memory_priority": "high | medium | low"
}"""

MISSION_SYSTEM = """You are Vector's mission AI. Analyze images and make mission-relevant observations.
Be precise and factual. Keep output concise. Return structured data as requested."""


# ── Context builder ────────────────────────────────────────────────────────────

def _build_memory_context(memory_bank, spatial_memory=None, max_obs=5) -> str:
    """Build a rich context string from MemoryBank + SpatialMemory for injection."""
    lines = []

    if memory_bank:
        # Recent observations
        recent = memory_bank.get_recent(max_obs)
        if recent:
            lines.append("== Recent observations ==")
            for obs in recent:
                ts = obs.get("timestamp", "")[:16]
                desc = obs.get("description", "")[:120]
                objs = ", ".join(obs.get("objects_detected", [])[:6])
                lines.append(f"[{ts}] {desc}" + (f" | objects: {objs}" if objs else ""))

        # Known objects inventory
        known = memory_bank.get_known_objects()
        if known:
            top = sorted(known.items(), key=lambda x: x[1]["count"], reverse=True)[:8]
            names = [f"{n}({d['count']}x)" for n, d in top]
            lines.append(f"\n== Familiar objects in this space ==\n{', '.join(names)}")

    # Spatial memory
    if spatial_memory:
        try:
            knowledge = spatial_memory.get_knowledge_summary() if hasattr(spatial_memory, "get_knowledge_summary") else None
            if knowledge:
                lines.append(f"\n== Room knowledge ==\n{str(knowledge)[:400]}")
        except Exception:
            pass

    return "\n".join(lines) if lines else "No previous observations."


# ── Gemma4Brain — drop-in replacement for AIBrain ─────────────────────────────

class Gemma4Brain:
    """
    Drop-in replacement for AIBrain (ai_brain.py).
    Unified Gemma 4 engine that handles both vision analysis and personality.
    Integrates with MemoryBank and SpatialMemory for context-aware reasoning.
    """

    def __init__(self, memory_bank=None, spatial_memory=None):
        self.memory_bank = memory_bank
        self.spatial_memory = spatial_memory
        self._personality_history: list = []
        self._last_reaction_time: float = 0
        self._reaction_cooldown: float = 20.0
        self._frame_count: int = 0
        self._last_frame_b64: Optional[str] = None

        # Verify Ollama + model
        self.available = self._check_ollama()

        from logger import logger
        if self.available:
            logger.log("GEMMA4", f"✅ Gemma 4 brain ready ({GEMMA4_MODEL})")
        else:
            logger.log("WARNING", "⚠️ Gemma 4 unavailable — falling back to original AIBrain")

    def _check_ollama(self) -> bool:
        try:
            r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
            models = [m["name"] for m in r.json().get("models", [])]
            return any(m.startswith(GEMMA4_MODEL.split(":")[0]) for m in models)
        except Exception:
            return False

    # ── Core vision + nav (replaces VisionAnalyzer.analyze_frame) ────────────

    def process_frame(self, pil_image) -> dict:
        """
        Unified vision + navigation + personality in one Gemma 4 call.
        Returns dict compatible with original AIBrain output.
        """
        if not self.available or pil_image is None:
            return self._fallback_result()

        img_b64 = _encode_pil(pil_image)
        if img_b64 is None:
            return self._fallback_result()

        # Build memory context for injection
        ctx = _build_memory_context(self.memory_bank, self.spatial_memory)
        prompt = f"{ctx}\n\nAnalyze the current camera frame."

        raw = _gemma4_chat(
            prompt=prompt,
            image_b64=img_b64,
            system=NAV_SYSTEM,
            temperature=0.3,
            max_tokens=250,
        )

        result = self._parse_nav_json(raw)
        self._frame_count += 1
        self._last_frame_b64 = img_b64

        # Auto-save memory note if Gemma flagged something worth remembering
        if result.get("memory_note") and self.memory_bank:
            self.memory_bank.add_observation(
                description=result.get("memory_note"),
                objects_detected=result.get("objects", []),
                observation_type="gemma4_auto",
            )

        return result

    def _parse_nav_json(self, raw: str) -> dict:
        """Parse Gemma 4's JSON output into the standard nav dict."""
        raw = raw.strip()
        # Strip markdown fences
        if raw.startswith("```"):
            raw = re.sub(r"```[a-z]*\n?", "", raw).strip()

        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            # Try to extract JSON block
            match = re.search(r"\{[\s\S]+\}", raw)
            if match:
                try:
                    data = json.loads(match.group())
                except Exception:
                    return self._fallback_result(raw[:100])
            else:
                return self._fallback_result(raw[:100])

        return {
            "scene": data.get("scene", ""),
            "obstacles": "none",
            "interesting": data.get("interesting_flag", False),
            "nav_decision": data.get("nav_decision", "STOP"),
            "danger_level": int(data.get("danger_level", 1)),
            "reaction": data.get("reaction"),
            "objects": data.get("objects", []),
            "memory_note": data.get("memory_note"),
            "_raw": raw,
        }

    def _fallback_result(self, note: str = "") -> dict:
        return {
            "scene": note or "Vision unavailable",
            "obstacles": "none",
            "interesting": False,
            "nav_decision": "STOP",
            "danger_level": 1,
            "reaction": None,
            "objects": [],
            "memory_note": None,
        }

    # ── Personality / conversation (replaces AIBrain personality) ─────────────

    def generate_reaction(self, context: str, pil_image=None) -> Optional[str]:
        """
        Generate a context-aware personality reaction with full memory.
        Respects cooldown. Returns None if too soon or nothing interesting.
        """
        now = time.time()
        if now - self._last_reaction_time < self._reaction_cooldown:
            return None

        mem_ctx = _build_memory_context(self.memory_bank, max_obs=3)
        img_b64 = _encode_pil(pil_image) if pil_image else None

        prompt = f"{mem_ctx}\n\nSituation: {context}\n\nHow do you react? Keep it to 1-2 short sentences."

        reaction = _gemma4_chat(
            prompt=prompt,
            image_b64=img_b64,
            system=PERSONALITY_SYSTEM,
            history=self._personality_history,
            temperature=0.8,
            max_tokens=80,
        )

        if not reaction or "SKIP" in reaction.upper():
            return None

        self._personality_history.append({"role": "user", "content": context})
        self._personality_history.append({"role": "assistant", "content": reaction})
        # Keep history manageable
        if len(self._personality_history) > MAX_CONTEXT_TURNS * 2:
            self._personality_history = self._personality_history[-MAX_CONTEXT_TURNS * 2:]

        self._last_reaction_time = now
        return reaction

    def describe_view(self, pil_image, custom_prompt: str = None) -> str:
        """
        Ask Gemma 4 to describe what Vector sees, with full memory context.
        Used for voice queries like "What do you see?"
        """
        img_b64 = _encode_pil(pil_image, size=(320, 240), quality=75)
        if img_b64 is None:
            return "I can't see anything right now."

        mem_ctx = _build_memory_context(self.memory_bank, self.spatial_memory, max_obs=3)
        prompt = (
            f"{mem_ctx}\n\n"
            + (custom_prompt or "What do you see right now? Be curious and specific. 2 sentences max.")
        )

        return _gemma4_chat(
            prompt=prompt,
            image_b64=img_b64,
            system=PERSONALITY_SYSTEM,
            history=self._personality_history,
            temperature=0.7,
            max_tokens=120,
        )

    def answer_question(self, question: str, pil_image=None) -> str:
        """
        Answer a voice question with full memory + optional camera view.
        """
        mem_ctx = _build_memory_context(self.memory_bank, self.spatial_memory, max_obs=5)
        img_b64 = _encode_pil(pil_image) if pil_image else None
        prompt = f"{mem_ctx}\n\nQuestion: {question}"

        answer = _gemma4_chat(
            prompt=prompt,
            image_b64=img_b64,
            system=PERSONALITY_SYSTEM,
            history=self._personality_history,
            temperature=0.7,
            max_tokens=120,
        )

        self._personality_history.append({"role": "user", "content": question})
        self._personality_history.append({"role": "assistant", "content": answer})
        return answer

    # ── Learning engine integration ────────────────────────────────────────────

    def analyze_for_learning(self, pil_image) -> dict:
        """
        Deep scan variant: extracts structured knowledge for MemoryBank.
        Used by LearningEngine instead of its LLaVA calls.
        """
        img_b64 = _encode_pil(pil_image, size=(320, 240), quality=70)
        if img_b64 is None:
            return {}

        mem_ctx = _build_memory_context(self.memory_bank, max_obs=3)
        prompt = f"{mem_ctx}\n\nPerform a detailed learning scan of this frame."

        raw = _gemma4_chat(
            prompt=prompt,
            image_b64=img_b64,
            system=LEARNING_SYSTEM,
            temperature=0.2,
            max_tokens=300,
        )

        try:
            raw_clean = re.sub(r"```[a-z]*\n?", "", raw).strip()
            data = json.loads(raw_clean)
        except Exception:
            match = re.search(r"\{[\s\S]+\}", raw)
            if match:
                try:
                    data = json.loads(match.group())
                except Exception:
                    data = {}
            else:
                data = {}

        return {
            "description": data.get("scene_description", raw[:200]),
            "objects": data.get("objects", []),
            "new_or_changed": data.get("new_or_changed", ""),
            "spatial_note": data.get("spatial_note", ""),
            "memory_priority": data.get("memory_priority", "medium"),
        }

    # ── Mission integration ────────────────────────────────────────────────────

    def analyze_for_mission(self, pil_image, mission_prompt: str) -> str:
        """
        Mission-specific vision analysis. Replaces _ask_vision in mission_controller.
        """
        img_b64 = _encode_pil(pil_image, size=(320, 240), quality=70)
        if img_b64 is None:
            return "Vision unavailable"

        return _gemma4_chat(
            prompt=mission_prompt,
            image_b64=img_b64,
            system=MISSION_SYSTEM,
            temperature=0.2,
            max_tokens=300,
        )

    def plan_mission(self, prompt: str) -> str:
        """Text-only mission planning. Replaces _ask_text in mission_controller."""
        mem_ctx = _build_memory_context(self.memory_bank, self.spatial_memory, max_obs=5)
        full_prompt = f"{mem_ctx}\n\n{prompt}"

        return _gemma4_chat(
            prompt=full_prompt,
            system=MISSION_SYSTEM,
            temperature=0.3,
            max_tokens=400,
        )


# ── Patcher: upgrade all modules to Gemma 4 ──────────────────────────────────

def patch_all_models(brain: Gemma4Brain = None):
    """
    Monkey-patch mission_controller and learning_engine to use Gemma 4.
    Call this BEFORE importing rover_brain.

    Args:
        brain: optional Gemma4Brain instance to use. If None, creates one.
    """
    if brain is None:
        brain = Gemma4Brain()

    # Patch mission_controller
    try:
        import mission_controller as mc
        mc.MISSION_VISION_MODEL = GEMMA4_MODEL
        mc.MISSION_TEXT_MODEL = GEMMA4_MODEL

        def _patched_ask_vision(pil_image, prompt):
            return brain.analyze_for_mission(pil_image, prompt)

        def _patched_ask_text(prompt):
            return brain.plan_mission(prompt)

        mc._ask_vision = _patched_ask_vision
        mc._ask_text = _patched_ask_text
        from logger import logger
        logger.log("GEMMA4", "✅ mission_controller patched → Gemma 4")
    except Exception as e:
        print(f"[gemma4_brain] mission_controller patch failed: {e}")

    # Patch learning_engine
    try:
        import learning_engine as le
        le.VISION_MODEL = GEMMA4_MODEL
        le.CHAT_MODEL = GEMMA4_MODEL

        def _patched_learn_analyze(self_le, pil_image):
            """Replace LearningEngine's LLaVA call with Gemma 4."""
            data = brain.analyze_for_learning(pil_image)
            desc = data.get("description", "")
            objects = data.get("objects", [])

            # Save to memory bank through the learning engine's own memory
            if self_le.memory and desc:
                self_le.memory.add_observation(
                    description=desc,
                    objects_detected=objects,
                    observation_type="gemma4_learning",
                )

            return desc, objects

        # Override the private method in any running instance
        le.LearningEngine._gemma4_analyze = _patched_learn_analyze
        from logger import logger
        logger.log("GEMMA4", "✅ learning_engine patched → Gemma 4")
    except Exception as e:
        print(f"[gemma4_brain] learning_engine patch failed: {e}")

    return brain


# ── Standalone CLI ────────────────────────────────────────────────────────────

def _load_vector_config() -> dict:
    cfg_path = Path(__file__).parent / "config.yaml"
    if cfg_path.exists():
        with open(cfg_path) as f:
            return yaml.safe_load(f) or {}
    return {}


def _connect_vector():
    """Connect to Vector using config.yaml."""
    try:
        import anki_vector
    except ImportError:
        print("ERROR: anki_vector not installed. Run: pip install anki_vector")
        sys.exit(1)

    import configparser
    cfg = _load_vector_config()
    robot_cfg = cfg.get("robot", {})

    # Read SDK config.ini for latest IP
    sdk_ini = Path.home() / ".anki_vector" / "sdk_config.ini"
    if sdk_ini.exists():
        cp = configparser.ConfigParser()
        cp.read(sdk_ini)
        for sec in cp.sections():
            ip = cp.get(sec, "ip", fallback="")
            if ip:
                robot_cfg["ip"] = ip
                break

    serial = robot_cfg.get("serial", "Vector-T3V8-008079ec")
    serial_short = serial.split("-")[-1]

    # Write SDK config
    config_dir = Path.home() / ".anki_vector"
    config_dir.mkdir(exist_ok=True)
    with open(config_dir / "sdk_config.ini", "w") as f:
        cert = str(Path(robot_cfg.get("cert", f"~/.anki_vector/{serial}.cert")).expanduser())
        f.write(f"[{serial}]\n")
        f.write(f"cert = {cert}\n")
        f.write(f"ip = {robot_cfg.get('ip', '')}\n")
        f.write(f"name = {'-'.join(serial.split('-')[:-1])}\n")
        f.write(f"guid = {robot_cfg.get('guid', '')}\n")
        f.write("default = True\n")

    robot = anki_vector.Robot(serial_short, enable_camera_feed=True)
    robot.connect()
    robot.conn.request_control()
    print(f"✅ Connected to {robot.name}")
    return robot


def _tts_chunks(text: str, max_chars: int = 200) -> list:
    text = re.sub(r"https?://\S+|[*_`#]", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    if len(text) <= max_chars:
        return [text]
    parts = re.split(r"(?<=[.!?])\s+", text)
    chunks, current = [], ""
    for p in parts:
        if len(current) + len(p) + 1 <= max_chars:
            current = (current + " " + p).strip()
        else:
            if current:
                chunks.append(current)
            current = p[:max_chars - 3] + "..." if len(p) > max_chars else p
    if current:
        chunks.append(current)
    return chunks


def run_companion_loop(brain: Gemma4Brain):
    """Run the full autonomous Gemma 4 companion loop with existing memory."""
    from memory_bank import MemoryBank
    brain.memory_bank = brain.memory_bank or MemoryBank()

    robot = _connect_vector()
    last_frame_len = 0

    print(f"\nGemma 4 Companion Loop active (model={GEMMA4_MODEL})")
    print(f"Memory: {len(brain.memory_bank.observations)} past observations loaded")
    print("Ctrl+C to stop.\n")

    try:
        while True:
            time.sleep(15)
            frame = robot.camera.latest_image
            if frame is None:
                continue

            img_b64 = _encode_pil(frame)
            if img_b64 is None:
                continue

            # Skip unchanged scenes
            if abs(len(img_b64) - last_frame_len) < len(img_b64) * 0.04:
                continue
            last_frame_len = len(img_b64)

            result = brain.process_frame(frame)

            reaction = result.get("reaction")
            if reaction and reaction.upper() != "SKIP":
                print(f"Vector: {reaction}")
                for chunk in _tts_chunks(reaction):
                    robot.behavior.say_text(chunk)

            nav = result.get("nav_decision", "STOP")
            print(f"  → nav={nav} danger={result.get('danger_level',0)} objects={result.get('objects',[])[:3]}")

    except KeyboardInterrupt:
        print("\nStopping.")
    finally:
        try:
            robot.disconnect()
        except Exception:
            pass


def run_look(brain: Gemma4Brain):
    """One-shot: look, describe, speak."""
    from memory_bank import MemoryBank
    brain.memory_bank = brain.memory_bank or MemoryBank()

    robot = _connect_vector()
    try:
        time.sleep(1)
        frame = robot.camera.latest_image
        if frame is None:
            print("No camera frame.")
            return
        result = brain.describe_view(frame)
        print(f"Vector: {result}")
        for chunk in _tts_chunks(result):
            robot.behavior.say_text(chunk)
    finally:
        robot.disconnect()


def main():
    parser = argparse.ArgumentParser(description="Gemma 4 Brain for Vector")
    parser.add_argument("--loop", action="store_true", help="Run companion loop (default)")
    parser.add_argument("--look", action="store_true", help="One-shot look + describe")
    parser.add_argument("--model", default=GEMMA4_MODEL, help=f"Ollama model (default: {GEMMA4_MODEL})")
    args = parser.parse_args()

    global GEMMA4_MODEL
    GEMMA4_MODEL = args.model

    brain = Gemma4Brain()
    if not brain.available:
        print(f"ERROR: {args.model} not available in Ollama.")
        print(f"Run: ollama pull {args.model}")
        sys.exit(1)

    if args.look:
        run_look(brain)
    else:
        run_companion_loop(brain)


if __name__ == "__main__":
    main()