bot.py

"""
groq-bot · bot.py
=================
A lightweight RAG-powered WhatsApp auto-reply server built for Termux.

Architecture
------------
  WhatsApp → WhatsAuto → POST /reply → TF-IDF retrieval → Groq LLM → reply

Features
--------
  • Pure-Python TF-IDF similarity (no numpy / sentence-transformers)
  • Per-sender conversation memory with configurable TTL
  • Relevance threshold: bot stays silent when no context matches
  • Response cache keyed by content hash (skips duplicate Groq calls)
  • Named Cloudflare tunnel support for a permanent public URL
  • Zero required environment variables – guided first-run config
"""

from __future__ import annotations

import hashlib
import json
import logging
import math
import os
import re
import sys
import time
from collections import defaultdict
from typing import Any

import requests
from flask import Flask, request, jsonify

# ─────────────────────────────────────────────────────────────────────────────
# Logging
# ─────────────────────────────────────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-7s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("groq-bot")

# ─────────────────────────────────────────────────────────────────────────────
# Paths
# ─────────────────────────────────────────────────────────────────────────────

BASE    = os.path.expanduser(os.environ.get("BOT_BASE", "~/groq-bot"))
DOCS    = os.path.join(BASE, "docs")
VECTORS = os.path.join(BASE, "vectors.json")
CACHE   = os.path.join(BASE, "cache.json")
CFG     = os.path.join(BASE, "config.json")

# ─────────────────────────────────────────────────────────────────────────────
# Runtime config (populated by load_config() at startup)
# ─────────────────────────────────────────────────────────────────────────────

_cfg: dict[str, Any] = {}

GROQ_URL = "https://api.groq.com/openai/v1/chat/completions"

# Defaults – all overridable via config.json or environment variables
_DEFAULTS: dict[str, Any] = {
    "api_key":           "",
    "model":             "llama-3.1-8b-instant",
    "port":              5000,
    # RAG settings
    "top_k":             5,       # number of chunks retrieved per query
    "chunk_words":       400,     # words per document chunk
    "chunk_overlap":     100,     # word overlap between consecutive chunks
    "relevance_threshold": 0.08,  # cosine score below which the bot stays silent
    # Session memory
    "session_ttl":       1800,    # seconds of inactivity before session expires (30 min)
    "session_max_turns": 8,       # maximum conversation turns kept in memory
    # System prompt
    "system_prompt": (
        "You are a professional cabin sales assistant.\n"
        "Answer ONLY using the provided context.\n"
        "Be concise, friendly, and never repeat information already given in this conversation.\n"
        "If the answer is not in the context, respond ONLY with the exact phrase:\n"
        "I will confirm that information with our team."
    ),
}


def cfg(key: str) -> Any:
    """Return a runtime config value, falling back to defaults."""
    return _cfg.get(key, _DEFAULTS[key])


# ─────────────────────────────────────────────────────────────────────────────
# Persistence helpers
# ─────────────────────────────────────────────────────────────────────────────

def load_json(path: str, default: Any) -> Any:
    if os.path.exists(path):
        try:
            with open(path, "r", encoding="utf-8") as fh:
                return json.load(fh)
        except Exception as exc:
            log.warning("Could not load %s: %s", path, exc)
    return default


def save_json(path: str, data: Any) -> None:
    try:
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(data, fh, ensure_ascii=False, indent=2)
    except Exception as exc:
        log.warning("Could not save %s: %s", path, exc)


# ─────────────────────────────────────────────────────────────────────────────
# Config management
# ─────────────────────────────────────────────────────────────────────────────

def load_config() -> None:
    """Load config.json and merge with environment overrides into _cfg."""
    global _cfg
    _cfg = load_json(CFG, {})
    # Environment variables take highest priority
    if os.environ.get("GROQ_API_KEY"):
        _cfg["api_key"] = os.environ["GROQ_API_KEY"]
    if os.environ.get("GROQ_MODEL"):
        _cfg["model"] = os.environ["GROQ_MODEL"]
    if os.environ.get("BOT_PORT"):
        _cfg["port"] = int(os.environ["BOT_PORT"])


