src/observer_llm.py

"""
Observer LLM (LLM-A) — watches sensory state and produces compact reasoning.

The observer compresses; the builder compiles.
"""

import os
import time
import json
import requests
from .stream_state import SessionState


GROQ_API_KEY = os.getenv("GROQ_API_KEY", os.getenv("GROK_API_KEY", ""))
GROQ_BASE_URL = "https://api.groq.com/openai/v1"
GROQ_VISION_MODEL = os.getenv("GROQ_VISION_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct")
GROQ_TEXT_MODEL = os.getenv("GROQ_TEXT_MODEL", "llama-3.3-70b-versatile")


OBSERVER_PROMPT = """You are the Observer in a Continuity Sensory Code Engine.

You are receiving a LIVE feed from a camera and microphone. The image attached is the current camera frame RIGHT NOW. The compressed state below contains the latest audio transcripts and motion metrics.

Your job is to produce a LIVE reasoning trace about what is happening RIGHT NOW, then a compact observation:

1. REASONING: Look at the attached camera frame. What do you see RIGHT NOW? Describe specific objects, screens, code, people, movements. Then read the audio transcripts — what was said? Connect what you see with what you hear. This must be specific to the current frame and audio, not generic.
2. SCENE: Specific description of what is visible in the current frame (not generic — name actual objects, text, colors, positions)
3. INFERRED_INTENT: Based on what you see and hear RIGHT NOW, what does the user want? Infer from the evidence. If they're showing code on screen, they want code help. If they're moving objects, they want something about those objects.
4. SIGNALS: What audio/speech signals are present? Quote the actual transcripts. Separate by speaker.
5. CANDIDATE_TASK: What specific code should the Builder generate based on this exact scene + audio?
6. UNCERTAINTY: What can't you determine from the current evidence?

Rules:
- You MUST reference specific things visible in the current frame. Do not give generic answers.
- If the frame is dark/blank, say so. If there's motion, describe the direction and speed.
- Quote actual audio transcript text in SIGNALS.
- Do not generate code. That is the Builder's job.
- Every observation must be tied to the current frame and current audio.
- If there is camera movement or sound, that IS the signal — infer what the user wants from it.
"""


def call_ollama_vision(frame_b64: str) -> str:
    """Call local Ollama vision model to describe a single frame.

    Returns a text description of what's visible in the frame.
    This runs on the user's local machine, tunneled via OLLAMA_HOST.
    """
    prompt = "Describe what you see in this camera frame in 2-3 sentences. Focus on: objects, screens, code, people, environment, and any visible text or error messages."
    payload = {"model": os.getenv("OLLAMA_MODEL", "llava"), "prompt": prompt, "stream": False, "images": [frame_b64]}
    r = requests.post(f"{os.getenv('OLLAMA_HOST', 'http://localhost:11434')}/api/generate", json=payload, timeout=60)
    if r.status_code >= 400:
        raise RuntimeError(f"Ollama vision error: {r.text[:500]}")
    return r.json().get("response", "")


def call_ollama(prompt: str, frame_b64: str | None = None) -> str:
    model = os.getenv("OLLAMA_MODEL", "llava")
    host = os.getenv("OLLAMA_HOST", "http://localhost:11434")
    payload = {"model": model, "prompt": prompt, "stream": False}
    if frame_b64:
        payload["images"] = [frame_b64]
    r = requests.post(f"{host}/api/generate", json=payload, timeout=120)
    if r.status_code >= 400:
        raise RuntimeError(f"Ollama error: {r.text[:500]}")
    return r.json().get("response", "")


def call_openai(prompt: str, frame_b64: str | None = None) -> str:
    from openai import OpenAI
    model = os.getenv("OPENAI_MODEL", "gpt-4o")
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    content = [{"type": "text", "text": prompt}]
    if frame_b64:
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{frame_b64}"},
        })
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": content}],
        max_tokens=800,
    )
    return response.choices[0].message.content


def call_grok(prompt: str, frame_b64: str | None = None) -> str:
    """Call Groq API (OpenAI-compatible) with vision support via llama-4-scout."""
    from openai import OpenAI
    client = OpenAI(api_key=GROQ_API_KEY, base_url=GROQ_BASE_URL)
    content = [{"type": "text", "text": prompt}]
    if frame_b64:
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{frame_b64}"},
        })
    response = client.chat.completions.create(
        model=GROQ_VISION_MODEL if frame_b64 else GROQ_TEXT_MODEL,
        messages=[{"role": "user", "content": content}],
        max_tokens=1000,
    )
    return response.choices[0].message.content


def call_hf_inference(prompt: str, frame_b64: str | None = None) -> str:
    """Call HF Inference API for text reasoning.

    In hybrid mode, vision is handled by local Ollama, and the frame
    description is injected into the text prompt.
    """
    from huggingface_hub import InferenceClient

    token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
    client = InferenceClient(token=token)
    model = os.getenv("HF_TEXT_MODEL", "Qwen/Qwen2.5-7B-Instruct")
    response = client.chat_completion(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=800,
    )
    return response.choices[0].message.content


def call_hybrid(prompt: str, frame_b64: str | None = None) -> str:
    """Hybrid mode: Ollama for vision, HF Inference for reasoning.

    1. If frame available, call local Ollama vision model to describe the scene
    2. Inject the visual description into the text prompt
    3. Call HF Inference for reasoning + structured observation
    """
    vision_description = ""
    if frame_b64:
        try:
            vision_description = call_ollama_vision(frame_b64)
        except Exception as e:
            vision_description = f"[Vision unavailable: {str(e)[:100]}]"

    enhanced_prompt = prompt
    if vision_description:
        enhanced_prompt = f"{prompt}\n\nVISION MODEL OUTPUT (from local Ollama):\n{vision_description}\n\nUse this visual description along with the compressed state to produce your observation."

    return call_hf_inference(enhanced_prompt)


def observe(state: SessionState, compact_state: dict) -> dict:
    """Run observer LLM on compressed state.

    Returns structured observation with scene, intent, signals,
    candidate_task, and uncertainty.
    """
    frames = list(state.frames)
    last_frame = frames[-1] if frames else None
    frame_b64 = last_frame.jpeg_b64 if last_frame else None

    prompt = f"""{OBSERVER_PROMPT}

COMPRESSED SENSORY STATE:
{json.dumps(compact_state, indent=2)}

Produce your observation now."""

    provider = os.getenv("PROVIDER", "groq").lower().strip()
    if provider == "ollama":
        output = call_ollama(prompt, frame_b64)
    elif provider == "openai":
        output = call_openai(prompt, frame_b64)
    elif provider == "grok" or provider == "groq":
        output = call_grok(prompt, frame_b64)
    elif provider == "huggingface":
        output = call_hf_inference(prompt, frame_b64)
    elif provider == "hybrid":
        output = call_hybrid(prompt, frame_b64)
    else:
        raise RuntimeError(f"PROVIDER must be 'groq', 'hybrid', 'ollama', 'openai', or 'huggingface', got: {provider!r}.")

    return {
        "observer_output": output,
        "state_hash": compact_state.get("state_hash", ""),
        "frame_hash": state.last_frame_hash,
        "timestamp": time.time(),
    }