Initial Commit
This commit is contained in:
@@ -0,0 +1,225 @@
|
||||
"""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."
|
||||
),
|
||||
)
|
||||
Reference in New Issue
Block a user