handler.py

"""Your bot lives here.

Inbound X Chat messages arrive as `Message`. Return a `Reply`.
The default implementation is a Grok pass-through:

  user text  → Grok chat
  /imagine   → Grok Imagine (uses an attached photo as an edit reference)
  /reset     → clears that user's Grok history

Replace `GrokPassthrough.handle` — or swap in any class with the same
`handle(Message) -> Reply` shape — to build a different bot.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any, Protocol

from xchat_bot.config import Settings
from xchat_bot.grok import Grok, _data_uri
from xchat_bot.store import Store

logger = logging.getLogger("xchat_bot.handler")


@dataclass(frozen=True)
class Message:
    user_id: str
    conversation_id: str
    text: str
    images: tuple[bytes, ...] = ()


@dataclass(frozen=True)
class Reply:
    text: str = ""
    image: bytes | None = None


class Handler(Protocol):
    def handle(self, message: Message) -> Reply: ...


class GrokPassthrough:
    """Grok in, Grok out. Customize `handle` or `system_prompt` to change behavior."""

    def __init__(self, grok: Grok, store: Store, settings: Settings) -> None:
        self.grok = grok
        self.store = store
        self.settings = settings

    def handle(self, message: Message) -> Reply:
        text = message.text.strip()
        lowered = text.lower()

        if lowered in {"/reset", "reset"}:
            self.store.clear_history(message.user_id)
            return Reply("Starting fresh.")

        if lowered.startswith("/imagine"):
            prompt = text.split(None, 1)[1].strip() if " " in text else ""
            if not prompt:
                return Reply("Send `/imagine` plus what you want to see.")
            try:
                image = self.grok.imagine(prompt, message.images[0] if message.images else None)
            except Exception:
                logger.exception("imagine_failed user=%s", message.user_id)
                return Reply("I couldn't generate that image. Try a different prompt.")
            self.store.append(message.user_id, "user", f"/imagine {prompt}")
            self.store.append(message.user_id, "assistant", "Here's what I imagined.")
            return Reply("Here's what I imagined.", image)

        return self._chat(message, text)

    def _chat(self, message: Message, text: str) -> Reply:
        history = self.store.history(message.user_id, self.settings.history_limit)
        user_content: str | list[dict[str, Any]] = text
        if message.images:
            parts: list[dict[str, Any]] = [
                {"type": "image_url", "image_url": {"url": _data_uri(img)}}
                for img in message.images[:3]
            ]
            parts.append({"type": "text", "text": text or "What is in this image?"})
            user_content = parts
        messages: list[dict[str, Any]] = [
            {"role": "system", "content": self.settings.system_prompt},
            *history,
            {"role": "user", "content": user_content},
        ]
        try:
            answer = self.grok.chat(messages)
        except Exception:
            logger.exception("grok_failed user=%s", message.user_id)
            return Reply("Something glitched on my side. Try that again in a second.")
        if not answer:
            answer = "…"
        stored_user = text if text else "[image]"
        self.store.append(message.user_id, "user", stored_user)
        self.store.append(message.user_id, "assistant", answer)
        return Reply(answer)