path: root/main.py

#!/usr/bin/env python3
"""
slopbot - Zulip bot that replies with stock prices when it sees $TICKER patterns,
and delegates @mention responses to a one-shot agent per message.

Usage:
    python main.py

Requires a .zuliprc file in the same directory (see .zuliprc.example).
"""

import html
import logging
import os
import re
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from threading import Semaphore

import yfinance as yf
import zulip

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("slopbot")

# Regex to match stock tickers like $NVDA, $AAPL, $BRK
TICKER_RE = re.compile(r"\$([A-Z]{1,5})")

# Regex to detect @mention of slop / slopbot
MENTION_RE = re.compile(r"@\*\*slop[^*]*\*\*", re.IGNORECASE)

BASE_DIR = Path.cwd()

AGENT_TIMEOUT = int(os.environ.get("SLOPBOT_AGENT_TIMEOUT", "120"))
AGENT_MODEL = os.environ.get("SLOPBOT_AGENT_MODEL", "parasail/kimi-k26")
AGENT_BIN = os.environ.get("SLOPBOT_AGENT", "agent")
AGENT_PATH = os.environ.get("SLOPBOT_AGENT_PATH", os.environ.get("PATH", ""))
WORKERS = int(os.environ.get("SLOPBOT_WORKERS", "4"))
AGENT_WORKERS = int(os.environ.get("SLOPBOT_AGENT_WORKERS", "1"))
HISTORY_LIMIT = int(os.environ.get("SLOPBOT_HISTORY_LIMIT", "20"))
SYSTEM_PROMPT_PATH = Path(os.environ.get("SLOPBOT_SYSTEM_PROMPT", "prompts/system.md"))
if not SYSTEM_PROMPT_PATH.is_absolute():
    SYSTEM_PROMPT_PATH = BASE_DIR / SYSTEM_PROMPT_PATH

SYSTEM_PROMPT = SYSTEM_PROMPT_PATH.read_text().strip()

log.info("Connecting to Zulip...")
client = zulip.Client(config_file=".zuliprc")
log.info("Connected as %s", client.email)

executor = ThreadPoolExecutor(max_workers=WORKERS, thread_name_prefix="slopbot")
agent_slots = Semaphore(AGENT_WORKERS)


# ---------------------------------------------------------------------------
# Zulip context
# ---------------------------------------------------------------------------

def fetch_topic_history(msg: dict, limit: int = HISTORY_LIMIT) -> str:
    """Fetch recent messages from the same Zulip stream/topic for LLM context."""
    if msg.get("type") != "stream":
        return ""

    stream = msg.get("display_recipient")
    topic = msg.get("subject") or msg.get("topic")
    if not stream or not topic:
        return ""

    try:
        resp = client.get_messages({
            "anchor": msg.get("id", "newest"),
            "num_before": limit,
            "num_after": 0,
            "narrow": [
                {"operator": "stream", "operand": stream},
                {"operator": "topic", "operand": topic},
            ],
        })
    except Exception as e:
        log.warning("Failed to fetch Zulip topic history: %s", e)
        return ""

    if resp.get("result") != "success":
        log.warning("Zulip history request failed: %s", resp)
        return ""

    lines = []
    for item in resp.get("messages", [])[-limit:]:
        if item.get("id") == msg.get("id"):
            continue
        sender = item.get("sender_full_name") or item.get("sender_email") or "someone"
        content = html.unescape(re.sub(r"<[^>]+>", " ", item.get("content", "")))
        content = re.sub(r"\s+", " ", content).strip()
        if content:
            lines.append(f"{sender} said: {content}")

    if not lines:
        return ""
    return "Recent topic history:\n" + "\n".join(lines)


# ---------------------------------------------------------------------------
# agent one-shot helper
# ---------------------------------------------------------------------------

def _agent(prompt: str) -> tuple[int, str, str]:
    """Run agent once; return (returncode, stdout, stderr)."""
    env = os.environ.copy()
    env["PATH"] = AGENT_PATH
    result = subprocess.run(
        [
            AGENT_BIN,
            "--model", AGENT_MODEL,
            "--eval-cwd", str(BASE_DIR),
            prompt,
        ],
        capture_output=True, text=True, timeout=AGENT_TIMEOUT, env=env,
    )
    return result.returncode, result.stdout.strip(), result.stderr.strip()


