"""Per-user budgets for the tools that reach out and touch a named host. dns_lookup, inspect_tls, http_headers and lookup_ip_or_asn all make this server connect to a host of the user's choosing. squid already guarantees that host is on the public internet, so this is not what keeps the operator's network safe - it is what stops one of the five guest accounts from turning a shared assistant into someone else's traffic source. Keyed on X-OpenWebUI-User-Id, which arrives because the open-webui container runs with ENABLE_FORWARD_USER_INFO_HEADERS=True. """ from __future__ import annotations import logging import time from fastapi import Header, HTTPException, status from .cache import get_redis from .settings import settings log = logging.getLogger(__name__) # Distinct pages per conversation before the tool starts saying so in its own # output. A first, gentler note lands earlier: by five separate routes to one # fact, the useful decision is usually "answer with what you have", and saying # it once at five is cheaper than saying it loudly at nine. FETCH_SOFT_LIMIT = 8 FETCH_FIRST_NOTE = 5 class Identity: __slots__ = ("user_id", "email", "role", "chat_id") def __init__(self, user_id: str | None, email: str | None, role: str | None, chat_id: str | None): self.user_id = user_id or "anonymous" self.email = email self.role = role self.chat_id = chat_id def __str__(self) -> str: return f"{self.email or self.user_id}({self.role or 'unknown'})" async def identity( x_openwebui_user_id: str | None = Header(default=None, include_in_schema=False), x_openwebui_user_email: str | None = Header(default=None, include_in_schema=False), x_openwebui_user_role: str | None = Header(default=None, include_in_schema=False), x_openwebui_chat_id: str | None = Header(default=None, include_in_schema=False), ) -> Identity: return Identity(x_openwebui_user_id, x_openwebui_user_email, x_openwebui_user_role, x_openwebui_chat_id) def _decode(value) -> str: return value.decode() if isinstance(value, bytes) else str(value) async def note_fetch(chat_id: str | None, url: str | None = None) -> str: """Count distinct pages opened in one conversation and say so as they add up. A system-prompt rule telling the model to stop after ~six fetches did not hold: asked for a specific number out of a paywalled paper, it made 33 tool calls across three "let me try another route" rounds and never produced an answer at all. The instruction loses to an explicit "give me the numbers". A count in the tool's own output does not lose, because the model reads it as a fact about the world rather than as advice. Keyed on the chat id that Open WebUI forwards, so it is per-conversation and resets on its own. What changed since: it counted *calls*, so paging through one long article with `offset` spent the whole budget on a single document and hit STOP FETCHING on page nine of one Wikipedia page - the precise opposite of the instruction that page needed, which was to keep reading. Counting distinct URLs instead means continuing inside a document you already opened is free, and only genuinely new routes to the answer are charged. Cache hits stop costing anything too, which they never should have. """ if not chat_id: return "" client = await get_redis() if client is None: return "" key = f"owuitools:fetchurls:{chat_id}" try: # No URL means the caller cannot tell pages apart; fall back to the old # per-call behaviour rather than losing the budget entirely. added = await client.sadd(key, url or f"call:{await client.incr(key + ':seq')}") if not added: return "" n = await client.scard(key) if n == 1: await client.expire(key, 1800) except Exception: return "" if n == FETCH_FIRST_NOTE: return (f"[Page {n} of this conversation. If the specific fact is still not in front " "of you after the next one, plan to answer with what you have and name what " "is missing. Reading further into a page you have already opened does not " "count towards this.]") if n == FETCH_SOFT_LIMIT: return (f"[You have now opened {n} different pages in this conversation. If the " "specific fact you are chasing is still not in front of you, it is probably " "not openly available. Consider answering with what you have and naming what " "is missing.]") if n > FETCH_SOFT_LIMIT: return (f"[STOP FETCHING. This is the {n}th different page in this conversation, past " "the point where more retrieval helps. The fact you are looking for is not " "openly available. Answer now with what you established, say plainly which " "part you could not verify and why, and let the user decide. Do not keep " "trying new routes to the same source. Continuing inside a page you already " "opened, with offset= or search=, does not count and is still fine.]") return "" async def recall_failure(chat_id: str | None, url: str) -> str | None: """Why this URL already failed in this conversation, if it did. Retrying a URL that has already 403'd or timed out produces the identical error, costs a call, and - measured - is what a stuck model does instead of answering: one graded run retried a 504'ing mirror and a timed-out vendor URL after both had already failed, and stopped only when the budget forced it. The tool can simply decline, which is both cheaper and a clearer signal than a second identical failure. """ if not chat_id or not url: return None client = await get_redis() if client is None: return None try: reason = await client.hget(f"owuitools:fetchfail:{chat_id}", url) except Exception: return None return _decode(reason) if reason is not None else None async def note_failure(chat_id: str | None, url: str, reason: str) -> None: """Remember that this URL failed, so the next attempt can be declined.""" if not chat_id or not url: return client = await get_redis() if client is None: return key = f"owuitools:fetchfail:{chat_id}" try: await client.hset(key, url, " ".join(str(reason).split())[:200]) await client.expire(key, 1800) except Exception: pass async def note_success(chat_id: str | None, url: str) -> None: """Remember that this URL was read, so search results can say so.""" if not chat_id or not url: return client = await get_redis() if client is None: return key = f"owuitools:fetchok:{chat_id}" try: await client.sadd(key, url) await client.expire(key, 1800) except Exception: pass async def fetch_history(chat_id: str | None) -> tuple[set[str], dict[str, str]]: """URLs already read, and URLs that already failed, in this conversation. `web_search` annotates its results with these. A result the model has already read is not a new lead, and one that already failed is not worth a second call - but from inside a search result they look identical to fresh ones, and the model re-fetched both in graded runs. """ if not chat_id: return set(), {} client = await get_redis() if client is None: return set(), {} try: read = {_decode(u) for u in await client.smembers(f"owuitools:fetchok:{chat_id}")} failed = {_decode(k): _decode(v) for k, v in (await client.hgetall(f"owuitools:fetchfail:{chat_id}")).items()} except Exception: return set(), {} return read, failed async def spend_probe(who: Identity, tool: str) -> None: """Consume one unit of the caller's hourly probe budget. Fails open on cache trouble. A rate limiter that takes the tools down when valkey hiccups causes more damage than the abuse it prevents. """ limit = settings.probe_rate_per_hour if limit <= 0: return client = await get_redis() if client is None: return window = int(time.time()) // 3600 key = f"owuitools:probe:{who.user_id}:{window}" try: used = await client.incr(key) if used == 1: await client.expire(key, 3700) except Exception: return if used > limit: log.warning("probe budget exhausted for %s on %s (%d/%d)", who, tool, used, limit) raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=( f"This user has used all {limit} network probes allowed this hour. " "Do not retry. Tell the user the probe limit was reached and answer " "from what you already know, or say you cannot verify it." ), )