491 lines
16 KiB
Python
491 lines
16 KiB
Python
|
"""
|
|||
|
|
Image parser using multimodal LLM.
|
|||
|
|
Supports parsing images into structured text content for knowledge base indexing.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import base64
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
from dataclasses import dataclass, field
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
from app.services.document.base import (
|
|||
|
|
DocumentParseException,
|
|||
|
|
DocumentParser,
|
|||
|
|
PageText,
|
|||
|
|
ParseResult,
|
|||
|
|
)
|
|||
|
|
from app.services.llm.factory import LLMUsageType, get_llm_config_manager
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
IMAGE_SYSTEM_PROMPT = """你是一个专业的图像内容分析助手。你的任务是分析图片内容，并将其智能拆分为适合知识库检索的独立数据块。
|
|||
|
|
|
|||
|
|
## 分析要求
|
|||
|
|
1. 仔细分析图片内容，识别其中的文字、图表、数据等信息
|
|||
|
|
2. 根据内容的逻辑结构，智能判断如何拆分为独立的知识条目
|
|||
|
|
3. 每个条目应该是独立、完整、可检索的知识单元
|
|||
|
|
|
|||
|
|
## 输出格式
|
|||
|
|
请严格按照以下 JSON 格式输出，不要添加任何其他内容：
|
|||
|
|
|
|||
|
|
```json
|
|||
|
|
{
|
|||
|
|
"image_summary": "图片整体概述（一句话描述图片主题）",
|
|||
|
|
"total_chunks": <分块总数>,
|
|||
|
|
"chunks": [
|
|||
|
|
{
|
|||
|
|
"chunk_index": 0,
|
|||
|
|
"content": "该分块的完整内容文字",
|
|||
|
|
"chunk_type": "text|table|list|diagram|chart|mixed",
|
|||
|
|
"keywords": ["关键词1", "关键词2"]
|
|||
|
|
}
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
## 分块策略
|
|||
|
|
- **单一内容**: 如果图片只有一段完整的文字/信息，可以只输出1个分块
|
|||
|
|
- **多段落内容**: 按段落或逻辑单元拆分，每个段落作为独立分块
|
|||
|
|
- **表格数据**: 将表格内容转换为结构化文字，作为一个分块
|
|||
|
|
- **图表数据**: 描述图表内容和数据，作为一个分块
|
|||
|
|
- **列表内容**: 每个列表项可作为独立分块，或合并为相关的一组
|
|||
|
|
- **混合内容**: 根据内容类型分别处理，确保每个分块主题明确
|
|||
|
|
|
|||
|
|
## 注意事项
|
|||
|
|
1. 每个分块的 content 必须是完整、可独立理解的文字
|
|||
|
|
2. chunk_type 用于标识内容类型，便于后续处理
|
|||
|
|
3. keywords 提取该分块的核心关键词，便于检索
|
|||
|
|
4. 确保输出的 JSON 格式正确，可以被解析"""
|
|||
|
|
|
|||
|
|
IMAGE_USER_PROMPT = "请分析这张图片，按照要求的 JSON 格式输出分块结果。"
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class ImageChunk:
|
|||
|
|
"""智能分块结果"""
|
|||
|
|
chunk_index: int
|
|||
|
|
content: str
|
|||
|
|
chunk_type: str = "text"
|
|||
|
|
keywords: list[str] = field(default_factory=list)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class ImageParseResult:
|
|||
|
|
"""图片解析结果（包含智能分块）"""
|
|||
|
|
image_summary: str
|
|||
|
|
chunks: list[ImageChunk]
|
|||
|
|
raw_text: str
|
|||
|
|
source_path: str
|
|||
|
|
file_size: int
|
|||
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ImageParser(DocumentParser):
|
|||
|
|
"""
|
|||
|
|
Image parser using multimodal LLM.
|
|||
|
|
|
|||
|
|
Supports common image formats and extracts text content using
|
|||
|
|
vision-capable LLM models (GPT-4V, GPT-4o, etc.).
|
|||
|
|
|
|||
|
|
Features:
|
|||
|
|
- Intelligent chunking based on content structure
|
|||
|
|
- Structured output with keywords and chunk types
|
|||
|
|
- Support for various content types (text, table, chart, etc.)
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
SUPPORTED_EXTENSIONS = [
|
|||
|
|
".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff", ".tif"
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
model: str | None = None,
|
|||
|
|
max_tokens: int = 4096,
|
|||
|
|
timeout_seconds: int = 120,
|
|||
|
|
):
|
|||
|
|
self._model = model
|
|||
|
|
self._max_tokens = max_tokens
|
|||
|
|
self._timeout_seconds = timeout_seconds
|
|||
|
|
|
|||
|
|
def parse(self, file_path: str | Path) -> ParseResult:
|
|||
|
|
"""
|
|||
|
|
Parse an image file and extract text content using multimodal LLM.
|
|||
|
|
|
|||
|
|
Note: This method is synchronous but internally uses async operations.
|
|||
|
|
For async contexts, use parse_async() instead.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
file_path: Path to the image file.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
ParseResult with extracted text content.
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
DocumentParseException: If parsing fails.
|
|||
|
|
"""
|
|||
|
|
path = Path(file_path)
|
|||
|
|
if not path.exists():
|
|||
|
|
raise DocumentParseException(
|
|||
|
|
f"Image file not found: {file_path}",
|
|||
|
|
file_path=str(path),
|
|||
|
|
parser="image",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
file_size = path.stat().st_size
|
|||
|
|
extension = path.suffix.lower()
|
|||
|
|
|
|||
|
|
if extension not in self.SUPPORTED_EXTENSIONS:
|
|||
|
|
raise DocumentParseException(
|
|||
|
|
f"Unsupported image format: {extension}",
|
|||
|
|
file_path=str(path),
|
|||
|
|
parser="image",
|
|||
|
|
details={"supported_formats": self.SUPPORTED_EXTENSIONS},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
with open(path, "rb") as f:
|
|||
|
|
image_data = f.read()
|
|||
|
|
|
|||
|
|
image_base64 = base64.b64encode(image_data).decode("utf-8")
|
|||
|
|
mime_type = self._get_mime_type(extension)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
loop = asyncio.get_running_loop()
|
|||
|
|
import concurrent.futures
|
|||
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|||
|
|
future = executor.submit(
|
|||
|
|
asyncio.run,
|
|||
|
|
self._analyze_image_async(image_base64, mime_type)
|
|||
|
|
)
|
|||
|
|
result = future.result()
|
|||
|
|
except RuntimeError:
|
|||
|
|
result = asyncio.run(self._analyze_image_async(image_base64, mime_type))
|
|||
|
|
|
|||
|
|
logger.info(
|
|||
|
|
f"[IMAGE-PARSER] Successfully parsed image: {path.name}, "
|
|||
|
|
f"size={file_size}, chunks={len(result.chunks)}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return ParseResult(
|
|||
|
|
text=result.raw_text,
|
|||
|
|
source_path=str(path),
|
|||
|
|
file_size=file_size,
|
|||
|
|
page_count=1,
|
|||
|
|
metadata={
|
|||
|
|
"format": extension,
|
|||
|
|
"parser": "image",
|
|||
|
|
"mime_type": mime_type,
|
|||
|
|
"image_summary": result.image_summary,
|
|||
|
|
"chunk_count": len(result.chunks),
|
|||
|
|
"chunks": [
|
|||
|
|
{
|
|||
|
|
"chunk_index": c.chunk_index,
|
|||
|
|
"content": c.content,
|
|||
|
|
"chunk_type": c.chunk_type,
|
|||
|
|
"keywords": c.keywords,
|
|||
|
|
}
|
|||
|
|
for c in result.chunks
|
|||
|
|
],
|
|||
|
|
},
|
|||
|
|
pages=[PageText(page=1, text=result.raw_text)],
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"[IMAGE-PARSER] Failed to parse image {path}: {e}")
|
|||
|
|
raise DocumentParseException(
|
|||
|
|
f"Failed to parse image: {str(e)}",
|
|||
|
|
file_path=str(path),
|
|||
|
|
parser="image",
|
|||
|
|
details={"error": str(e)},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
async def parse_async(self, file_path: str | Path) -> ParseResult:
|
|||
|
|
"""
|
|||
|
|
Async version of parse method for use in async contexts.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
file_path: Path to the image file.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
ParseResult with extracted text content.
|
|||
|
|
"""
|
|||
|
|
path = Path(file_path)
|
|||
|
|
if not path.exists():
|
|||
|
|
raise DocumentParseException(
|
|||
|
|
f"Image file not found: {file_path}",
|
|||
|
|
file_path=str(path),
|
|||
|
|
parser="image",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
file_size = path.stat().st_size
|
|||
|
|
extension = path.suffix.lower()
|
|||
|
|
|
|||
|
|
if extension not in self.SUPPORTED_EXTENSIONS:
|
|||
|
|
raise DocumentParseException(
|
|||
|
|
f"Unsupported image format: {extension}",
|
|||
|
|
file_path=str(path),
|
|||
|
|
parser="image",
|
|||
|
|
details={"supported_formats": self.SUPPORTED_EXTENSIONS},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
with open(path, "rb") as f:
|
|||
|
|
image_data = f.read()
|
|||
|
|
|
|||
|
|
image_base64 = base64.b64encode(image_data).decode("utf-8")
|
|||
|
|
mime_type = self._get_mime_type(extension)
|
|||
|
|
|
|||
|
|
result = await self._analyze_image_async(image_base64, mime_type)
|
|||
|
|
|
|||
|
|
logger.info(
|
|||
|
|
f"[IMAGE-PARSER] Successfully parsed image (async): {path.name}, "
|
|||
|
|
f"size={file_size}, chunks={len(result.chunks)}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return ParseResult(
|
|||
|
|
text=result.raw_text,
|
|||
|
|
source_path=str(path),
|
|||
|
|
file_size=file_size,
|
|||
|
|
page_count=1,
|
|||
|
|
metadata={
|
|||
|
|
"format": extension,
|
|||
|
|
"parser": "image",
|
|||
|
|
"mime_type": mime_type,
|
|||
|
|
"image_summary": result.image_summary,
|
|||
|
|
"chunk_count": len(result.chunks),
|
|||
|
|
"chunks": [
|
|||
|
|
{
|
|||
|
|
"chunk_index": c.chunk_index,
|
|||
|
|
"content": c.content,
|
|||
|
|
"chunk_type": c.chunk_type,
|
|||
|
|
"keywords": c.keywords,
|
|||
|
|
}
|
|||
|
|
for c in result.chunks
|
|||
|
|
],
|
|||
|
|
},
|
|||
|
|
pages=[PageText(page=1, text=result.raw_text)],
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"[IMAGE-PARSER] Failed to parse image {path}: {e}")
|
|||
|
|
raise DocumentParseException(
|
|||
|
|
f"Failed to parse image: {str(e)}",
|
|||
|
|
file_path=str(path),
|
|||
|
|
parser="image",
|
|||
|
|
details={"error": str(e)},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
async def parse_with_chunks(self, file_path: str | Path) -> ImageParseResult:
|
|||
|
|
"""
|
|||
|
|
Parse image and return structured result with intelligent chunks.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
file_path: Path to the image file.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
ImageParseResult with intelligent chunks.
|
|||
|
|
"""
|
|||
|
|
path = Path(file_path)
|
|||
|
|
if not path.exists():
|
|||
|
|
raise DocumentParseException(
|
|||
|
|
f"Image file not found: {file_path}",
|
|||
|
|
file_path=str(path),
|
|||
|
|
parser="image",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
file_size = path.stat().st_size
|
|||
|
|
extension = path.suffix.lower()
|
|||
|
|
|
|||
|
|
if extension not in self.SUPPORTED_EXTENSIONS:
|
|||
|
|
raise DocumentParseException(
|
|||
|
|
f"Unsupported image format: {extension}",
|
|||
|
|
file_path=str(path),
|
|||
|
|
parser="image",
|
|||
|
|
details={"supported_formats": self.SUPPORTED_EXTENSIONS},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
with open(path, "rb") as f:
|
|||
|
|
image_data = f.read()
|
|||
|
|
|
|||
|
|
image_base64 = base64.b64encode(image_data).decode("utf-8")
|
|||
|
|
mime_type = self._get_mime_type(extension)
|
|||
|
|
|
|||
|
|
result = await self._analyze_image_async(image_base64, mime_type)
|
|||
|
|
result.source_path = str(path)
|
|||
|
|
result.file_size = file_size
|
|||
|
|
result.metadata = {
|
|||
|
|
"format": extension,
|
|||
|
|
"parser": "image",
|
|||
|
|
"mime_type": mime_type,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
async def _analyze_image_async(self, image_base64: str, mime_type: str) -> ImageParseResult:
|
|||
|
|
"""
|
|||
|
|
Analyze image using multimodal LLM and return structured chunks.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
image_base64: Base64 encoded image data.
|
|||
|
|
mime_type: MIME type of the image.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
ImageParseResult with intelligent chunks.
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
manager = get_llm_config_manager()
|
|||
|
|
client = manager.get_kb_processing_client()
|
|||
|
|
|
|||
|
|
config = manager.kb_processing_config
|
|||
|
|
model = self._model or config.get("model", "gpt-4o-mini")
|
|||
|
|
|
|||
|
|
messages = [
|
|||
|
|
{
|
|||
|
|
"role": "system",
|
|||
|
|
"content": IMAGE_SYSTEM_PROMPT,
|
|||
|
|
},
|
|||
|
|
{
|
|||
|
|
"role": "user",
|
|||
|
|
"content": [
|
|||
|
|
{
|
|||
|
|
"type": "text",
|
|||
|
|
"text": IMAGE_USER_PROMPT,
|
|||
|
|
},
|
|||
|
|
{
|
|||
|
|
"type": "image_url",
|
|||
|
|
"image_url": {
|
|||
|
|
"url": f"data:{mime_type};base64,{image_base64}",
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
],
|
|||
|
|
},
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
from app.services.llm.base import LLMConfig
|
|||
|
|
|
|||
|
|
llm_config = LLMConfig(
|
|||
|
|
model=model,
|
|||
|
|
max_tokens=self._max_tokens,
|
|||
|
|
temperature=0.3,
|
|||
|
|
timeout_seconds=self._timeout_seconds,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
response = await client.generate(messages=messages, config=llm_config)
|
|||
|
|
|
|||
|
|
if not response.content:
|
|||
|
|
raise DocumentParseException(
|
|||
|
|
"LLM returned empty response for image analysis",
|
|||
|
|
parser="image",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return self._parse_llm_response(response.content)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"[IMAGE-PARSER] LLM analysis failed: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
def _parse_llm_response(self, response_content: str) -> ImageParseResult:
|
|||
|
|
"""
|
|||
|
|
Parse LLM response into structured ImageParseResult.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
response_content: Raw LLM response content.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
ImageParseResult with parsed chunks.
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
json_str = self._extract_json(response_content)
|
|||
|
|
data = json.loads(json_str)
|
|||
|
|
|
|||
|
|
image_summary = data.get("image_summary", "")
|
|||
|
|
chunks_data = data.get("chunks", [])
|
|||
|
|
|
|||
|
|
chunks = []
|
|||
|
|
for chunk_data in chunks_data:
|
|||
|
|
chunk = ImageChunk(
|
|||
|
|
chunk_index=chunk_data.get("chunk_index", len(chunks)),
|
|||
|
|
content=chunk_data.get("content", ""),
|
|||
|
|
chunk_type=chunk_data.get("chunk_type", "text"),
|
|||
|
|
keywords=chunk_data.get("keywords", []),
|
|||
|
|
)
|
|||
|
|
if chunk.content.strip():
|
|||
|
|
chunks.append(chunk)
|
|||
|
|
|
|||
|
|
if not chunks:
|
|||
|
|
chunks.append(ImageChunk(
|
|||
|
|
chunk_index=0,
|
|||
|
|
content=response_content,
|
|||
|
|
chunk_type="text",
|
|||
|
|
keywords=[],
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
raw_text = "\n\n".join([c.content for c in chunks])
|
|||
|
|
|
|||
|
|
return ImageParseResult(
|
|||
|
|
image_summary=image_summary,
|
|||
|
|
chunks=chunks,
|
|||
|
|
raw_text=raw_text,
|
|||
|
|
source_path="",
|
|||
|
|
file_size=0,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
except json.JSONDecodeError as e:
|
|||
|
|
logger.warning(f"[IMAGE-PARSER] Failed to parse JSON response: {e}, using fallback")
|
|||
|
|
return ImageParseResult(
|
|||
|
|
image_summary="图片内容",
|
|||
|
|
chunks=[ImageChunk(
|
|||
|
|
chunk_index=0,
|
|||
|
|
content=response_content,
|
|||
|
|
chunk_type="text",
|
|||
|
|
keywords=[],
|
|||
|
|
)],
|
|||
|
|
raw_text=response_content,
|
|||
|
|
source_path="",
|
|||
|
|
file_size=0,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def _extract_json(self, content: str) -> str:
|
|||
|
|
"""
|
|||
|
|
Extract JSON from LLM response content.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
content: Raw response content that may contain JSON.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Extracted JSON string.
|
|||
|
|
"""
|
|||
|
|
content = content.strip()
|
|||
|
|
|
|||
|
|
if content.startswith("{") and content.endswith("}"):
|
|||
|
|
return content
|
|||
|
|
|
|||
|
|
json_start = content.find("{")
|
|||
|
|
json_end = content.rfind("}")
|
|||
|
|
|
|||
|
|
if json_start != -1 and json_end != -1 and json_end > json_start:
|
|||
|
|
return content[json_start:json_end + 1]
|
|||
|
|
|
|||
|
|
return content
|
|||
|
|
|
|||
|
|
def _get_mime_type(self, extension: str) -> str:
|
|||
|
|
"""Get MIME type for image extension."""
|
|||
|
|
mime_types = {
|
|||
|
|
".jpg": "image/jpeg",
|
|||
|
|
".jpeg": "image/jpeg",
|
|||
|
|
".png": "image/png",
|
|||
|
|
".gif": "image/gif",
|
|||
|
|
".webp": "image/webp",
|
|||
|
|
".bmp": "image/bmp",
|
|||
|
|
".tiff": "image/tiff",
|
|||
|
|
".tif": "image/tiff",
|
|||
|
|
}
|
|||
|
|
return mime_types.get(extension.lower(), "image/jpeg")
|
|||
|
|
|
|||
|
|
def get_supported_extensions(self) -> list[str]:
|
|||
|
|
"""Get list of supported image extensions."""
|
|||
|
|
return ImageParser.SUPPORTED_EXTENSIONS
|