mirror of
https://github.com/DrizzleTime/Foxel.git
synced 2026-09-05 23:57:19 +08:00
Initial commit
This commit is contained in:
shiyu
@@ -0,0 +1,251 @@
|
||||
import base64
|
||||
import mimetypes
|
||||
import os
|
||||
from io import BytesIO
|
||||
from typing import Dict, Any, List, Tuple
|
||||
|
||||
from fastapi.responses import Response
|
||||
from PIL import Image
|
||||
|
||||
from ..base import BaseProcessor
|
||||
from domain.ai import (
|
||||
DEFAULT_VECTOR_DIMENSION,
|
||||
FILE_COLLECTION_NAME,
|
||||
VECTOR_COLLECTION_NAME,
|
||||
VectorDBService,
|
||||
describe_image_base64,
|
||||
get_text_embedding,
|
||||
provider_service,
|
||||
)
|
||||
|
||||
|
||||
CHUNK_SIZE = 800
|
||||
CHUNK_OVERLAP = 200
|
||||
MAX_IMAGE_EDGE = 1600
|
||||
JPEG_QUALITY = 85
|
||||
|
||||
|
||||
def _chunk_text(content: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> List[Tuple[int, str, int, int]]:
|
||||
"""按固定窗口拆分文本，返回(chunk_id, chunk_text, start, end)。"""
|
||||
if chunk_size <= 0:
|
||||
chunk_size = CHUNK_SIZE
|
||||
if overlap >= chunk_size:
|
||||
overlap = max(chunk_size // 4, 1)
|
||||
|
||||
chunks: List[Tuple[int, str, int, int]] = []
|
||||
step = chunk_size - overlap
|
||||
idx = 0
|
||||
start = 0
|
||||
length = len(content)
|
||||
|
||||
while start < length:
|
||||
end = min(length, start + chunk_size)
|
||||
chunk = content[start:end].strip()
|
||||
if chunk:
|
||||
chunks.append((idx, chunk, start, end))
|
||||
idx += 1
|
||||
if end >= length:
|
||||
break
|
||||
start += step
|
||||
return chunks
|
||||
|
||||
|
||||
def _guess_mime(path: str) -> str:
|
||||
mime, _ = mimetypes.guess_type(path)
|
||||
return mime or "application/octet-stream"
|
||||
|
||||
|
||||
def _chunk_key(path: str, chunk_id: str) -> str:
|
||||
return f"{path}#chunk={chunk_id}"
|
||||
|
||||
|
||||
def _compress_image_for_embedding(input_bytes: bytes) -> Tuple[bytes, Dict[str, Any] | None]:
|
||||
"""压缩图片，降低发送到视觉模型的体积。"""
|
||||
if Image is None:
|
||||
return input_bytes, None
|
||||
|
||||
try:
|
||||
with Image.open(BytesIO(input_bytes)) as img:
|
||||
img = img.convert("RGB")
|
||||
width, height = img.size
|
||||
longest_edge = max(width, height)
|
||||
scale = 1.0
|
||||
if longest_edge > MAX_IMAGE_EDGE:
|
||||
scale = MAX_IMAGE_EDGE / float(longest_edge)
|
||||
new_size = (max(int(width * scale), 1), max(int(height * scale), 1))
|
||||
resample_mode = getattr(getattr(Image, "Resampling", Image), "LANCZOS")
|
||||
img = img.resize(new_size, resample=resample_mode)
|
||||
|
||||
buffer = BytesIO()
|
||||
img.save(buffer, format="JPEG", quality=JPEG_QUALITY, optimize=True)
|
||||
compressed = buffer.getvalue()
|
||||
|
||||
if len(compressed) < len(input_bytes):
|
||||
return compressed, {
|
||||
"original_bytes": len(input_bytes),
|
||||
"compressed_bytes": len(compressed),
|
||||
"scaled": scale < 1.0,
|
||||
"width": img.width,
|
||||
"height": img.height,
|
||||
}
|
||||
except Exception: # pragma: no cover - 任意图像处理异常时回退
|
||||
return input_bytes, None
|
||||
|
||||
return input_bytes, None
|
||||
|
||||
|
||||
class VectorIndexProcessor:
|
||||
name = "向量索引"
|
||||
supported_exts: List[str] = [] # 留空表示不限扩展名
|
||||
config_schema = [
|
||||
{
|
||||
"key": "action", "label": "操作", "type": "select", "required": True, "default": "create",
|
||||
"options": [
|
||||
{"value": "create", "label": "创建索引"},
|
||||
{"value": "destroy", "label": "销毁索引"},
|
||||
]
|
||||
},
|
||||
{
|
||||
"key": "index_type", "label": "索引类型", "type": "select", "required": True, "default": "vector",
|
||||
"options": [
|
||||
{"value": "vector", "label": "向量索引"},
|
||||
{"value": "simple", "label": "普通索引"},
|
||||
]
|
||||
}
|
||||
]
|
||||
produces_file = False
|
||||
requires_input_bytes = False
|
||||
|
||||
async def process(self, input_bytes: bytes, path: str, config: Dict[str, Any]) -> Response:
|
||||
async def ensure_input_bytes() -> bytes:
|
||||
if input_bytes:
|
||||
return input_bytes
|
||||
from domain.virtual_fs import VirtualFSService
|
||||
return await VirtualFSService.read_file(path)
|
||||
|
||||
action = config.get("action", "create")
|
||||
index_type = config.get("index_type", "vector")
|
||||
vector_db = VectorDBService()
|
||||
vector_collection = VECTOR_COLLECTION_NAME
|
||||
file_collection = FILE_COLLECTION_NAME
|
||||
|
||||
if action == "destroy":
|
||||
target_collection = file_collection if index_type == "simple" else vector_collection
|
||||
await vector_db.delete_vector(target_collection, path)
|
||||
return Response(content=f"文件 {path} 的 {index_type} 索引已销毁", media_type="text/plain")
|
||||
|
||||
mime_type = _guess_mime(path)
|
||||
|
||||
if index_type == "simple":
|
||||
await vector_db.ensure_collection(file_collection, vector=False)
|
||||
await vector_db.delete_vector(file_collection, path)
|
||||
await vector_db.upsert_vector(file_collection, {
|
||||
"path": path,
|
||||
"source_path": path,
|
||||
"chunk_id": "filename",
|
||||
"mime": mime_type,
|
||||
"type": "filename",
|
||||
"name": os.path.basename(path),
|
||||
})
|
||||
return Response(content=f"文件 {path} 的普通索引已创建", media_type="text/plain")
|
||||
|
||||
file_ext = path.split('.')[-1].lower()
|
||||
details: Dict[str, Any] = {"path": path, "action": "create", "index_type": "vector"}
|
||||
|
||||
embedding_model = await provider_service.get_default_model("embedding")
|
||||
vector_dim = DEFAULT_VECTOR_DIMENSION
|
||||
if embedding_model and getattr(embedding_model, "embedding_dimensions", None):
|
||||
try:
|
||||
vector_dim = int(embedding_model.embedding_dimensions)
|
||||
except (TypeError, ValueError):
|
||||
vector_dim = DEFAULT_VECTOR_DIMENSION
|
||||
if vector_dim <= 0:
|
||||
vector_dim = DEFAULT_VECTOR_DIMENSION
|
||||
|
||||
await vector_db.ensure_collection(vector_collection, vector=True, dim=vector_dim)
|
||||
await vector_db.delete_vector(vector_collection, path)
|
||||
|
||||
if file_ext in ["jpg", "jpeg", "png", "bmp"]:
|
||||
file_bytes = await ensure_input_bytes()
|
||||
processed_bytes, compression = _compress_image_for_embedding(file_bytes)
|
||||
base64_image = base64.b64encode(processed_bytes).decode("utf-8")
|
||||
description = await describe_image_base64(base64_image)
|
||||
embedding = await get_text_embedding(description)
|
||||
image_mime = "image/jpeg" if compression else mime_type
|
||||
await vector_db.upsert_vector(vector_collection, {
|
||||
"path": _chunk_key(path, "image"),
|
||||
"source_path": path,
|
||||
"chunk_id": "image",
|
||||
"embedding": embedding,
|
||||
"text": description,
|
||||
"mime": image_mime,
|
||||
"type": "image",
|
||||
})
|
||||
details["description"] = description
|
||||
if compression:
|
||||
details["image_compression"] = compression
|
||||
return Response(content=f"图片已索引，描述：{description}", media_type="text/plain")
|
||||
|
||||
if file_ext in ["txt", "md"]:
|
||||
try:
|
||||
file_bytes = await ensure_input_bytes()
|
||||
text = file_bytes.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return Response(content="文本文件解码失败", status_code=400)
|
||||
|
||||
chunks = _chunk_text(text)
|
||||
if not chunks:
|
||||
await vector_db.upsert_vector(vector_collection, {
|
||||
"path": _chunk_key(path, "0"),
|
||||
"source_path": path,
|
||||
"chunk_id": "0",
|
||||
"embedding": await get_text_embedding(text or path),
|
||||
"text": text,
|
||||
"mime": mime_type,
|
||||
"type": "text",
|
||||
"start_offset": 0,
|
||||
"end_offset": len(text),
|
||||
})
|
||||
details["chunks"] = 1
|
||||
return Response(content="文本文件已索引", media_type="text/plain")
|
||||
|
||||
chunk_count = 0
|
||||
for chunk_id, chunk_text, start, end in chunks:
|
||||
embedding = await get_text_embedding(chunk_text)
|
||||
await vector_db.upsert_vector(vector_collection, {
|
||||
"path": _chunk_key(path, str(chunk_id)),
|
||||
"source_path": path,
|
||||
"chunk_id": str(chunk_id),
|
||||
"embedding": embedding,
|
||||
"text": chunk_text,
|
||||
"mime": mime_type,
|
||||
"type": "text",
|
||||
"start_offset": start,
|
||||
"end_offset": end,
|
||||
})
|
||||
chunk_count += 1
|
||||
|
||||
details["chunks"] = chunk_count
|
||||
sample = chunks[0][1]
|
||||
details["sample"] = sample[:120]
|
||||
return Response(content="文本文件已索引", media_type="text/plain")
|
||||
|
||||
# 其他类型暂未支持向量索引，回退为文件名索引
|
||||
await vector_db.ensure_collection(file_collection, vector=False)
|
||||
await vector_db.delete_vector(file_collection, path)
|
||||
await vector_db.upsert_vector(file_collection, {
|
||||
"path": path,
|
||||
"source_path": path,
|
||||
"chunk_id": "filename",
|
||||
"mime": mime_type,
|
||||
"type": "filename",
|
||||
"name": os.path.basename(path),
|
||||
})
|
||||
return Response(content="暂不支持该类型的向量索引，已创建文件名索引", media_type="text/plain")
|
||||
|
||||
|
||||
PROCESSOR_TYPE = "vector_index"
|
||||
PROCESSOR_NAME = VectorIndexProcessor.name
|
||||
SUPPORTED_EXTS = VectorIndexProcessor.supported_exts
|
||||
CONFIG_SCHEMA = VectorIndexProcessor.config_schema
|
||||
def PROCESSOR_FACTORY(): return VectorIndexProcessor()
|
||||
Reference in New Issue
Block a user