def get_api_key(interactive: bool = False) -> str:
    """
    Return the configured Groq API key.

    Parameters
    ----------
    interactive : bool
        When True (CLI only), prompt the user if no key is set.
        When False (server mode), raise RuntimeError instead of prompting,
        since stdin is not available in a background nohup process.
    """
    key = cfg("api_key")
    if key:
        return key

    if not interactive:
        raise RuntimeError(
            "No Groq API key configured. "
            "Run  bot config  to set it, then  bot start  again."
        )

    key = input("Enter your Groq API key: ").strip()
    if not key:
        raise RuntimeError("API key cannot be empty.")
    _cfg["api_key"] = key
    save_json(CFG, _cfg)
    return key


def config_cmd() -> None:
    """Interactive config editor (bot config)."""
    load_config()
    fields = [
        ("api_key",             "Groq API key"),
        ("model",               "Model name"),
        ("relevance_threshold", "Relevance threshold (0.0–1.0, default 0.08)"),
        ("session_ttl",         "Session TTL in seconds (default 1800)"),
        ("session_max_turns",   "Max conversation turns (default 8)"),
    ]
    print("\n  Leave blank to keep current value.\n")
    for key, label in fields:
        current = _cfg.get(key, _DEFAULTS.get(key, ""))
        val = input(f"  {label} [{current}]: ").strip()
        if val:
            # cast to original type
            orig = _DEFAULTS.get(key, "")
            try:
                _cfg[key] = type(orig)(val) if orig != "" else val
            except (ValueError, TypeError):
                _cfg[key] = val
    save_json(CFG, _cfg)
    print("\n  Config saved.\n")


# ─────────────────────────────────────────────────────────────────────────────
# TF-IDF retrieval  (zero external dependencies)
# ─────────────────────────────────────────────────────────────────────────────

def _tokenize(text: str) -> list[str]:
    return re.findall(r"[a-z0-9]+", text.lower())


def _build_tfidf(corpus: list[str]) -> tuple[list[dict[str, float]], dict[str, float]]:
    """
    Compute L2-normalised TF-IDF vectors for every document in *corpus*.

    Returns
    -------
    vecs : list of sparse dicts  {token: tfidf_weight}
    idf  : global IDF lookup     {token: idf_weight}
    """
    N  = len(corpus)
    df: dict[str, int] = defaultdict(int)
    tfs: list[dict[str, float]] = []

    for doc in corpus:
        tokens = _tokenize(doc)
        tf: dict[str, int] = defaultdict(int)
        for t in tokens:
            tf[t] += 1
        total = len(tokens) or 1
        tfs.append({t: c / total for t, c in tf.items()})
        for t in tf:
            df[t] += 1

    idf = {t: math.log((N + 1) / (cnt + 1)) + 1 for t, cnt in df.items()}

    vecs: list[dict[str, float]] = []
    for tf_doc in tfs:
        v = {t: tf_doc[t] * idf.get(t, 1.0) for t in tf_doc}
        norm = math.sqrt(sum(x * x for x in v.values())) or 1.0
        vecs.append({t: x / norm for t, x in v.items()})

    return vecs, idf


def _cosine(a: dict[str, float], b: dict[str, float]) -> float:
    return sum(a[t] * b[t] for t in set(a) & set(b))


def _query_vec(query: str, idf: dict[str, float]) -> dict[str, float]:
    tokens = _tokenize(query)
    tf: dict[str, int] = defaultdict(int)
    for t in tokens:
        tf[t] += 1
    total = len(tokens) or 1
    v = {t: (c / total) * idf.get(t, 1.0) for t, c in tf.items()}
    norm = math.sqrt(sum(x * x for x in v.values())) or 1.0
    return {t: x / norm for t, x in v.items()}


# ─────────────────────────────────────────────────────────────────────────────
# Document index
# ─────────────────────────────────────────────────────────────────────────────

