Refine AI image context handling
This commit is contained in:
Codex
297
app/ai/image_understand.py
Normal file
297
app/ai/image_understand.py
Normal file
@@ -0,0 +1,297 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from app.config import AIImageUnderstandConfig
|
||||
|
||||
|
||||
SUPPORTED_IMAGE_MIME_PREFIXES = ("image/",)
|
||||
SUPPORTED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp"}
|
||||
MAX_IMAGE_BYTES = 20 * 1024 * 1024
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class TelegramImageSource:
|
||||
message_id: int
|
||||
index: int
|
||||
kind: str
|
||||
media: Any
|
||||
unique_id: str
|
||||
mime_type: str | None = None
|
||||
file_name: str | None = None
|
||||
text_context: str = ""
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return f"message {self.message_id} #{self.index + 1} ({self.kind})"
|
||||
|
||||
|
||||
def _text_context(message: Any) -> str:
|
||||
text = getattr(message, "text", None) or getattr(message, "caption", None)
|
||||
if isinstance(text, str):
|
||||
return text.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def _is_supported_image(media: Any) -> bool:
|
||||
mime_type = getattr(media, "mime_type", None)
|
||||
if isinstance(mime_type, str) and mime_type.startswith(SUPPORTED_IMAGE_MIME_PREFIXES):
|
||||
return True
|
||||
|
||||
file_name = getattr(media, "file_name", None)
|
||||
if isinstance(file_name, str):
|
||||
suffix = Path(file_name).suffix.lower()
|
||||
if suffix in SUPPORTED_IMAGE_EXTENSIONS:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def extract_image_sources(message: Any) -> list[TelegramImageSource]:
|
||||
message_id = int(getattr(message, "id", 0) or 0)
|
||||
text_context = _text_context(message)
|
||||
sources: list[TelegramImageSource] = []
|
||||
|
||||
candidates: list[tuple[str, Any]] = []
|
||||
photo = getattr(message, "photo", None)
|
||||
if photo is not None:
|
||||
candidates.append(("photo", photo))
|
||||
|
||||
document = getattr(message, "document", None)
|
||||
if document is not None and _is_supported_image(document):
|
||||
candidates.append(("document", document))
|
||||
|
||||
animation = getattr(message, "animation", None)
|
||||
if animation is not None and _is_supported_image(animation):
|
||||
candidates.append(("animation", animation))
|
||||
|
||||
sticker = getattr(message, "sticker", None)
|
||||
if sticker is not None and _is_supported_image(sticker):
|
||||
candidates.append(("sticker", sticker))
|
||||
|
||||
for index, (kind, media) in enumerate(candidates):
|
||||
unique_id = (
|
||||
str(getattr(media, "file_unique_id", "") or "")
|
||||
or str(getattr(media, "file_id", "") or "")
|
||||
or f"{message_id}:{kind}:{index}"
|
||||
)
|
||||
sources.append(
|
||||
TelegramImageSource(
|
||||
message_id=message_id,
|
||||
index=index,
|
||||
kind=kind,
|
||||
media=media,
|
||||
unique_id=unique_id,
|
||||
mime_type=getattr(media, "mime_type", None),
|
||||
file_name=getattr(media, "file_name", None),
|
||||
text_context=text_context,
|
||||
)
|
||||
)
|
||||
|
||||
return sources
|
||||
|
||||
|
||||
class MCPError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class MinimaxImageUnderstandClient:
|
||||
def __init__(self, config: AIImageUnderstandConfig) -> None:
|
||||
self.config = config
|
||||
self._process: asyncio.subprocess.Process | None = None
|
||||
self._request_id = 0
|
||||
self._stdout: asyncio.StreamReader | None = None
|
||||
self._stdin: asyncio.StreamWriter | None = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def __aenter__(self) -> "MinimaxImageUnderstandClient":
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> None:
|
||||
await self.close()
|
||||
|
||||
async def start(self) -> None:
|
||||
if self._process is not None:
|
||||
return
|
||||
|
||||
env = os.environ.copy()
|
||||
env["MINIMAX_API_KEY"] = self.config.api_key
|
||||
env["MINIMAX_API_HOST"] = self.config.api_host
|
||||
|
||||
try:
|
||||
self._process = await asyncio.create_subprocess_exec(
|
||||
*self.config.command,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.DEVNULL,
|
||||
env=env,
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
raise MCPError(
|
||||
f"Unable to start MiniMax MCP command: {self.config.command[0]}"
|
||||
) from exc
|
||||
|
||||
if self._process.stdin is None or self._process.stdout is None:
|
||||
raise MCPError("MiniMax MCP process did not expose stdio pipes")
|
||||
|
||||
self._stdin = self._process.stdin
|
||||
self._stdout = self._process.stdout
|
||||
request_timeout = self.config.timeout if self.config.timeout is not None else 120.0
|
||||
|
||||
await self._request(
|
||||
"initialize",
|
||||
{
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {},
|
||||
"clientInfo": {
|
||||
"name": "multi-bot-tg",
|
||||
"version": "1.0.0",
|
||||
},
|
||||
},
|
||||
timeout=request_timeout,
|
||||
)
|
||||
await self._notify("notifications/initialized", {})
|
||||
|
||||
async def close(self) -> None:
|
||||
process = self._process
|
||||
self._process = None
|
||||
|
||||
if process is None:
|
||||
return
|
||||
|
||||
try:
|
||||
if self._stdin is not None:
|
||||
try:
|
||||
await asyncio.wait_for(self._request("shutdown", {}, timeout=None), timeout=5)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
process.terminate()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=5)
|
||||
except Exception:
|
||||
try:
|
||||
process.kill()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
|
||||
self._stdin = None
|
||||
self._stdout = None
|
||||
|
||||
async def understand_image(self, *, image_path: Path, prompt: str) -> str:
|
||||
if self._process is None:
|
||||
await self.start()
|
||||
|
||||
size = image_path.stat().st_size
|
||||
if size > MAX_IMAGE_BYTES:
|
||||
raise MCPError(
|
||||
f"Image {image_path.name} is too large for understand_image ({size} bytes)"
|
||||
)
|
||||
request_timeout = self.config.timeout if self.config.timeout is not None else 120.0
|
||||
|
||||
last_error: MCPError | None = None
|
||||
for image_field in ("image_url", "image_source"):
|
||||
try:
|
||||
result = await self._request(
|
||||
"tools/call",
|
||||
{
|
||||
"name": "understand_image",
|
||||
"arguments": {
|
||||
"prompt": prompt,
|
||||
image_field: str(image_path),
|
||||
},
|
||||
},
|
||||
timeout=request_timeout,
|
||||
)
|
||||
break
|
||||
except MCPError as exc:
|
||||
last_error = exc
|
||||
else:
|
||||
raise last_error or MCPError("MiniMax understand_image call failed")
|
||||
|
||||
text = self._extract_text(result)
|
||||
if not text:
|
||||
raise MCPError("MiniMax understand_image returned an empty response")
|
||||
return text
|
||||
|
||||
async def _notify(self, method: str, params: dict[str, Any]) -> None:
|
||||
await self._write_message({"jsonrpc": "2.0", "method": method, "params": params})
|
||||
|
||||
async def _request(self, method: str, params: dict[str, Any], *, timeout: float | None = None) -> Any:
|
||||
async with self._lock:
|
||||
self._request_id += 1
|
||||
request_id = self._request_id
|
||||
await self._write_message(
|
||||
{"jsonrpc": "2.0", "id": request_id, "method": method, "params": params}
|
||||
)
|
||||
|
||||
async def _wait_for_response() -> Any:
|
||||
while True:
|
||||
message = await self._read_message()
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
if message.get("id") != request_id:
|
||||
continue
|
||||
if "error" in message:
|
||||
error = message["error"]
|
||||
message_text = error.get("message") if isinstance(error, dict) else None
|
||||
raise MCPError(message_text or f"MCP request failed for {method}")
|
||||
return message.get("result")
|
||||
|
||||
if timeout is None:
|
||||
return await _wait_for_response()
|
||||
return await asyncio.wait_for(_wait_for_response(), timeout=timeout)
|
||||
|
||||
async def _write_message(self, payload: dict[str, Any]) -> None:
|
||||
if self._stdin is None:
|
||||
raise MCPError("MiniMax MCP stdin is not available")
|
||||
|
||||
data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||
header = f"Content-Length: {len(data)}\r\n\r\n".encode("ascii")
|
||||
self._stdin.write(header + data)
|
||||
await self._stdin.drain()
|
||||
|
||||
async def _read_message(self) -> Any:
|
||||
if self._stdout is None:
|
||||
raise MCPError("MiniMax MCP stdout is not available")
|
||||
|
||||
header_block = await self._stdout.readuntil(b"\r\n\r\n")
|
||||
headers = header_block.decode("ascii", errors="ignore").split("\r\n")
|
||||
content_length = None
|
||||
for header in headers:
|
||||
if header.lower().startswith("content-length:"):
|
||||
content_length = int(header.split(":", 1)[1].strip())
|
||||
break
|
||||
if content_length is None:
|
||||
raise MCPError("MiniMax MCP response is missing Content-Length header")
|
||||
|
||||
payload = await self._stdout.readexactly(content_length)
|
||||
return json.loads(payload.decode("utf-8"))
|
||||
|
||||
def _extract_text(self, result: Any) -> str:
|
||||
if isinstance(result, dict):
|
||||
content = result.get("content")
|
||||
if isinstance(content, list):
|
||||
parts: list[str] = []
|
||||
for item in content:
|
||||
if isinstance(item, dict):
|
||||
text = item.get("text")
|
||||
if isinstance(text, str):
|
||||
parts.append(text.strip())
|
||||
return "\n".join(part for part in parts if part).strip()
|
||||
if isinstance(content, str):
|
||||
return content.strip()
|
||||
if isinstance(result.get("text"), str):
|
||||
return str(result["text"]).strip()
|
||||
if isinstance(result, str):
|
||||
return result.strip()
|
||||
return ""
|
||||
@@ -8,6 +8,7 @@ import shlex
|
||||
class AIInvocation:
|
||||
context_messages: int
|
||||
model_flags: list[str]
|
||||
use_image_understanding: bool
|
||||
prompt: str
|
||||
|
||||
|
||||
@@ -19,6 +20,7 @@ def parse_ai_invocation(text: str, *, default_context_messages: int) -> AIInvoca
|
||||
|
||||
context_messages = default_context_messages
|
||||
model_flags: list[str] = []
|
||||
use_image_understanding = False
|
||||
|
||||
if tokens:
|
||||
first = tokens[0]
|
||||
@@ -31,14 +33,18 @@ def parse_ai_invocation(text: str, *, default_context_messages: int) -> AIInvoca
|
||||
flag = flag.split("=", 1)[0]
|
||||
flag = flag.strip().lower()
|
||||
if flag:
|
||||
model_flags.append(flag)
|
||||
if flag == "img":
|
||||
use_image_understanding = True
|
||||
else:
|
||||
model_flags.append(flag)
|
||||
|
||||
prompt = " ".join(tokens).strip()
|
||||
if not prompt:
|
||||
raise ValueError("Usage: ..ai [N] [--model] prompt")
|
||||
raise ValueError("Usage: ..ai [N] [--img] [--model] prompt")
|
||||
|
||||
return AIInvocation(
|
||||
context_messages=context_messages,
|
||||
model_flags=model_flags,
|
||||
use_image_understanding=use_image_understanding,
|
||||
prompt=prompt,
|
||||
)
|
||||
|
||||
@@ -1,11 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import inspect
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from app.ai.image_understand import (
|
||||
MCPError,
|
||||
MinimaxImageUnderstandClient,
|
||||
TelegramImageSource,
|
||||
extract_image_sources,
|
||||
)
|
||||
from app.ai.parser import AIInvocation
|
||||
from app.ai.prompts import PromptBook
|
||||
from app.config import AIConfig, AIModelConfig, AIProviderConfig
|
||||
@@ -16,6 +25,8 @@ if TYPE_CHECKING:
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
ProgressCallback = Callable[[str], Awaitable[None] | None]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AIResponse:
|
||||
@@ -31,6 +42,14 @@ class ResolvedAIModel:
|
||||
model: AIModelConfig
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class HistoryMessage:
|
||||
message: Any
|
||||
role: str
|
||||
text: str
|
||||
images: list[TelegramImageSource]
|
||||
|
||||
|
||||
class AIService:
|
||||
def __init__(self, config: AIConfig) -> None:
|
||||
self.config = config
|
||||
@@ -51,6 +70,8 @@ class AIService:
|
||||
model = self.config.models[alias]
|
||||
provider = self.config.providers[model.provider]
|
||||
lines.append(f"{alias} -> {provider.name}:{model.model}")
|
||||
if self.config.image_understand is not None:
|
||||
lines.append("image-understand=enabled")
|
||||
return "; ".join(lines)
|
||||
|
||||
def select_model_alias(self, model_flags: list[str]) -> str:
|
||||
@@ -91,47 +112,12 @@ class AIService:
|
||||
self._system_prompt = self.prompt_book.load_system_prompt()
|
||||
return self._system_prompt
|
||||
|
||||
async def build_messages(
|
||||
self,
|
||||
ctx: "CommandContext",
|
||||
*,
|
||||
prompt: str,
|
||||
context_messages: int,
|
||||
) -> list[dict[str, str]]:
|
||||
messages: list[dict[str, str]] = [{"role": "system", "content": self.load_system_prompt()}]
|
||||
history = await self.collect_history(ctx, limit=context_messages)
|
||||
messages.extend(history)
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
return messages
|
||||
|
||||
async def collect_history(self, ctx: "CommandContext", *, limit: int) -> list[dict[str, str]]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
collected: list[dict[str, str]] = []
|
||||
current_message_id = getattr(ctx.message, "id", 0)
|
||||
owner_id = ctx.owner_id
|
||||
|
||||
async for message in ctx.client.get_chat_history(
|
||||
ctx.message.chat.id,
|
||||
limit=limit + 5,
|
||||
offset_id=current_message_id,
|
||||
):
|
||||
message_id = getattr(message, "id", 0)
|
||||
if message_id >= current_message_id:
|
||||
continue
|
||||
|
||||
text = self._extract_message_text(message)
|
||||
if not text:
|
||||
continue
|
||||
|
||||
role = "assistant" if self._is_own_message(message, owner_id) else "user"
|
||||
collected.append({"role": role, "content": text})
|
||||
if len(collected) >= limit:
|
||||
break
|
||||
|
||||
collected.reverse()
|
||||
return collected
|
||||
async def _maybe_call_progress(self, progress: ProgressCallback | None, text: str) -> None:
|
||||
if progress is None:
|
||||
return
|
||||
result = progress(text)
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
|
||||
def _extract_message_text(self, message: Any) -> str:
|
||||
text = getattr(message, "text", None) or getattr(message, "caption", None)
|
||||
@@ -146,14 +132,186 @@ class AIService:
|
||||
return bool(getattr(message, "outgoing", False))
|
||||
return user_id == owner_id or bool(getattr(message, "outgoing", False))
|
||||
|
||||
async def generate(self, ctx: "CommandContext", invocation: AIInvocation) -> AIResponse:
|
||||
async def collect_history(self, ctx: "CommandContext", *, limit: int) -> list[HistoryMessage]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
collected: list[HistoryMessage] = []
|
||||
current_message_id = getattr(ctx.message, "id", 0)
|
||||
owner_id = ctx.owner_id
|
||||
|
||||
async for message in ctx.client.get_chat_history(
|
||||
ctx.message.chat.id,
|
||||
limit=limit,
|
||||
offset_id=current_message_id,
|
||||
):
|
||||
message_id = getattr(message, "id", 0)
|
||||
if message_id >= current_message_id:
|
||||
continue
|
||||
|
||||
text = self._extract_message_text(message)
|
||||
role = "assistant" if self._is_own_message(message, owner_id) else "user"
|
||||
images = extract_image_sources(message)
|
||||
collected.append(HistoryMessage(message=message, role=role, text=text, images=images))
|
||||
if len(collected) >= limit:
|
||||
break
|
||||
|
||||
collected.reverse()
|
||||
return collected
|
||||
|
||||
async def _download_image_source(
|
||||
self,
|
||||
ctx: "CommandContext",
|
||||
source: TelegramImageSource,
|
||||
directory: Path,
|
||||
) -> Path:
|
||||
path = await ctx.client.download_media(source.media, file_name=str(directory))
|
||||
if not isinstance(path, str) or not path:
|
||||
raise RuntimeError(f"Failed to download {source.label}")
|
||||
downloaded = Path(path)
|
||||
if not downloaded.exists():
|
||||
raise RuntimeError(f"Downloaded file for {source.label} does not exist")
|
||||
return downloaded
|
||||
|
||||
def _build_image_prompt(self, invocation: AIInvocation, source: TelegramImageSource) -> str:
|
||||
context_text = source.text_context or "none"
|
||||
return "\n".join(
|
||||
[
|
||||
"Analyze the attached Telegram image for a user-bot.",
|
||||
f"User request: {invocation.prompt}",
|
||||
f"Source: {source.label}",
|
||||
f"Message text or caption: {context_text}",
|
||||
"Return a concise structured description of visible objects, readable text, notable actions, and anything relevant to the user request.",
|
||||
]
|
||||
)
|
||||
|
||||
async def _analyze_images(
|
||||
self,
|
||||
ctx: "CommandContext",
|
||||
invocation: AIInvocation,
|
||||
history: list[HistoryMessage],
|
||||
*,
|
||||
progress: ProgressCallback | None,
|
||||
) -> list[str]:
|
||||
if self.config.image_understand is None:
|
||||
raise ValueError("Image understanding is not configured")
|
||||
|
||||
sources: list[TelegramImageSource] = []
|
||||
seen: set[str] = set()
|
||||
for item in history:
|
||||
for source in item.images:
|
||||
if source.unique_id in seen:
|
||||
continue
|
||||
seen.add(source.unique_id)
|
||||
sources.append(source)
|
||||
|
||||
if not sources:
|
||||
await self._maybe_call_progress(
|
||||
progress,
|
||||
"🛈 No images found in the selected context; continuing with text only",
|
||||
)
|
||||
return []
|
||||
|
||||
analyses: list[str] = []
|
||||
await self._maybe_call_progress(
|
||||
progress,
|
||||
f"🛠 Found {len(sources)} image(s) in context, running understand_image",
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="multi-bot-tg-ai-") as tmpdir:
|
||||
temp_root = Path(tmpdir)
|
||||
try:
|
||||
async with MinimaxImageUnderstandClient(self.config.image_understand) as image_client:
|
||||
for index, source in enumerate(sources, start=1):
|
||||
await self._maybe_call_progress(
|
||||
progress,
|
||||
f"🛠 Understanding image {index}/{len(sources)} from message {source.message_id}",
|
||||
)
|
||||
downloaded = await self._download_image_source(ctx, source, temp_root)
|
||||
if downloaded.stat().st_size > 20 * 1024 * 1024:
|
||||
raise RuntimeError(
|
||||
f"Image {source.label} exceeds the 20MB understand_image limit"
|
||||
)
|
||||
|
||||
try:
|
||||
text = await image_client.understand_image(
|
||||
image_path=downloaded,
|
||||
prompt=self._build_image_prompt(invocation, source),
|
||||
)
|
||||
except MCPError as exc:
|
||||
raise RuntimeError(
|
||||
f"Image analysis failed for {source.label}: {exc}"
|
||||
) from exc
|
||||
|
||||
analyses.append(
|
||||
"\n".join(
|
||||
[
|
||||
f"[Image {index}/{len(sources)}]",
|
||||
f"Source: {source.label}",
|
||||
f"File: {downloaded.name}",
|
||||
f"Analysis:\n{text}",
|
||||
]
|
||||
)
|
||||
)
|
||||
except MCPError as exc:
|
||||
raise RuntimeError(f"Image understanding tool failed: {exc}") from exc
|
||||
|
||||
await self._maybe_call_progress(progress, "🛠 Image analysis complete")
|
||||
return analyses
|
||||
|
||||
def _build_image_summary_message(self, image_notes: list[str]) -> str:
|
||||
return "\n\n".join(
|
||||
[
|
||||
"Visual context extracted from images:",
|
||||
*image_notes,
|
||||
]
|
||||
).strip()
|
||||
|
||||
async def build_messages(
|
||||
self,
|
||||
ctx: "CommandContext",
|
||||
*,
|
||||
prompt: str,
|
||||
history: list[HistoryMessage],
|
||||
image_notes: list[str],
|
||||
) -> list[dict[str, str]]:
|
||||
messages: list[dict[str, str]] = [{"role": "system", "content": self.load_system_prompt()}]
|
||||
if image_notes:
|
||||
messages.append({"role": "system", "content": self._build_image_summary_message(image_notes)})
|
||||
for item in history:
|
||||
if not item.text:
|
||||
continue
|
||||
messages.append({"role": item.role, "content": item.text})
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
return messages
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
ctx: "CommandContext",
|
||||
invocation: AIInvocation,
|
||||
*,
|
||||
progress: ProgressCallback | None = None,
|
||||
) -> AIResponse:
|
||||
model = self.resolve_model(invocation.model_flags[0] if invocation.model_flags else None)
|
||||
client = self._get_client(model.provider)
|
||||
|
||||
history = await self.collect_history(ctx, limit=invocation.context_messages)
|
||||
image_notes: list[str] = []
|
||||
|
||||
if invocation.use_image_understanding:
|
||||
image_notes = await self._analyze_images(
|
||||
ctx,
|
||||
invocation,
|
||||
history,
|
||||
progress=progress,
|
||||
)
|
||||
|
||||
await self._maybe_call_progress(progress, "🛠 Generating final answer")
|
||||
messages = await self.build_messages(
|
||||
ctx,
|
||||
prompt=invocation.prompt,
|
||||
context_messages=invocation.context_messages,
|
||||
history=history,
|
||||
image_notes=image_notes,
|
||||
)
|
||||
|
||||
try:
|
||||
|
||||
Reference in New Issue
Block a user