def _unique_session_name(msg: dict) -> str:
    """Build a unique, short-lived session name for a single message."""
    ts = int(time.time() * 1000)
    if msg.get("type") == "private":
        sender_id = str(msg.get("sender_id") or "dm")
        safe_id = re.sub(r"[^a-z0-9-]", "-", sender_id.lower())[:20].strip("-")
        return f"slop-dm-{safe_id}-{ts}"
    stream_id = msg.get("stream_id", "0")
    topic = msg.get("subject") or msg.get("topic") or "general"
    slug = re.sub(r"[^a-z0-9-]", "-", topic.lower())[:20].strip("-")
    return f"slop-{stream_id}-{slug}-{ts}"


def ask_agent(msg: dict, content: str) -> str:
    """Run a one-shot agent for a single mention."""
    name = _unique_session_name(msg)
    sender = msg.get("sender_full_name") or msg.get("sender_email") or "user"

    topic_history = fetch_topic_history(msg)
    prompt_parts = [SYSTEM_PROMPT]
    if topic_history:
        prompt_parts.append(f"\n\n{topic_history}")
    prompt_parts.append("\nYou have access to the agent shell tools if you need them.")
    prompt_parts.append("\nRespond to the latest Zulip message.")
    prompt_parts.append(f"\nLatest message from {sender}:\n{content}")
    prompt = "\n".join(prompt_parts)

    log.info("Running one-shot agent: %s", name)
    try:
        with agent_slots:
            rc, reply, err = _agent(prompt)
    except subprocess.TimeoutExpired:
        log.warning("agent timed out for %s", name)
        return "sorry, agent timed out"
    except Exception as e:
        log.warning("agent failed for %s: %s", name, e)
        return "sorry, couldn't reach the agent"

    if rc != 0 or not reply:
        log.warning("agent failed for %s (rc=%d): %s", name, rc, err)
        return "sorry, agent timed out"

    return reply


# ---------------------------------------------------------------------------
# Price fetcher
# ---------------------------------------------------------------------------

def fetch_price(ticker: str) -> str:
    """Return a formatted price string for a ticker, or an error message."""
    log.info("Fetching price for %s", ticker)
    try:
        price = yf.Ticker(ticker).fast_info["lastPrice"]
        result = f"${ticker}: ${price:,.2f}"
        log.info("  -> %s", result)
        return result
    except Exception as e:
        log.warning("  -> %s unavailable: %s", ticker, e)
        return f"${ticker}: unavailable ({e})"


# ---------------------------------------------------------------------------
# Event handler
# ---------------------------------------------------------------------------

def handle_message(msg: dict) -> None:
    """Process a Zulip message and reply to mentions and/or tickers."""
    try:
        if msg["sender_email"] == client.email:
            return

        content = msg.get("content", "")
        has_mention = MENTION_RE.search(content)
        tickers = TICKER_RE.findall(content.upper())

        if not has_mention and not tickers:
            return

        if has_mention:
            cleaned = MENTION_RE.sub("", content).strip()
            reply = ask_agent(msg, cleaned)
            if tickers:
                prices = [fetch_price(t) for t in dict.fromkeys(tickers)]
                reply += "\n\n" + "\n".join(prices)
        else:
            lines = [fetch_price(t) for t in dict.fromkeys(tickers)]
            reply = "\n".join(lines)

        client.send_message({
            "type": msg["type"],
            "to": msg["stream_id"] if msg["type"] == "stream" else msg["sender_email"],
            "topic": msg.get("subject", ""),
            "content": reply,
        })
        log.info("Reply sent.")
    except Exception:
        log.exception("Failed to handle message %s", msg.get("id"))


def handle_event(event: dict) -> None:
    """Queue Zulip message work without blocking the event loop."""
    if event.get("type") != "message":
        return
    executor.submit(handle_message, event["message"])


def main() -> None:
    log.info("Listening for messages...")
    client.call_on_each_event(handle_event, event_types=["message"])


if __name__ == "__main__":
    main()