class Index:
    """Manages document ingestion, chunking, and TF-IDF retrieval."""

    def __init__(self) -> None:
        self.chunks: list[str]            = []
        self.vecs:   list[dict[str, float]] = []
        self.idf:    dict[str, float]     = {}

    # ── persistence ──────────────────────────────────────────────────────────

    def load(self) -> None:
        data = load_json(VECTORS, {"chunks": []})
        self.chunks = data.get("chunks", [])
        if self.chunks:
            self.vecs, self.idf = _build_tfidf(self.chunks)
        log.info("Index loaded: %d chunks", len(self.chunks))

    def save(self) -> None:
        save_json(VECTORS, {"chunks": self.chunks})

    # ── build ─────────────────────────────────────────────────────────────────

    def build(self, force: bool = False) -> None:
        """Scan DOCS directory and (re)build the index."""
        if self.chunks and not force:
            return
        if not os.path.isdir(DOCS) or not os.listdir(DOCS):
            log.info("No docs to index yet. Drop .txt files into %s", DOCS)
            return

        log.info("Building index from docs...")
        self.chunks = []

        for fname in sorted(os.listdir(DOCS)):
            fpath = os.path.join(DOCS, fname)
            if not os.path.isfile(fpath):
                continue
            try:
                with open(fpath, "r", encoding="utf-8", errors="ignore") as fh:
                    raw = fh.read()
                for chunk in self._chunk(raw):
                    self.chunks.append(chunk)
            except Exception as exc:
                log.warning("Skipping %s: %s", fname, exc)

        self.vecs, self.idf = _build_tfidf(self.chunks)
        self.save()
        log.info("Indexed %d chunks", len(self.chunks))

    def _chunk(self, text: str) -> list[str]:
        size    = cfg("chunk_words")
        overlap = cfg("chunk_overlap")
        words   = text.split()
        out: list[str] = []
        i = 0
        while i < len(words):
            out.append(" ".join(words[i : i + size]))
            i += size - overlap
        return out

    # ── search ────────────────────────────────────────────────────────────────

    def search(self, query: str, k: int | None = None) -> tuple[list[str], float]:
        """
        Return the top-k most relevant chunks and the best score found.

        Returns
        -------
        chunks     : list[str]  – up to k text chunks
        best_score : float      – highest cosine similarity (0.0 if no index)
        """
        if not self.chunks:
            return [], 0.0

        k = k or cfg("top_k")
        qv     = _query_vec(query, self.idf)
        scores = sorted(
            ((score, i) for i, v in enumerate(self.vecs) if (score := _cosine(qv, v)) > 0),
            reverse=True,
        )
        best   = scores[0][0] if scores else 0.0
        chunks = [self.chunks[i] for _, i in scores[:k]]
        return chunks, best


index = Index()


# ─────────────────────────────────────────────────────────────────────────────
# Session memory
# ─────────────────────────────────────────────────────────────────────────────

class Session:
    """Conversation history for a single sender."""

    __slots__ = ("turns", "last_active")

    def __init__(self) -> None:
        self.turns:       list[dict[str, str]] = []
        self.last_active: float                = time.time()

    def add(self, role: str, content: str) -> None:
        self.turns.append({"role": role, "content": content})
        self.last_active = time.time()
        # Keep only the most recent N turns (each turn = user + assistant)
        max_msgs = cfg("session_max_turns") * 2
        if len(self.turns) > max_msgs:
            self.turns = self.turns[-max_msgs:]

    def is_expired(self) -> bool:
        return (time.time() - self.last_active) > cfg("session_ttl")

    def messages(self) -> list[dict[str, str]]:
        return list(self.turns)


class SessionStore:
    """Thread-safe in-memory store for per-sender sessions."""

    def __init__(self) -> None:
        self._store: dict[str, Session] = {}

    def get(self, sender: str) -> Session:
        self._evict_expired()
        if sender not in self._store:
            self._store[sender] = Session()
        return self._store[sender]

    def _evict_expired(self) -> None:
        expired = [k for k, s in self._store.items() if s.is_expired()]
        for k in expired:
            del self._store[k]
            log.debug("Session expired: %s", k)

    def clear(self, sender: str) -> None:
        self._store.pop(sender, None)

    @property
    def active_count(self) -> int:
        self._evict_expired()
        return len(self._store)


sessions = SessionStore()


# ─────────────────────────────────────────────────────────────────────────────
# Response cache
# ─────────────────────────────────────────────────────────────────────────────

_cache: dict[str, str] = {}


def _cache_key(text: str) -> str:
    return hashlib.md5(text.lower().strip().encode()).hexdigest()


def cache_get(text: str) -> str | None:
    return _cache.get(_cache_key(text))


def cache_set(text: str, answer: str) -> None:
    _cache[_cache_key(text)] = answer
    save_json(CACHE, _cache)


