diff options
| -rw-r--r-- | .envrc | 2 | ||||
| -rw-r--r-- | main.py | 208 | ||||
| -rw-r--r-- | requirements.txt | 2 | ||||
| -rw-r--r-- | shell.nix | 2 |
4 files changed, 208 insertions, 6 deletions
@@ -0,0 +1,2 @@ +export PARASAIL_API_KEY="psk-parasailzmJ0-Yj3t3mofugwBetQifS0r" +export KAGI_API_KEY="9Iyt2Wy_ZQFJzkFPL7t6AbvlrnlY8uTBcxC8fhqakqk.MxMiJl0SPTbtjRuZmaSTGi8LQBvq8dajYN29qdOtJ8Y" @@ -1,6 +1,7 @@ #!/usr/bin/env python3 """ -slopbot - Zulip bot that replies with stock prices when it sees $TICKER patterns. +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 @@ -8,10 +9,15 @@ Usage: 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, @@ -23,11 +29,188 @@ 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 +# --------------------------------------------------------------------------- + +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("Calling Parasail for LLM response") + # Prefer a small/fast model; fall back to Llama if unavailable + try: + r = requests.get( + "https://api.parasail.io/v1/models", + headers={"Authorization": f"Bearer {parasail_key}"}, + timeout=8, + ) + model_ids = [m["id"] for m in r.json().get("data", [])] + log.info("Parasail available models: %s", model_ids) + 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" + + c = OpenAI(api_key=parasail_key, base_url="https://api.parasail.io/v1") + return c, model + + log.info("Calling Ollama for LLM response") + c = OpenAI(api_key="ollama", base_url="http://localhost:11434/v1") + return c, "qwen2.5:3b" + + +# --------------------------------------------------------------------------- +# 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) @@ -41,8 +224,12 @@ def fetch_price(ticker: str) -> str: return f"${ticker}: unavailable ({e})" +# --------------------------------------------------------------------------- +# Event handler +# --------------------------------------------------------------------------- + def handle_event(event: dict) -> None: - """Process a Zulip event and reply if tickers are found.""" + """Process a Zulip event and reply to mentions and/or tickers.""" if event.get("type") != "message": return @@ -53,21 +240,30 @@ def handle_event(event: dict) -> None: return content = msg.get("content", "") + has_mention = MENTION_RE.search(content) tickers = TICKER_RE.findall(content.upper()) - if not tickers: + if not has_mention and not tickers: return log.info( - "Message from %s in %s/%s: found tickers %s", + "Message from %s in %s/%s: mention=%s tickers=%s", msg["sender_email"], msg.get("display_recipient", "?"), msg.get("subject", "?"), + bool(has_mention), tickers, ) - lines = [fetch_price(t) for t in dict.fromkeys(tickers)] # dedup, preserve order - reply = "\n".join(lines) + 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"], diff --git a/requirements.txt b/requirements.txt index 9184f43..437e60e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,4 @@ zulip>=0.9.0 yfinance>=0.2.0 +openai>=1.0.0 +requests>=2.28.0 @@ -5,6 +5,8 @@ pkgs.mkShell { (pkgs.python3.withPackages (ps: [ ps.zulip ps.yfinance + ps.openai + ps.requests ])) ]; |