mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-22 16:35:20 +08:00
fix(storage): unify CAS persistence for inline images
- Externalize inline images in messages and run events before persistence - Keep snapshots, transcripts, and metadata free of inline payloads - Serialize messages after externalization to reduce peak memory usage - Preserve live messages, URL-only images, and SDK image retrieval behavior - Add regression tests for deduplication, redaction, storage failures, and serialization order
This commit is contained in:
@@ -794,7 +794,9 @@ async def execute_platform_tool(
|
||||
bot_name=str(getattr(getattr(bot, 'bot_entity', None), 'name', bot_id)),
|
||||
pipeline_id=str(authorization.get('processor_id') or ''),
|
||||
pipeline_name=str(authorization.get('processor_id') or authorization.get('runner_id') or ''),
|
||||
message_content=json.dumps(content, ensure_ascii=False),
|
||||
message_content=content
|
||||
if isinstance(content, (list, dict))
|
||||
else json.dumps(content, ensure_ascii=False),
|
||||
session_id=str(authorization.get('conversation_id') or ''),
|
||||
role='assistant',
|
||||
status='error' if delivery_error else 'success',
|
||||
|
||||
@@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
import typing
|
||||
|
||||
from ...core import app
|
||||
from ...storage.media import MediaCache
|
||||
from .descriptor import RunnerDescriptor
|
||||
from .errors import RunnerProtocolError
|
||||
from .host_models import AgentBinding, AgentEventEnvelope
|
||||
@@ -44,10 +45,10 @@ class AgentRunJournal:
|
||||
if not item:
|
||||
return value
|
||||
item_type = item.get('type')
|
||||
if item_type == 'image_base64' and item.get('image_base64'):
|
||||
item['image_base64'] = None
|
||||
if MediaCache.has_inline_images(item):
|
||||
item = MediaCache.redact_inline_images(item)
|
||||
item['content_redacted'] = True
|
||||
elif item_type == 'file_base64' and item.get('file_base64'):
|
||||
if item_type == 'file_base64' and item.get('file_base64'):
|
||||
item['file_base64'] = None
|
||||
item['content_redacted'] = True
|
||||
return item
|
||||
@@ -77,6 +78,8 @@ class AgentRunJournal:
|
||||
return [cls._sanitize_event_data(item) for item in value]
|
||||
if not isinstance(value, dict):
|
||||
return value
|
||||
if value.get('type') in ('Image', 'image', 'image_url', 'image_base64'):
|
||||
value = cls._sanitize_content_item(value)
|
||||
result = {}
|
||||
for key, item in value.items():
|
||||
if key in ('image_base64', 'file_base64', 'base64') and item:
|
||||
@@ -157,11 +160,16 @@ class AgentRunJournal:
|
||||
usage = result_dict.get('usage')
|
||||
if hasattr(usage, 'model_dump'):
|
||||
usage = usage.model_dump(mode='json')
|
||||
data = result_dict.get('data') if isinstance(result_dict.get('data'), dict) else {}
|
||||
if MediaCache.has_inline_images(data):
|
||||
data = await self.ap.storage_mgr.media_cache.externalize_chain_dump(data)
|
||||
if MediaCache.has_inline_images(metadata):
|
||||
metadata = MediaCache.redact_inline_images(metadata)
|
||||
return await self._get_run_ledger_store().append_event(
|
||||
run_id=run_id,
|
||||
sequence=sequence,
|
||||
event_type=str(result_dict.get('type') or 'unknown'),
|
||||
data=result_dict.get('data') if isinstance(result_dict.get('data'), dict) else {},
|
||||
data=data,
|
||||
usage=usage if isinstance(usage, dict) else None,
|
||||
source=source,
|
||||
metadata=metadata,
|
||||
|
||||
@@ -15,6 +15,7 @@ from ....entity.persistence import monitoring as persistence_monitoring
|
||||
from ....entity.persistence import agent_run as persistence_agent_run
|
||||
from ....entity.persistence import agent_interaction as persistence_agent_interaction
|
||||
from ....entity.persistence.event_log import EventLog
|
||||
from ....storage.media import MediaCache
|
||||
from ..authz import WorkspaceRequiredError
|
||||
from ..context import ExecutionContext
|
||||
from .tenant import TenantContext, require_workspace_uuid
|
||||
@@ -504,30 +505,31 @@ class MonitoringService(ExecutionDetailsMixin):
|
||||
|
||||
def _sanitize_message_content(self, content: str) -> str:
|
||||
"""Strip raw base64 data to protect database storage from unbounded bloating."""
|
||||
if not content or len(content) < 10000 or (';base64,' not in content and 'data:image/' not in content):
|
||||
if not content or ('base64' not in content and 'data:image/' not in content):
|
||||
return content
|
||||
try:
|
||||
data = json.loads(content)
|
||||
|
||||
def _strip_node(node):
|
||||
if isinstance(node, list):
|
||||
return [_strip_node(x) for x in node]
|
||||
if isinstance(node, dict):
|
||||
res = dict(node)
|
||||
if res.get('type') == 'Image' and res.get('base64'):
|
||||
res['base64'] = None
|
||||
for k, v in list(res.items()):
|
||||
if isinstance(v, (list, dict)):
|
||||
res[k] = _strip_node(v)
|
||||
return res
|
||||
return node
|
||||
|
||||
return json.dumps(_strip_node(data), ensure_ascii=False)
|
||||
return json.dumps(MediaCache.redact_inline_images(data), ensure_ascii=False)
|
||||
except Exception:
|
||||
return re.sub(
|
||||
r'data:image/[a-zA-Z0-9.+_-]+;base64,[\sA-Za-z0-9+/=]{1000,}', '[base64 image omitted]', content
|
||||
)
|
||||
|
||||
async def _externalize_message_content(self, content: str | list | dict) -> str:
|
||||
"""Externalize structured images before allocating their JSON representation."""
|
||||
if isinstance(content, str):
|
||||
if not content or ('base64' not in content and 'data:image/' not in content):
|
||||
return content
|
||||
try:
|
||||
payload = json.loads(content)
|
||||
except (TypeError, ValueError):
|
||||
return self._sanitize_message_content(content)
|
||||
else:
|
||||
payload = content
|
||||
if MediaCache.has_inline_images(payload):
|
||||
payload = await self.ap.storage_mgr.media_cache.externalize_chain_dump(payload)
|
||||
return json.dumps(MediaCache.redact_inline_images(payload), ensure_ascii=False)
|
||||
|
||||
@_workspace_transaction
|
||||
async def record_message(
|
||||
self,
|
||||
@@ -536,7 +538,7 @@ class MonitoringService(ExecutionDetailsMixin):
|
||||
bot_name: str,
|
||||
pipeline_id: str,
|
||||
pipeline_name: str,
|
||||
message_content: str,
|
||||
message_content: str | list | dict,
|
||||
session_id: str,
|
||||
status: str = 'success',
|
||||
level: str = 'info',
|
||||
@@ -558,7 +560,7 @@ class MonitoringService(ExecutionDetailsMixin):
|
||||
event_id = event_id or (trace_state.event_id if trace_state else None)
|
||||
run_id = run_id or trace.current_run() or None
|
||||
message_id = str(uuid.uuid4())
|
||||
message_content = self._sanitize_message_content(message_content)
|
||||
message_content = await self._externalize_message_content(message_content)
|
||||
from ....utils.execution_lease import host_execution_lease
|
||||
|
||||
message_data = {
|
||||
|
||||
@@ -143,12 +143,9 @@ class MonitoringHelper:
|
||||
sender_name = query.message_event.sender.member_name
|
||||
|
||||
# Try to record message
|
||||
# Use JSON serialization to preserve message chain structure (including image URLs, etc.)
|
||||
# Let persistence externalize images before serializing the message chain.
|
||||
if hasattr(query, 'message_chain') and hasattr(query.message_chain, 'model_dump'):
|
||||
chain_dump = query.message_chain.model_dump()
|
||||
if hasattr(ap, 'storage_mgr') and hasattr(ap.storage_mgr, 'media_cache'):
|
||||
chain_dump = await ap.storage_mgr.media_cache.externalize_chain_dump(chain_dump)
|
||||
message_content = json.dumps(chain_dump, ensure_ascii=False)
|
||||
message_content = query.message_chain.model_dump()
|
||||
else:
|
||||
message_content = str(query)
|
||||
|
||||
@@ -257,10 +254,10 @@ class MonitoringHelper:
|
||||
|
||||
# Extract response content from resp_message_chain
|
||||
if hasattr(query, 'resp_message_chain') and query.resp_message_chain:
|
||||
# Serialize the last response message chain
|
||||
# Preserve the structure until persistence has externalized images.
|
||||
last_resp = query.resp_message_chain[-1]
|
||||
if hasattr(last_resp, 'model_dump'):
|
||||
message_content = json.dumps(last_resp.model_dump(), ensure_ascii=False)
|
||||
message_content = last_resp.model_dump()
|
||||
else:
|
||||
message_content = str(last_resp)
|
||||
elif hasattr(query, 'resp_messages') and query.resp_messages:
|
||||
@@ -268,10 +265,7 @@ class MonitoringHelper:
|
||||
if hasattr(last_resp, 'get_content_platform_message_chain'):
|
||||
chain = last_resp.get_content_platform_message_chain()
|
||||
if hasattr(chain, 'model_dump'):
|
||||
chain_dump = chain.model_dump()
|
||||
if hasattr(ap, 'storage_mgr') and hasattr(ap.storage_mgr, 'media_cache'):
|
||||
chain_dump = await ap.storage_mgr.media_cache.externalize_chain_dump(chain_dump)
|
||||
message_content = json.dumps(chain_dump, ensure_ascii=False)
|
||||
message_content = chain.model_dump()
|
||||
else:
|
||||
message_content = str(chain)
|
||||
else:
|
||||
|
||||
@@ -1356,10 +1356,7 @@ class RuntimeBot:
|
||||
"""Record a discarded message in the monitoring system."""
|
||||
try:
|
||||
if hasattr(message_chain, 'model_dump'):
|
||||
chain_dump = message_chain.model_dump()
|
||||
if hasattr(self.ap, 'storage_mgr') and hasattr(self.ap.storage_mgr, 'media_cache'):
|
||||
chain_dump = await self.ap.storage_mgr.media_cache.externalize_chain_dump(chain_dump)
|
||||
message_content = json.dumps(chain_dump, ensure_ascii=False)
|
||||
message_content = message_chain.model_dump()
|
||||
else:
|
||||
message_content = str(message_chain)
|
||||
|
||||
|
||||
@@ -11,6 +11,8 @@ import xxhash
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..utils.image import DEFAULT_BASE64_MEDIA_LIMIT
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...core import app
|
||||
from . import mgr as storage_mgr
|
||||
@@ -49,14 +51,80 @@ class MediaCache:
|
||||
if split_index != -1:
|
||||
mime_type = data_url[5:split_index]
|
||||
b64_data = data_url[split_index + 8 :]
|
||||
return base64.b64decode(b64_data), mime_type
|
||||
# Try raw base64 if sufficiently long
|
||||
if len(data_url) > 20 and not data_url.startswith(('http://', 'https://', '/')):
|
||||
return base64.b64decode(data_url), 'application/octet-stream'
|
||||
return MediaCache._decode_base64(b64_data), mime_type
|
||||
if not data_url.startswith(('http://', 'https://', '/')):
|
||||
return MediaCache._decode_base64(data_url), 'application/octet-stream'
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _decode_base64(value: str) -> bytes:
|
||||
"""Bound decoding before allocating a binary payload for persistence."""
|
||||
if len(value) > 4 * ((DEFAULT_BASE64_MEDIA_LIMIT + 2) // 3):
|
||||
raise ValueError('Inline image exceeds the media size limit')
|
||||
decoded = base64.b64decode(value, validate=True)
|
||||
if len(decoded) > DEFAULT_BASE64_MEDIA_LIMIT:
|
||||
raise ValueError('Inline image exceeds the media size limit')
|
||||
return decoded
|
||||
|
||||
@staticmethod
|
||||
def _inline_image_payload(node: dict) -> str | None:
|
||||
"""Recognize platform and provider image payloads without fetching URLs."""
|
||||
node_type = node.get('type')
|
||||
if node_type == 'Image':
|
||||
payload = node.get('base64')
|
||||
url = node.get('url')
|
||||
elif node_type in ('image', 'image_url', 'image_base64'):
|
||||
payload = node.get('image_base64')
|
||||
url = node.get('image_url')
|
||||
if isinstance(url, dict):
|
||||
url = url.get('url')
|
||||
else:
|
||||
return None
|
||||
if isinstance(payload, str) and payload:
|
||||
return payload
|
||||
if isinstance(url, str) and url.startswith('data:image/'):
|
||||
return url
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def has_inline_images(cls, value: Any) -> bool:
|
||||
if isinstance(value, list):
|
||||
return any(cls.has_inline_images(item) for item in value)
|
||||
if isinstance(value, dict):
|
||||
return cls._inline_image_payload(value) is not None or any(
|
||||
cls.has_inline_images(item) for item in value.values() if isinstance(item, (list, dict))
|
||||
)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _clear_inline_image(node: dict) -> None:
|
||||
field = 'base64' if node.get('type') == 'Image' else 'image_base64'
|
||||
if node.get(field):
|
||||
node[field] = None
|
||||
url_field = 'url' if node.get('type') == 'Image' else 'image_url'
|
||||
url = node.get(url_field)
|
||||
if isinstance(url, str) and url.startswith('data:image/'):
|
||||
node[url_field] = None
|
||||
elif isinstance(url, dict) and str(url.get('url', '')).startswith('data:image/'):
|
||||
node[url_field] = {**url, 'url': None}
|
||||
|
||||
@classmethod
|
||||
def redact_inline_images(cls, value: Any) -> Any:
|
||||
"""Keep metadata-only records binary-free, including short raw base64."""
|
||||
if isinstance(value, list):
|
||||
return [cls.redact_inline_images(item) for item in value]
|
||||
if isinstance(value, dict):
|
||||
node = dict(value)
|
||||
if cls._inline_image_payload(node) is not None:
|
||||
cls._clear_inline_image(node)
|
||||
return {
|
||||
key: cls.redact_inline_images(item) if isinstance(item, (list, dict)) else item
|
||||
for key, item in node.items()
|
||||
}
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def guess_extension(mime_type: str | None, default: str = '.jpg') -> str:
|
||||
"""Guess appropriate file extension from MIME type."""
|
||||
@@ -201,29 +269,33 @@ class MediaCache:
|
||||
return [await self.externalize_chain_dump(item) for item in chain_dump]
|
||||
if isinstance(chain_dump, dict):
|
||||
node = copy.copy(chain_dump)
|
||||
node_type = node.get('type')
|
||||
if node_type == 'Image':
|
||||
b64 = node.get('base64')
|
||||
if b64 and isinstance(b64, str):
|
||||
try:
|
||||
parsed = self.parse_data_url(b64)
|
||||
if parsed is not None:
|
||||
raw_bytes, mime_type = parsed
|
||||
hash_str, storage_key, size = await self.save_media(raw_bytes, mime_type)
|
||||
filename = os.path.basename(storage_key)
|
||||
current_url = node.get('url') or ''
|
||||
if current_url and not current_url.startswith('data:'):
|
||||
node['original_url'] = current_url
|
||||
node['url'] = f'/api/v1/files/media/{filename}'
|
||||
node['base64'] = None
|
||||
node['hash'] = hash_str
|
||||
node['storage_key'] = storage_key
|
||||
node['size'] = size
|
||||
node['mime_type'] = mime_type
|
||||
except Exception as e:
|
||||
if hasattr(self.ap, 'logger') and self.ap.logger:
|
||||
self.ap.logger.warning(f'Failed to externalize image to media cache: {e}')
|
||||
node['base64'] = None
|
||||
payload = self._inline_image_payload(node)
|
||||
if payload is not None:
|
||||
try:
|
||||
parsed = await asyncio.to_thread(self.parse_data_url, payload)
|
||||
if parsed is None:
|
||||
raise ValueError('Invalid inline image payload')
|
||||
raw_bytes, mime_type = parsed
|
||||
hash_str, storage_key, size = await self.save_media(raw_bytes, mime_type)
|
||||
url = f'/api/v1/files/media/{os.path.basename(storage_key)}'
|
||||
if node.get('type') == 'Image':
|
||||
current_url = node.get('url') or ''
|
||||
if current_url and not current_url.startswith('data:'):
|
||||
node['original_url'] = current_url
|
||||
node['url'] = url
|
||||
else:
|
||||
current_url = node.get('image_url')
|
||||
if isinstance(current_url, dict):
|
||||
current_url = current_url.get('url')
|
||||
if current_url and not current_url.startswith('data:'):
|
||||
node['original_url'] = current_url
|
||||
image_url = node.get('image_url')
|
||||
node['image_url'] = {**image_url, 'url': url} if isinstance(image_url, dict) else {'url': url}
|
||||
node.update(hash=hash_str, storage_key=storage_key, size=size, mime_type=mime_type)
|
||||
except Exception as e:
|
||||
if hasattr(self.ap, 'logger') and self.ap.logger:
|
||||
self.ap.logger.warning(f'Failed to externalize image to media cache: {e}')
|
||||
self._clear_inline_image(node)
|
||||
for k, v in list(node.items()):
|
||||
if isinstance(v, (list, dict)):
|
||||
node[k] = await self.externalize_chain_dump(v)
|
||||
|
||||
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.agent.runner.run_journal import AgentRunJournal
|
||||
from langbot.pkg.storage.media import MediaCache
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -68,3 +69,21 @@ def test_event_snapshot_preserves_custom_values_and_redacts_nested_binary():
|
||||
assert saved['custom']['contents'][0]['image_base64'] is None
|
||||
assert saved['attachments'][0]['content'] is None
|
||||
assert data['custom']['contents'][0]['image_base64'] == 'secret'
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'image',
|
||||
[
|
||||
{'type': 'Image', 'base64': 'eA=='},
|
||||
{'type': 'image_base64', 'image_base64': 'eA=='},
|
||||
{'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,eA=='}},
|
||||
],
|
||||
)
|
||||
def test_metadata_and_transcript_redaction_cover_all_image_encodings(image):
|
||||
content = AgentRunJournal._sanitize_content_item(image)
|
||||
event = AgentRunJournal._sanitize_event_data({'nested': [image]})
|
||||
assert not MediaCache.has_inline_images(content)
|
||||
assert not MediaCache.has_inline_images(event)
|
||||
assert content['content_redacted'] is True
|
||||
assert event['nested'][0]['content_redacted'] is True
|
||||
assert MediaCache.has_inline_images(image)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import base64
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
@@ -127,6 +128,51 @@ class TestMediaCache:
|
||||
# Should not raise; base64 should be stripped as fallback
|
||||
assert result[0]['base64'] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'image',
|
||||
[
|
||||
{'type': 'Image', 'base64': 'eA=='},
|
||||
{'type': 'image_base64', 'image_base64': 'eA=='},
|
||||
{'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,eA==', 'detail': 'low'}},
|
||||
{'type': 'image', 'image_url': 'data:image/png;base64,eA=='},
|
||||
{'type': 'Image', 'url': 'data:image/png;base64,eA=='},
|
||||
],
|
||||
)
|
||||
async def test_externalize_image_encodings_without_mutating_runtime_input(self, image):
|
||||
original = json.dumps(image)
|
||||
result = await self.media_cache.externalize_chain_dump({'nested': [image]})
|
||||
saved = result['nested'][0]
|
||||
assert saved['storage_key'].startswith('media_cache/')
|
||||
self.mock_provider.save.assert_awaited_once_with(saved['storage_key'], b'x')
|
||||
assert not MediaCache.has_inline_images(result)
|
||||
assert json.dumps(image) == original
|
||||
if image.get('type') == 'image_url':
|
||||
assert saved['image_url']['detail'] == 'low'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('payload', ['invalid!!!', 'https://example.com/image.png'])
|
||||
async def test_invalid_base64_is_removed_without_saving_or_downloading(self, payload):
|
||||
result = await self.media_cache.externalize_chain_dump({'type': 'Image', 'base64': payload})
|
||||
assert result['base64'] is None
|
||||
self.mock_provider.save.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_url_only_images_do_not_enter_cas(self):
|
||||
image = {'type': 'Image', 'url': 'https://example.com/image.png'}
|
||||
assert not MediaCache.has_inline_images(image)
|
||||
assert await self.media_cache.externalize_chain_dump(image) == image
|
||||
self.mock_provider.exists.assert_not_awaited()
|
||||
self.mock_provider.save.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oversized_inline_image_is_removed_before_decode(self):
|
||||
image = {'type': 'Image', 'base64': 'data:image/png;base64,' + 'A' * 100}
|
||||
with patch('langbot.pkg.storage.media.DEFAULT_BASE64_MEDIA_LIMIT', 16):
|
||||
result = await self.media_cache.externalize_chain_dump(image)
|
||||
assert result['base64'] is None
|
||||
self.mock_provider.save.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_retention_and_max_size(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
@@ -197,8 +243,8 @@ class TestMonitoringServiceSanitization:
|
||||
def test_sanitize_oversized_base64_payload(self):
|
||||
svc = MonitoringService.__new__(MonitoringService)
|
||||
small_content = '{"type": "Image", "base64": "data:image/png;base64,tiny"}'
|
||||
# Should leave small contents untouched
|
||||
assert svc._sanitize_message_content(small_content) == small_content
|
||||
# Metadata-only records must not retain even small inline images.
|
||||
assert json.loads(svc._sanitize_message_content(small_content))['base64'] is None
|
||||
|
||||
# Large content with base64 data URL
|
||||
huge_b64 = 'A' * 60000
|
||||
@@ -213,3 +259,21 @@ class TestMonitoringServiceSanitization:
|
||||
sanitized_raw = svc._sanitize_message_content(raw_corrupted)
|
||||
assert '[base64 image omitted]' in sanitized_raw
|
||||
assert multiline_b64 not in sanitized_raw
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'image',
|
||||
[
|
||||
{'type': 'Image', 'base64': 'eA=='},
|
||||
{'type': 'Image', 'base64': 'A' * 20000},
|
||||
{'type': 'image_base64', 'image_base64': 'eA=='},
|
||||
{'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,eA=='}},
|
||||
],
|
||||
)
|
||||
def test_redaction_covers_raw_and_provider_images_at_any_size(self, image):
|
||||
svc = MonitoringService.__new__(MonitoringService)
|
||||
payload = {'event': {'images': [image]}, 'count': 0, 'enabled': False}
|
||||
saved = json.loads(svc._sanitize_message_content(json.dumps(payload)))
|
||||
assert not MediaCache.has_inline_images(saved)
|
||||
assert saved['count'] == 0
|
||||
assert saved['enabled'] is False
|
||||
assert MediaCache.has_inline_images(payload)
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
"""Persist media references without changing live messages or metadata-only policy."""
|
||||
|
||||
import base64
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
|
||||
from langbot.pkg.agent.runner.run_journal import AgentRunJournal
|
||||
from langbot.pkg.entity.persistence.agent_run import AgentRunEvent
|
||||
from langbot.pkg.entity.persistence.event_log import EventLog
|
||||
from langbot.pkg.entity.persistence.monitoring import MonitoringMessage
|
||||
from langbot.pkg.entity.persistence.transcript import Transcript
|
||||
from langbot.pkg.pipeline.monitoring_helper import MonitoringHelper
|
||||
from langbot.pkg.storage.media import MediaCache
|
||||
from tests.unit_tests.api.service.test_monitoring_tenancy import (
|
||||
WORKSPACE_A,
|
||||
_context,
|
||||
_record_message,
|
||||
service as base_service, # noqa: F401
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def media_service(base_service): # noqa: F811
|
||||
service = base_service
|
||||
objects = {}
|
||||
|
||||
async def save(key, data):
|
||||
objects[key] = data
|
||||
|
||||
provider = SimpleNamespace(
|
||||
exists=AsyncMock(side_effect=lambda key: key in objects),
|
||||
save=AsyncMock(side_effect=save),
|
||||
)
|
||||
service.ap.logger = logging.getLogger(__name__)
|
||||
service.ap.storage_mgr = SimpleNamespace(storage_provider=provider)
|
||||
cache = MediaCache(service.ap, service.ap.storage_mgr)
|
||||
service.ap.storage_mgr.media_cache = cache
|
||||
service.ap.monitoring_service = service
|
||||
return service, objects, cache
|
||||
|
||||
|
||||
def image(raw=b'tiny adapter image', *, encoded=True):
|
||||
payload = base64.b64encode(raw).decode()
|
||||
return {'type': 'Image', 'base64': f'data:image/png;base64,{payload}' if encoded else payload}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'response_kind',
|
||||
['input', 'resp_message_chain', 'resp_messages', 'direct', 'structured', 'structured_dict', 'discarded'],
|
||||
)
|
||||
async def test_all_retained_monitoring_messages_use_cas(media_service, response_kind, monkeypatch):
|
||||
service, objects, _ = media_service
|
||||
context = _context(WORKSPACE_A)
|
||||
chain = [image()]
|
||||
original = copy.deepcopy(chain)
|
||||
legacy_content = json.dumps(chain)
|
||||
json_dumps = json.dumps
|
||||
|
||||
def serialize_references(value, *args, **kwargs):
|
||||
assert not MediaCache.has_inline_images(value), 'Inline images must reach CAS before JSON serialization'
|
||||
return json_dumps(value, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(json, 'dumps', serialize_references)
|
||||
query = SimpleNamespace(
|
||||
_execution_context=context,
|
||||
launcher_type='person',
|
||||
launcher_id='user',
|
||||
sender_id='user',
|
||||
message_chain=SimpleNamespace(model_dump=lambda: chain),
|
||||
variables={},
|
||||
)
|
||||
resource = dict(bot_id='same-bot', bot_name='Bot', pipeline_id='pipeline', pipeline_name='Pipeline')
|
||||
if response_kind == 'input':
|
||||
await MonitoringHelper.record_query_start(service.ap, query, **resource)
|
||||
elif response_kind == 'direct':
|
||||
await _record_message(service, context, legacy_content)
|
||||
elif response_kind in ('structured', 'structured_dict'):
|
||||
content = {'chain': chain} if response_kind == 'structured_dict' else chain
|
||||
await _record_message(service, context, content)
|
||||
elif response_kind == 'discarded':
|
||||
from langbot.pkg.platform.botmgr import RuntimeBot
|
||||
|
||||
bot = SimpleNamespace(
|
||||
ap=service.ap,
|
||||
execution_context=context,
|
||||
bot_entity=SimpleNamespace(uuid='same-bot', name='Bot'),
|
||||
PIPELINE_DISCARD=RuntimeBot.PIPELINE_DISCARD,
|
||||
PIPELINE_DISCARD_DISPLAY_NAME=RuntimeBot.PIPELINE_DISCARD_DISPLAY_NAME,
|
||||
logger=SimpleNamespace(error=AsyncMock()),
|
||||
)
|
||||
await RuntimeBot._record_discarded_message(
|
||||
bot, 'person', 'user', 'user', SimpleNamespace(), query.message_chain
|
||||
)
|
||||
else:
|
||||
response = SimpleNamespace(model_dump=lambda: chain)
|
||||
if response_kind == 'resp_message_chain':
|
||||
query.resp_message_chain = [response]
|
||||
else:
|
||||
query.resp_messages = [SimpleNamespace(get_content_platform_message_chain=lambda: response)]
|
||||
await MonitoringHelper.record_query_response(service.ap, query, **resource)
|
||||
|
||||
async with service.ap.persistence_mgr.get_db_engine().connect() as conn:
|
||||
stored = (await conn.execute(sa.select(MonitoringMessage.message_content))).scalar_one()
|
||||
saved_chain = json.loads(stored)
|
||||
saved = (saved_chain['chain'] if isinstance(saved_chain, dict) else saved_chain)[0]
|
||||
assert saved['base64'] is None
|
||||
assert saved['url'].startswith('/api/v1/files/media/')
|
||||
assert objects[saved['storage_key']] == b'tiny adapter image'
|
||||
assert chain == original
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_only_ingress_discards_images_without_populating_cas(media_service):
|
||||
service, objects, _ = media_service
|
||||
payload = {'type': 'message.received', 'message_chain': [image(encoded=False)], 'count': 0}
|
||||
original = copy.deepcopy(payload)
|
||||
await service.record_ingress_event(
|
||||
_context(WORKSPACE_A),
|
||||
event_id='ingress',
|
||||
event_type='message.received',
|
||||
bot_id='same-bot',
|
||||
bot_name='Bot',
|
||||
platform='slack',
|
||||
payload=payload,
|
||||
)
|
||||
async with service.ap.persistence_mgr.get_db_engine().connect() as conn:
|
||||
stored = (await conn.execute(sa.select(EventLog.input_json))).scalar_one()
|
||||
assert json.loads(stored)['message_chain'][0]['base64'] is None
|
||||
assert not objects
|
||||
assert payload == original
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeated_images_deduplicate_across_messages_and_run_events(media_service):
|
||||
service, objects, cache = media_service
|
||||
raw = b'shared image'
|
||||
payload = image(raw)
|
||||
await _record_message(service, _context(WORKSPACE_A), json.dumps([payload]))
|
||||
journal = AgentRunJournal(service.ap)
|
||||
result = {
|
||||
'type': 'message.completed',
|
||||
'data': {
|
||||
'message': {'content': [payload]},
|
||||
'provider_image': {'type': 'image_base64', 'image_base64': payload['base64']},
|
||||
},
|
||||
}
|
||||
original = copy.deepcopy(result)
|
||||
metadata = {'diagnostic_image': image(b'metadata-only image')}
|
||||
await journal.append_run_result(result_dict=result, run_id='run', sequence=1, metadata=metadata)
|
||||
async with service.ap.persistence_mgr.get_db_engine().connect() as conn:
|
||||
stored = (await conn.execute(sa.select(AgentRunEvent.data_json))).scalar_one()
|
||||
saved_metadata = (await conn.execute(sa.select(AgentRunEvent.metadata_json))).scalar_one()
|
||||
saved = json.loads(stored)
|
||||
assert saved['message']['content'][0]['hash'] == cache.hash_bytes(raw)
|
||||
assert saved['provider_image']['image_base64'] is None
|
||||
assert saved['provider_image']['image_url']['url'].startswith('/api/v1/files/media/')
|
||||
assert len(objects) == 1
|
||||
service.ap.storage_mgr.storage_provider.save.assert_awaited_once()
|
||||
assert result == original
|
||||
assert not MediaCache.has_inline_images(json.loads(saved_metadata))
|
||||
assert MediaCache.has_inline_images(metadata)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cas_failure_does_not_write_inline_payload_to_database(media_service):
|
||||
service, objects, _ = media_service
|
||||
service.ap.storage_mgr.storage_provider.save.side_effect = OSError('Disk full')
|
||||
message_id = await _record_message(service, _context(WORKSPACE_A), json.dumps([image()]))
|
||||
async with service.ap.persistence_mgr.get_db_engine().connect() as conn:
|
||||
stored = (
|
||||
await conn.execute(sa.select(MonitoringMessage.message_content).where(MonitoringMessage.id == message_id))
|
||||
).scalar_one()
|
||||
assert json.loads(stored)[0]['base64'] is None
|
||||
assert not objects
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('role', ['user', 'assistant'])
|
||||
async def test_transcript_remains_metadata_only_without_creating_cas_objects(media_service, role):
|
||||
service, objects, _ = media_service
|
||||
contents = [image(), {'type': 'image_base64', 'image_base64': 'eA=='}]
|
||||
original = copy.deepcopy(contents)
|
||||
event = SimpleNamespace(
|
||||
event_id='event',
|
||||
conversation_id='conversation',
|
||||
bot_id='same-bot',
|
||||
workspace_id=WORKSPACE_A,
|
||||
thread_id=None,
|
||||
actor=None,
|
||||
input=SimpleNamespace(text='image', contents=contents, attachments=[]),
|
||||
)
|
||||
journal = AgentRunJournal(service.ap)
|
||||
if role == 'user':
|
||||
await journal.write_user_transcript(event, 'event')
|
||||
else:
|
||||
await journal.write_assistant_transcript(
|
||||
result_dict={'data': {'message': {'role': 'assistant', 'content': contents}}},
|
||||
event=event,
|
||||
run_id='run',
|
||||
runner_id='runner',
|
||||
)
|
||||
async with service.ap.persistence_mgr.get_db_engine().connect() as conn:
|
||||
stored = (await conn.execute(sa.select(Transcript.content_json))).scalar_one()
|
||||
assert not MediaCache.has_inline_images(json.loads(stored))
|
||||
assert not objects
|
||||
assert contents == original
|
||||
Reference in New Issue
Block a user