# ─────────────────────────────────────────────────────────────────────────────
# Groq API
# ─────────────────────────────────────────────────────────────────────────────

def call_groq(messages: list[dict[str, str]], api_key: str) -> str:
    """
    Send *messages* to the Groq chat completions endpoint.

    Returns the assistant's reply or an error string prefixed with "Error:".
    """
    try:
        resp = requests.post(
            GROQ_URL,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type":  "application/json",
            },
            json={
                "model":       cfg("model"),
                "messages":    messages,
                "temperature": 0.3,
                "max_tokens":  1024,
            },
            timeout=30,
        )
        resp.raise_for_status()
        return resp.json()["choices"][0]["message"]["content"].strip()

    except requests.exceptions.Timeout:
        return "Error: request timed out – please try again."
    except requests.exceptions.ConnectionError:
        return "Error: no internet connection."
    except requests.exceptions.HTTPError as exc:
        code = exc.response.status_code
        if code == 401:
            return "Error: invalid API key. Run `bot config` to update it."
        if code == 429:
            return "Error: rate limit reached – please wait a moment."
        return f"Error: Groq API returned HTTP {code}."
    except (KeyError, IndexError, ValueError):
        return "Error: unexpected response from Groq API."


# ─────────────────────────────────────────────────────────────────────────────
# Core answer logic
# ─────────────────────────────────────────────────────────────────────────────

# Sentinel returned when relevance is too low – caller decides what to do
NO_CONTEXT = "__NO_CONTEXT__"


def answer(question: str, sender: str = "cli") -> str:
    """
    Produce a reply to *question* from *sender*.

    Steps
    -----
    1. Retrieve relevant context chunks (TF-IDF).
    2. If best relevance score < threshold  → return NO_CONTEXT.
    3. Check single-question cache (no session involved).
    4. Build message list: system + session history + new user turn.
    5. Call Groq and update session memory.

    Returns
    -------
    str  – the bot's reply, or NO_CONTEXT if nothing is relevant enough.
    """
    try:
        api_key = get_api_key(interactive=False)
    except RuntimeError:
        return "Error: no API key configured. Run `bot config` then restart the bot."

    # ── 1. Retrieve context ───────────────────────────────────────────────
    chunks, best_score = index.search(question)
    threshold = cfg("relevance_threshold")

    log.info(
        "Query: %r | best_score=%.3f | threshold=%.3f | sender=%s",
        question[:60], best_score, threshold, sender,
    )

    # ── 2. Silence if no relevant context ────────────────────────────────
    if best_score < threshold:
        log.info("Score below threshold – staying silent.")
        return NO_CONTEXT

    context = "\n\n".join(chunks)

    # ── 3. Cache lookup (stateless queries only) ──────────────────────────
    session = sessions.get(sender)
    if not session.turns:
        cached = cache_get(question)
        if cached:
            session.add("user",      question)
            session.add("assistant", cached)
            return cached

    # ── 4. Build messages ─────────────────────────────────────────────────
    messages: list[dict[str, str]] = [
        {"role": "system", "content": cfg("system_prompt")},
        *session.messages(),
        {
            "role":    "user",
            "content": f"Context:\n{context}\n\nQuestion: {question}",
        },
    ]

    # ── 5. Call Groq and persist ──────────────────────────────────────────
    reply = call_groq(messages, api_key)

    session.add("user",      question)
    session.add("assistant", reply)

    # Only cache when session was empty (first interaction)
    if len(session.turns) == 2:
        cache_set(question, reply)

    return reply


# ─────────────────────────────────────────────────────────────────────────────
# Flask application
# ─────────────────────────────────────────────────────────────────────────────

app = Flask(__name__)


def _parse_body() -> dict[str, str]:
    """
    Extract request fields in order of preference:
      1. JSON body            (application/json)
      2. Form data            (application/x-www-form-urlencoded / multipart)
      3. URL query parameters (?message=...)
      4. Raw text body        (plain text fallback)
    """
    body = request.get_json(silent=True, force=True) or {}
    if not body:
        body = request.form.to_dict()
    if not body:
        body = request.args.to_dict()
    if not body:
        raw = request.get_data(as_text=True).strip()
        if raw:
            body = {"message": raw}
    return body


