mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-27 02:45:24 +08:00
feat(media): centralize resource URLs and persist refreshed image URLs
- Deduplicate message and history resource URLs using stable catalog references - Persist renewed URLs with concurrent refresh protection - Refresh OneBot images via get_msg intervals or NapCat rkey TTL - Migrate historical payloads and cover SQLite and PostgreSQL behavior
This commit is contained in:
@@ -7,6 +7,7 @@ import datetime
|
||||
import sqlalchemy
|
||||
|
||||
from .base import Base
|
||||
from ...persistence.media import MediaReferenceText
|
||||
|
||||
|
||||
class AgentInteraction(Base):
|
||||
@@ -29,15 +30,15 @@ class AgentInteraction(Base):
|
||||
actor_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True)
|
||||
|
||||
status = sqlalchemy.Column(sqlalchemy.String(50), nullable=False, index=True)
|
||||
request_json = sqlalchemy.Column(sqlalchemy.Text, nullable=False)
|
||||
request_json = sqlalchemy.Column(MediaReferenceText, nullable=False)
|
||||
delivery_target_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
|
||||
delivery_result_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
|
||||
delivery_result_json = sqlalchemy.Column(MediaReferenceText, nullable=True)
|
||||
replaces_interaction_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True)
|
||||
callback_token_hash = sqlalchemy.Column(sqlalchemy.String(64), nullable=False, unique=True, index=True)
|
||||
|
||||
expires_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True, index=True)
|
||||
submitted_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True)
|
||||
submission_json = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
|
||||
submission_json = sqlalchemy.Column(MediaReferenceText, nullable=True)
|
||||
status_reason = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
|
||||
|
||||
created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, default=datetime.datetime.utcnow)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import sqlalchemy as sa
|
||||
|
||||
from .base import Base
|
||||
from ...persistence.media import MediaReferenceJSON
|
||||
|
||||
|
||||
class AssistantConversation(Base):
|
||||
@@ -13,7 +14,7 @@ class AssistantConversation(Base):
|
||||
account_uuid = sa.Column(sa.String(36), sa.ForeignKey('users.uuid', ondelete='CASCADE'), nullable=False)
|
||||
revision = sa.Column(sa.Integer, nullable=False, default=0)
|
||||
status = sa.Column(sa.String(20), nullable=False, default='ready')
|
||||
messages = sa.Column(sa.JSON, nullable=False, default=list)
|
||||
messages = sa.Column(MediaReferenceJSON, nullable=False, default=list)
|
||||
error = sa.Column(sa.Text, nullable=True)
|
||||
model_name = sa.Column(sa.String(255), nullable=True)
|
||||
model_uuid = sa.Column(sa.String(36), nullable=True)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import sqlalchemy
|
||||
|
||||
from ...persistence.media import MediaReferenceText
|
||||
from ...persistence.media import MediaReferenceText, MediaReferenceJSON
|
||||
from .base import Base
|
||||
|
||||
|
||||
@@ -107,7 +107,7 @@ class MonitoringLLMCall(Base):
|
||||
error_message = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
|
||||
message_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True) # Associated message ID
|
||||
|
||||
details = sqlalchemy.Column(sqlalchemy.JSON, nullable=True)
|
||||
details = sqlalchemy.Column(MediaReferenceJSON, nullable=True)
|
||||
|
||||
__table_args__ = (
|
||||
sqlalchemy.Index('ix_monitoring_llm_calls_workspace_timestamp', 'workspace_uuid', 'timestamp'),
|
||||
@@ -137,8 +137,8 @@ class MonitoringToolCall(Base):
|
||||
pipeline_name = sqlalchemy.Column(sqlalchemy.String(255), nullable=False)
|
||||
session_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True)
|
||||
message_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True, index=True)
|
||||
arguments = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
|
||||
result = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
|
||||
arguments = sqlalchemy.Column(MediaReferenceText, nullable=True)
|
||||
result = sqlalchemy.Column(MediaReferenceText, nullable=True)
|
||||
error_message = sqlalchemy.Column(sqlalchemy.Text, nullable=True)
|
||||
|
||||
__table_args__ = (
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import sqlalchemy
|
||||
|
||||
from .base import Base
|
||||
from ...persistence.media import MediaReferenceJSON
|
||||
|
||||
|
||||
class LegacyPipeline(Base):
|
||||
@@ -74,7 +75,7 @@ class PipelineRunRecord(Base):
|
||||
)
|
||||
started_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
|
||||
finished_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
|
||||
result = sqlalchemy.Column(sqlalchemy.JSON, nullable=False)
|
||||
result = sqlalchemy.Column(MediaReferenceJSON, nullable=False)
|
||||
knowledge_base_uuid = sqlalchemy.Column(sqlalchemy.String(255), nullable=True)
|
||||
|
||||
__table_args__ = (
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Shared, stable identities for resource URLs embedded in conversation payloads."""
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from .base import Base
|
||||
|
||||
|
||||
class ResourceURL(Base):
|
||||
__tablename__ = 'resource_urls'
|
||||
|
||||
id = sa.Column(sa.String(36), primary_key=True)
|
||||
url = sa.Column(sa.Text, nullable=True)
|
||||
url_hash = sa.Column(sa.String(64), nullable=True, unique=True)
|
||||
canonical_id = sa.Column(sa.String(36), sa.ForeignKey('resource_urls.id'), nullable=True, index=True)
|
||||
created_at = sa.Column(sa.DateTime, nullable=False)
|
||||
last_checked_at = sa.Column(sa.DateTime, nullable=True)
|
||||
updated_at = sa.Column(sa.DateTime, nullable=False)
|
||||
|
||||
__table_args__ = (
|
||||
sa.CheckConstraint(
|
||||
'(canonical_id IS NULL AND url IS NOT NULL AND url_hash IS NOT NULL) '
|
||||
'OR (canonical_id IS NOT NULL AND url IS NULL AND url_hash IS NULL)',
|
||||
name='ck_resource_url_canonical',
|
||||
),
|
||||
sa.Index('ix_resource_urls_last_checked', 'last_checked_at'),
|
||||
)
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Normalize retained resource URLs behind stable shared references."""
|
||||
|
||||
import datetime
|
||||
import json
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = '0041_resource_urls'
|
||||
down_revision = '0040_session_summaries'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.resource_url import ResourceURL
|
||||
from langbot.pkg.persistence.media import MediaReferenceText, MediaReferenceJSON
|
||||
from langbot.pkg.persistence.resource_urls import create_postgres_expansion, intern, install
|
||||
|
||||
connection = op.get_bind()
|
||||
ResourceURL.__table__.create(connection, checkfirst=True)
|
||||
install(connection.engine)
|
||||
if connection.dialect.name == 'postgresql':
|
||||
create_postgres_expansion(connection)
|
||||
inspector = sa.inspect(connection)
|
||||
existing = set(inspector.get_table_names())
|
||||
for table in sorted(Base.metadata.tables.values(), key=lambda table: table.name):
|
||||
if table.name not in existing:
|
||||
continue
|
||||
retained = [
|
||||
column.name for column in table.columns if isinstance(column.type, (MediaReferenceText, MediaReferenceJSON))
|
||||
]
|
||||
json_columns = {column.name for column in table.columns if isinstance(column.type, MediaReferenceJSON)}
|
||||
if not retained:
|
||||
continue
|
||||
# Reflection intentionally reads literal stored values without expansion.
|
||||
source = sa.Table(table.name, sa.MetaData(), autoload_with=connection)
|
||||
retained = [name for name in retained if name in source.c]
|
||||
keys = list(source.primary_key.columns)
|
||||
if len(keys) != 1:
|
||||
raise RuntimeError(f'Resource URL migration requires one primary key: {table.name}')
|
||||
key = keys[0]
|
||||
cursor = None
|
||||
stamp = source.c.get('created_at')
|
||||
if stamp is None:
|
||||
stamp = source.c.get('timestamp')
|
||||
while True:
|
||||
query = sa.select(source).order_by(key).limit(100)
|
||||
if cursor is not None:
|
||||
query = query.where(key > cursor)
|
||||
rows = connection.execute(query).mappings().all()
|
||||
if not rows:
|
||||
break
|
||||
for row in rows:
|
||||
created = row[stamp.name] if stamp is not None else None
|
||||
if created is not None and created.tzinfo is not None:
|
||||
created = created.astimezone(datetime.timezone.utc).replace(tzinfo=None)
|
||||
updates = {}
|
||||
for name in retained:
|
||||
raw = json.dumps(row[name]) if name in json_columns else row[name]
|
||||
converted = intern(connection, raw, created)
|
||||
updates[name] = json.loads(converted) if name in json_columns else converted
|
||||
updates = {name: value for name, value in updates.items() if value != row[name]}
|
||||
if updates:
|
||||
connection.execute(source.update().where(key == row[key.name]).values(**updates))
|
||||
cursor = rows[-1][key.name]
|
||||
|
||||
|
||||
def downgrade():
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.persistence.media import MediaReferenceText, MediaReferenceJSON
|
||||
from langbot.pkg.persistence.resource_urls import uninstall
|
||||
|
||||
connection = op.get_bind()
|
||||
existing = set(sa.inspect(connection).get_table_names())
|
||||
for table in sorted(Base.metadata.tables.values(), key=lambda table: table.name):
|
||||
if table.name not in existing:
|
||||
continue
|
||||
columns = {column['name'] for column in sa.inspect(connection).get_columns(table.name)}
|
||||
for column in table.columns:
|
||||
if column.name in columns and isinstance(column.type, (MediaReferenceText, MediaReferenceJSON)):
|
||||
name = connection.dialect.identifier_preparer.quote(table.name)
|
||||
field = connection.dialect.identifier_preparer.quote(column.name)
|
||||
expanded = f'langbot_expand_resource_urls(CAST({field} AS TEXT))'
|
||||
if isinstance(column.type, MediaReferenceJSON) and connection.dialect.name == 'postgresql':
|
||||
expanded = f'CAST({expanded} AS JSON)'
|
||||
connection.exec_driver_sql(f'UPDATE {name} SET {field} = {expanded}')
|
||||
op.drop_table('resource_urls')
|
||||
if connection.dialect.name == 'postgresql':
|
||||
connection.exec_driver_sql('DROP FUNCTION langbot_expand_resource_urls(text)')
|
||||
uninstall(connection.engine)
|
||||
@@ -5,7 +5,8 @@ from __future__ import annotations
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import Text
|
||||
from sqlalchemy import Text, JSON, cast
|
||||
from sqlalchemy.sql.elements import ColumnElement
|
||||
from sqlalchemy.engine import Dialect
|
||||
from sqlalchemy.sql.operators import OperatorType
|
||||
from sqlalchemy.types import TypeDecorator
|
||||
@@ -23,9 +24,20 @@ class MediaReferenceText(TypeDecorator[str]):
|
||||
impl = Text
|
||||
cache_ok = True
|
||||
|
||||
class comparator_factory(Text.Comparator):
|
||||
def operate(self, op, *other, **kw):
|
||||
from .resource_urls import ExpandedResourceURLs
|
||||
|
||||
# ORM identity maps compare annotated copies of mapped columns.
|
||||
if other and isinstance(other[0], ColumnElement) and self.expr.shares_lineage(other[0]):
|
||||
return super().operate(op, *other, **kw)
|
||||
# CAS and textual filters compare the same expanded value readers see.
|
||||
return op(ExpandedResourceURLs(self.expr), *other, **kw)
|
||||
|
||||
def process_bind_param(self, value: str | None, dialect: Dialect) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
reject_internal_references(value)
|
||||
try:
|
||||
payload = json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
@@ -44,3 +56,37 @@ class MediaReferenceText(TypeDecorator[str]):
|
||||
def coerce_compared_value(self, op: OperatorType | None, value: Any) -> Text:
|
||||
# Legacy payloads must remain usable in searches and compare-and-swap updates.
|
||||
return self.impl
|
||||
|
||||
def column_expression(self, column):
|
||||
from .resource_urls import ExpandedResourceURLs
|
||||
|
||||
return ExpandedResourceURLs(column)
|
||||
|
||||
|
||||
class MediaReferenceJSON(TypeDecorator):
|
||||
"""Preserve native JSON API values while interning structured resource fields."""
|
||||
|
||||
impl = JSON
|
||||
cache_ok = True
|
||||
|
||||
def process_bind_param(self, value, dialect):
|
||||
if value is not None:
|
||||
reject_internal_references(json.dumps(value))
|
||||
return value
|
||||
|
||||
def column_expression(self, column):
|
||||
from .resource_urls import ExpandedResourceJSON
|
||||
|
||||
return ExpandedResourceJSON(cast(column, Text))
|
||||
|
||||
|
||||
def reject_internal_references(value):
|
||||
"""Only the transactional normalizer may introduce catalog references."""
|
||||
from .resource_urls import REFERENCE
|
||||
|
||||
if REFERENCE.search(value):
|
||||
raise ValueError('Internal resource references cannot be supplied as conversation content')
|
||||
|
||||
|
||||
# Register the catalog and its database hooks before engines are opened.
|
||||
from . import resource_urls as _resource_urls # noqa: E402,F401
|
||||
|
||||
@@ -183,16 +183,25 @@ class PersistenceManager:
|
||||
|
||||
if self.mode == PersistenceMode.CLOUD_RUNTIME:
|
||||
await self._validate_cloud_runtime()
|
||||
from .resource_urls import install
|
||||
|
||||
install(engine.sync_engine)
|
||||
return
|
||||
|
||||
if self.mode == PersistenceMode.RELEASE_MIGRATION:
|
||||
async with self._release_migration_lock():
|
||||
await self._initialize_managed_schema()
|
||||
await self._validate_release_schema()
|
||||
from .resource_urls import install
|
||||
|
||||
install(engine.sync_engine)
|
||||
return
|
||||
|
||||
await self._initialize_managed_schema()
|
||||
await self._enable_sqlite_foreign_keys_after_migration()
|
||||
from .resource_urls import install
|
||||
|
||||
install(engine.sync_engine)
|
||||
|
||||
if self.mode == PersistenceMode.OSS_COMPAT:
|
||||
await self.write_space_model_providers()
|
||||
|
||||
@@ -0,0 +1,391 @@
|
||||
"""Normalize resource URLs at the typed SQL boundary, within the source transaction.
|
||||
|
||||
Only typed conversation payload writes are interned. Reads expand references in SQL so
|
||||
ORM and Core callers retain their existing string/JSON contracts. The catalog
|
||||
is shared; authorization continues to belong to the enclosing business record.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
import sqlite3
|
||||
from collections import OrderedDict
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql, sqlite
|
||||
from sqlalchemy.ext.compiler import compiles
|
||||
from sqlalchemy.sql.functions import FunctionElement
|
||||
|
||||
from ..entity.persistence.resource_url import ResourceURL
|
||||
|
||||
|
||||
PREFIX = 'langbot-resource://'
|
||||
REFERENCE = re.compile(r'langbot-resource://([a-f0-9-]{36})')
|
||||
FIELDS = frozenset({'url', 'image_url', 'audio_url', 'file_url', 'video_url', 'original_url'})
|
||||
|
||||
|
||||
def now():
|
||||
return datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
|
||||
|
||||
|
||||
def digest(url):
|
||||
return hashlib.sha256(url.encode()).hexdigest()
|
||||
|
||||
|
||||
def fields(value):
|
||||
if isinstance(value, list):
|
||||
for child in value:
|
||||
yield from fields(child)
|
||||
elif isinstance(value, dict):
|
||||
for key, child in value.items():
|
||||
if key in FIELDS and isinstance(child, str) and child.startswith(('http://', 'https://', '/api/v1/files/')):
|
||||
yield value, key, child
|
||||
else:
|
||||
yield from fields(child)
|
||||
|
||||
|
||||
def intern(connection, raw, created_at=None, identities_out=None):
|
||||
"""Intern resource fields synchronously in the current SQLAlchemy transaction."""
|
||||
if not isinstance(raw, str):
|
||||
return raw
|
||||
try:
|
||||
payload = json.loads(raw)
|
||||
except ValueError:
|
||||
return raw
|
||||
nodes = list(fields(payload))
|
||||
if not nodes:
|
||||
return raw
|
||||
table = ResourceURL.__table__
|
||||
stamp = now()
|
||||
urls = {url for _, _, url in nodes}
|
||||
aliases = connection.engine._langbot_resource_aliases
|
||||
known = {url: aliases[url] for url in urls if url in aliases}
|
||||
valid = (
|
||||
set(connection.execute(sa.select(table.c.id).where(table.c.id.in_(set(known.values())))).scalars())
|
||||
if known
|
||||
else set()
|
||||
)
|
||||
identities = {url: identity for url, identity in known.items() if identity in valid}
|
||||
urls -= identities.keys()
|
||||
hashes = dict(sorted((digest(url), url) for url in urls))
|
||||
insert = {'sqlite': sqlite.insert, 'postgresql': postgresql.insert}[connection.dialect.name]
|
||||
if hashes:
|
||||
connection.execute(
|
||||
insert(table)
|
||||
.values(
|
||||
[
|
||||
dict(id=str(uuid.uuid4()), url=url, url_hash=key, created_at=created_at or stamp, updated_at=stamp)
|
||||
for key, url in hashes.items()
|
||||
]
|
||||
)
|
||||
.on_conflict_do_nothing(index_elements=['url_hash'])
|
||||
)
|
||||
for identity, url, key in connection.execute(
|
||||
sa.select(table.c.id, table.c.url, table.c.url_hash).where(table.c.url_hash.in_(hashes))
|
||||
):
|
||||
if hashes[key] != url:
|
||||
raise ValueError('Resource URL hash collision')
|
||||
identities[url] = identity
|
||||
if created_at is not None and identities:
|
||||
# Migration order must not make an old shared URL appear newly created.
|
||||
connection.execute(
|
||||
table.update()
|
||||
.where(table.c.id.in_(set(identities.values())), table.c.created_at > created_at)
|
||||
.values(created_at=created_at)
|
||||
)
|
||||
for node, key, url in nodes:
|
||||
node[key] = PREFIX + identities[url]
|
||||
aliases[url] = identities[url]
|
||||
aliases.move_to_end(url)
|
||||
while len(aliases) > 8192:
|
||||
aliases.popitem(last=False)
|
||||
if identities_out is not None:
|
||||
identities_out.update(identities)
|
||||
return json.dumps(payload, ensure_ascii=False)
|
||||
|
||||
|
||||
class ExpandedResourceURLs(FunctionElement):
|
||||
type = sa.Text()
|
||||
inherit_cache = True
|
||||
|
||||
|
||||
class ExpandedResourceJSON(FunctionElement):
|
||||
type = sa.JSON()
|
||||
inherit_cache = True
|
||||
|
||||
|
||||
@compiles(ExpandedResourceJSON)
|
||||
def compile_json_expansion(element, compiler, **kw):
|
||||
expression = compile_expansion(element, compiler, **kw)
|
||||
if compiler.dialect.name == 'postgresql':
|
||||
return f'CAST({expression} AS JSON)'
|
||||
return expression
|
||||
|
||||
|
||||
@compiles(ExpandedResourceURLs)
|
||||
def compile_expansion(element, compiler, **kw):
|
||||
source = compiler.process(list(element.clauses)[0], **kw)
|
||||
if not getattr(compiler.dialect, '_langbot_resource_urls', False):
|
||||
return source
|
||||
return f'langbot_expand_resource_urls({source})'
|
||||
|
||||
|
||||
def install(engine):
|
||||
"""Enable normalization after schema initialization, before producers start."""
|
||||
engine.dialect._langbot_resource_urls = True
|
||||
if not hasattr(engine, '_langbot_resource_aliases'):
|
||||
engine._langbot_resource_aliases = OrderedDict()
|
||||
engine.clear_compiled_cache()
|
||||
if not sa.event.contains(engine, 'before_cursor_execute', normalize_parameters):
|
||||
sa.event.listen(engine, 'before_cursor_execute', normalize_parameters, retval=True)
|
||||
|
||||
|
||||
def uninstall(engine):
|
||||
"""Restore literal payload reads and writes after catalog downgrade."""
|
||||
if sa.event.contains(engine, 'before_cursor_execute', normalize_parameters):
|
||||
sa.event.remove(engine, 'before_cursor_execute', normalize_parameters)
|
||||
engine.dialect._langbot_resource_urls = False
|
||||
engine.clear_compiled_cache()
|
||||
engine._langbot_resource_aliases.clear()
|
||||
|
||||
|
||||
def normalize_parameters(connection, cursor, statement, parameters, context, executemany):
|
||||
"""Replace only typed payload binds; never inspect arbitrary SQL or configuration."""
|
||||
from .media import MediaReferenceText, MediaReferenceJSON
|
||||
|
||||
compiled = context.compiled
|
||||
if compiled is None or not (context.isinsert or context.isupdate):
|
||||
return statement, parameters
|
||||
names = {
|
||||
name
|
||||
for bind, name in compiled.bind_names.items()
|
||||
if isinstance(bind.type, (MediaReferenceText, MediaReferenceJSON))
|
||||
}
|
||||
if not names:
|
||||
return statement, parameters
|
||||
|
||||
def normalize(row):
|
||||
if isinstance(row, dict):
|
||||
return {key: intern(connection, value) if key in names else value for key, value in row.items()}
|
||||
positions = compiled.positiontup or []
|
||||
if not positions:
|
||||
return row
|
||||
return tuple(
|
||||
intern(connection, value) if positions[index % len(positions)] in names else value
|
||||
for index, value in enumerate(row)
|
||||
)
|
||||
|
||||
batched = executemany and getattr(compiled, '_insertmanyvalues', None) is None
|
||||
return statement, [normalize(row) for row in parameters] if batched else normalize(parameters)
|
||||
|
||||
|
||||
def register_sqlite(dbapi, record):
|
||||
"""The UDF runs on aiosqlite's worker, including its catalog SELECTs."""
|
||||
driver = getattr(dbapi, 'driver_connection', None)
|
||||
raw_connection = driver._conn if hasattr(driver, '_conn') else dbapi
|
||||
if not isinstance(raw_connection, sqlite3.Connection):
|
||||
return
|
||||
|
||||
def expand(raw):
|
||||
if not isinstance(raw, str) or PREFIX not in raw:
|
||||
return raw
|
||||
|
||||
resolved = {}
|
||||
|
||||
def lookup(match):
|
||||
identity = match.group(1)
|
||||
if identity in resolved:
|
||||
return resolved[identity]
|
||||
visited = set()
|
||||
while identity not in visited:
|
||||
visited.add(identity)
|
||||
row = raw_connection.execute(
|
||||
'SELECT url, canonical_id FROM resource_urls WHERE id=?', (identity,)
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise ValueError('Missing resource URL')
|
||||
if row[1] is None:
|
||||
value = json.dumps(row[0], ensure_ascii=False)[1:-1]
|
||||
for reference in visited:
|
||||
resolved[reference] = value
|
||||
return value
|
||||
identity = row[1]
|
||||
raise ValueError('Cyclic resource URL reference')
|
||||
|
||||
return REFERENCE.sub(lookup, raw)
|
||||
|
||||
dbapi.create_function('langbot_expand_resource_urls', 1, expand)
|
||||
|
||||
|
||||
POSTGRES_EXPANSION = r"""
|
||||
CREATE OR REPLACE FUNCTION public.langbot_expand_resource_urls(payload text)
|
||||
RETURNS text LANGUAGE plpgsql STABLE SECURITY INVOKER SET search_path = public, pg_temp AS $$
|
||||
DECLARE token text; identity text; target text; parent text; result text := payload; visited text[];
|
||||
BEGIN
|
||||
IF payload IS NULL OR position('langbot-resource://' in payload) = 0 THEN RETURN payload; END IF;
|
||||
FOR token IN SELECT DISTINCT m[1] FROM regexp_matches(payload, 'langbot-resource://([a-f0-9-]{36})', 'g') m LOOP
|
||||
identity := token; visited := ARRAY[]::text[];
|
||||
LOOP
|
||||
IF identity = ANY(visited) THEN RAISE EXCEPTION 'Cyclic resource URL reference'; END IF;
|
||||
visited := array_append(visited, identity);
|
||||
SELECT url, canonical_id INTO target, parent FROM resource_urls WHERE id = identity;
|
||||
IF NOT FOUND THEN RAISE EXCEPTION 'Missing resource URL'; END IF;
|
||||
EXIT WHEN parent IS NULL;
|
||||
identity := parent;
|
||||
END LOOP;
|
||||
result := replace(result, 'langbot-resource://' || token, substr(to_json(target)::text, 2, length(to_json(target)::text)-2));
|
||||
END LOOP;
|
||||
RETURN result;
|
||||
END $$;
|
||||
"""
|
||||
|
||||
|
||||
def catalog_created(table, connection, **kw):
|
||||
if connection.dialect.name == 'postgresql':
|
||||
create_postgres_expansion(connection)
|
||||
install(connection.engine)
|
||||
|
||||
|
||||
def create_postgres_expansion(connection):
|
||||
schema = connection.scalar(sa.select(sa.func.current_schema()))
|
||||
quoted = connection.dialect.identifier_preparer.quote(schema)
|
||||
connection.exec_driver_sql(POSTGRES_EXPANSION.replace('public', quoted))
|
||||
|
||||
|
||||
async def refresh_images(engine, peer, payload, *, is_current=None):
|
||||
"""Serialize refreshes on catalog identities and overwrite current URLs atomically."""
|
||||
from ..platform.adapters.aiocqhttp.image_urls import image_nodes
|
||||
|
||||
nodes = list(image_nodes(payload))
|
||||
if not nodes:
|
||||
return
|
||||
|
||||
def connected():
|
||||
return not peer.closed and (is_current is None or is_current())
|
||||
|
||||
table = ResourceURL.__table__
|
||||
# Establish catalog identities for runtime-only inputs too. Existing history
|
||||
# already has identities created in its original write or migration.
|
||||
async with engine.begin() as db:
|
||||
references = {}
|
||||
await db.run_sync(lambda sync: intern(sync, json.dumps(payload), identities_out=references))
|
||||
identities = set(references.values())
|
||||
async with engine.connect() as db:
|
||||
if engine.dialect.name == 'sqlite':
|
||||
await db.exec_driver_sql('BEGIN IMMEDIATE')
|
||||
else:
|
||||
await db.begin()
|
||||
try:
|
||||
rows = (
|
||||
(
|
||||
await db.execute(
|
||||
sa.select(table).where(table.c.id.in_(identities)).order_by(table.c.id).with_for_update()
|
||||
)
|
||||
)
|
||||
.mappings()
|
||||
.all()
|
||||
)
|
||||
by_id = {row['id']: row for row in rows}
|
||||
for identity in list(by_id):
|
||||
row = by_id[identity]
|
||||
while row['canonical_id'] is not None:
|
||||
row = (
|
||||
(await db.execute(sa.select(table).where(table.c.id == row['canonical_id']).with_for_update()))
|
||||
.mappings()
|
||||
.one()
|
||||
)
|
||||
by_id[identity] = row
|
||||
by_url = {node['url']: by_id[references[node['url']]] for node, _, _ in nodes}
|
||||
# Capture origins even when renewal is not due or identification fails.
|
||||
for node, message_id, image_id in nodes:
|
||||
peer.image_urls.remember(by_url[node['url']]['url'], message_id, image_id)
|
||||
# Identification and upstream credential handling remain connection-local.
|
||||
if not connected() or not await peer.image_urls.identify():
|
||||
return
|
||||
napcat = peer.image_urls.app_name == 'NapCat.Onebot'
|
||||
due = []
|
||||
stamp = now()
|
||||
for node, message_id, image_id in nodes:
|
||||
row = by_url.get(node['url'])
|
||||
if row is None:
|
||||
continue
|
||||
node['url'] = row['url']
|
||||
checked = row['last_checked_at'] or row['created_at']
|
||||
if napcat or (stamp - checked).total_seconds() >= 1800:
|
||||
due.append((node, message_id, image_id, row))
|
||||
if due:
|
||||
# Keep Source/Image identities, including those recovered from events.
|
||||
batch = [
|
||||
dict(message_id=message_id, image={'type': 'Image', 'image_id': image_id, 'url': node['url']})
|
||||
for node, message_id, image_id, _ in due
|
||||
]
|
||||
old = [entry['image']['url'] for entry in batch]
|
||||
completed = await peer.image_urls.refresh(batch)
|
||||
if not completed:
|
||||
return
|
||||
if not connected():
|
||||
return
|
||||
stamp = now()
|
||||
updated = set()
|
||||
for (node, _, _, row), entry, previous in zip(due, batch, old):
|
||||
if previous not in completed:
|
||||
continue
|
||||
fresh = entry['image']['url']
|
||||
node['url'] = fresh
|
||||
if napcat and fresh == previous:
|
||||
continue
|
||||
if row['id'] in updated:
|
||||
continue
|
||||
updated.add(row['id'])
|
||||
values = dict(last_checked_at=stamp)
|
||||
if fresh != previous:
|
||||
fresh_hash = digest(fresh)
|
||||
# Lock the destination before merging identical current URLs.
|
||||
destination = (
|
||||
await db.execute(
|
||||
sa.select(table.c.id, table.c.url)
|
||||
.where(table.c.url_hash == fresh_hash)
|
||||
.with_for_update()
|
||||
)
|
||||
).first()
|
||||
if destination is None:
|
||||
try:
|
||||
# Another source may discover the same renewed URL
|
||||
# concurrently. Let the unique index arbitrate it.
|
||||
async with db.begin_nested():
|
||||
await db.execute(
|
||||
table.update()
|
||||
.where(table.c.id == row['id'])
|
||||
.values(url=fresh, url_hash=fresh_hash, updated_at=stamp, **values)
|
||||
)
|
||||
continue
|
||||
except sa.exc.IntegrityError as exc:
|
||||
if getattr(exc.orig, 'sqlstate', None) != '23505':
|
||||
raise
|
||||
destination = (
|
||||
await db.execute(
|
||||
sa.select(table.c.id, table.c.url)
|
||||
.where(table.c.url_hash == fresh_hash)
|
||||
.with_for_update()
|
||||
)
|
||||
).one()
|
||||
if destination is not None and destination.id != row['id']:
|
||||
if destination.url != fresh:
|
||||
raise ValueError('Resource URL hash collision')
|
||||
values.update(url=None, url_hash=None, canonical_id=destination.id, updated_at=stamp)
|
||||
await db.execute(
|
||||
table.update().where(table.c.id == destination.id).values(last_checked_at=stamp)
|
||||
)
|
||||
else:
|
||||
values.update(url=fresh, url_hash=fresh_hash, updated_at=stamp)
|
||||
await db.execute(table.update().where(table.c.id == row['id']).values(**values))
|
||||
await db.commit()
|
||||
except BaseException:
|
||||
await db.rollback()
|
||||
raise
|
||||
|
||||
|
||||
sa.event.listen(sa.engine.Engine, 'connect', register_sqlite)
|
||||
sa.event.listen(ResourceURL.__table__, 'after_create', catalog_created)
|
||||
@@ -19,7 +19,7 @@ from sqlalchemy.dialects.postgresql.dml import OnConflictDoUpdate as PostgreSQLO
|
||||
from sqlalchemy.dialects.sqlite.dml import OnConflictDoNothing as SQLiteOnConflictDoNothing
|
||||
from sqlalchemy.dialects.sqlite.dml import OnConflictDoUpdate as SQLiteOnConflictDoUpdate
|
||||
|
||||
from .media import MediaReferenceText
|
||||
from .media import MediaReferenceText, MediaReferenceJSON
|
||||
|
||||
|
||||
TENANT_SETTING = 'langbot.workspace_uuid'
|
||||
@@ -295,7 +295,9 @@ def _validate_scoped_sql_type(
|
||||
|
||||
if type(sql_type) in {Vector, HALFVEC}:
|
||||
return
|
||||
if type(sql_type) is not MediaReferenceText and not type(sql_type).__module__.startswith('sqlalchemy.'):
|
||||
if type(sql_type) not in {MediaReferenceText, MediaReferenceJSON} and not type(sql_type).__module__.startswith(
|
||||
'sqlalchemy.'
|
||||
):
|
||||
raise ScopedSessionTransactionError('TenantUnitOfWork does not allow custom SQL types in public statements')
|
||||
|
||||
for identifier_attribute in ('name', 'schema', 'collation'):
|
||||
@@ -501,6 +503,14 @@ def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[s
|
||||
)
|
||||
|
||||
if isinstance(element, sqlalchemy.sql.functions.FunctionElement):
|
||||
from .resource_urls import ExpandedResourceURLs
|
||||
|
||||
resource_expansion_allowed = (
|
||||
type(element) is ExpandedResourceURLs
|
||||
and len(list(element.clauses)) == 1
|
||||
and isinstance(list(element.clauses)[0], sqlalchemy.Column)
|
||||
and type(list(element.clauses)[0].type) is MediaReferenceText
|
||||
)
|
||||
function_name = str(getattr(element, 'name', '')).casefold()
|
||||
package_names = tuple(getattr(element, 'packagenames', ()))
|
||||
generic_function_allowed = (
|
||||
@@ -508,7 +518,9 @@ def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[s
|
||||
and function_name in _ALLOWED_SCOPED_GENERIC_FUNCTIONS
|
||||
)
|
||||
builtin_function_allowed = _ALLOWED_SCOPED_BUILTIN_FUNCTION_TYPES.get(function_name) is type(element)
|
||||
if package_names or not (generic_function_allowed or builtin_function_allowed):
|
||||
if package_names or not (
|
||||
generic_function_allowed or builtin_function_allowed or resource_expansion_allowed
|
||||
):
|
||||
raise ScopedSessionTransactionError(
|
||||
f'TenantUnitOfWork does not allow SQL function {function_name or "<unknown>"!r}'
|
||||
)
|
||||
|
||||
@@ -77,6 +77,7 @@ class OneBotImageURLs:
|
||||
self.keys = {}
|
||||
self.references = OrderedDict()
|
||||
self.messages = OrderedDict()
|
||||
self.urls = OrderedDict()
|
||||
self.retry_after = 0
|
||||
|
||||
def remember(self, url, message_id, image_id):
|
||||
@@ -87,23 +88,37 @@ class OneBotImageURLs:
|
||||
while len(self.references) > 2048:
|
||||
self.references.popitem(last=False)
|
||||
|
||||
async def identify(self):
|
||||
"""Identify lazily with the same backoff as resource renewal."""
|
||||
async with self.lock:
|
||||
if self.connection.closed or time.monotonic() < self.retry_after:
|
||||
return False
|
||||
try:
|
||||
await self._identify()
|
||||
return True
|
||||
except Exception:
|
||||
self.retry_after = time.monotonic() + 30
|
||||
raise
|
||||
|
||||
async def _identify(self):
|
||||
if self.app_name is None:
|
||||
info = await self.connection.call_action('get_version_info')
|
||||
self.app_name = info.get('app_name', '')
|
||||
|
||||
async def refresh(self, payload):
|
||||
nodes = list(image_nodes(payload))
|
||||
if not nodes or self.connection.closed:
|
||||
return
|
||||
return False
|
||||
for node, message_id, image_id in nodes:
|
||||
self.remember(node['url'], message_id, image_id)
|
||||
async with self.lock:
|
||||
if self.connection.closed or time.monotonic() < self.retry_after:
|
||||
return
|
||||
return False
|
||||
try:
|
||||
if self.app_name is None:
|
||||
info = await self.connection.call_action('get_version_info')
|
||||
self.app_name = info.get('app_name', '')
|
||||
await self._identify()
|
||||
if self.app_name == 'NapCat.Onebot':
|
||||
await self._napcat(nodes)
|
||||
else:
|
||||
await self._standard(nodes)
|
||||
return await self._napcat(nodes)
|
||||
return await self._standard(nodes)
|
||||
except Exception:
|
||||
# Failed/unsupported APIs must not break history reads or hammer the peer.
|
||||
self.retry_after = time.monotonic() + 30
|
||||
@@ -112,7 +127,7 @@ class OneBotImageURLs:
|
||||
async def _napcat(self, nodes):
|
||||
wanted = {kind for node, _, _ in nodes if (kind := nt_image_type(node['url'])) is not None}
|
||||
if not wanted:
|
||||
return
|
||||
return set()
|
||||
if any(self.keys.get(kind, ('', 0))[1] <= time.time() for kind in wanted):
|
||||
result = await self.connection.call_action('nc_get_rkey')
|
||||
keys = {}
|
||||
@@ -129,10 +144,12 @@ class OneBotImageURLs:
|
||||
if not wanted.issubset(keys):
|
||||
raise ValueError('NapCat returned no valid image credentials')
|
||||
self.keys.update(keys)
|
||||
completed = set()
|
||||
for node, _, _ in nodes:
|
||||
kind = nt_image_type(node['url'])
|
||||
if kind not in wanted:
|
||||
continue
|
||||
completed.add(node['url'])
|
||||
parsed = urlsplit(node['url'])
|
||||
parts = parsed.query.split('&')
|
||||
# Preserve every other parameter byte-for-byte, including signed resource IDs.
|
||||
@@ -143,10 +160,17 @@ class OneBotImageURLs:
|
||||
for part in parts
|
||||
]
|
||||
node['url'] = urlunsplit(parsed._replace(query='&'.join(parts)))
|
||||
return completed
|
||||
|
||||
async def _standard(self, nodes):
|
||||
completed = set()
|
||||
for node, _, _ in nodes:
|
||||
old_url = node['url']
|
||||
refreshed = self.urls.get(old_url)
|
||||
if refreshed is not None and refreshed[1] > time.monotonic():
|
||||
node['url'] = refreshed[0]
|
||||
completed.add(old_url)
|
||||
continue
|
||||
ref = self.references.get(old_url)
|
||||
if ref is None:
|
||||
continue
|
||||
@@ -160,12 +184,19 @@ class OneBotImageURLs:
|
||||
data = segment.data
|
||||
if data.get('file') and remote_image_url(data.get('url')):
|
||||
images[str(data['file'])] = data['url']
|
||||
cached = (images, time.monotonic() + 60)
|
||||
cached = (images, time.monotonic() + 1800)
|
||||
self.messages[str(message_id)] = cached
|
||||
self.messages.move_to_end(str(message_id))
|
||||
while len(self.messages) > 256:
|
||||
self.messages.popitem(last=False)
|
||||
new_url = cached[0].get(image_id)
|
||||
if new_url:
|
||||
completed.add(old_url)
|
||||
self.remember(new_url, message_id, image_id)
|
||||
deadline = cached[1]
|
||||
self.urls[old_url] = (new_url, deadline)
|
||||
self.urls[new_url] = (new_url, deadline)
|
||||
while len(self.urls) > 2048:
|
||||
self.urls.popitem(last=False)
|
||||
node['url'] = new_url
|
||||
return completed
|
||||
|
||||
@@ -102,6 +102,7 @@ class OneBotConnection:
|
||||
self.image_urls.keys.clear()
|
||||
self.image_urls.references.clear()
|
||||
self.image_urls.messages.clear()
|
||||
self.image_urls.urls.clear()
|
||||
for future in self.pending.values():
|
||||
if not future.done():
|
||||
future.set_exception(ApiNotAvailable('OneBot connection closed'))
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Refresh image references at authorized read/use boundaries, without persistence."""
|
||||
"""Refresh image references at authorized read/use boundaries."""
|
||||
|
||||
import copy
|
||||
import json
|
||||
@@ -21,7 +21,14 @@ async def refresh_image_payload(ap, workspace_uuid, bot_id, payload):
|
||||
return payload
|
||||
result = copy.deepcopy(payload)
|
||||
try:
|
||||
await connection.image_urls.refresh(result)
|
||||
from ..persistence.resource_urls import refresh_images
|
||||
|
||||
await refresh_images(
|
||||
ap.persistence_mgr.get_db_engine(),
|
||||
connection,
|
||||
result,
|
||||
is_current=lambda: bot.adapter.bot.connection is connection,
|
||||
)
|
||||
if bot.adapter.bot.connection is not connection or connection.closed:
|
||||
return payload
|
||||
return result
|
||||
|
||||
Reference in New Issue
Block a user