226 lines
8.8 KiB
Python
226 lines
8.8 KiB
Python
"""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."
|
|
),
|
|
)
|