@app.route("/reply", methods=["GET", "POST"])
def route_reply():
    """
    Main webhook consumed by WhatsAuto (and any HTTP client).

    Request fields (any format)
    ---------------------------
    message     : str  – the incoming message text          (required)
    sender      : str  – sender name or phone number        (optional)
    phone       : str  – sender phone, used as session key  (optional)
    app         : str  – source app name, e.g. "WhatsAuto"  (optional)
    group_name  : str  – WhatsApp group name if applicable  (optional)

    Response
    --------
    { "reply": "<text>" }
    An empty "reply" string means the bot chose to stay silent.
    HTTP 200 is always returned so WhatsAuto does not retry.
    """
    body = _parse_body()

    question = (
        body.get("message") or body.get("text") or
        body.get("msg")     or body.get("content") or ""
    ).strip()

    # Use phone number as session key when available, fall back to sender name
    sender = (body.get("phone") or body.get("sender") or "unknown").strip()

    log.info(
        "← %s | sender=%s | group=%s | msg=%r",
        body.get("app", "unknown"),
        sender,
        body.get("group_name", "—"),
        question[:80],
    )

    if not question:
        return jsonify({"reply": ""}), 200

    result = answer(question, sender)

    # NO_CONTEXT → stay silent (empty reply so WhatsAuto sends nothing)
    if result == NO_CONTEXT:
        log.info("→ silent (no relevant context)")
        return jsonify({"reply": ""}), 200

    log.info("→ %r", result[:80])
    return jsonify({"reply": result}), 200


@app.route("/health", methods=["GET"])
def route_health():
    """Health-check endpoint – useful for monitoring."""
    return jsonify({
        "status":   "ok",
        "chunks":   len(index.chunks),
        "model":    cfg("model"),
        "sessions": sessions.active_count,
    }), 200


@app.route("/reindex", methods=["POST"])
def route_reindex():
    """Trigger a live reindex without restarting the server."""
    index.build(force=True)
    return jsonify({"chunks": len(index.chunks)}), 200


@app.route("/session/<sender>", methods=["DELETE"])
def route_clear_session(sender: str):
    """Clear the conversation history for a specific sender."""
    sessions.clear(sender)
    return jsonify({"cleared": sender}), 200


# ─────────────────────────────────────────────────────────────────────────────
# CLI
# ─────────────────────────────────────────────────────────────────────────────

def cli_chat() -> None:
    """Interactive REPL for local testing."""
    index.build()
    print("\n  Groq Bot  –  type 'exit' to quit, 'reindex' to rebuild docs\n")

    while True:
        try:
            question = input("You: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nGoodbye.")
            break

        if not question:
            continue
        if question.lower() == "exit":
            break
        if question.lower() == "reindex":
            index.build(force=True)
            print("Bot: Index rebuilt.\n")
            continue

        result = answer(question, sender="cli")
        if result == NO_CONTEXT:
            print("Bot: [silent – question outside knowledge base]\n")
        else:
            print(f"Bot: {result}\n")


def cli_start_server() -> None:
    """Start the Flask development server."""
    os.makedirs(DOCS, exist_ok=True)
    try:
        get_api_key(interactive=False)
    except RuntimeError as exc:
        log.error("%s", exc)
        sys.exit(1)
    index.load()
    index.build()
    port = cfg("port")
    log.info("Server starting on port %d", port)
    app.run(host="0.0.0.0", port=port, use_reloader=False)


# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

def main() -> None:
    load_config()
    _cache.update(load_json(CACHE, {}))

    command = sys.argv[1] if len(sys.argv) > 1 else "server"

    # Internal helper used by bot.sh to verify a key is set before
    # launching the server in the background.
    if command == "config" and len(sys.argv) > 2 and sys.argv[2] == "--check-key":
        key = cfg("api_key")
        if not key:
            sys.exit(1)   # non-zero → bot.sh shows the friendly error
        sys.exit(0)

    commands = {
        "server":  cli_start_server,
        "chat":    cli_chat,
        "config":  config_cmd,
        "reindex": lambda: (index.load(), index.build(force=True)),
    }

    handler = commands.get(command)
    if handler is None:
        print(f"Unknown command: {command}")
        print(f"Valid commands: {', '.join(commands)}")
        sys.exit(1)

    handler()


if __name__ == "__main__":
    main()