diff --git a/CLAUDE.md b/CLAUDE.md index fac1272..2079044 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -15,7 +15,7 @@ MCP data tools and web search. User-facing docs are in `README.md`; this file is - `config.py`: `Settings` from env (`load()`), `ENV_VARS` registry, tuning constants. The only place that knows the provider. Every env var must appear in README.md (a test enforces it). - `store.py`: SQLite `Store`: history (`messages` shared by groups, `dm_messages` per bot), users, kv, - holding-news state, `usage` ledger. Thread-safe; call through `asyncio.to_thread`. + holding-news state, `usage` ledger, `period_cache` (finished Sharesight periods, kept until flushed). Thread-safe; call through `asyncio.to_thread`. - `msgtext.py` sender_name/describe. `history.py` transcript lines. `textfmt.py` HTML helpers (md_to_html, split_html, plain_text, TG_TAG_RE). `tickers.py` Yahoo links (+ `tickers_data.py`). - `mcp/`: `server.py` (MCPServer lifecycle, TTL cache, `Registry`), `schema.py` (neutral `ToolDef`, @@ -24,7 +24,7 @@ MCP data tools and web search. User-facing docs are in `README.md`; this file is `openrouter.py` + `zai.py` (the only provider-specific code), `runner.py` (THE tool loop), `policy.py` (THE retry ladder), `gate.py` (which tools a message needs). - `prompts.py` all prompts. `ask.py` one entry point: route -> Request -> policy.ask -> usage -> NEEDS_TOOLS rerun. - `telegram/`: `handlers.py` (trigger -> request -> reply), `send.py`, `draft.py` (typing + streaming - drafts), `commands.py` (/credits), `access.py` (who is an admin). `features/`: `holding_news.py`, `movers.py`, `dm_buttons.py`, `post.py`. + drafts), `commands.py` (/credits), `access.py` (who is an admin). `features/`: `holding_news.py`, `movers.py`, `dm_buttons.py`, `post.py`, `periods/` (local `period_returns` tool, cache-first; `windows.py`, `rows.py`, `backfill.py` = PERIOD_BACKFILL loop). - `app.py` wiring/startup; `context.py` `Ctx` (settings, store, backend, registry, bot). - `lib/` holds only what the bot needs to run. Developer helpers live in `tools/`: `report.py` (`make prompt-report`), `capture.py` (real MCP output for fixtures); the bot never imports them. @@ -57,6 +57,7 @@ MCP data tools and web search. User-facing docs are in `README.md`; this file is - z.ai: `tool_choice` is ignored (even "none", so the backend omits tools instead), its built-in `web_search` is ignored whenever function tools are present (so search is a function tool run through its search API), thinking can't be disabled (default `reasoning_effort` low), an invalid search recency is accepted silently (so the tool schema fixes the list), and usage has no cost (price table). +- Finished Sharesight periods are cached forever (`period_cache`); after a backdated change run `/flushcache [portfolio] [year|month|day]` or `python bot.py --flush-cache ...`. The unfinished period is never cached. `_keep_sales` in tool args is reserved for the bot's own tools (the runner strips it from the model's). - `must_search` requests are never retried without tools (answering from memory would invent news). - Telegram HTML: only b/i/u/s/code/pre/a/blockquote/tg-spoiler. `Draft` strips tags (half-written HTML is rejected). - No live Telegram/xAI/OpenRouter/MCP access in tests or in this sandbox; say so when behaviour can't be checked. diff --git a/README.md b/README.md index d637524..d8cceb2 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,23 @@ The optional `label` is ignored; it only appears in `ps`/`top` so you can tell i A systemd unit just needs `ExecStart=/usr/bin/python3 /path/to/bot.py mimo` and an `Environment=` or `EnvironmentFile=` for the variables below. +### Running under systemd +With `StateDirectory=`, systemd sets `STATE_DIRECTORY` and the bot keeps its database there (`DATA_DIR` +defaults to it), so the code directory can stay read-only. `mcp_servers.json` is read from the checkout +unless `MCP_CONFIG` says otherwise. Copy an existing `chat_log.db` (and `-wal`/`-shm`) into the state +directory before the first start. + +```ini +[Service] +DynamicUser=yes +WorkingDirectory=/opt/grokbot +EnvironmentFile=/opt/grokbot.env +StateDirectory=grokbot +Environment=PYTHONUNBUFFERED=1 +ExecStart=/opt/grokbot/.venv/bin/python bot.py +Restart=on-failure +``` + ## Providers - `OPENROUTER_API_KEY` set -> OpenRouter (default model `z-ai/glm-5.3-flash`). - `XAI_API_KEY` set -> xAI (default model `grok-4.3`). @@ -161,6 +178,7 @@ Each is off by default. The holding-news DM and the movers reply are switched on | `/start` | DM | Greeting (and the buttons if enabled) | | `/holdingnews` | DM, if holding news is on | Run the holding-news check now | | `/credits` | an admin (see below), in a group or a DM | Provider balance (OpenRouter; xAI and z.ai have none to show), then requests, tokens (and % cached) and cost per model for the last 24 h and 7 days | +| `/flushcache [portfolio_id] [year or month or day]` | an admin | Forget cached Sharesight period returns: all, or of one portfolio, or those overlapping a year (`2024`), month (`2024-05`) or day (`2024-05-17`), e.g. `/flushcache 684141 2024`; see [Caching Sharesight history](#caching-sharesight-history) | **Admins** are the user IDs in `ADMIN_CHAT_IDS` plus anyone who administers a group the bot is in (checked with Telegram, cached for five minutes). The bot only knows a group once someone has written in it, or it was added, since @@ -170,6 +188,23 @@ gets "Portfolio data is for admins only." Holding-news DMs, buttons and down-ser unaffected. With nobody in `ADMIN_CHAT_IDS` and no known group, nobody can DM the bot or reach Sharesight: set `ADMIN_ONLY=off` to allow everyone. +## Caching Sharesight history +Questions like "best day / week / month / quarter / year since inception" are answered by the `period_returns` tool +(admins only), which asks Sharesight for one report per period (about 2.4 s each). Every **finished** period is +stored in the SQLite database (next to the chat history, so it survives upgrades with `DATA_DIR`) and **never +expires**: the current, unfinished period is always fetched fresh. Past history rarely changes; after a backdated +trade or correction in Sharesight, flush the cache with `/flushcache [portfolio_id] [year|month|day]` (an admin, in +Telegram) or `python bot.py --flush-cache [portfolio_id] [year|month|day]` (needs only `DB_PATH`/`DATA_DIR`, not the +bot's keys). With a year, month or day, only the periods overlapping it go (every day, week, month, quarter and year +that contains it), and the background fill re-fetches just those. + +Without the cache, a request may fetch at most 64 days, 26 weeks or 120 months/quarters/years that aren't cached; a +longer range is refused with the earliest start that fits. Set `PERIOD_BACKFILL=[,...]` (or `all`) to +fill the cache in the background: months first, then quarters, years, weeks and days (each type for every portfolio before the next), up to two Sharesight calls at a time, +never while a user's request is using the server (about 1,650 calls per portfolio, an hour or two the first time; progress is logged and it resumes after a restart). Once filled, any range is instant and the +limits above no longer matter. Every six hours it fetches whatever newly finished. If five fetches in a row fail (Sharesight down or throttling), +the pass stops and tries again after ten minutes; windows Sharesight simply has no data for are skipped until a restart. + ## Caching and token saving **Provider prompt caching** (xAI and OpenRouter keep a conversation on the server holding its cached prompt; z.ai caches automatically and gets no key): xAI gets `prompt_cache_key` and the `x-grok-conv-id` header; OpenRouter gets `session_id`, @@ -246,12 +281,14 @@ Set `TELEGRAM_BOT_TOKEN` and exactly one of the three LLM keys. |---|---| | `ADMIN_CHAT_IDS` | Telegram IDs of the bot's admins (comma/space separated). Your user ID gets a message when an MCP server goes down and may use `/credits`; a group's chat ID gets the down-server messages only. Admins of the groups the bot is in count as admins too. | | `ADMIN_ONLY` | on|off. Only admins may DM the bot or get Sharesight data (default on). | +| `DATA_DIR` | Directory for the bot's data (the database); a relative `DB_PATH` resolves inside it, an absolute one is used as given. Default: systemd's `STATE_DIRECTORY` (see [Running under systemd](#running-under-systemd)), else the working directory. In Docker, set `DATA_DIR=/data` and bind-mount a host directory there so history survives upgrades. | | `DB_PATH` | SQLite file for chat history (default chat_log.db). Instances may share it. | | `HISTORY_LIMIT` | Messages of chat history in the prompt; the window is HISTORY_LIMIT to 1.5x (default 20). | | `MAX_TOKENS` | Reply cap in tokens, reasoning included (default 3000). | -| `MCP_CONFIG` | MCP server config file (default mcp_servers.json; missing = no MCP tools). | +| `MCP_CONFIG` | MCP server config file (default mcp_servers.json next to bot.py, whatever `DATA_DIR` is; missing = no MCP tools). | | `MOVERS_BOTS` | Usernames of bots whose end-of-day big-movers lists get explained. Empty (default) = feature off. | | `PORTFOLIO_NAMES` | Comma list of Sharesight portfolio names; a message naming one is treated as a portfolio question (default: the recipients' portfolio names). | +| `PERIOD_BACKFILL` | Comma list of Sharesight portfolio names (or `all`) whose returns per day, week, month, quarter and year are fetched into the local cache in the background (about 1,650 Sharesight calls per portfolio the first time, a few hours), so "best day since inception" answers instantly. Default `all`; list names to limit it, or `off` to disable. See [Caching Sharesight history](#caching-sharesight-history). | | `POST_TO_GROUPS_FROM_DM` | on|off. A group admin can have the bot post in that group from a DM (default off). | | `SEARCH_MODEL` | OpenRouter only: model that runs the searches (default xiaomi/mimo-v2.6-flash:online). | | `SHARESIGHT_HOLDING_NEWS_RECIPIENTS` | Comma list of portfolio:telegram_username pairs to notify. Empty (default) = daily holding-news DM off. | diff --git a/lib/app.py b/lib/app.py index d8b6bd5..96f977d 100644 --- a/lib/app.py +++ b/lib/app.py @@ -8,7 +8,7 @@ import json import logging import os -import subprocess +import re import sys from collections.abc import Sequence from pathlib import Path @@ -29,6 +29,8 @@ from . import config from .context import Ctx from .features import dm_buttons, holding_news, post +from .features.periods import backfill +from .features.periods.windows import describe_scope, flush_scope from .llm.base import Backend from .llm.openrouter import OpenRouterBackend from .llm.xai import XaiBackend @@ -52,6 +54,14 @@ def setup_logging() -> None: format="%(levelname)s %(message)s" if os.getenv("JOURNAL_STREAM") else "%(asctime)s %(levelname)s %(message)s") logging.getLogger("httpx").setLevel(logging.WARNING) + for handler in logging.getLogger().handlers: + handler.addFilter(_drop_protocol_line) + + +def _drop_protocol_line(record: logging.LogRecord) -> bool: + """The MCP library's own "Negotiated protocol version" line doesn't say which server it is about; the + bot logs "MCP : connected, protocol ..." instead.""" + return not record.getMessage().startswith("Negotiated protocol version") async def alert_down(ctx: Ctx, server: MCPServer) -> None: @@ -101,6 +111,8 @@ async def post_init(app: Application) -> None: background.append(asyncio.create_task(holding_news.daily_loop(ctx))) if st.dm_buttons: background.append(asyncio.create_task(dm_buttons.push_on_startup(ctx))) + if st.period_backfill: + background.append(asyncio.create_task(backfill.loop(ctx))) async def post_shutdown(app: Application) -> None: for task in background: @@ -124,6 +136,7 @@ async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Commands first: the first matching handler wins, and the catch-all would swallow them. app.add_handler(CommandHandler("credits", command(commands.credits))) + app.add_handler(CommandHandler("flushcache", command(commands.flush_cache))) app.add_handler(MessageHandler( (filters.ChatType.GROUPS | filters.ChatType.PRIVATE) & (filters.UpdateType.MESSAGE | filters.UpdateType.EDITED_MESSAGE) @@ -141,14 +154,38 @@ def allowed_updates(st: config.Settings) -> list[str]: + (["my_chat_member"] if st.post_to_groups or st.admin_only else [])) -def git_hash() -> str: - """Short hash of the commit this checkout is on, or "" if it can't be read (no git, no checkout).""" +def git_hash(root: Path | None = None) -> str: + """Short hash of the commit this checkout is on, or "" if it can't be read. Read from .git directly: the + git command refuses a checkout owned by another user (a systemd DynamicUser runs as one).""" + root = root or Path(__file__).resolve().parent.parent try: - done = subprocess.run(["git", "rev-parse", "--short", "HEAD"], cwd=Path(__file__).resolve().parent.parent, - capture_output=True, text=True, timeout=2) - except (OSError, subprocess.SubprocessError): + git = root / ".git" + if git.is_file(): # a linked worktree: ".git" holds "gitdir: " + git = (root / git.read_text().split(":", 1)[1].strip()).resolve() + common = git / (git / "commondir").read_text().strip() if (git / "commondir").exists() else git + head = (git / "HEAD").read_text().strip() + if head.startswith("ref:"): + ref = head.split(None, 1)[1] + try: + head = (common / ref).read_text().strip() + except FileNotFoundError: # packed by `git gc` + packed = (common / "packed-refs").read_text().splitlines() + head = next((line.split()[0] for line in packed if line.endswith(" " + ref)), "") + return head[:7] if re.fullmatch(r"[0-9a-f]{40,64}", head) else "" + except (OSError, IndexError): return "" - return done.stdout.strip() if done.returncode == 0 else "" + + +def flush_cache(words: list[str]) -> str: + """--flush-cache [portfolio ID] [year|month|day]: needs only the database path, not the bot's keys. + Returns the message to print.""" + scope = flush_scope(words) + if scope is None: + return "Usage: --flush-cache [portfolio ID] [2024 | 2024-05 | 2024-05-17]" + pid, first, last = scope + path = config.data_path(os.environ, "DB_PATH", "chat_log.db") + n = Store(path).period_flush(pid, first and first.isoformat(), last and last.isoformat()) + return f"Flushed {n} cached periods ({describe_scope(pid, first, last)})." def main(argv: Sequence[str] | None = None) -> None: @@ -156,7 +193,14 @@ def main(argv: Sequence[str] | None = None) -> None: # side by side with different environments (python bot.py mimo). parser = argparse.ArgumentParser(description="Telegram group bot that answers @mentions with an LLM.") parser.add_argument("label", nargs="?", help="ignored; identifies this instance in ps/top") + parser.add_argument("--flush-cache", nargs="*", metavar="WHAT", + help="delete the cached Sharesight period returns and exit: all, or of a portfolio ID " + "and/or a year, month or day (2024, 2024-05, 2024-05-17)") args, _ = parser.parse_known_args(argv) + if args.flush_cache is not None: + message = flush_cache(args.flush_cache) + print(message) + sys.exit(0 if message.startswith("Flushed") else 1) try: st = config.load() except config.ConfigError as e: @@ -168,4 +212,6 @@ def main(argv: Sequence[str] | None = None) -> None: log.info("Starting%s%s: %s on %s, reasoning %s, MCP servers: %s", f" instance {args.label}" if args.label else "", f" ({commit})" if commit else "", st.model, st.provider, st.reasoning or "default", ", ".join(ctx.registry.servers) or "none") + if st.period_backfill: + log.info("Period backfill: on for %s", ", ".join(sorted(st.period_backfill))) app.run_polling(allowed_updates=allowed_updates(st)) diff --git a/lib/config.py b/lib/config.py index 3368092..fdf6509 100644 --- a/lib/config.py +++ b/lib/config.py @@ -6,6 +6,7 @@ import re from collections.abc import Mapping from dataclasses import dataclass +from pathlib import Path from zoneinfo import ZoneInfo # Every environment variable the bot reads, with a one-line description. README.md must @@ -25,12 +26,14 @@ # Advanced and optional "ADMIN_CHAT_IDS": "Telegram IDs of the bot's admins (comma/space separated). Your user ID gets MCP-down alerts and may use /credits; a group's chat ID gets the alerts only. Admins of the groups the bot is in count as admins too.", "ADMIN_ONLY": "on|off. Only admins may DM the bot or get Sharesight data (default on).", + "DATA_DIR": "Directory for the bot's data (the database); a relative DB_PATH resolves inside it. Default: systemd's STATE_DIRECTORY, else the working directory.", "DB_PATH": "SQLite file for chat history (default chat_log.db). Instances may share it.", "HISTORY_LIMIT": "Messages of chat history in the prompt; the window is HISTORY_LIMIT to 1.5x (default 20).", "MAX_TOKENS": "Reply cap in tokens, reasoning included (default 3000).", - "MCP_CONFIG": "MCP server config file (default mcp_servers.json; missing = no MCP tools).", + "MCP_CONFIG": "MCP server config file (default mcp_servers.json next to bot.py; missing = no MCP tools).", "MOVERS_BOTS": "Usernames of bots whose end-of-day big-movers lists get explained. Empty (default) = feature off.", "PORTFOLIO_NAMES": "Comma list of Sharesight portfolio names; a message naming one is treated as a portfolio question (default: the recipients' portfolio names).", + "PERIOD_BACKFILL": "Comma list of Sharesight portfolio names (or all) whose daily/weekly/monthly/quarterly/yearly returns are fetched into the local cache in the background, so period questions answer instantly. Default all; list names to limit it, or off to disable.", "POST_TO_GROUPS_FROM_DM": "on|off. Lets an admin of a group the bot is in make it post there from a DM (default off).", "SEARCH_MODEL": "OpenRouter only: model that runs the searches (default xiaomi/mimo-v2.6-flash:online).", "SHARESIGHT_HOLDING_NEWS_RECIPIENTS": "Comma list of portfolio:telegram_username pairs to notify. Empty (default) = daily holding-news DM off.", @@ -81,6 +84,7 @@ class Settings: holding_news_time: str holding_news_recipients: dict[str, str] portfolio_names: frozenset[str] + period_backfill: frozenset[str] token_saver: bool @property @@ -108,6 +112,22 @@ def history_own_line_max(self) -> int | None: return 160 if self.token_saver else None +ROOT = Path(__file__).resolve().parent.parent # the checkout, where bot.py is + + +def data_path(env: Mapping[str, str], var: str, default: str) -> str: + """A data file named by env var `var` (or `default`), inside DATA_DIR, which defaults to systemd's + STATE_DIRECTORY (first entry if it lists several). An absolute path is used as given.""" + name = env.get(var, "").strip() or default + data_dir = env.get("DATA_DIR", "").strip() or env.get("STATE_DIRECTORY", "").split(":")[0].strip() + return os.path.join(data_dir, name) if data_dir else name # join keeps an absolute `name` as is + + +def mcp_config_path(env: Mapping[str, str]) -> str: + """MCP_CONFIG, default mcp_servers.json next to bot.py. Deployment config, so DATA_DIR doesn't apply.""" + return env.get("MCP_CONFIG", "").strip() or str(ROOT / "mcp_servers.json") + + def _flag(env: Mapping[str, str], name: str, default: bool = False) -> bool: value = env.get(name, "").strip().lower() if not value: @@ -130,6 +150,11 @@ def _names(text: str) -> frozenset[str]: return frozenset(n.strip().lower() for n in re.split(r"[,\n]", text) if n.strip()) +def _backfill(text: str) -> frozenset[str]: + names = _names(text) + return frozenset() if names & {"off", "none", "false", "0"} else names + + def load(env: Mapping[str, str] | None = None) -> Settings: """Build Settings from `env` (default: os.environ). Raises ConfigError if unusable.""" env = os.environ if env is None else env @@ -164,9 +189,9 @@ def load(env: Mapping[str, str] | None = None) -> Settings: search_model=search_model, max_tokens=int(env.get("MAX_TOKENS") or 3000), history_limit=int(env.get("HISTORY_LIMIT") or 20), - db_path=env.get("DB_PATH", "chat_log.db"), + db_path=data_path(env, "DB_PATH", "chat_log.db"), tz=ZoneInfo(env.get("BOT_TZ", "UTC")), - mcp_config=env.get("MCP_CONFIG", "mcp_servers.json"), + mcp_config=mcp_config_path(env), admin_chats=_ids(env.get("ADMIN_CHAT_IDS", "")), admin_only=_flag(env, "ADMIN_ONLY", default=True), movers_bots=frozenset( @@ -178,4 +203,5 @@ def load(env: Mapping[str, str] | None = None) -> Settings: holding_news_recipients=recipients, portfolio_names=_names(env.get("PORTFOLIO_NAMES", "")) or frozenset(recipients), token_saver=saver, + period_backfill=_backfill(env.get("PERIOD_BACKFILL", "all")), ) diff --git a/lib/features/periods/__init__.py b/lib/features/periods/__init__.py new file mode 100644 index 0000000..b93afcf --- /dev/null +++ b/lib/features/periods/__init__.py @@ -0,0 +1,125 @@ +"""period_returns: a Sharesight portfolio's return per day, week, month, quarter or year, in one tool call +instead of one model-driven report per period. Finished periods are cached in SQLite (see backfill.py).""" + +from __future__ import annotations + +import asyncio +import json +import logging +from datetime import date, datetime + +from ...context import Ctx +from ...mcp.schema import ToolDef +from .rows import drivers, fetch_row +from .windows import PERIODS, key, parse_date, period_windows, refusal + +log = logging.getLogger("bot") + +TOOL = ToolDef( + "period_returns", + "Total return of a Sharesight portfolio for each day, week, month, quarter or year, as a percentage and a " + "money gain, with the best and worst one by each measure and the holdings behind them (a long range lists " + "only the top and bottom five, so one call answers 'best day since inception'). Call it once per portfolio " + "and cover every portfolio asked about. Use it for best/worst day, week, month, quarter or year, or any " + "period-by-period question, instead of fetching one get_performance_report per period. Sold holdings are " + "included. Omit start for 'since inception'. History that isn't cached yet is fetched live, which is slow " + "for days and weeks: if the tool refuses a range, tell the user why and offer the start it suggests.", + {"type": "object", "required": ["portfolio_id"], "properties": { + "portfolio_id": {"type": "number", "description": "From list_portfolios"}, + "period": {"type": "string", "enum": list(PERIODS), "description": "Default month"}, + "start": {"type": "string", "description": "First day: YYYY-MM-DD, YYYY-MM or YYYY (default: inception)"}, + "end": {"type": "string", "description": "Last day, same formats (default: today)"}}}) + + +async def inception(server, portfolio_id: int) -> date | None: + try: + portfolios = json.loads(await server.call("list_portfolios", {})).get("portfolios") or [] + except (ValueError, AttributeError): + return None + for p in portfolios: + if isinstance(p, dict) and p.get("id") == portfolio_id: + return parse_date(p.get("inception_date") or "") + return None + + +def sharesight(ctx: Ctx): + """The MCP server that has a get_performance_report tool, if it is up.""" + return next((s for s in ctx.registry.up() if "get_performance_report" in s.fn_names.values()), None) + + +MAX_TABLE_ROWS = 40 # longer results list only the extremes +TOP_N = 5 + + +def pct_of(row) -> float: + return row[1].pct + + +def gain_of(row) -> float: + return row[1].gain + + +def tool(ctx: Ctx): + """(definition, function) of period_returns.""" + + async def periods(args: dict) -> str: + server = sharesight(ctx) + if server is None: + return "Error: the portfolio data source isn't connected." + try: + pid = int(args.get("portfolio_id")) + except (TypeError, ValueError): + return "Error: give the portfolio_id (call list_portfolios first)." + period = str(args.get("period") or "month").lower() + if period not in PERIODS: + return f"Error: period must be one of {', '.join(PERIODS)}." + today = datetime.now(ctx.st.tz).date() + last = min(parse_date(args["end"], end=True) or today, today) if args.get("end") else today + first = parse_date(args["start"]) if args.get("start") else await inception(server, pid) + if first is None: + return "Error: couldn't tell when to start; pass start as YYYY-MM-DD." + if first > last: + return "Error: start is after end." + windows = period_windows(period, first, last) + if not windows: + return "Error: no trading days in that range." + cached = await asyncio.to_thread(ctx.store.period_get, pid, first.isoformat(), last.isoformat()) + have = {key(w): cached[key(w)] for w in windows if key(w) in cached} + if refused := refusal(period, windows, set(have)): + return "Error: " + refused + todo = [w for w in windows if key(w) not in have] + fetched = await asyncio.gather(*(fetch_row(server, pid, a, b) for _, a, b in todo)) + fresh = [r for r in fetched if r is not None] + have.update({(r.start, r.end): r for r in fresh}) + if finished := [r for r in fresh if date.fromisoformat(r.end) < today]: # never the unfinished period + await asyncio.to_thread(ctx.store.period_put, finished) + rows = [(label, have[key((label, a, b))]) for label, a, b in windows if key((label, a, b)) in have] + missing = len(windows) - len(rows) + if not rows: + return "Error: no reports came back." + best = max(rows, key=lambda r: r[1].pct) + worst = min(rows, key=lambda r: r[1].pct) + lines = [f"Total return per {period} ({rows[0][1].currency}), {rows[0][0]} to {rows[-1][0]}, " + "sold holdings included:"] + rich, poor = max(rows, key=lambda r: r[1].gain), min(rows, key=lambda r: r[1].gain) + both = lambda row: f"{row[0]} {row[1].pct:+.2f}% ({row[1].gain:+,.0f})" # noqa: E731 + if len(rows) <= MAX_TABLE_ROWS: + lines += [f"{label} {r.pct:+.2f}% {r.gain:+,.0f}" for label, r in rows] + else: # thousands of rows would cost tens of thousands of tokens; the extremes are what gets asked + lines[0] = lines[0][:-1] + f" ({len(rows)} {period}s; the full list is left out, ask for a shorter range to see it):" + for title, key_fn, newest_first in (("Top", pct_of, True), ("Bottom", pct_of, False), + ("Top", gain_of, True), ("Bottom", gain_of, False)): + picked = sorted(rows, key=key_fn, reverse=newest_first)[:TOP_N] + lines.append(f"{title} {TOP_N} by {'%' if key_fn is pct_of else 'money'}: " + ", ".join(both(r) for r in picked)) + lines.append(f"Best by %: {both(best)}. Worst by %: {both(worst)}. " + f"Best by money: {both(rich)}. Worst by money: {both(poor)}." + + (f" {missing} {period}s had no data." if missing else "")) + # What moved the extremes (best/worst by percent and by money), so the answer can say why. + extremes = dict.fromkeys(label for label, _ in (best, worst, rich, poor)) + for label, r in rows: + if label in extremes and (text := drivers(r.movers)): + lines.append(f"Holdings behind {label}: {text}") + log.info("period_returns: %d %ss for portfolio %s (%d fetched)", len(rows), period, pid, len(todo)) + return "\n".join(lines) + + return TOOL, periods diff --git a/lib/features/periods/backfill.py b/lib/features/periods/backfill.py new file mode 100644 index 0000000..c9e3c50 --- /dev/null +++ b/lib/features/periods/backfill.py @@ -0,0 +1,182 @@ +"""PERIOD_BACKFILL: fetch finished days, weeks, months, quarters and years of selected Sharesight portfolios into +the period cache in the background, politely (two calls at a time, never while a user's request is using the server).""" + +from __future__ import annotations + +import asyncio +import json +import logging +import time +from datetime import date, datetime, timedelta + +from ... import config +from ...context import Ctx +from . import sharesight +from .rows import fetch +from .windows import Window, key, parse_date, period_windows + +log = logging.getLogger("bot") + +ORDER = ("month", "quarter", "year", "week", "day") # what pays off first: few calls, most questions +BUSY_POLL = 2.0 +RECHECK = 6 * 3600 # seconds between passes (a pass fetches only what newly finished) +RETRY_AFTER_FAILURE = 600 # seconds before trying again after a pass stopped by MAX_FAILS failures in a row +WORKERS = 2 # fetches at a time (needs max_concurrent >= 2 on the Sharesight server in mcp_servers.json) +MAX_FAILS = 5 # consecutive failed fetches that end a pass (Sharesight down or throttling) +PROGRESS_EVERY = 50 +WAIT_LOG = 30.0 # seconds held back by a user's request before the journal says so +SAMPLE_NODATA = 3 # "no data" windows logged individually per step +SLOW_FETCH = 20.0 # seconds one Sharesight fetch may take before the journal says so + +_wake: asyncio.Event | None = None # set to start a pass now (after a flush) +BAD: set[tuple[int, str, str]] = set() # windows Sharesight had no data for; not retried until restart or reset() + + +def wake_event() -> asyncio.Event: + global _wake + if _wake is None: + _wake = asyncio.Event() + return _wake + + +def reset() -> None: + """After a cache flush: forget what failed and run a pass now.""" + BAD.clear() + wake_event().set() + + +def natural_end(period: str, w: Window) -> bool: + """Does the window end where its period ends (so it is a finished, canonical window)?""" + end, nxt = w[2], w[2] + timedelta(days=1) + return {"day": True, "week": end.weekday() == 6, "month": nxt.day == 1, + "quarter": nxt.day == 1 and nxt.month in (1, 4, 7, 10), + "year": end.month == 12 and end.day == 31}[period] + + +def wanted(first: date, today: date, periods: tuple[str, ...] = ORDER) -> list[Window]: + """Finished windows of these period types from `first` up to yesterday, in the order worth fetching them.""" + yesterday = today - timedelta(days=1) + return [w for p in periods for w in period_windows(p, first, yesterday) if natural_end(p, w)] \ + if first <= yesterday else [] + + +def selected(st: config.Settings, listing: str) -> list[tuple[int, str, date]]: + """(id, name, inception) of the portfolios PERIOD_BACKFILL names (or all of them).""" + try: + portfolios = json.loads(listing).get("portfolios") or [] + except (ValueError, AttributeError): + return [] + out = [] + for p in portfolios: + if not isinstance(p, dict) or not isinstance(p.get("id"), int): + continue + name = str(p.get("name") or "") + start = parse_date(p.get("inception_date") or "") + if start and ("all" in st.period_backfill or name.lower() in st.period_backfill): + out.append((p["id"], name, start)) + return out + + +async def fill_portfolio(ctx: Ctx, server, pid: int, name: str, todo: list[Window], + streak: list[int] | None = None) -> tuple[int, bool]: + """Fetch these windows with WORKERS fetches at a time. Returns (how many were stored, whether the pass + was stopped by MAX_FAILS transient failures in a row). A window Sharesight has no data for is skipped + (until the next restart or flush) without counting as a failure: early days before any trade are normal.""" + pending = iter(todo) # shared: each window goes to exactly one worker + streak = streak if streak is not None else [0] # failures in a row; shared across a pass's steps + stored = done = nodata = transient = 0 + stopped = False + + async def worker() -> None: + nonlocal stored, done, nodata, transient, stopped + for w in pending: + waited, said = time.monotonic(), False + while (users := server.user_calls()) and not stopped: # a user's request comes first + await asyncio.sleep(BUSY_POLL) + if not said and time.monotonic() - waited >= WAIT_LOG: + said = True + log.info("Period backfill: %s held back for %.0f s by: %s", name, WAIT_LOG, "; ".join(users[:3])) + if stopped: + return + began = time.monotonic() + got = await fetch(server, pid, w[1], w[2], background=True) + if time.monotonic() - began >= SLOW_FETCH: + log.warning("Period backfill: %s %s..%s was slow: %.0f s", name, w[1], w[2], time.monotonic() - began) + if got.row is not None: + streak[0] = 0 + await asyncio.to_thread(ctx.store.period_put, [got.row]) + stored += 1 + elif got.kind == "nodata": + BAD.add((pid, *key(w))) + nodata += 1 + if nodata <= SAMPLE_NODATA: # enough of them to see what Sharesight says, not one line per empty day + log.info("Period backfill: %s %s..%s has no data: %s", name, w[1], w[2], got.text[:150]) + else: # a blip says nothing about the window itself, so it is tried again next pass + streak[0] += 1 + transient += 1 + log.warning("Period backfill: %s %s..%s failed: %s", name, w[1], w[2], got.text[:150]) + if streak[0] >= MAX_FAILS and not stopped: + stopped = True + log.warning("Period backfill: %d fetches in a row failed; stopping this pass", streak[0]) + done += 1 + if done % PROGRESS_EVERY == 0: + log.info("Period backfill: %s %d/%d", name, done, len(todo)) + + await asyncio.gather(*(worker() for _ in range(WORKERS))) + log.info("Period backfill: %s: %d stored, %d had no data, %d failed", name, stored, nodata, transient) + return stored, stopped + + +async def fill_pass(ctx: Ctx, server, portfolios: list[tuple[int, str, date]], today: date) -> tuple[int, bool]: + """One pass: fetch every wanted window that isn't cached yet, one period type at a time across all the + portfolios (everyone's months, then quarters, years, weeks, and the many days last). Returns (how many + were stored, whether failures in a row cut it short).""" + cached = {} + for pid, _, first in portfolios: + cached[pid] = set(await asyncio.to_thread(ctx.store.period_get, pid, first.isoformat(), today.isoformat())) + stored, streak = 0, [0] + for period in ORDER: + for pid, name, first in portfolios: + todo = [w for w in wanted(first, today, (period,)) + if key(w) not in cached[pid] and (pid, *key(w)) not in BAD] + if not todo: + continue + log.info("Period backfill: %d %ss to fetch for %s", len(todo), period, name) + n, stopped = await fill_portfolio(ctx, server, pid, f"{name} {period}s", todo, streak) + stored += n + if stopped: + return stored, True + return stored, False + + +async def fill_once(ctx: Ctx, server, portfolios: list[tuple[int, str, date]], today: date) -> int: + return (await fill_pass(ctx, server, portfolios, today))[0] + + +async def loop(ctx: Ctx) -> None: + await ctx.registry.wait_started(config.MCP_STARTUP_WAIT) + warned = False + while True: + wait = RECHECK + server = sharesight(ctx) + if server is None: # on by default, so a bot without Sharesight must not nag every 10 minutes + (log.debug if warned else log.warning)("Period backfill: Sharesight isn't connected; trying again in 10 minutes") + warned = True + wait = 600 + else: + try: + portfolios = selected(ctx.st, await server.call("list_portfolios", {})) + if not portfolios: + log.warning("Period backfill: no Sharesight portfolio matches PERIOD_BACKFILL") + n, stopped = await fill_pass(ctx, server, portfolios, datetime.now(ctx.st.tz).date()) + log.info("Period backfill: pass %s, %d new windows cached%s", "stopped" if stopped else "done", n, + f"; trying again in {RETRY_AFTER_FAILURE // 60} minutes" if stopped else "") + if stopped: + wait = RETRY_AFTER_FAILURE + except Exception: + log.exception("Period backfill pass failed") + try: + await asyncio.wait_for(wake_event().wait(), wait) + except TimeoutError: + pass + wake_event().clear() diff --git a/lib/features/periods/rows.py b/lib/features/periods/rows.py new file mode 100644 index 0000000..3edea76 --- /dev/null +++ b/lib/features/periods/rows.py @@ -0,0 +1,71 @@ +"""Turning Sharesight performance reports into cached period rows, and the "what drove it" text.""" + +from __future__ import annotations + +import json +import re +from datetime import date +from typing import NamedTuple + +from ...mcp.results import table_records +from ...store import PeriodRow + +MAX_MOVERS = 12 # holdings kept per period +Mover = tuple[str, float, object] # code, total gain, gain % + + +def parse_report(text: str) -> dict | None: + try: + report = json.loads(text).get("report") + except (ValueError, AttributeError): + return None + return report if isinstance(report, dict) and report.get("total_gain_percent") is not None else None + + +def movers(report: dict) -> list[list]: + """The holdings with the largest gain or loss in a report, sold ones included: [[code, gain, gain %], ...].""" + rows = [[h["code"], float(h["total_gain"]), h.get("total_gain_percent")] + for h in table_records(report.get("holdings")) + if h.get("code") and isinstance(h.get("total_gain"), (int, float)) and h["total_gain"]] + return sorted(rows, key=lambda m: -abs(m[1]))[:MAX_MOVERS] + + +def drivers(rows: list, n: int = 3) -> str: + """The holdings that added most and lost most: "up: CRWD +9,975 (+64.5%); down: ...".""" + def show(m): + return f"{m[0]} {m[1]:+,.0f}" + (f" ({m[2]:+.1f}%)" if isinstance(m[2], (int, float)) else "") + ups = sorted((m for m in rows if m[1] > 0), key=lambda m: -m[1])[:n] + downs = sorted((m for m in rows if m[1] < 0), key=lambda m: m[1])[:n] + return "; ".join(f"{name}: {', '.join(map(show, group))}" for name, group in (("up", ups), ("down", downs)) if group) + + +THROTTLE = re.compile(r"429|rate.?limit|too many|throttl|timed? ?out|unavailable|temporar|overload|\b5\d\d\b", re.I) + + +class Fetched(NamedTuple): + row: PeriodRow | None + kind: str # "ok", "nodata" (Sharesight has nothing for that window) or "transient" (outage, timeout, throttling) + text: str # what came back, for the log + + +async def fetch(server, portfolio_id: int, first: date, last: date, background: bool = False) -> Fetched: + """One window's return from Sharesight (sold holdings included). `background`: the backfill, which yields to + users' calls.""" + try: + text = await server.call("get_performance_report", { + "portfolio_id": portfolio_id, "start_date": first.isoformat(), "end_date": last.isoformat(), + "_keep_sales": True, **({"_bg": True} if background else {})}) + except Exception as e: # a timeout or a dropped connection + return Fetched(None, "transient", f"{type(e).__name__}: {e}") + report = parse_report(text) + if report is not None: + row = PeriodRow(portfolio_id, first.isoformat(), last.isoformat(), float(report["total_gain_percent"]), + float(report.get("total_gain") or 0), report.get("currency") or "", movers(report)) + return Fetched(row, "ok", "") + transient = text.startswith("Error:") or bool(THROTTLE.search(text[:300])) + return Fetched(None, "transient" if transient else "nodata", text) + + +async def fetch_row(server, portfolio_id: int, first: date, last: date) -> PeriodRow | None: + """fetch(), keeping only the row (None if no report came back).""" + return (await fetch(server, portfolio_id, first, last)).row diff --git a/lib/features/periods/windows.py b/lib/features/periods/windows.py new file mode 100644 index 0000000..c0fa320 --- /dev/null +++ b/lib/features/periods/windows.py @@ -0,0 +1,105 @@ +"""Period windows for period_returns: parsing dates, cutting a date range into days/weeks/months/quarters/years, +and the limits on how many uncached ones one request may fetch.""" + +from __future__ import annotations + +from datetime import date, datetime, timedelta + +PERIODS = ("day", "week", "month", "quarter", "year") +# Uncached reports one call may fetch (about 2.4 s each); cached ones are free. +LIVE_LIMIT = {"day": 64, "week": 26, "month": 120, "quarter": 120, "year": 120} + +Window = tuple[str, date, date] # label, first day, last day + + +def parse_date(text: str, *, end: bool = False) -> date | None: + """A date from YYYY-MM-DD, YYYY-MM or YYYY; a partial one means its first day, or its last for `end`.""" + text = str(text).strip() + for fmt, size in (("%Y-%m-%d", 10), ("%Y-%m", 7), ("%Y", 4)): + try: + d = datetime.strptime(text, fmt).date() + except ValueError: + continue + if end and size < 10: + d = (date(d.year + 1, 1, 1) if size == 4 else month_add(d, 1)) - timedelta(days=1) + return d + return None + + +def month_add(d: date, months: int) -> date: + n = d.year * 12 + d.month - 1 + months + return date(n // 12, n % 12 + 1, 1) + + +def period_windows(period: str, first: date, last: date) -> list[Window]: + """(label, first day, last day) of each period between `first` and `last`; the ends are clipped to them.""" + out: list[Window] = [] + if period == "day": + d = first + while d <= last: + if d.weekday() < 5: # no trading at weekends + out.append((d.isoformat(), d, d)) + d += timedelta(days=1) + elif period == "week": + monday = first - timedelta(days=first.weekday()) + while monday <= last: + iso = monday.isocalendar() + out.append((f"{iso[0]}-W{iso[1]:02d}", max(monday, first), min(monday + timedelta(days=6), last))) + monday += timedelta(days=7) + else: + step = {"month": 1, "quarter": 3, "year": 12}[period] + cur = date(first.year, ((first.month - 1) // step) * step + 1, 1) + while cur <= last: + nxt = month_add(cur, step) + label = {"month": f"{cur.year}-{cur.month:02d}", "quarter": f"{cur.year}-Q{(cur.month - 1) // 3 + 1}", + "year": str(cur.year)}[period] + out.append((label, max(cur, first), min(nxt - timedelta(days=1), last))) + cur = nxt + return out + + +def key(w: Window) -> tuple[str, str]: + return w[1].isoformat(), w[2].isoformat() + + +def refusal(period: str, windows: list[Window], cached: set[tuple[str, str]]) -> str | None: + """Why a request needs too many uncached reports (with the earliest start that fits), or None.""" + limit = LIVE_LIMIT[period] + missing = [w for w in windows if key(w) not in cached] + if len(missing) <= limit: + return None + seen, start = 0, windows[-1][1] + for w in reversed(windows): # walk back from the end until the budget of uncached reports is spent + if key(w) not in cached: + seen += 1 + if seen > limit: + break + start = w[1] + return (f"{len(windows) - len(missing)} of {len(windows)} {period}s are cached and the other {len(missing)} " + f"would take about {len(missing) * 2.4 / 60:.0f} minutes to fetch (the background fill may still be " + f"running or off). Ask again with start={start} or later, or try again later") + + +def flush_scope(words: list[str]) -> tuple[int | None, date | None, date | None] | None: + """What /flushcache and --flush-cache were asked to forget: any of a portfolio ID (5+ digits), a year + (2024), a month (2024-05) or a day (2024-05-17), in any order, plus an optional "all". Returns + (portfolio ID, first day, last day), each None when not given, or None if a word isn't understood.""" + pid, first, last = None, None, None + for word in words: + word = word.strip().lower() + if word in ("", "all"): + continue + if word.isdigit() and len(word) >= 5 and pid is None: + pid = int(word) + elif first is None and (d := parse_date(word)) and 1990 <= d.year <= 2100 and word[:4].isdigit(): + first, last = d, parse_date(word, end=True) + else: + return None + return pid, first, last + + +def describe_scope(pid: int | None, first: date | None, last: date | None) -> str: + parts = ["all portfolios" if pid is None else f"portfolio {pid}"] + if first: + parts.append(f"{first}" if first == last else f"{first} to {last}") + return ", ".join(parts) diff --git a/lib/llm/gate.py b/lib/llm/gate.py index 53c7432..b80dbc0 100644 --- a/lib/llm/gate.py +++ b/lib/llm/gate.py @@ -20,7 +20,7 @@ # A portfolio question: strong phrases always count; "my " phrases count; weak words # count only when no ticker is named ("performance of NVDA" is a market question). PORTFOLIO_STRONG = re.compile( - r"\b(portfolios?|holdings?|sharesight|smsf|superannuation|super fund|net worth|" + r"\b(portfolios?|holdings?|sharesight|\w*smsf|superannuation|super fund|net worth|" r"what do i (?:own|hold)|how am i doing|am i (?:up|down)|" r"what (?:did|have) i (?:make|made|lose|lost)|how did i do|" r"(?:i|we) (?:out|under)perform\w*|(?:out|under)perform\w* (?:me|my|us|our)|" diff --git a/lib/llm/runner.py b/lib/llm/runner.py index 9a3670e..cfefe87 100644 --- a/lib/llm/runner.py +++ b/lib/llm/runner.py @@ -71,6 +71,8 @@ async def _run_mcp(call: Call, registry: Registry) -> str: args = json.loads(call.arguments or "{}") if not isinstance(args, dict): return f"Error: arguments for {tool} must be a JSON object." + args.pop("_keep_sales", None) # reserved for the bot's own tools + args.pop("_bg", None) log.info("MCP %s.%s %s", server.label, tool, brief_args(args)) out = await server.call(tool, args) # A tiny result is usually an empty list or an error: show it, so "why did it say diff --git a/lib/mcp/server.py b/lib/mcp/server.py index 9707020..871cc9b 100644 --- a/lib/mcp/server.py +++ b/lib/mcp/server.py @@ -4,6 +4,7 @@ import asyncio import difflib +import itertools import json import logging import re @@ -54,6 +55,8 @@ def __init__(self, label: str, cfg: dict, *, timeout: float, max_output: int, self.fn_names: dict[str, str] = {} # function name sent to the model -> MCP tool name # Cap simultaneous calls: a 20-stock request fired all at once gets rate-limited. self._sem = asyncio.Semaphore(int(cfg.get("max_concurrent", 4))) + self._active: dict[int, tuple[str, bool]] = {} # calls running or waiting for a slot: text, background? + self._call_ids = itertools.count() self._ready = asyncio.Event() self._stop = asyncio.Event() self._task: asyncio.Task | None = None @@ -107,7 +110,8 @@ async def _run(self) -> None: env=expand_env(self.cfg.get("env"), self.label)) read, write = await stack.enter_async_context(stdio_client(params, errlog=errlog)) session = await stack.enter_async_context(ClientSession(read, write)) - await session.initialize() + init = await session.initialize() + log.info("MCP %s: connected, protocol %s", self.label, getattr(init, "protocolVersion", "?")) await self._load_tools(session) self.session, self.error = session, None self._ready.set() @@ -230,12 +234,32 @@ async def _warm_once(self) -> None: else: fresh.used = False + @property + def in_flight(self) -> int: + """Calls running or waiting for a slot.""" + return len(self._active) + + def user_calls(self) -> list[str]: + """What the calls that are not background work (a chat request, a period_returns question) are doing. + Background work (the period backfill) waits while there are any.""" + return [text for text, background in self._active.values() if not background] + async def _call(self, tool: str, args: dict) -> str: holdings_only = bool(self.cfg.get("current_holdings_only")) + keep_sales = args.get("_keep_sales") # set by the bot's own tools (history), never by the model + background = bool(args.get("_bg")) # the period backfill, which must not hold up a user's request + args = {k: v for k, v in args.items() if k not in ("_keep_sales", "_bg")} if holdings_only and tool == "get_performance_report": - args = {**args, "include_sales": False} # Sharesight's own switch for sold holdings - async with self._sem: - result = await asyncio.wait_for(self.session.call_tool(tool, args), self.timeout) + args = {**args, "include_sales": bool(keep_sales)} # Sharesight's own switch for sold holdings + if keep_sales: + holdings_only = False # closed holdings stay in the report too + token = next(self._call_ids) + self._active[token] = (describe_call(tool, args, bool(keep_sales)), background) + try: + async with self._sem: + result = await asyncio.wait_for(self.session.call_tool(tool, args), self.timeout) + finally: + del self._active[token] parts = [c.text if c.type == "text" else f"[{c.type} content omitted]" for c in result.content] if not parts and getattr(result, "structuredContent", None): parts.append(json.dumps(result.structuredContent)) @@ -244,12 +268,20 @@ async def _call(self, tool: str, args: dict) -> str: out = "Tool error: " + out else: out = slim_result(out, self.cfg.get("slim", self.label), drop_closed=holdings_only) - if len(out) > self.max_output: + if len(out) > self.max_output and not keep_sales: # the cap protects the model's context; the bot's own tools aren't sent to it log.warning("MCP %s.%s result truncated (%d chars)", self.label, tool, len(out)) out = out[:self.max_output] + f"\n...[truncated; {len(out)} chars total]" return out +def describe_call(tool: str, args: dict, bot_own: bool = False) -> str: + """A call in a few words, for the log: "get_performance_report portfolio_id=7 2024-01-01..2024-01-31".""" + shown = [f"{k}={v}" for k, v in args.items() if k in ("portfolio_id", "symbol", "ticker", "query")] + if args.get("start_date") or args.get("end_date"): + shown.append(f"{args.get('start_date', '')}..{args.get('end_date', '')}") + return " ".join([tool, *shown]) + (" (period_returns)" if bot_own else "") + + class Registry: """All configured MCP servers.""" diff --git a/lib/prompts.py b/lib/prompts.py index 3c58121..d9cb3da 100644 --- a/lib/prompts.py +++ b/lib/prompts.py @@ -85,9 +85,10 @@ def chat_prompt(*, transcript: str, sender: str, private: bool, reply_quote: str "earlier replies, and follow the formatting rules in your instructions even " "where your earlier replies didn't.") if fresh: # a portfolio question: stale numbers and refusals in the history must not be reused - tail += (" Get current figures with your portfolio tools before answering; don't reuse " - "numbers or refusals from earlier replies. For a comparison or a \"why\", write 2-3 " - "short paragraphs on the tailwinds and headwinds of each portfolio, not a summary.") + tail += (" Get current figures with your portfolio tools before answering (and your market tools " + "for anything else asked, such as prices or a comparison with market data); don't reuse " + "numbers or refusals from earlier replies. When comparing portfolios, or asked \"why\", " + "write 2-3 short paragraphs on the tailwinds and headwinds of each portfolio, not a summary.") return head + tail + down diff --git a/lib/store.py b/lib/store.py index ea39c34..b58850d 100644 --- a/lib/store.py +++ b/lib/store.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import sqlite3 import threading import time @@ -33,7 +34,24 @@ CREATE TABLE IF NOT EXISTS usage ( ts INTEGER, bot_id INTEGER, chat_id INTEGER, model TEXT, kind TEXT, rounds INTEGER, tokens_in INTEGER, tokens_cached INTEGER, tokens_out INTEGER, cost REAL); +-- Finished-period Sharesight returns (sold holdings included), kept until flushed. pct NULL = no data +-- for that window. movers = the holdings that moved most, as JSON [[code, gain, gain_pct], ...]. +CREATE TABLE IF NOT EXISTS period_cache ( + portfolio_id INTEGER, start TEXT, end TEXT, pct REAL, gain REAL, currency TEXT, movers TEXT, + v INTEGER, PRIMARY KEY (portfolio_id, start, end)); """ +PERIOD_CACHE_VERSION = 1 # bump when the stored shape changes: rows of another version are dropped + + +class PeriodRow(NamedTuple): + """One cached window of a portfolio's return (start/end are ISO dates).""" + portfolio_id: int + start: str + end: str + pct: float | None + gain: float | None + currency: str + movers: list class HistoryRow(NamedTuple): @@ -65,6 +83,7 @@ def __init__(self, path: str): self._conn.execute( "DELETE FROM holding_news WHERE length(text) < 40 AND upper(text) LIKE '%NOTHING%'" ) + self._conn.execute("DELETE FROM period_cache WHERE v != ?", (PERIOD_CACHE_VERSION,)) def _run(self, sql: str, params=()) -> list[tuple]: with self._lock: @@ -181,3 +200,36 @@ def usage_summary(self, since_ts: int) -> list[tuple]: "SELECT model, COUNT(*), SUM(tokens_in), SUM(tokens_cached), SUM(tokens_out), SUM(cost) " "FROM usage WHERE bot_id = ? AND ts >= ? GROUP BY model ORDER BY SUM(cost) DESC", (self.bot_id or 0, since_ts)) + + # -- Sharesight period cache ------------------------------------------------------- + def period_get(self, portfolio_id: int, first: str, last: str) -> dict[tuple[str, str], PeriodRow]: + """Cached windows of a portfolio lying within first..last (ISO dates), by (start, end).""" + rows = self._run("SELECT start, end, pct, gain, currency, movers FROM period_cache " + "WHERE portfolio_id = ? AND start >= ? AND end <= ?", (portfolio_id, first, last)) + return {(r[0], r[1]): PeriodRow(portfolio_id, r[0], r[1], r[2], r[3], r[4] or "", json.loads(r[5] or "[]")) + for r in rows} + + def period_put(self, rows: list[PeriodRow]) -> None: + with self._lock: + self._conn.executemany( + "INSERT OR REPLACE INTO period_cache VALUES (?,?,?,?,?,?,?,?)", + [(r.portfolio_id, r.start, r.end, r.pct, r.gain, r.currency, json.dumps(r.movers), + PERIOD_CACHE_VERSION) for r in rows]) + + def period_count(self, portfolio_id: int, first: str, last: str) -> int: + return self._run("SELECT COUNT(*) FROM period_cache WHERE portfolio_id = ? AND start >= ? AND end <= ?", + (portfolio_id, first, last))[0][0] + + def period_flush(self, portfolio_id: int | None = None, first: str | None = None, last: str | None = None) -> int: + """Delete cached windows, of one portfolio or all, and (if first/last ISO dates are given) only those + overlapping that date range. Returns how many went.""" + where, args = [], [] + if portfolio_id is not None: + where.append("portfolio_id = ?") + args.append(portfolio_id) + if first and last: + where.append("start <= ? AND end >= ?") + args += [last, first] + with self._lock: + sql = "DELETE FROM period_cache" + (" WHERE " + " AND ".join(where) if where else "") + return self._conn.execute(sql, args).rowcount diff --git a/lib/telegram/commands.py b/lib/telegram/commands.py index f2ac424..a4130a2 100644 --- a/lib/telegram/commands.py +++ b/lib/telegram/commands.py @@ -10,6 +10,8 @@ from telegram.ext import ContextTypes from ..context import Ctx +from ..features.periods import backfill +from ..features.periods.windows import describe_scope, flush_scope from .access import is_admin log = logging.getLogger("bot") @@ -58,3 +60,22 @@ async def credits(ctx: Ctx, update: Update, context: ContextTypes.DEFAULT_TYPE) rows = await asyncio.to_thread(ctx.store.usage_summary, now - secs) parts.append(usage_text(rows, label)) await _reply(ctx, update, "\n\n".join(parts)) + + +async def flush_cache(ctx: Ctx, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """/flushcache [portfolio_id] [year | month | day] - admins only: forget the cached Sharesight period + returns (all, or of one portfolio, or the ones overlapping a year/month/day), e.g. after a backdated + change. The background fill, if on, re-fetches them.""" + if not await _gate(ctx, update): + return + scope = flush_scope(context.args or []) + if scope is None: + await _reply(ctx, update, "Usage: /flushcache [portfolio_id] [2024 | 2024-05 | 2024-05-17]") + return + pid, first, last = scope + n = await asyncio.to_thread(ctx.store.period_flush, pid, first and first.isoformat(), last and last.isoformat()) + backfill.reset() + what = describe_scope(pid, first, last) + log.info("Flushed %d cached period returns (%s)", n, what) + await _reply(ctx, update, f"Flushed {n} cached periods ({what})." + + (" Refilling in the background." if ctx.st.period_backfill else "")) diff --git a/lib/telegram/handlers.py b/lib/telegram/handlers.py index d7d8009..fd6ba8a 100644 --- a/lib/telegram/handlers.py +++ b/lib/telegram/handlers.py @@ -15,7 +15,7 @@ from .. import config from ..ask import ask from ..context import Ctx -from ..features import dm_buttons, holding_news, movers, post +from ..features import dm_buttons, holding_news, movers, periods, post from ..history import compact, format_rows from ..llm.gate import Route, is_portfolio_question, route from ..msgtext import describe, sender_name @@ -29,6 +29,9 @@ log = logging.getLogger("bot") PORTFOLIO_ADMINS_ONLY = "Portfolio data is for admins only." +FOLLOW_UP_WORDS = 12 # a message this short right after the bot's reply may be a follow-up to it +FOLLOW_UP_SECONDS = 30 * 60 +FOLLOW_UP_CHARS = 400 # of the bot's reply that the tool gate also reads READ_IMAGE_WORDS = re.compile(r"\b(read|text|says?|chart|table|number|screenshot|ocr|zoom|detail|label)\b", re.I) @@ -162,6 +165,19 @@ def _quote(self, trig: Trigger) -> str | None: body = describe(t) return f"[#{t.message_id}] {who}: {compact(body) if self.ctx.st.token_saver else body}" + async def _gate_text(self, msg: Message, trig: Trigger) -> str: + """What the tool gate reads: the message and what it replies to. A short follow-up ("what about in dollar + terms?") sent right after the bot's own reply also gets the start of that reply, so it counts as part of + the same (portfolio, market...) conversation instead of reaching the model with no tools.""" + text = trig.text + " " + (describe(trig.reply_target) if trig.reply_target else "") + if trig.reply_target is None and len(trig.text.split()) <= FOLLOW_UP_WORDS: + rows = await asyncio.to_thread(self.ctx.store.history, msg.chat_id, 2) + if len(rows) >= 2 and rows[-1].message_id == msg.message_id: + prev = rows[-2] + if prev.sender == self.ctx.self_name and int(msg.date.timestamp()) - prev.ts <= FOLLOW_UP_SECONDS: + text += " " + compact(prev.text)[:FOLLOW_UP_CHARS] + return text + async def _chat_text(self, msg: Message, trig: Trigger, middle: str = "", fresh: bool = False) -> str: ctx, st = self.ctx, self.ctx.st rows = await asyncio.to_thread(ctx.store.history, msg.chat_id, st.history_limit) @@ -222,13 +238,15 @@ async def _generate(self, msg: Message, trig: Trigger, draft: Draft | None) -> s must_search=True, cache_key=f"chat-{chat_id}", on_text=on_text, kind="movers", chat_id=chat_id) else: - context = trig.text + " " + (describe(trig.reply_target) if trig.reply_target else "") + context = await self._gate_text(msg, trig) chosen = route(context, st, ctx.registry) admin = not st.admin_only or await self._is_admin(msg) if not admin and is_portfolio_question(context, st) and any( s.cfg.get("gate") == "portfolio" for s in ctx.registry.up()): return PORTFOLIO_ADMINS_ONLY extra = [] + if admin and any(s.cfg.get("gate") == "portfolio" for s in chosen.servers): + extra.append(periods.tool(ctx)) if st.post_to_groups and trig.private and post.wants_post(trig.text): extra = [post.tool(ctx, msg.from_user.id, msg.from_user.full_name or "")] if chosen.simple: # the model must not be told it has no tools diff --git a/mcp_servers.json b/mcp_servers.json index e34c573..8ae5a81 100644 --- a/mcp_servers.json +++ b/mcp_servers.json @@ -39,7 +39,7 @@ ], "slim": "sharesight", "gate": "portfolio", - "max_concurrent": 1, + "max_concurrent": 2, "cache_ttl": 1800, "keep_warm": {"list_portfolios": null, "get_performance_report": 0} } diff --git a/tests/test_app.py b/tests/test_app.py index 1152f29..f8f84e7 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -1,3 +1,6 @@ +import asyncio +from types import SimpleNamespace + import pytest from telegram.ext import CallbackQueryHandler, CommandHandler, MessageHandler @@ -30,12 +33,12 @@ def build(**extra): return app.build_app(Ctx(st, store, ScriptedBackend([]), registry())), st plain, st = build() - assert handler_kinds(plain) == ["CommandHandler", "MessageHandler"] and app.allowed_updates(st) == ["message", "edited_message"] + assert handler_kinds(plain) == ["CommandHandler", "CommandHandler", "MessageHandler"] and app.allowed_updates(st) == ["message", "edited_message"] full, st = build(ADMIN_CHAT_IDS="5", SHARESIGHT_HOLDING_NEWS_RECIPIENTS="Pf:alice") kinds = handler_kinds(full) - assert kinds == ["CommandHandler", "MessageHandler", "CallbackQueryHandler"] + assert kinds == ["CommandHandler", "CommandHandler", "MessageHandler", "CallbackQueryHandler"] assert "callback_query" in app.allowed_updates(st) - assert {c for h in full.handlers[0] if isinstance(h, CommandHandler) for c in h.commands} == {"credits"} + assert {c for h in full.handlers[0] if isinstance(h, CommandHandler) for c in h.commands} == {"credits", "flushcache"} assert any(isinstance(h, MessageHandler) for h in full.handlers[0]) locked, st = build(ADMIN_ONLY="on") # remembers the groups it is in, to know who their admins are assert "ChatMemberHandler" in handler_kinds(locked) and "my_chat_member" in app.allowed_updates(st) @@ -70,7 +73,7 @@ def test_main_accepts_and_ignores_a_label_and_unknown_args(monkeypatch, tmp_path assert started -def test_startup_logs_the_git_hash_and_survives_without_git(monkeypatch, tmp_path, caplog): +def test_startup_logs_the_git_hash_and_whether_the_backfill_is_on(monkeypatch, tmp_path, caplog): monkeypatch.setenv("TELEGRAM_BOT_TOKEN", "1:x") monkeypatch.setenv("OPENROUTER_API_KEY", "b") monkeypatch.delenv("XAI_API_KEY", raising=False) @@ -78,18 +81,48 @@ def test_startup_logs_the_git_hash_and_survives_without_git(monkeypatch, tmp_pat monkeypatch.setenv("MCP_CONFIG", str(tmp_path / "none.json")) monkeypatch.setattr(app.Application, "run_polling", lambda self, **kw: None) caplog.set_level("INFO", logger="bot") - monkeypatch.setattr(app.subprocess, "run", lambda *a, **kw: type("R", (), {"returncode": 0, "stdout": "abc1234\n"})()) + monkeypatch.setattr(app, "git_hash", lambda: "abc1234") + monkeypatch.setenv("PERIOD_BACKFILL", "off") app.main(["grokbot"]) - assert "Starting instance grokbot (abc1234): " in caplog.text + assert "Starting instance grokbot (abc1234): " in caplog.text and "Period backfill" not in caplog.text caplog.clear() - - def no_git(*a, **kw): - raise FileNotFoundError("git") - - monkeypatch.setattr(app.subprocess, "run", no_git) - assert app.git_hash() == "" + monkeypatch.setattr(app, "git_hash", lambda: "") + monkeypatch.setenv("PERIOD_BACKFILL", "RobSMSF, other") app.main(["grokbot"]) assert "Starting instance grokbot: " in caplog.text + assert "Period backfill: on for other, robsmsf" in caplog.text + + +def test_git_hash_is_read_from_the_git_directory_without_running_git(tmp_path): + sha = "0123456789abcdef0123456789abcdef01234567" + assert app.git_hash(tmp_path) == "" # no checkout + git = tmp_path / ".git" + (git / "refs" / "heads").mkdir(parents=True) + (git / "HEAD").write_text("ref: refs/heads/main\n") + assert app.git_hash(tmp_path) == "" # the branch has no commit yet + (git / "refs" / "heads" / "main").write_text(sha + "\n") + assert app.git_hash(tmp_path) == "0123456" + (git / "refs" / "heads" / "main").unlink() # packed by gc + (git / "packed-refs").write_text(f"# pack-refs\n{sha} refs/heads/main\n") + assert app.git_hash(tmp_path) == "0123456" + (git / "HEAD").write_text(sha + "\n") # detached + assert app.git_hash(tmp_path) == "0123456" + (git / "HEAD").write_text("garbage") + assert app.git_hash(tmp_path) == "" + + +def test_git_hash_follows_a_linked_worktree(tmp_path): + sha = "fedcba9876543210fedcba9876543210fedcba98" + main, tree = tmp_path / "main" / ".git", tmp_path / "wt" + (main / "refs" / "heads").mkdir(parents=True) + (main / "refs" / "heads" / "feat").write_text(sha + "\n") + admin = main / "worktrees" / "wt" + admin.mkdir(parents=True) + (admin / "HEAD").write_text("ref: refs/heads/feat\n") + (admin / "commondir").write_text("../..\n") + tree.mkdir() + (tree / ".git").write_text(f"gitdir: {admin}\n") + assert app.git_hash(tree) == "fedcba9" async def test_alert_down_notifies_admin_chats_only(env, store): @@ -104,3 +137,63 @@ async def test_alert_down_notifies_admin_chats_only(env, store): quiet = Ctx(config.load(env), store, ScriptedBackend([]), registry(), FakeBot()) await app.alert_down(quiet, server) assert quiet.bot.sent == [] + + +def test_flush_cache_cli_needs_no_keys_and_flushes_by_portfolio_and_period(monkeypatch, tmp_path, capsys): + from lib.store import PeriodRow, Store + + db = str(tmp_path / "x.db") + for var in ("TELEGRAM_BOT_TOKEN", "OPENROUTER_API_KEY", "XAI_API_KEY", "ZAI_API_KEY"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("DB_PATH", db) + s = Store(db) + row = lambda pid, a, b: PeriodRow(pid, a, b, 1.0, 1.0, "AUD", []) # noqa: E731 + s.period_put([row(70007, "2023-12-26", "2024-01-01"), row(70007, "2024-03-01", "2024-03-31"), + row(70007, "2025-01-01", "2025-12-31"), row(80008, "2024-03-01", "2024-03-31")]) + + def run(*words): + with pytest.raises(SystemExit) as e: + app.main(["--flush-cache", *words]) + return e.value.code, capsys.readouterr().out + + assert run("70007", "2024") == (0, "Flushed 2 cached periods (portfolio 70007, 2024-01-01 to 2024-12-31).\n") + assert run("2025-06") == (0, "Flushed 1 cached periods (all portfolios, 2025-06-01 to 2025-06-30).\n") + assert run("80008") == (0, "Flushed 1 cached periods (portfolio 80008).\n") + assert run() == (0, "Flushed 0 cached periods (all portfolios).\n") + code, out = run("abc") + assert code == 1 and "Usage" in out + assert run("70007", "70008")[0] == 1 # one portfolio at a time + + +async def test_backfill_task_starts_unless_period_backfill_is_off(env, store, monkeypatch): + started = [] + + async def fake_loop(ctx): + started.append(ctx) + + monkeypatch.setattr(app.backfill, "loop", fake_loop) + for extra, expected in (({"PERIOD_BACKFILL": "off"}, 0), ({}, 1), ({"PERIOD_BACKFILL": "x"}, 2)): + st = config.load({**env, **extra}) + ctx = Ctx(st, store, ScriptedBackend([]), registry(), FakeBot()) + built = app.build_app(ctx) + await built.post_init(SimpleNamespace(bot=SimpleNamespace(id=1))) + await asyncio.sleep(0) + assert len(started) == expected + + +def test_the_mcp_librarys_unlabelled_protocol_line_is_dropped(monkeypatch): + import logging + + monkeypatch.delenv("JOURNAL_STREAM", raising=False) + root = logging.getLogger() + saved = list(root.handlers), root.level + root.handlers.clear() + try: + app.setup_logging() + keep = [h for h in root.handlers] + mk = lambda msg: logging.LogRecord("mcp.client", logging.INFO, "f", 1, msg, None, None) # noqa: E731 + assert all(not h.filter(mk("Negotiated protocol version: 2025-11-25")) for h in keep) + assert all(h.filter(mk("MCP yahoo: connected, protocol 2025-11-25")) for h in keep) + finally: + root.handlers[:] = saved[0] + root.setLevel(saved[1]) diff --git a/tests/test_backfill.py b/tests/test_backfill.py new file mode 100644 index 0000000..279f872 --- /dev/null +++ b/tests/test_backfill.py @@ -0,0 +1,240 @@ +import asyncio +import json +from datetime import date, timedelta + +import pytest + +from lib import config +from lib.context import Ctx +from lib.features.periods import backfill +from lib.features.periods.windows import key, period_windows + +from .conftest import FakeBot +from .fakes import FakeMcp, ScriptedBackend, registry + + +class Sharesight(FakeMcp): + def __init__(self, *a, fail=False, inception="2026-01-15", **kw): + super().__init__(*a, **kw) + self.fail = fail + self.fail_text = "Error: timed out" + self.data_from = "" + self.inception = inception + self.busy_for = 0 # polls during which the server counts as busy + self.polls = 0 + self.delay = 0 + self.running = self.peak = 0 + self.order: list[tuple[str, str]] = [] + + def user_calls(self): # a user's request is running for the first few polls + self.polls += 1 + return ["get_performance_report portfolio_id=7 2024-01-01..2024-12-31"] if self.polls <= self.busy_for else [] + + async def call(self, tool, args): + self.calls.append((tool, args)) + if tool == "list_portfolios": + return json.dumps({"portfolios": [{"id": 7, "name": "RobSMSF", "inception_date": self.inception}, + {"id": 8, "name": "Other", "inception_date": "2026-01-15"}]}) + self.order.append((args["start_date"], args["end_date"])) + self.running += 1 + self.peak = max(self.peak, self.running) + await asyncio.sleep(self.delay) + self.running -= 1 + if self.fail: + return self.fail_text + if args["end_date"] < self.data_from: # nothing in Sharesight before this date + return "Tool error: no data for this period" + return json.dumps({"report": {"total_gain_percent": 1.0, "total_gain": 10.0, "currency": "AUD"}}) + + +@pytest.fixture(autouse=True) +def fast(monkeypatch): + monkeypatch.setattr(backfill, "BUSY_POLL", 0) + monkeypatch.setattr(backfill, "_wake", None) + backfill.BAD.clear() + + +def make(env, store, ss, **extra): + st = config.load({**env, "PERIOD_BACKFILL": "robsmsf", **extra}) + return Ctx(st, store, ScriptedBackend([]), registry(ss), FakeBot()), st + + +def test_only_finished_canonical_windows_are_wanted_months_first(): + w = backfill.wanted(date(2026, 1, 15), date(2026, 3, 3)) # yesterday = 2026-03-02 + labels = [x[0] for x in w] + assert labels[:2] == ["2026-01", "2026-02"] and "2026-03" not in labels # March isn't over + assert labels.index("2026-W03") < labels.index("2026-01-15") # weeks before days + assert [x[0] for x in w if x[0].startswith("2026-W")][-1] == "2026-W09" # week ending Sun 1 Mar; not the next + assert w[0][1] == date(2026, 1, 15) # the first window starts at inception + assert not any(x[0] == "2026-Q1" or x[0] == "2026" for x in w) # quarter and year not finished + assert backfill.wanted(date(2026, 3, 5), date(2026, 3, 3)) == [] + + +def test_portfolios_are_selected_by_name_or_all(env): + listing = json.dumps({"portfolios": [{"id": 7, "name": "RobSMSF", "inception_date": "2026-01-15"}, + {"id": 8, "name": "Other", "inception_date": "2026-01-15"}, + {"id": 9, "name": "NoDate"}]}) + pick = lambda **e: [p[1] for p in backfill.selected(config.load({**env, "TELEGRAM_BOT_TOKEN": "1:x", **e}), listing)] # noqa: E731 + assert pick(PERIOD_BACKFILL="robsmsf") == ["RobSMSF"] + assert pick(PERIOD_BACKFILL="all") == ["RobSMSF", "Other"] # a portfolio without an inception date is skipped + assert pick() == ["RobSMSF", "Other"] # on by default + assert pick(PERIOD_BACKFILL="off") == [] + + +async def test_a_pass_fetches_what_is_missing_then_only_what_is_new(env, store): + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report")) + ctx, st = make(env, store, ss) + portfolios = backfill.selected(st, await ss.call("list_portfolios", {})) + want = backfill.wanted(date(2026, 1, 15), date(2026, 3, 3)) + assert await backfill.fill_once(ctx, ss, portfolios, date(2026, 3, 3)) == len(want) + assert ss.order[0] == (want[0][1].isoformat(), want[0][2].isoformat()) # the first month, from inception + assert store.period_count(7, "2026-01-01", "2026-03-31") == len(want) and store.period_count(8, "2000-01-01", "2100-01-01") == 0 + ss.order.clear() + assert await backfill.fill_once(ctx, ss, portfolios, date(2026, 3, 3)) == 0 and ss.order == [] + assert await backfill.fill_once(ctx, ss, portfolios, date(2026, 3, 4)) == 1 # yesterday, Tue 3 Mar, newly finished + assert ss.order == [("2026-03-03", "2026-03-03")] + + +async def test_every_portfolios_cheap_periods_come_before_anyones_days(env, store): + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report")) + ctx, _ = make(env, store, ss) + seen = [] + real = ss.call + + async def spy(tool, args): + if tool == "get_performance_report": + seen.append((args["portfolio_id"], args["start_date"], args["end_date"])) + return await real(tool, args) + + ss.call = spy + folios = [(7, "RobSMSF", date(2026, 1, 15)), (8, "Other", date(2026, 1, 15))] + today = date(2026, 3, 3) + total = await backfill.fill_once(ctx, ss, folios, today) + assert total == 2 * len(backfill.wanted(date(2026, 1, 15), today)) and len(set(seen)) == len(seen) # once each + is_day = lambda x: x[1] == x[2] # noqa: E731 + last_cheap = max(i for i, x in enumerate(seen) if not is_day(x)) + first_day = min(i for i, x in enumerate(seen) if is_day(x)) + assert last_cheap < first_day # all weeks/months of both portfolios, then days + assert {x[0] for x in seen[:last_cheap + 1]} == {7, 8} + assert await backfill.fill_once(ctx, ss, folios, today) == 0 + + +async def test_it_waits_while_a_user_request_is_using_the_server(env, store): + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report")) + ss.busy_for = 3 + ctx, st = make(env, store, ss) + await backfill.fill_once(ctx, ss, [(7, "RobSMSF", date(2026, 3, 1))], date(2026, 3, 3)) + assert ss.polls > 3 and ss.order # polled until free, then fetched + + +async def test_two_fetches_run_at_a_time_and_each_window_is_fetched_once(env, store): + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report")) + ss.delay = 0.01 + ctx, _ = make(env, store, ss) + want = backfill.wanted(date(2026, 1, 15), date(2026, 3, 3)) + assert await backfill.fill_once(ctx, ss, [(7, "RobSMSF", date(2026, 1, 15))], date(2026, 3, 3)) == len(want) + assert ss.peak == backfill.WORKERS == 2 + assert len(ss.order) == len(set(ss.order)) == len(want) # none twice + assert store.period_count(7, "2026-01-01", "2026-03-31") == len(want) + + +async def test_failures_end_the_pass_and_the_same_windows_are_tried_again_next_time(env, store): + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report"), fail=True) + ctx, _ = make(env, store, ss) + args = ctx, ss, [(7, "RobSMSF", date(2026, 1, 15))], date(2026, 3, 3) + assert await backfill.fill_once(*args) == 0 + assert backfill.MAX_FAILS <= len(ss.order) <= backfill.MAX_FAILS + backfill.WORKERS - 1 # gave up after five + assert not backfill.BAD # a blip doesn't mark the windows + first_pass = list(ss.order) + ss.order.clear() + await backfill.fill_once(*args) + assert ss.order[:backfill.MAX_FAILS] == first_pass[:backfill.MAX_FAILS] # same windows again + backfill.reset() + assert not backfill.BAD and backfill.wake_event().is_set() + + +async def test_a_stopped_pass_is_retried_after_ten_minutes_not_six_hours(env, store, monkeypatch): + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report"), fail=True) + ctx, _ = make(env, store, ss) + monkeypatch.setattr(backfill, "RECHECK", 3600) + monkeypatch.setattr(backfill, "RETRY_AFTER_FAILURE", 0.05) + task = asyncio.create_task(backfill.loop(ctx)) + for _ in range(200): + await asyncio.sleep(0.02) + if len(ss.order) >= 3 * backfill.MAX_FAILS: # several passes, each stopped by failures + break + task.cancel() + assert len(ss.order) >= 3 * backfill.MAX_FAILS + + +async def test_windows_with_no_data_do_not_stop_the_pass_and_are_not_retried(env, store, caplog): + caplog.set_level("INFO", logger="bot") + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report")) + ss.data_from = "2026-02-15" # everything ending before this has no data: dozens of empty windows in a row + ctx, _ = make(env, store, ss) + args = ctx, ss, [(7, "RobSMSF", date(2026, 1, 15))], date(2026, 3, 3) + stored = await backfill.fill_once(*args) + assert stored > 0 and store.period_count(7, "2026-01-01", "2026-03-31") == stored # later windows still got stored + assert "had no data" in caplog.text and "fetches in a row failed" not in caplog.text + shown = [r.message for r in caplog.records if "has no data:" in r.message] + assert 0 < len(shown) <= 3 * 5 and "no data for this period" in shown[0] # a few samples per step, with Sharesight's words + ss.order.clear() + assert await backfill.fill_once(*args) == 0 and ss.order == [] # the empty ones are skipped until reset() + + +async def test_throttling_text_counts_as_a_transient_failure_and_is_logged(env, store, caplog): + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report"), fail=True) + ss.fail_text = "Tool error: HTTP 429 Too Many Requests" + ctx, _ = make(env, store, ss) + assert await backfill.fill_once(ctx, ss, [(7, "RobSMSF", date(2026, 1, 15))], date(2026, 3, 3)) == 0 + assert "fetches in a row failed" in caplog.text and "429 Too Many Requests" in caplog.text + + +async def test_the_journal_says_when_it_is_held_back_or_a_fetch_is_slow(env, store, monkeypatch, caplog): + caplog.set_level("INFO", logger="bot") + monkeypatch.setattr(backfill, "WAIT_LOG", 0) + monkeypatch.setattr(backfill, "SLOW_FETCH", 0.005) + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report")) + ss.busy_for = 3 + ss.delay = 0.01 + ctx, _ = make(env, store, ss) + await backfill.fill_once(ctx, ss, [(7, "RobSMSF", date(2026, 3, 1))], date(2026, 3, 3)) + assert "held back for 0 s by: get_performance_report portfolio_id=7" in caplog.text and "was slow" in caplog.text + + +async def test_progress_is_logged_every_n_fetches(env, store, monkeypatch, caplog): + caplog.set_level("INFO", logger="bot") + monkeypatch.setattr(backfill, "PROGRESS_EVERY", 5) + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report")) + ctx, _ = make(env, store, ss) + await backfill.fill_once(ctx, ss, [(7, "RobSMSF", date(2026, 1, 15))], date(2026, 3, 3)) + assert "RobSMSF days 5/" in caplog.text and "RobSMSF days 10/" in caplog.text + + +async def test_the_loop_runs_a_pass_then_sleeps_until_woken(env, store, monkeypatch): + recent = (date.today() - timedelta(days=4)).isoformat() + ss = Sharesight("sharesight", tools=("list_portfolios", "get_performance_report"), inception=recent) + ctx, _ = make(env, store, ss) + monkeypatch.setattr(backfill, "RECHECK", 3600) + task = asyncio.create_task(backfill.loop(ctx)) + expected = len(backfill.wanted(date.today() - timedelta(days=4), date.today())) + for _ in range(500): # the store runs in a thread, so give it real time + await asyncio.sleep(0.01) + if store.period_count(7, recent, "2100-01-01") == expected: + break + assert expected and store.period_count(7, recent, "2100-01-01") == expected + assert store.period_count(8, "2000-01-01", "2100-01-01") == 0 + n = len(ss.order) + backfill.reset() # as after a flush: runs again straight away (nothing new, so nothing fetched) + for _ in range(50): + await asyncio.sleep(0.01) + if not backfill.wake_event().is_set(): + break + await asyncio.sleep(0.05) + assert len(ss.order) == n and not backfill.wake_event().is_set() + task.cancel() + + +def test_keys_match_the_tool_windows(): + w = period_windows("month", date(2026, 1, 15), date(2026, 3, 2))[0] + assert key(w) == ("2026-01-15", "2026-01-31") diff --git a/tests/test_config.py b/tests/test_config.py index 2f975d4..d8cc75e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -6,6 +6,34 @@ from lib import config +def test_data_dir_holds_the_data_files(): + base = {"TELEGRAM_BOT_TOKEN": "1:x", "XAI_API_KEY": "a"} + mcp = str(config.ROOT / "mcp_servers.json") # next to bot.py, whatever DATA_DIR is + plain = config.load(base) + assert (plain.db_path, plain.mcp_config) == ("chat_log.db", mcp) + s = config.load({**base, "DATA_DIR": "/data"}) + assert (s.db_path, s.mcp_config) == ("/data/chat_log.db", mcp) + s = config.load({**base, "DATA_DIR": "/data", "DB_PATH": "other.db", "MCP_CONFIG": "/etc/mcp.json"}) + assert (s.db_path, s.mcp_config) == ("/data/other.db", "/etc/mcp.json") # relative joins, absolute stays + + +def test_data_dir_defaults_to_systemd_state_directory(): + base = {"TELEGRAM_BOT_TOKEN": "1:x", "XAI_API_KEY": "a"} + assert config.load({**base, "STATE_DIRECTORY": "/var/lib/x"}).db_path == "/var/lib/x/chat_log.db" + assert config.load({**base, "STATE_DIRECTORY": "/a:/b"}).db_path == "/a/chat_log.db" # first of several + assert config.load({**base, "STATE_DIRECTORY": "/a", "DATA_DIR": "/data"}).db_path == "/data/chat_log.db" + assert config.load({**base, "STATE_DIRECTORY": "/a", "DB_PATH": "/x/y.db"}).db_path == "/x/y.db" + + +def test_period_backfill_is_all_by_default_and_can_be_listed_or_off(): + base = {"TELEGRAM_BOT_TOKEN": "1:x", "XAI_API_KEY": "a"} + assert config.load(base).period_backfill == {"all"} + assert config.load({**base, "PERIOD_BACKFILL": "off"}).period_backfill == frozenset() + assert config.load({**base, "PERIOD_BACKFILL": ""}).period_backfill == frozenset() + assert config.load({**base, "PERIOD_BACKFILL": "RobSMSF, bobsmsf"}).period_backfill == {"robsmsf", "bobsmsf"} + assert "all" in config.load({**base, "PERIOD_BACKFILL": "all"}).period_backfill + + def test_both_keys_refused(): with pytest.raises(config.ConfigError, match="only one"): config.load({"TELEGRAM_BOT_TOKEN": "1:x", "XAI_API_KEY": "a", "OPENROUTER_API_KEY": "b"}) diff --git a/tests/test_features.py b/tests/test_features.py index 7d3f2e9..afc607a 100644 --- a/tests/test_features.py +++ b/tests/test_features.py @@ -207,3 +207,30 @@ async def test_group_admin_can_run_credits_in_the_group_and_in_a_dm(env, store): @pytest.mark.parametrize("cached,tin,pct", [(0, 0, 0), (50, 100, 50)]) def test_usage_text_percentages(cached, tin, pct): assert f"({pct}% cached)" in commands.usage_text([("m", 1, tin, cached, 1, 0.0)], "x") + + +async def test_flushcache_command_is_admin_only_and_flushes_by_portfolio_and_period(env, store): + from lib.features.periods import backfill + from lib.store import PeriodRow + + st = config.load({**env, "ADMIN_CHAT_IDS": "5", "PERIOD_BACKFILL": "x"}) + ctx = Ctx(st, store, ScriptedBackend([]), registry(), FakeBot()) + store.period_put([PeriodRow(p, f"{y}-0{m}-01", f"{y}-0{m}-28", 1.0, 1.0, "AUD", []) + for p in (70007, 80008) for y in (2024, 2025) for m in (1, 2)]) + count = lambda p: store.period_count(p, "2000-01-01", "2100-01-01") # noqa: E731 + stranger = user_msg(ctx.bot, "/flushcache", user_id=9, mid=1) + await commands.flush_cache(ctx, update_for(stranger), N(args=[])) + assert stranger.replies == [] and count(70007) == 4 + year = user_msg(ctx.bot, "/flushcache 70007 2024", user_id=5, mid=2) + await commands.flush_cache(ctx, update_for(year), N(args=["70007", "2024"])) + assert "Flushed 2 cached periods (portfolio 70007, 2024-01-01 to 2024-12-31). Refilling" in year.replies[0]["text"] + assert count(70007) == 2 and count(80008) == 4 and backfill.wake_event().is_set() + month = user_msg(ctx.bot, "/flushcache 2025-01", user_id=5, mid=3) + await commands.flush_cache(ctx, update_for(month), N(args=["2025-01"])) + assert "Flushed 2 cached periods (all portfolios, 2025-01-01 to 2025-01-31)" in month.replies[0]["text"] + bad = user_msg(ctx.bot, "/flushcache x", user_id=5, mid=4) + await commands.flush_cache(ctx, update_for(bad), N(args=["x"])) + assert "Usage" in bad.replies[0]["text"] + every = user_msg(ctx.bot, "/flushcache", user_id=5, mid=5) + await commands.flush_cache(ctx, update_for(every), N(args=[])) + assert "Flushed 4 cached periods (all portfolios)" in every.replies[0]["text"] diff --git a/tests/test_handlers.py b/tests/test_handlers.py index f3ae829..1b2c487 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -97,6 +97,22 @@ async def test_market_question_offers_tools_up_front(env, store): assert backend.seen[0]["tools"] == ["yahoo__get_quote"] and backend.seen[0]["search"] is True +async def test_a_short_follow_up_to_a_portfolio_reply_still_gets_the_portfolio_tools(env, store): + sharesight = FakeMcp("sharesight", gate="portfolio", tools=("list_portfolios",)) + ctx, backend, h = make_ctx(env, store, [step("Your portfolio best month was May"), step("In dollars: ..."), step("c"), step("d")], + servers=[sharesight]) + await run(h, user_msg(ctx.bot, "@stockbot how is my portfolio doing?", mid=1)) + await run(h, user_msg(ctx.bot, "@stockbot what about in dollar terms?", mid=2000)) + assert "sharesight__list_portfolios" in backend.seen[1]["tools"] and "period_returns" in backend.seen[1]["tools"] + # too long after the reply: a plain message again + await run(h, user_msg(ctx.bot, "@stockbot and what about the other one?", mid=3000, ts=100_000)) + assert backend.seen[2]["tools"] == [] + # a long unrelated message right after a reply is not read as a follow-up + long = "@stockbot " + " ".join(["tell me a long story about the history of the word mortgage"] * 3) + await run(h, user_msg(ctx.bot, long, mid=4000)) + assert backend.seen[3]["tools"] == [] + + async def test_portfolio_question_reminds_the_model_to_fetch_fresh_figures(env, store): sharesight = FakeMcp("sharesight", gate="portfolio", tools=("list_portfolios",)) ctx, backend, h = make_ctx(env, store, [step("a"), step("b")], servers=[sharesight]) @@ -104,6 +120,7 @@ async def test_portfolio_question_reminds_the_model_to_fetch_fresh_figures(env, await run(h, user_msg(ctx.bot, "@stockbot how is NVDA looking?", mid=2)) assert "portfolio tools before answering" in backend.seen[0]["prompt"] assert "tailwinds and headwinds" in backend.seen[0]["prompt"] + assert "your market tools" in backend.seen[0]["prompt"] # not portfolio tools only (Yahoo comparisons) assert "portfolio tools before answering" not in backend.seen[1]["prompt"] @@ -252,7 +269,7 @@ async def test_non_admin_gets_no_portfolio_data_in_a_group(env, store): assert asked.replies[0]["text"] == "Portfolio data is for admins only." and backend.seen == [] admin = user_msg(ctx.bot, "@stockbot how is my portfolio doing?", user_id=5, mid=2) await run(h, admin) - assert backend.seen[0]["tools"] == ["sharesight__list_portfolios"] + assert backend.seen[0]["tools"] == ["sharesight__list_portfolios", "period_returns"] # admins also get the period tool async def test_non_admin_never_gets_portfolio_tools_even_on_the_rerun(env, store): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 42ec112..24cc860 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -224,6 +224,68 @@ async def test_sharesight_holdings_only_flag_sets_include_sales(): assert s.session.calls[0][1]["include_sales"] is False +async def test_bots_own_tools_can_keep_sold_holdings(): + s = MCPServer("sharesight", {"current_holdings_only": True}, timeout=5, max_output=10_000) + s.session = FakeSession() + await s.call("get_performance_report", {"portfolio_id": 1, "_keep_sales": True}) + sent = s.session.calls[0][1] + assert sent["include_sales"] is True and "_keep_sales" not in sent + + +async def test_bots_own_tools_are_not_truncated_but_the_models_are(): + s = MCPServer("sharesight", {}, timeout=5, max_output=20) + s.session = FakeSession() + assert "truncated" in await s.call("get_performance_report", {"portfolio_id": 1}) + whole = await s.call("get_performance_report", {"portfolio_id": 1, "_keep_sales": True}) + assert "truncated" not in whole and len(whole) > 20 + + +async def test_user_calls_describe_what_is_running_but_not_background_work(): + s = MCPServer("sharesight", {"max_concurrent": 2}, timeout=5, max_output=1000) + gate = asyncio.Event() + + class Held(FakeSession): + async def call_tool(self, name, args): + await gate.wait() + return await super().call_tool(name, args) + + s.session = Held() + chat = asyncio.create_task(s.call("get_performance_report", {"portfolio_id": 7, "start_date": "2024-01-01", "end_date": "2024-01-31"})) + own = asyncio.create_task(s.call("get_performance_report", {"portfolio_id": 7, "start_date": "2024-02-01", "end_date": "2024-02-29", "_keep_sales": True})) + bg = asyncio.create_task(s.call("get_performance_report", {"portfolio_id": 7, "start_date": "2024-03-01", "end_date": "2024-03-31", "_keep_sales": True, "_bg": True})) + await asyncio.sleep(0.01) + assert s.in_flight == 3 + assert s.user_calls() == ["get_performance_report portfolio_id=7 2024-01-01..2024-01-31", + "get_performance_report portfolio_id=7 2024-02-01..2024-02-29 (period_returns)"] + gate.set() + await asyncio.gather(chat, own, bg) + assert s.user_calls() == [] and s.in_flight == 0 + + +async def test_in_flight_counts_running_calls_and_returns_to_zero(): + s = MCPServer("y", {"max_concurrent": 1}, timeout=0.05, max_output=1000) + gate = asyncio.Event() + + class Slow(FakeSession): + async def call_tool(self, name, args): + if name == "hang": + await asyncio.sleep(1) # longer than the timeout + await gate.wait() + return await super().call_tool(name, args) + + s.session = Slow() + first = asyncio.create_task(s.call("get_a", {})) + second = asyncio.create_task(s.call("get_b", {})) # waits for the only slot, still counted + await asyncio.sleep(0.01) + assert s.in_flight == 2 + gate.set() + await asyncio.gather(first, second) + assert s.in_flight == 0 + with pytest.raises(TimeoutError): + await s.call("hang", {}) + assert s.in_flight == 0 # also after a timeout + + async def test_output_cap(): s = MCPServer("y", {}, timeout=5, max_output=20) s.session = FakeSession() diff --git a/tests/test_periods.py b/tests/test_periods.py new file mode 100644 index 0000000..84e8899 --- /dev/null +++ b/tests/test_periods.py @@ -0,0 +1,197 @@ +import json +from datetime import date, datetime + +import pytest + +from lib import config +from lib.context import Ctx +from lib.features import periods +from lib.features.periods import windows +from lib.store import PeriodRow + +from .conftest import FakeBot +from .fakes import FakeMcp, ScriptedBackend, registry + + +class Sharesight(FakeMcp): + """Reports by start date: +1% in 2026-01, +5% in 2026-02, -2% in 2025-12; no data otherwise.""" + + async def call(self, tool, args): + self.calls.append((tool, args)) + if tool == "list_portfolios": + return json.dumps({"portfolios": [{"id": 7, "inception_date": "2025-12-20"}]}) + pct = {"2025-12": -2.0, "2026-01": 1.0, "2026-02": 5.0}.get(args["start_date"][:7]) + if pct is None: + return "Tool error: no data" + rows = [["AAA", pct * 150, 40.0], ["BBB", -pct * 50, -10.0], ["CCC", 0, 0]] # a sold holding could be among them + return json.dumps({"report": {"total_gain_percent": pct, "total_gain": pct * 100, "currency": "AUD", + "holdings": {"columns": ["code", "total_gain", "total_gain_percent"], "rows": rows}}}) + + +class Always(Sharesight): + async def call(self, tool, args): + self.calls.append((tool, args)) + return json.dumps({"report": {"total_gain_percent": 1.0, "total_gain": 10.0, "currency": "AUD"}}) + + +def make(env, store, ss=None): + ss = ss or Sharesight("sharesight", tools=("list_portfolios", "get_performance_report")) + return Ctx(config.load(env), store, ScriptedBackend([]), registry(ss), FakeBot()), ss + + +def reports(ss): + return [a for t, a in ss.calls if t == "get_performance_report"] + + +async def test_months_run_from_inception_with_sold_holdings_and_best_worst(env, store): + ctx, ss = make(env, store) + defn, fn = periods.tool(ctx) + out = await fn({"portfolio_id": 7, "end": "2026-03"}) + assert defn.name == "period_returns" and "Total return per month (AUD)" in out + assert "2025-12 -2.00% -200" in out and "2026-02 +5.00% +500" in out + assert "Best by %: 2026-02 +5.00% (+500). Worst by %: 2025-12 -2.00% (-200)." in out + assert "Best by money: 2026-02 +5.00% (+500). Worst by money: 2025-12 -2.00% (-200). 1 months had no data." in out # March: no data + first = reports(ss)[0] + assert (first["start_date"], first["end_date"]) == ("2025-12-20", "2025-12-31") # clipped to inception + assert all(r["_keep_sales"] is True for r in reports(ss)) + + +async def test_best_and_worst_periods_say_which_holdings_drove_them(env, store): + ctx, ss = make(env, store) + _, fn = periods.tool(ctx) + out = await fn({"portfolio_id": 7, "end": "2026-03"}) + assert "Holdings behind 2026-02: up: AAA +750 (+40.0%); down: BBB -250 (-10.0%)" in out + assert "Holdings behind 2025-12: up: BBB +100 (-10.0%); down: AAA -300 (+40.0%)" in out # worst month + assert "Holdings behind 2026-01" not in out and "CCC" not in out # only the extremes, only movers + + +async def test_finished_periods_are_cached_forever_and_served_without_sharesight(env, store): + ctx, ss = make(env, store) + _, fn = periods.tool(ctx) + first = await fn({"portfolio_id": 7, "start": "2025-12", "end": "2026-02"}) + n = len(reports(ss)) + assert n == 3 and store.period_count(7, "2000-01-01", "2100-01-01") == 3 + assert await fn({"portfolio_id": 7, "start": "2025-12", "end": "2026-02"}) == first # same text, from the cache + assert len(reports(ss)) == n + assert store.period_flush(7) == 3 and store.period_flush() == 0 + await fn({"portfolio_id": 7, "start": "2025-12", "end": "2026-02"}) + assert len(reports(ss)) == 2 * n # fetched again after the flush + + +async def test_the_unfinished_period_is_never_cached(env, store): + ctx, ss = make(env, store, Always("sharesight", tools=("list_portfolios", "get_performance_report"))) + _, fn = periods.tool(ctx) + today = datetime.now(ctx.st.tz).date() + await fn({"portfolio_id": 7, "start": today.strftime("%Y-%m")}) # this month, ends today + assert len(reports(ss)) == 1 and store.period_count(7, "2000-01-01", "2100-01-01") == 0 + + +async def test_quarters_and_years_are_asked_for_whole_windows(env, store): + ctx, ss = make(env, store) + _, fn = periods.tool(ctx) + out = await fn({"portfolio_id": 7, "period": "quarter", "start": "2025-10", "end": "2026-06"}) + assert [(r["start_date"], r["end_date"]) for r in reports(ss)] == [ + ("2025-10-01", "2025-12-31"), ("2026-01-01", "2026-03-31"), ("2026-04-01", "2026-06-30")] + assert "2026-Q1 +1.00%" in out and "2 quarters had no data" in out + ctx, ss = make(env, store) + _, fn = periods.tool(ctx) + await fn({"portfolio_id": 7, "period": "year", "start": "2024", "end": "2025-06"}) + assert [(r["start_date"], r["end_date"]) for r in reports(ss)] == [ + ("2024-01-01", "2024-12-31"), ("2025-01-01", "2025-06-30")] + + +async def test_start_given_skips_the_inception_lookup_and_bad_input_is_reported(env, store): + ctx, ss = make(env, store) + _, fn = periods.tool(ctx) + out = await fn({"portfolio_id": 7, "start": "2026-01", "end": "2026-02"}) + assert out.splitlines()[1].startswith("2026-01 +1.00%") and not any(t == "list_portfolios" for t, _ in ss.calls) + assert "portfolio_id" in await fn({}) + assert "period must be" in await fn({"portfolio_id": 7, "period": "decade"}) + assert "start as YYYY-MM-DD" in await fn({"portfolio_id": 99}) # unknown portfolio: no inception date + assert "after end" in await fn({"portfolio_id": 7, "start": "2026-03", "end": "2026-01"}) + none_ctx = Ctx(config.load(env), store, ScriptedBackend([]), registry(), FakeBot()) + assert "isn't connected" in await periods.tool(none_ctx)[1]({"portfolio_id": 7}) + + +async def test_only_uncached_windows_count_towards_the_live_limit(env, store): + ctx, ss = make(env, store) + _, fn = periods.tool(ctx) + refused = await fn({"portfolio_id": 7, "period": "day", "start": "2020-01-01"}) + assert "0 of " in refused and "days are cached" in refused and "start=" in refused and reports(ss) == [] + # Cache every weekday from 2025-01-01 to 2026-03-31 except the last ten: the rest fits the live budget. + days = windows.period_windows("day", date(2025, 1, 1), date(2026, 3, 31)) + store.period_put([PeriodRow(7, a.isoformat(), b.isoformat(), 0.5, 5.0, "AUD", []) for _, a, b in days[:-10]]) + out = await fn({"portfolio_id": 7, "period": "day", "start": "2025-01-01", "end": "2026-03-31"}) + assert not out.startswith("Error") and len(reports(ss)) == 10 # only the ten uncached days were fetched + + +@pytest.mark.parametrize("period, first, last, labels", [ + ("month", date(2025, 11, 15), date(2026, 2, 10), ["2025-11", "2025-12", "2026-01", "2026-02"]), + ("quarter", date(2025, 8, 1), date(2026, 2, 10), ["2025-Q3", "2025-Q4", "2026-Q1"]), + ("year", date(2024, 6, 1), date(2026, 2, 10), ["2024", "2025", "2026"]), + ("week", date(2026, 3, 4), date(2026, 3, 17), ["2026-W10", "2026-W11", "2026-W12"]), + ("day", date(2026, 3, 6), date(2026, 3, 10), ["2026-03-06", "2026-03-09", "2026-03-10"]), # Fri, Mon, Tue +]) +def test_period_windows_labels_and_weekends(period, first, last, labels): + w = windows.period_windows(period, first, last) + assert [x[0] for x in w] == labels + assert w[0][1] == first and w[-1][2] == last # clipped to the requested range + + +def test_dates_accept_day_month_or_year_and_end_means_the_period_end(): + d = windows.parse_date + assert d("2026-02-03") == date(2026, 2, 3) + assert d("2026-02") == date(2026, 2, 1) and d("2026-02", end=True) == date(2026, 2, 28) + assert d("2025-12", end=True) == date(2025, 12, 31) and d("2025", end=True) == date(2025, 12, 31) + assert d("nope") is None + + +def test_refusal_suggests_the_earliest_start_that_fits_and_ignores_cached_windows(): + w = windows.period_windows("week", date(2026, 1, 5), date(2027, 1, 3)) # 52 weeks + msg = windows.refusal("week", w, set()) + assert "0 of 52 weeks are cached" in msg and f"start={w[-26][1]}" in msg + assert windows.refusal("week", w, {windows.key(x) for x in w[:30]}) is None # 22 uncached: fits + + +def test_flush_scope_understands_portfolio_year_month_day_and_all(): + f = windows.flush_scope + assert f([]) == (None, None, None) and f(["all"]) == (None, None, None) + assert f(["684141"]) == (684141, None, None) + assert f(["2024"]) == (None, date(2024, 1, 1), date(2024, 12, 31)) + assert f(["2024-05", "684141"]) == (684141, date(2024, 5, 1), date(2024, 5, 31)) + assert f(["2024-05-17"]) == (None, date(2024, 5, 17), date(2024, 5, 17)) + assert f(["nope"]) is None and f(["684141", "702349"]) is None and f(["2024", "2025"]) is None + assert f(["1234"]) is None # 4 digits that aren't a year in range + + +async def test_best_by_money_can_differ_from_best_by_percent(env, store): + class Mixed(Sharesight): + async def call(self, tool, args): + self.calls.append((tool, args)) + pct, gain = {"2026-01": (10.0, 100.0), "2026-02": (2.0, 900.0)}[args["start_date"][:7]] + return json.dumps({"report": {"total_gain_percent": pct, "total_gain": gain, "currency": "AUD"}}) + + ctx, _ = make(env, store, Mixed("sharesight", tools=("list_portfolios", "get_performance_report"))) + _, fn = periods.tool(ctx) + out = await fn({"portfolio_id": 7, "start": "2026-01", "end": "2026-02"}) + assert "Best by %: 2026-01 +10.00% (+100)" in out and "Best by money: 2026-02 +2.00% (+900)" in out + assert "Worst by money: 2026-01 +10.00% (+100)" in out + + +async def test_long_results_list_only_the_extremes_not_every_row(env, store): + class Spread(Sharesight): + async def call(self, tool, args): + self.calls.append((tool, args)) + day = int(args["start_date"][-2:]) + 31 * int(args["start_date"][5:7]) # a different number every day + return json.dumps({"report": {"total_gain_percent": day / 10, "total_gain": -day * 10.0, "currency": "AUD"}}) + + ctx, _ = make(env, store, Spread("sharesight", tools=("list_portfolios", "get_performance_report"))) + _, fn = periods.tool(ctx) + out = await fn({"portfolio_id": 7, "period": "day", "start": "2026-01-01", "end": "2026-03-31"}) + lines = out.splitlines() + assert len(lines) < 15 and "full list is left out" in lines[0] and "64 days" in lines[0] + assert lines[1].startswith("Top 5 by %: 2026-03-31 ") and lines[2].startswith("Bottom 5 by %: 2026-01-01 ") + assert lines[3].startswith("Top 5 by money: 2026-01-01 ") and lines[4].startswith("Bottom 5 by money: 2026-03-31 ") + assert "Best by %:" in out + short = await fn({"portfolio_id": 7, "period": "day", "start": "2026-03-02", "end": "2026-03-13"}) + assert "Top 5" not in short and "2026-03-02 " in short # a short range still lists every row diff --git a/tests/test_runner.py b/tests/test_runner.py index 6ba5e2d..9e775b5 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -35,6 +35,14 @@ async def test_tool_round_trip_runs_mcp_and_returns_answer(): assert ans.usage.tokens_in == 200 and ans.usage.cached == 80 +async def test_the_model_cannot_ask_for_sold_holdings_to_be_kept(): + from lib.llm.base import Call + mcp = FakeMcp("sharesight", tools=("get_performance_report",)) + await run_calls([Call("1", "sharesight__get_performance_report", '{"portfolio_id":1,"_keep_sales":true,"_bg":true}')], + ScriptedBackend([]), registry(mcp), with_tools(mcp), {}, True, Usage()) + assert mcp.calls == [("get_performance_report", {"portfolio_id": 1})] + + async def test_final_round_forces_answer_without_tools(): mcp = FakeMcp() loop = [step(calls=[(str(i), "yahoo__get_quote", f'{{"n":{i}}}')], finish="tool_calls") @@ -246,6 +254,7 @@ def test_gate_full_when_forced_or_token_saver_off(settings, env): "my position in CBA", "am I up or down today", "show performance this year", "bob smsf vs bob personal", "are my dividends coming", "why did I outperform Alice over the past month?", "how did I do last month", "did I beat Alice this year", "why did Alice underperform me", "who's ahead of who this quarter, me or Alice", + "what was the best month for robsmsf since inception?", "how is BobSMSF doing", ] diff --git a/tests/test_store.py b/tests/test_store.py index f04f7a1..397dde5 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -80,3 +80,32 @@ def test_writes_are_committed_immediately(tmp_path): second = Store(path) # another process sharing the file sees the writes without any explicit commit second.bot_id = 7 assert second.kv_get("k") == "v" and second.chats() == [(-1, "Group")] + + +def test_period_cache_round_trip_flush_and_version(tmp_path): + from lib.store import PeriodRow + + path = str(tmp_path / "p.db") + s = Store(path) + s.period_put([PeriodRow(7, "2026-01-01", "2026-01-31", 1.5, 150.0, "AUD", [["AAA", 90.0, 12.0]]), + PeriodRow(7, "2026-02-01", "2026-02-28", -2.0, -200.0, "AUD", []), + PeriodRow(8, "2026-01-01", "2026-01-31", 3.0, 30.0, "USD", [])]) + got = s.period_get(7, "2026-01-01", "2026-02-15") # only windows lying within the range + assert list(got) == [("2026-01-01", "2026-01-31")] and got[("2026-01-01", "2026-01-31")].movers == [["AAA", 90.0, 12.0]] + assert s.period_count(7, "2000-01-01", "2100-01-01") == 2 + s.period_put([PeriodRow(7, "2026-01-01", "2026-01-31", 9.0, 1.0, "AUD", [])]) # replaces + assert s.period_get(7, "2026-01-01", "2026-01-31")[("2026-01-01", "2026-01-31")].pct == 9.0 + s._conn.execute("UPDATE period_cache SET v = 0 WHERE portfolio_id = 8") # an older format + assert Store(path).period_count(8, "2000-01-01", "2100-01-01") == 0 # dropped on open + assert s.period_flush(7) == 2 and s.period_flush() == 0 + + +def test_period_flush_by_date_range_removes_every_window_that_overlaps_it(tmp_path): + from lib.store import PeriodRow + + s = Store(str(tmp_path / "p.db")) + row = lambda a, b: PeriodRow(7, a, b, 1.0, 1.0, "AUD", []) # noqa: E731 + s.period_put([row("2023-12-26", "2024-01-01"), row("2024-12-31", "2025-01-06"), row("2024-06-01", "2024-06-30"), + row("2023-01-01", "2023-12-31"), row("2025-01-07", "2025-01-13")]) + assert s.period_flush(7, "2024-01-01", "2024-12-31") == 3 # the week spanning New Year both ways, and June + assert sorted(s.period_get(7, "2000-01-01", "2100-01-01")) == [("2023-01-01", "2023-12-31"), ("2025-01-07", "2025-01-13")] diff --git a/tools/capture.py b/tools/capture.py index 60137b0..6f9a916 100644 --- a/tools/capture.py +++ b/tools/capture.py @@ -106,7 +106,7 @@ async def capture(server: MCPServer, days: int) -> None: async def run(label: str, days: int) -> int: - registry = Registry.load(os.environ.get("MCP_CONFIG", "mcp_servers.json"), timeout=config.MCP_TIMEOUT, + registry = Registry.load(config.mcp_config_path(os.environ), timeout=config.MCP_TIMEOUT, max_output=10**7) server = registry.servers.get(label) if server is None: