path: root/main.py

#!/usr/bin/env python3
"""
slopbot - Zulip bot that replies with stock prices when it sees $TICKER patterns,
and responds to @mention with an agentic LLM (Parasail/Ollama fallback + Kagi search).

Usage:
    python main.py

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

import json
import logging
import os
import re

import requests
import zulip
import yfinance as yf
from openai import OpenAI

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)

MAX_ITER = 5

SYSTEM_PROMPT = (
    "You are slop-bot, a helpful assistant in the meshheads.org Zulip chat. "
    "You answer questions about stocks, finance, and general topics. Be concise and conversational. "
    "Use kagi_search to find current information when needed. "
    "Keep replies short (2-4 sentences max unless asked for more)."
)

KAGI_TOOL = {
    "type": "function",
    "function": {
        "name": "kagi_search",
        "description": "Search the web for current information. Use this for news, prices, or anything time-sensitive.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "Search query"}
            },
            "required": ["query"],
        },
    },
}

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


# ---------------------------------------------------------------------------
# LLM client setup: Parasail primary, Ollama fallback
# ---------------------------------------------------------------------------

# Model config: big Qwen3 on Parasail, smaller local Qwen2.5 as fallback
PARASAIL_MODEL = "parasail-qwen3-235b-a22b-instruct-2507"
OLLAMA_MODEL = "qwen2.5:14b-instruct-q4_K_M"


def _make_llm_client() -> tuple[OpenAI, str]:
    """Return (openai_client, model_name) for the best available backend."""
    parasail_key = os.environ.get("PARASAIL_API_KEY", "")
    if parasail_key:
        log.info("Using Parasail model: %s", PARASAIL_MODEL)
        c = OpenAI(api_key=parasail_key, base_url="https://api.parasail.io/v1")
        return c, PARASAIL_MODEL

    log.info("Using Ollama model: %s", OLLAMA_MODEL)
    c = OpenAI(api_key="ollama", base_url="http://localhost:11434/v1")
    return c, OLLAMA_MODEL


# ---------------------------------------------------------------------------
# Kagi search tool
# ---------------------------------------------------------------------------

def kagi_search(query: str, limit: int = 3) -> str:
    log.info("Tool call: kagi_search(%s)", query)
    kagi_key = os.environ.get("KAGI_API_KEY", "")
    if not kagi_key:
        log.warning("KAGI_API_KEY not set; skipping search")
        return ""
    try:
        resp = requests.get(
            "https://kagi.com/api/v0/search",
            headers={"Authorization": f"Bot {kagi_key}"},
            params={"q": query, "limit": limit},
            timeout=10,
        )
        results = resp.json().get("data", [])
        log.info("Kagi returned %d results", len(results[:limit]))
        return "\n".join(
            f"- {r['title']}: {r.get('snippet', '')}" for r in results[:limit]
        )
    except Exception as e:
        log.warning("kagi_search failed: %s", e)
        return ""


# ---------------------------------------------------------------------------
# Agentic LLM loop
# ---------------------------------------------------------------------------

def ask_llm(content: str) -> str:
    log.info("LLM mention detected: %s", content)
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": content},
    ]

    # Try Parasail first; fall back to Ollama on error
    backends = []
    parasail_key = os.environ.get("PARASAIL_API_KEY", "")
    if parasail_key:
        backends.append(("Parasail", lambda: _build_parasail(parasail_key)))
    backends.append(("Ollama", lambda: (
        OpenAI(api_key="ollama", base_url="http://localhost:11434/v1"),
        "qwen2.5:3b",
    )))

    last_error = None
    for backend_name, get_client in backends:
        log.info("Calling %s for LLM response", backend_name)
        try:
            llm, model = get_client()
            result = _run_agentic_loop(llm, model, messages)
            log.info("LLM response: %s", result)
            return result
        except Exception as e:
            log.warning("%s failed: %s", backend_name, e)
            last_error = e

    log.error("All LLM backends failed: %s", last_error)
    return "sorry, LLM unavailable right now"


def _build_parasail(key: str) -> tuple[OpenAI, str]:
    """Build Parasail client and resolve best model."""
    try:
        r = requests.get(
            "https://api.parasail.io/v1/models",
            headers={"Authorization": f"Bearer {key}"},
            timeout=8,
        )
        model_ids = [m["id"] for m in r.json().get("data", [])]
        preferred = ["qwen3-4b", "meta-llama/Llama-3.1-8B-Instruct"]
        model = next((m for m in preferred if m in model_ids), None)
        if model is None and model_ids:
            model = model_ids[0]
        model = model or "meta-llama/Llama-3.1-8B-Instruct"
    except Exception as e:
        log.warning("Could not fetch Parasail models (%s); using default", e)
        model = "meta-llama/Llama-3.1-8B-Instruct"
    return OpenAI(api_key=key, base_url="https://api.parasail.io/v1"), model


def _run_agentic_loop(llm: OpenAI, model: str, messages: list) -> str:
    msg = None
    for _ in range(MAX_ITER):
        resp = llm.chat.completions.create(
            model=model,
            messages=messages,
            tools=[KAGI_TOOL],
            tool_choice="auto",
        )
        msg = resp.choices[0].message
        if not msg.tool_calls:
            return msg.content or "sorry, I couldn't generate a response"

        # Execute tool calls and append results
        messages.append(msg)
        for tc in msg.tool_calls:
            args = json.loads(tc.function.arguments)
            result = kagi_search(args["query"])
            messages.append({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": result,
            })

    return (msg.content if msg else None) or "sorry, I couldn't generate a response"


# ---------------------------------------------------------------------------
# Price fetcher (unchanged)
# ---------------------------------------------------------------------------

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_event(event: dict) -> None:
    """Process a Zulip event and reply to mentions and/or tickers."""
    if event.get("type") != "message":
        return

    msg = event["message"]

    # Don't reply to ourselves
    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

    log.info(
        "Message from %s in %s/%s: mention=%s tickers=%s",
        msg["sender_email"],
        msg.get("display_recipient", "?"),
        msg.get("subject", "?"),
        bool(has_mention),
        tickers,
    )

    if has_mention:
        cleaned = MENTION_RE.sub("", content).strip()
        reply = ask_llm(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.")


log.info("Listening for messages...")
client.call_on_each_event(handle_event, event_types=["message"])