feat(avatar): understand BOXIM image messages
This commit is contained in:
@@ -45,6 +45,7 @@ def init_db():
|
||||
("token_account", "total_consumed", "BIGINT DEFAULT 0"),
|
||||
("token_account", "created_at", "TIMESTAMP"),
|
||||
("token_account", "updated_at", "TIMESTAMP"),
|
||||
("takeover_messages", "attachment_id", "VARCHAR DEFAULT NULL"),
|
||||
)
|
||||
_normalize_optional_unique_values()
|
||||
_normalize_takeover_delays()
|
||||
|
||||
@@ -120,6 +120,7 @@ class TakeoverMessage(Base):
|
||||
direction = Column(String, nullable=False) # incoming | outgoing
|
||||
message_type = Column(Integer, default=0)
|
||||
content = Column(Text, default="")
|
||||
attachment_id = Column(String, nullable=True)
|
||||
is_avatar = Column(Boolean, default=False)
|
||||
send_time = Column(DateTime, nullable=False)
|
||||
created_at = Column(DateTime, server_default=func.now())
|
||||
@@ -263,7 +264,7 @@ class ChatAttachment(Base):
|
||||
__tablename__ = "chat_attachments"
|
||||
id = Column(String, primary_key=True, default=lambda: uuid.uuid4().hex)
|
||||
avatar_id = Column(String, nullable=False, default="", index=True)
|
||||
uploader_kind = Column(String, default="owner") # owner | public
|
||||
uploader_kind = Column(String, default="owner") # owner | public | boxim
|
||||
filename = Column(String, default="")
|
||||
mime_type = Column(String, default="")
|
||||
file_size = Column(Integer, default=0)
|
||||
|
||||
@@ -236,11 +236,40 @@ async def _analyze_uploaded_image(
|
||||
max_bytes = max(1024, int(os.getenv("CHAT_IMAGE_MAX_BYTES", str(8 * 1024 * 1024))))
|
||||
content = await file.read(max_bytes + 1)
|
||||
filename = os.path.basename(file.filename or "图片")[:255]
|
||||
try:
|
||||
return _analyze_image_bytes(
|
||||
db,
|
||||
avatar,
|
||||
content,
|
||||
filename=filename,
|
||||
mime_type=file.content_type or "",
|
||||
uploader_kind=uploader_kind,
|
||||
)
|
||||
except ImageValidationError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
except InsufficientTokensError:
|
||||
raise
|
||||
except RuntimeError as exc:
|
||||
raise HTTPException(status_code=502, detail=str(exc)) from exc
|
||||
finally:
|
||||
content = b""
|
||||
|
||||
|
||||
def _analyze_image_bytes(
|
||||
db: Session,
|
||||
avatar: Avatar,
|
||||
content: bytes,
|
||||
*,
|
||||
filename: str,
|
||||
mime_type: str,
|
||||
uploader_kind: str,
|
||||
) -> ChatAttachment:
|
||||
"""Analyze image bytes from either HTTP upload or BOXIM without persisting raw data."""
|
||||
attachment = ChatAttachment(
|
||||
avatar_id=avatar.id,
|
||||
uploader_kind=uploader_kind,
|
||||
filename=filename,
|
||||
mime_type=(file.content_type or "")[:100],
|
||||
mime_type=(mime_type or "")[:100],
|
||||
file_size=len(content),
|
||||
status="processing",
|
||||
expires_at=_attachment_expiry(),
|
||||
@@ -312,7 +341,7 @@ async def _analyze_uploaded_image(
|
||||
attachment.status = "failed"
|
||||
attachment.warning = str(exc)
|
||||
db.commit()
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
raise
|
||||
except InsufficientTokensError:
|
||||
attachment.status = "failed"
|
||||
attachment.warning = "积分余额不足"
|
||||
@@ -328,9 +357,7 @@ async def _analyze_uploaded_image(
|
||||
avatar.id,
|
||||
type(exc).__name__,
|
||||
)
|
||||
raise HTTPException(status_code=502, detail=str(exc)) from exc
|
||||
finally:
|
||||
content = b""
|
||||
raise
|
||||
|
||||
|
||||
def _normalize_question(value: str) -> str:
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
"""Parse and safely download image payloads from BOXIM private messages."""
|
||||
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
from dataclasses import dataclass
|
||||
from pathlib import PurePosixPath
|
||||
from urllib.parse import unquote, urljoin, urlsplit
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
MAX_REDIRECTS = 3
|
||||
|
||||
|
||||
class BoxIMImageError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DownloadedBoxIMImage:
|
||||
content: bytes
|
||||
filename: str
|
||||
mime_type: str
|
||||
source_url: str
|
||||
|
||||
|
||||
def parse_boxim_image_url(content: str, *, base_url: str = "") -> str:
|
||||
try:
|
||||
payload = json.loads(content or "")
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise BoxIMImageError("BOXIM 图片消息格式无效") from exc
|
||||
if not isinstance(payload, dict):
|
||||
raise BoxIMImageError("BOXIM 图片消息格式无效")
|
||||
|
||||
value = payload.get("originUrl") or payload.get("thumbUrl") or payload.get("url")
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise BoxIMImageError("BOXIM 图片消息缺少图片地址")
|
||||
value = value.strip()
|
||||
if value.startswith("/"):
|
||||
if not base_url:
|
||||
raise BoxIMImageError("BOXIM 图片地址不完整")
|
||||
value = urljoin(f"{base_url.rstrip('/')}/", value)
|
||||
return value
|
||||
|
||||
|
||||
def _configured_hosts(name: str) -> set[str]:
|
||||
return {
|
||||
value.strip().lower().rstrip(".")
|
||||
for value in os.getenv(name, "").split(",")
|
||||
if value.strip()
|
||||
}
|
||||
|
||||
|
||||
def _host_matches(host: str, configured: set[str]) -> bool:
|
||||
return any(host == value or host.endswith(f".{value}") for value in configured)
|
||||
|
||||
|
||||
def _resolved_addresses(host: str, port: int) -> set[ipaddress.IPv4Address | ipaddress.IPv6Address]:
|
||||
try:
|
||||
return {
|
||||
ipaddress.ip_address(item[4][0])
|
||||
for item in socket.getaddrinfo(host, port, type=socket.SOCK_STREAM)
|
||||
}
|
||||
except (OSError, ValueError) as exc:
|
||||
raise BoxIMImageError("BOXIM 图片地址无法解析") from exc
|
||||
|
||||
|
||||
def _is_safe_remote_url(url: str) -> None:
|
||||
parsed = urlsplit(url)
|
||||
scheme = parsed.scheme.lower()
|
||||
allow_http = os.getenv("BOXIM_IMAGE_ALLOW_HTTP", "").lower() in {"1", "true", "yes"}
|
||||
if scheme not in ({"https", "http"} if allow_http else {"https"}):
|
||||
raise BoxIMImageError("BOXIM 图片地址必须使用 HTTPS")
|
||||
if parsed.username or parsed.password or not parsed.hostname:
|
||||
raise BoxIMImageError("BOXIM 图片地址无效")
|
||||
|
||||
host = parsed.hostname.lower().rstrip(".")
|
||||
allowed_hosts = _configured_hosts("BOXIM_IMAGE_ALLOWED_HOSTS")
|
||||
if allowed_hosts and not _host_matches(host, allowed_hosts):
|
||||
raise BoxIMImageError("BOXIM 图片地址不在允许的域名范围内")
|
||||
|
||||
private_hosts = _configured_hosts("BOXIM_IMAGE_PRIVATE_HOSTS")
|
||||
try:
|
||||
addresses = {ipaddress.ip_address(host)}
|
||||
except ValueError:
|
||||
addresses = _resolved_addresses(host, parsed.port or (443 if scheme == "https" else 80))
|
||||
if not addresses:
|
||||
raise BoxIMImageError("BOXIM 图片地址无法解析")
|
||||
if _host_matches(host, private_hosts):
|
||||
return
|
||||
if any(not address.is_global for address in addresses):
|
||||
raise BoxIMImageError("BOXIM 图片地址指向受限网络")
|
||||
|
||||
|
||||
def _filename_from_url(url: str) -> str:
|
||||
value = unquote(PurePosixPath(urlsplit(url).path).name).strip()
|
||||
value = value.replace("\x00", "")
|
||||
return (value or "boxim-image")[:255]
|
||||
|
||||
|
||||
def download_boxim_image(
|
||||
content: str,
|
||||
*,
|
||||
base_url: str = "",
|
||||
transport: httpx.BaseTransport | None = None,
|
||||
) -> DownloadedBoxIMImage:
|
||||
"""Download one BOXIM image without redirects or oversized responses escaping checks."""
|
||||
url = parse_boxim_image_url(content, base_url=base_url)
|
||||
max_bytes = max(1024, int(os.getenv("CHAT_IMAGE_MAX_BYTES", str(8 * 1024 * 1024))))
|
||||
timeout = max(1.0, min(float(os.getenv("BOXIM_IMAGE_TIMEOUT_SECONDS", "15")), 60.0))
|
||||
|
||||
with httpx.Client(
|
||||
timeout=timeout,
|
||||
follow_redirects=False,
|
||||
trust_env=False,
|
||||
transport=transport,
|
||||
) as client:
|
||||
for _ in range(MAX_REDIRECTS + 1):
|
||||
_is_safe_remote_url(url)
|
||||
try:
|
||||
with client.stream("GET", url, headers={"Accept": "image/*"}) as response:
|
||||
if response.status_code in {301, 302, 303, 307, 308}:
|
||||
location = response.headers.get("location", "").strip()
|
||||
if not location:
|
||||
raise BoxIMImageError("BOXIM 图片跳转地址无效")
|
||||
url = urljoin(url, location)
|
||||
continue
|
||||
response.raise_for_status()
|
||||
raw_length = response.headers.get("content-length", "")
|
||||
if raw_length.isdigit() and int(raw_length) > max_bytes:
|
||||
raise BoxIMImageError("BOXIM 图片超过大小限制")
|
||||
chunks = bytearray()
|
||||
for chunk in response.iter_bytes():
|
||||
chunks.extend(chunk)
|
||||
if len(chunks) > max_bytes:
|
||||
raise BoxIMImageError("BOXIM 图片超过大小限制")
|
||||
if not chunks:
|
||||
raise BoxIMImageError("BOXIM 图片内容为空")
|
||||
return DownloadedBoxIMImage(
|
||||
content=bytes(chunks),
|
||||
filename=_filename_from_url(url),
|
||||
mime_type=response.headers.get("content-type", "").split(";", 1)[0][:100],
|
||||
source_url=url,
|
||||
)
|
||||
except BoxIMImageError:
|
||||
raise
|
||||
except (httpx.HTTPError, OSError) as exc:
|
||||
raise BoxIMImageError("BOXIM 图片下载失败") from exc
|
||||
raise BoxIMImageError("BOXIM 图片跳转次数过多")
|
||||
@@ -3,6 +3,7 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
import time
|
||||
@@ -13,12 +14,19 @@ from sqlalchemy.orm import Session
|
||||
|
||||
from models import (
|
||||
Avatar,
|
||||
ChatAttachment,
|
||||
TakeoverCursor,
|
||||
TakeoverMessage,
|
||||
TakeoverReplyTask,
|
||||
User,
|
||||
)
|
||||
from services.boxim_client import BoxIMClient, BoxIMError
|
||||
from services.boxim_image_service import (
|
||||
BoxIMImageError,
|
||||
download_boxim_image,
|
||||
parse_boxim_image_url,
|
||||
)
|
||||
from services.vision_service import ImageValidationError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -37,6 +45,10 @@ HUMAN_PAUSE_SECONDS = 600
|
||||
RATE_LIMIT_WINDOW_SECONDS = 300
|
||||
RATE_LIMIT_MAX_REPLIES = 5
|
||||
AVATAR_LOCAL_ID_PREFIX = "880"
|
||||
BOXIM_TEXT_MESSAGE_TYPE = 0
|
||||
BOXIM_IMAGE_MESSAGE_TYPE = 1
|
||||
BOXIM_IMAGE_PROMPT = "请看看这张图片。"
|
||||
BOXIM_IMAGE_UNAVAILABLE_REPLY = "这张图片我暂时没看清,麻烦重新发送一张清晰的原图。"
|
||||
|
||||
|
||||
def _utcnow() -> datetime:
|
||||
@@ -107,6 +119,14 @@ def _configured_reply_delay(avatar: Avatar, fallback: int | None = None) -> int:
|
||||
return delay
|
||||
|
||||
|
||||
def _event_prompt(event: TakeoverMessage) -> str:
|
||||
if event.message_type == BOXIM_TEXT_MESSAGE_TYPE:
|
||||
return event.content.strip()
|
||||
if event.message_type == BOXIM_IMAGE_MESSAGE_TYPE:
|
||||
return BOXIM_IMAGE_PROMPT
|
||||
return ""
|
||||
|
||||
|
||||
class TakeoverService:
|
||||
"""Poll BOXIM, honor the owner grace period, then generate and send one reply."""
|
||||
|
||||
@@ -498,8 +518,27 @@ class TakeoverService:
|
||||
if not is_avatar:
|
||||
self._cancel_conversation(db, avatar.owner_id, peer_id, "owner_replied")
|
||||
return
|
||||
if not schedule_reply or event.message_type != 0 or not event.content.strip():
|
||||
if not schedule_reply or event.message_type not in {
|
||||
BOXIM_TEXT_MESSAGE_TYPE,
|
||||
BOXIM_IMAGE_MESSAGE_TYPE,
|
||||
}:
|
||||
return
|
||||
if event.message_type == BOXIM_TEXT_MESSAGE_TYPE and not event.content.strip():
|
||||
return
|
||||
if event.message_type == BOXIM_IMAGE_MESSAGE_TYPE:
|
||||
try:
|
||||
parse_boxim_image_url(
|
||||
event.content,
|
||||
base_url=getattr(self.boxim, "im_base_url", ""),
|
||||
)
|
||||
except BoxIMImageError as exc:
|
||||
logger.warning(
|
||||
"Ignored invalid BOXIM image message %s for avatar %s: %s",
|
||||
message_id,
|
||||
avatar.id,
|
||||
exc,
|
||||
)
|
||||
return
|
||||
if (now - send_time).total_seconds() > self.max_message_age_seconds:
|
||||
logger.info(
|
||||
"Ignored stale BOXIM message %s for avatar %s (age=%ss)",
|
||||
@@ -606,7 +645,7 @@ class TakeoverService:
|
||||
task.status = "cancelled"
|
||||
task.cancel_reason = "newer_incoming_message"
|
||||
task.locked_at = None
|
||||
prompt_parts.append(event.content.strip())
|
||||
prompt_parts.append(_event_prompt(event))
|
||||
source_ids.append(event.boxim_message_id)
|
||||
prompt = "\n".join(part for part in prompt_parts if part).strip()[-MAX_PROMPT_LENGTH:]
|
||||
due_at = max(
|
||||
@@ -665,6 +704,50 @@ class TakeoverService:
|
||||
results = await asyncio.gather(*(generate(task_id) for task_id in task_ids))
|
||||
return sum(bool(result) for result in results)
|
||||
|
||||
def _takeover_image_attachment(
|
||||
self,
|
||||
db: Session,
|
||||
avatar: Avatar,
|
||||
event: TakeoverMessage,
|
||||
) -> ChatAttachment:
|
||||
now = self.now()
|
||||
if event.attachment_id:
|
||||
cached = db.get(ChatAttachment, event.attachment_id)
|
||||
if cached and cached.status == "ready" and cached.expires_at > now:
|
||||
cached.used_at = now
|
||||
db.commit()
|
||||
return cached
|
||||
|
||||
downloaded = download_boxim_image(
|
||||
event.content,
|
||||
base_url=getattr(
|
||||
self.boxim,
|
||||
"im_base_url",
|
||||
os.getenv("BOXIM_API_BASE_URL", "https://im.99hui.com/api"),
|
||||
),
|
||||
)
|
||||
from routers.chat import _analyze_image_bytes
|
||||
|
||||
attachment = _analyze_image_bytes(
|
||||
db,
|
||||
avatar,
|
||||
downloaded.content,
|
||||
filename=downloaded.filename,
|
||||
mime_type=downloaded.mime_type,
|
||||
uploader_kind="boxim",
|
||||
)
|
||||
event.attachment_id = attachment.id
|
||||
attachment.used_at = now
|
||||
db.commit()
|
||||
logger.info(
|
||||
"BOXIM image analyzed message=%s attachment=%s avatar=%s category=%s",
|
||||
event.boxim_message_id,
|
||||
attachment.id,
|
||||
avatar.id,
|
||||
attachment.category,
|
||||
)
|
||||
return attachment
|
||||
|
||||
def _generate_reply(self, task_id: str) -> bool:
|
||||
db = self.session_factory()
|
||||
try:
|
||||
@@ -694,9 +777,36 @@ class TakeoverService:
|
||||
.limit(30)
|
||||
.all()
|
||||
)
|
||||
source_events = {
|
||||
event.boxim_message_id: event
|
||||
for event in events
|
||||
if event.boxim_message_id in excluded_ids
|
||||
}
|
||||
image_attachments = []
|
||||
image_failed = False
|
||||
for message_id in (task.source_message_ids or [])[-3:]:
|
||||
event = source_events.get(message_id)
|
||||
if not event or event.message_type != BOXIM_IMAGE_MESSAGE_TYPE:
|
||||
continue
|
||||
try:
|
||||
image_attachments.append(
|
||||
self._takeover_image_attachment(db, avatar, event)
|
||||
)
|
||||
except (BoxIMImageError, ImageValidationError) as exc:
|
||||
image_failed = True
|
||||
logger.warning(
|
||||
"BOXIM image unavailable message=%s avatar=%s: %s",
|
||||
event.boxim_message_id,
|
||||
avatar.id,
|
||||
exc,
|
||||
)
|
||||
history = []
|
||||
for event in reversed(events):
|
||||
if event.boxim_message_id in excluded_ids or not event.content.strip():
|
||||
if (
|
||||
event.boxim_message_id in excluded_ids
|
||||
or event.message_type != BOXIM_TEXT_MESSAGE_TYPE
|
||||
or not event.content.strip()
|
||||
):
|
||||
continue
|
||||
if event.direction == "incoming" and event.is_avatar:
|
||||
continue
|
||||
@@ -708,10 +818,25 @@ class TakeoverService:
|
||||
)
|
||||
history = history[-10:]
|
||||
|
||||
from routers.chat import _resolve_reply
|
||||
from routers.chat import _attachment_contexts, _resolve_reply
|
||||
|
||||
result = _resolve_reply(db, avatar, task.prompt, history, usage_source="takeover")
|
||||
answer = _plain_text_reply(result.get("answer", ""))
|
||||
image_contexts = _attachment_contexts(image_attachments)
|
||||
has_source_text = any(
|
||||
event.message_type == BOXIM_TEXT_MESSAGE_TYPE and event.content.strip()
|
||||
for event in source_events.values()
|
||||
)
|
||||
if image_failed and not image_contexts and not has_source_text:
|
||||
answer = BOXIM_IMAGE_UNAVAILABLE_REPLY
|
||||
else:
|
||||
result = _resolve_reply(
|
||||
db,
|
||||
avatar,
|
||||
task.prompt,
|
||||
history,
|
||||
usage_source="takeover",
|
||||
image_contexts=image_contexts,
|
||||
)
|
||||
answer = _plain_text_reply(result.get("answer", ""))
|
||||
db.refresh(task)
|
||||
if task.status != "generating":
|
||||
return False
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
import ipaddress
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from services.boxim_image_service import (
|
||||
BoxIMImageError,
|
||||
download_boxim_image,
|
||||
parse_boxim_image_url,
|
||||
)
|
||||
|
||||
|
||||
def test_parse_boxim_image_prefers_origin_and_supports_relative_url():
|
||||
content = json.dumps({"originUrl": "/files/original.png", "thumbUrl": "/thumb.png"})
|
||||
assert parse_boxim_image_url(content, base_url="https://im.example/api") == (
|
||||
"https://im.example/files/original.png"
|
||||
)
|
||||
|
||||
|
||||
def test_download_boxim_image_streams_public_https(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"services.boxim_image_service._resolved_addresses",
|
||||
lambda _host, _port: {ipaddress.ip_address("8.8.8.8")},
|
||||
)
|
||||
transport = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "image/png"},
|
||||
content=b"png-bytes",
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
image = download_boxim_image(
|
||||
json.dumps({"originUrl": "https://cdn.example/case%20photo.png"}),
|
||||
transport=transport,
|
||||
)
|
||||
|
||||
assert image.content == b"png-bytes"
|
||||
assert image.filename == "case photo.png"
|
||||
assert image.mime_type == "image/png"
|
||||
|
||||
|
||||
def test_download_boxim_image_rejects_private_network_url():
|
||||
with pytest.raises(BoxIMImageError, match="受限网络"):
|
||||
download_boxim_image(
|
||||
json.dumps({"originUrl": "https://127.0.0.1/private.png"}),
|
||||
transport=httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
|
||||
)
|
||||
|
||||
|
||||
def test_download_boxim_image_stops_oversized_stream(monkeypatch):
|
||||
monkeypatch.setenv("CHAT_IMAGE_MAX_BYTES", "1024")
|
||||
monkeypatch.setattr(
|
||||
"services.boxim_image_service._resolved_addresses",
|
||||
lambda _host, _port: {ipaddress.ip_address("8.8.8.8")},
|
||||
)
|
||||
transport = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200,
|
||||
headers={"content-length": "2048"},
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
with pytest.raises(BoxIMImageError, match="超过大小限制"):
|
||||
download_boxim_image(
|
||||
json.dumps({"originUrl": "https://cdn.example/large.png"}),
|
||||
transport=transport,
|
||||
)
|
||||
@@ -17,6 +17,7 @@ from routers.chat import (
|
||||
_resolve_reply,
|
||||
)
|
||||
from services.chat_attachment_service import purge_expired_chat_attachments
|
||||
from services.token_billing import InsufficientTokensError
|
||||
from services.vision_service import PreparedImage
|
||||
|
||||
|
||||
@@ -75,6 +76,22 @@ def test_non_owner_cannot_upload_chat_image(authorization_context):
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
def test_image_upload_preserves_insufficient_points_response(authorization_context):
|
||||
context = authorization_context
|
||||
with patch(
|
||||
"routers.chat._analyze_image_bytes",
|
||||
side_effect=InsufficientTokensError("积分余额不足"),
|
||||
):
|
||||
response = client.post(
|
||||
f"/api/avatar/{context['avatar'].id}/chat/images",
|
||||
headers=context["owner_headers"],
|
||||
files={"file": ("private.png", b"image-bytes", "image/png")},
|
||||
)
|
||||
|
||||
assert response.status_code == 402
|
||||
assert response.json()["detail"] == "积分余额不足"
|
||||
|
||||
|
||||
def test_public_share_can_upload_without_exposing_analysis_details(authorization_context):
|
||||
context = authorization_context
|
||||
db = SessionLocal()
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""End-to-end service tests for BOXIM takeover timing and human priority."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from threading import Barrier
|
||||
from unittest.mock import AsyncMock, patch
|
||||
@@ -10,8 +11,9 @@ from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from database import Base
|
||||
from models import Avatar, TakeoverCursor, TakeoverMessage, TakeoverReplyTask, User
|
||||
from models import Avatar, ChatAttachment, TakeoverCursor, TakeoverMessage, TakeoverReplyTask, User
|
||||
from services.boxim_client import BoxIMError
|
||||
from services.boxim_image_service import DownloadedBoxIMImage
|
||||
from services.takeover_service import (
|
||||
AVATAR_LOCAL_ID_PREFIX,
|
||||
TakeoverService,
|
||||
@@ -171,6 +173,121 @@ async def test_incoming_message_is_prepared_then_sent_at_three_seconds(service_c
|
||||
db.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_incoming_image_is_analyzed_and_used_in_takeover_reply(service_context):
|
||||
session_factory, service, boxim, clock = service_context
|
||||
await service.poll_and_process_messages()
|
||||
boxim.messages.append(
|
||||
{
|
||||
"id": 111,
|
||||
"localId": 111,
|
||||
"sendId": 200,
|
||||
"recvId": 100,
|
||||
"sendTime": clock.millis(),
|
||||
"type": 1,
|
||||
"content": json.dumps(
|
||||
{
|
||||
"originUrl": "https://cdn.example/case.png",
|
||||
"thumbUrl": "https://cdn.example/case-thumb.png",
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
await service.poll_and_process_messages()
|
||||
db = session_factory()
|
||||
try:
|
||||
scheduled = db.query(TakeoverReplyTask).filter_by(trigger_message_id="111").one()
|
||||
assert scheduled.status == "pending"
|
||||
assert scheduled.prompt == "请看看这张图片。"
|
||||
finally:
|
||||
db.close()
|
||||
clock.advance(3)
|
||||
|
||||
def analyze(db, avatar, content, **kwargs):
|
||||
assert content == b"image-content"
|
||||
attachment = ChatAttachment(
|
||||
avatar_id=avatar.id,
|
||||
uploader_kind=kwargs["uploader_kind"],
|
||||
filename=kwargs["filename"],
|
||||
mime_type="image/jpeg",
|
||||
file_size=len(content),
|
||||
status="ready",
|
||||
category="medical_document",
|
||||
summary="一张门诊病例",
|
||||
extracted_text="主诉:咳嗽三天",
|
||||
structured_data={"medical": {"chief_complaint": "咳嗽三天"}},
|
||||
warning="请核对原始资料",
|
||||
expires_at=clock.now() + timedelta(hours=24),
|
||||
)
|
||||
db.add(attachment)
|
||||
db.commit()
|
||||
db.refresh(attachment)
|
||||
return attachment
|
||||
|
||||
downloaded = DownloadedBoxIMImage(
|
||||
content=b"image-content",
|
||||
filename="case.png",
|
||||
mime_type="image/png",
|
||||
source_url="https://cdn.example/case.png",
|
||||
)
|
||||
with (
|
||||
patch("services.takeover_service.download_boxim_image", return_value=downloaded),
|
||||
patch("routers.chat._analyze_image_bytes", side_effect=analyze) as analyzer,
|
||||
patch("routers.chat._resolve_reply", return_value={"answer": "这份资料里写的是咳嗽三天。"}) as resolver,
|
||||
):
|
||||
await service.poll_and_process_messages()
|
||||
|
||||
analyzer.assert_called_once()
|
||||
assert resolver.call_args.args[2] == "请看看这张图片。"
|
||||
image_contexts = resolver.call_args.kwargs["image_contexts"]
|
||||
assert image_contexts[0]["summary"] == "一张门诊病例"
|
||||
assert image_contexts[0]["extractedText"] == "主诉:咳嗽三天"
|
||||
assert [item["content"] for item in boxim.sent] == ["这份资料里写的是咳嗽三天。"]
|
||||
|
||||
db = session_factory()
|
||||
try:
|
||||
event = db.query(TakeoverMessage).filter_by(boxim_message_id="111").one()
|
||||
task = db.query(TakeoverReplyTask).filter_by(trigger_message_id="111").one()
|
||||
assert event.attachment_id
|
||||
assert db.get(ChatAttachment, event.attachment_id).uploader_kind == "boxim"
|
||||
assert task.status == "sent"
|
||||
with patch(
|
||||
"services.takeover_service.download_boxim_image",
|
||||
side_effect=AssertionError("cached image must not be downloaded again"),
|
||||
):
|
||||
cached = service._takeover_image_attachment(db, db.get(Avatar, "avatar-1"), event)
|
||||
assert cached.id == event.attachment_id
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_image_message_is_recorded_but_not_scheduled(service_context):
|
||||
session_factory, service, boxim, clock = service_context
|
||||
await service.poll_and_process_messages()
|
||||
boxim.messages.append(
|
||||
{
|
||||
"id": 112,
|
||||
"localId": 112,
|
||||
"sendId": 200,
|
||||
"recvId": 100,
|
||||
"sendTime": clock.millis(),
|
||||
"type": 1,
|
||||
"content": json.dumps({"width": 100, "height": 100}),
|
||||
}
|
||||
)
|
||||
|
||||
await service.poll_and_process_messages()
|
||||
|
||||
db = session_factory()
|
||||
try:
|
||||
assert db.query(TakeoverMessage).filter_by(boxim_message_id="112").one()
|
||||
assert db.query(TakeoverReplyTask).count() == 0
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_reply_delay_is_three_minutes(service_context):
|
||||
session_factory, service, boxim, clock = service_context
|
||||
|
||||
Reference in New Issue
Block a user