Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1ab56ad0b1 | ||
|
|
8585d101d5 | ||
|
|
62eb9578fd | ||
|
|
848657219e |
@@ -5,6 +5,7 @@ pydantic
|
|||||||
python-multipart
|
python-multipart
|
||||||
httpx
|
httpx
|
||||||
pypdf
|
pypdf
|
||||||
|
PyMuPDF>=1.24,<2
|
||||||
python-docx
|
python-docx
|
||||||
openpyxl
|
openpyxl
|
||||||
apscheduler>=3.10
|
apscheduler>=3.10
|
||||||
|
|||||||
@@ -8,7 +8,8 @@ import threading
|
|||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
from database import SessionLocal
|
from database import SessionLocal
|
||||||
from models import KnowledgeChunk, KnowledgeDoc
|
from models import Avatar, KnowledgeChunk, KnowledgeDoc
|
||||||
|
from services.pdf_ocr_service import extract_scanned_pdf_text
|
||||||
import embeddings
|
import embeddings
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -76,7 +77,23 @@ class KnowledgeVectorizer:
|
|||||||
|
|
||||||
self._set_progress(db, doc, "extracting", 8)
|
self._set_progress(db, doc, "extracting", 8)
|
||||||
text = embeddings.extract_text(path, f".{doc.file_type}")
|
text = embeddings.extract_text(path, f".{doc.file_type}")
|
||||||
self._set_progress(db, doc, "chunking", 22)
|
if doc.file_type == "pdf" and not text.strip():
|
||||||
|
avatar = db.get(Avatar, doc.avatar_id)
|
||||||
|
if not avatar:
|
||||||
|
raise ValueError("文档所属分身不存在")
|
||||||
|
|
||||||
|
def ocr_progress(done: int, total: int):
|
||||||
|
percent = 8 + int((done / max(1, total)) * 20)
|
||||||
|
self._set_progress(db, doc, "ocr", min(percent, 28))
|
||||||
|
|
||||||
|
self._set_progress(db, doc, "ocr", 8)
|
||||||
|
text = extract_scanned_pdf_text(
|
||||||
|
db,
|
||||||
|
avatar,
|
||||||
|
path,
|
||||||
|
on_progress=ocr_progress,
|
||||||
|
)
|
||||||
|
self._set_progress(db, doc, "chunking", 29)
|
||||||
chunks = embeddings.chunk_text(text)
|
chunks = embeddings.chunk_text(text)
|
||||||
if not chunks:
|
if not chunks:
|
||||||
raise ValueError("文档没有可建立索引的文字内容")
|
raise ValueError("文档没有可建立索引的文字内容")
|
||||||
|
|||||||
@@ -0,0 +1,130 @@
|
|||||||
|
"""OCR fallback for image-only PDF knowledge documents."""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import time
|
||||||
|
from typing import Callable
|
||||||
|
|
||||||
|
from sqlalchemy.orm import Session
|
||||||
|
|
||||||
|
from models import Avatar
|
||||||
|
from services.chat_model_config import get_chat_model_config
|
||||||
|
from services.token_billing import (
|
||||||
|
estimate_fallback_usage,
|
||||||
|
release_reservation,
|
||||||
|
reserve_avatar_tokens,
|
||||||
|
settle_reservation,
|
||||||
|
)
|
||||||
|
from services.vision_service import call_vision_model, prepare_image
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
PDF_OCR_PROMPT = (
|
||||||
|
"请逐字转录这一页扫描文档中的全部可见文字和表格,只输出转录内容,不要解释,不要使用 Markdown 代码块。"
|
||||||
|
"保留标题、段落、项目编号、数值和自然换行;看不清的内容写作[无法辨认],不要猜测、纠错或补全。"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _positive_int(name: str, default: int, minimum: int, maximum: int) -> int:
|
||||||
|
try:
|
||||||
|
value = int(os.getenv(name, str(default)))
|
||||||
|
except ValueError:
|
||||||
|
value = default
|
||||||
|
return max(minimum, min(maximum, value))
|
||||||
|
|
||||||
|
|
||||||
|
def extract_scanned_pdf_text(
|
||||||
|
db: Session,
|
||||||
|
avatar: Avatar,
|
||||||
|
path: str,
|
||||||
|
*,
|
||||||
|
on_progress: Callable[[int, int], None] | None = None,
|
||||||
|
) -> str:
|
||||||
|
"""Render and OCR an image-only PDF while preserving page order."""
|
||||||
|
try:
|
||||||
|
import pymupdf
|
||||||
|
except ImportError as exc:
|
||||||
|
raise RuntimeError("扫描型 PDF 识别组件未安装") from exc
|
||||||
|
|
||||||
|
max_pages = _positive_int("KNOWLEDGE_PDF_OCR_MAX_PAGES", 80, 1, 300)
|
||||||
|
render_dpi = _positive_int("KNOWLEDGE_PDF_OCR_DPI", 144, 96, 200)
|
||||||
|
max_attempts = _positive_int("KNOWLEDGE_PDF_OCR_ATTEMPTS", 3, 1, 5)
|
||||||
|
model_config = get_chat_model_config()
|
||||||
|
model = model_config.ocr_model or model_config.vision_model
|
||||||
|
if not model_config.api_key or not model:
|
||||||
|
raise RuntimeError("扫描型 PDF 需要配置视觉 OCR 模型")
|
||||||
|
|
||||||
|
texts: list[str] = []
|
||||||
|
with pymupdf.open(path) as document:
|
||||||
|
total_pages = document.page_count
|
||||||
|
if total_pages <= 0:
|
||||||
|
raise ValueError("PDF 没有可识别页面")
|
||||||
|
if total_pages > max_pages:
|
||||||
|
raise ValueError(
|
||||||
|
f"扫描型 PDF 共 {total_pages} 页,超过单次 OCR 上限 {max_pages} 页,请拆分后上传"
|
||||||
|
)
|
||||||
|
|
||||||
|
scale = render_dpi / 72
|
||||||
|
for page_index in range(total_pages):
|
||||||
|
page = document.load_page(page_index)
|
||||||
|
pixmap = page.get_pixmap(
|
||||||
|
matrix=pymupdf.Matrix(scale, scale),
|
||||||
|
colorspace=pymupdf.csRGB,
|
||||||
|
alpha=False,
|
||||||
|
)
|
||||||
|
prepared = prepare_image(pixmap.tobytes("jpeg", jpg_quality=88))
|
||||||
|
estimate_messages = [{
|
||||||
|
"role": "user",
|
||||||
|
"content": f"[扫描 PDF 第 {page_index + 1}/{total_pages} 页]\n{PDF_OCR_PROMPT}",
|
||||||
|
}]
|
||||||
|
reservation = reserve_avatar_tokens(
|
||||||
|
db,
|
||||||
|
avatar,
|
||||||
|
"knowledge_pdf_ocr",
|
||||||
|
model,
|
||||||
|
estimate_messages,
|
||||||
|
model_config.vision_max_tokens,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
result = None
|
||||||
|
for attempt in range(1, max_attempts + 1):
|
||||||
|
try:
|
||||||
|
result = call_vision_model(
|
||||||
|
prepared,
|
||||||
|
model_config,
|
||||||
|
model=model,
|
||||||
|
prompt=PDF_OCR_PROMPT,
|
||||||
|
json_output=False,
|
||||||
|
)
|
||||||
|
break
|
||||||
|
except RuntimeError:
|
||||||
|
if attempt == max_attempts:
|
||||||
|
raise
|
||||||
|
time.sleep(min(4, attempt))
|
||||||
|
content = str((result or {}).get("content") or "").strip()
|
||||||
|
if not content:
|
||||||
|
raise RuntimeError("扫描型 PDF 页面识别结果为空")
|
||||||
|
settle_reservation(
|
||||||
|
db,
|
||||||
|
reservation,
|
||||||
|
(result or {}).get("usage"),
|
||||||
|
fallback_total=estimate_fallback_usage(estimate_messages, content),
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
release_reservation(db, reservation, str(exc))
|
||||||
|
raise RuntimeError(
|
||||||
|
f"扫描型 PDF 第 {page_index + 1}/{total_pages} 页识别失败:{exc}"
|
||||||
|
) from exc
|
||||||
|
|
||||||
|
texts.append(f"[第 {page_index + 1} 页]\n{content}")
|
||||||
|
if on_progress:
|
||||||
|
on_progress(page_index + 1, total_pages)
|
||||||
|
logger.info(
|
||||||
|
"Scanned PDF OCR completed avatar=%s page=%s/%s",
|
||||||
|
avatar.id,
|
||||||
|
page_index + 1,
|
||||||
|
total_pages,
|
||||||
|
)
|
||||||
|
|
||||||
|
return "\n\n".join(texts).strip()
|
||||||
@@ -238,6 +238,53 @@ def test_background_vectorizer_keeps_failure_reason_for_retry(
|
|||||||
db.close()
|
db.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_background_vectorizer_uses_ocr_for_image_only_pdf(
|
||||||
|
tmp_path: Path,
|
||||||
|
authorization_context,
|
||||||
|
):
|
||||||
|
context = authorization_context
|
||||||
|
with (
|
||||||
|
patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)),
|
||||||
|
patch("routers.knowledge.knowledge_vectorizer.enqueue"),
|
||||||
|
):
|
||||||
|
response = client.post(
|
||||||
|
f"/api/avatar/{context['avatar'].id}/knowledge/docs",
|
||||||
|
headers=context["owner_headers"],
|
||||||
|
files={"file": ("scanned.pdf", b"image-only-pdf", "application/pdf")},
|
||||||
|
)
|
||||||
|
|
||||||
|
payload = response.json()["data"]
|
||||||
|
progress = []
|
||||||
|
with (
|
||||||
|
patch("services.knowledge_vectorizer.UPLOAD_DIR", str(tmp_path)),
|
||||||
|
patch("services.knowledge_vectorizer.embeddings.extract_text", return_value=""),
|
||||||
|
patch(
|
||||||
|
"services.knowledge_vectorizer.extract_scanned_pdf_text",
|
||||||
|
side_effect=lambda _db, _avatar, _path, on_progress: (
|
||||||
|
on_progress(1, 2), on_progress(2, 2), "扫描页文字"
|
||||||
|
)[-1],
|
||||||
|
) as ocr,
|
||||||
|
patch("services.knowledge_vectorizer.embeddings.embed", return_value=[[1.0, 0.0]]),
|
||||||
|
patch.object(knowledge_vectorizer, "_set_progress", wraps=knowledge_vectorizer._set_progress) as set_progress,
|
||||||
|
):
|
||||||
|
knowledge_vectorizer.vectorize_document(payload["id"])
|
||||||
|
progress = [(call.args[2], call.args[3]) for call in set_progress.call_args_list]
|
||||||
|
|
||||||
|
db = SessionLocal()
|
||||||
|
try:
|
||||||
|
stored = db.query(KnowledgeDoc).filter(KnowledgeDoc.id == payload["id"]).one()
|
||||||
|
assert stored.status == "ready"
|
||||||
|
assert stored.chunk_count == 1
|
||||||
|
assert ("ocr", 18) in progress
|
||||||
|
assert ("ocr", 28) in progress
|
||||||
|
ocr.assert_called_once()
|
||||||
|
db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).delete()
|
||||||
|
db.delete(stored)
|
||||||
|
db.commit()
|
||||||
|
finally:
|
||||||
|
db.close()
|
||||||
|
|
||||||
|
|
||||||
def test_retry_queues_a_failed_document_again(
|
def test_retry_queues_a_failed_document_again(
|
||||||
tmp_path: Path,
|
tmp_path: Path,
|
||||||
authorization_context,
|
authorization_context,
|
||||||
|
|||||||
@@ -0,0 +1,104 @@
|
|||||||
|
from types import SimpleNamespace
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from services.pdf_ocr_service import extract_scanned_pdf_text
|
||||||
|
|
||||||
|
|
||||||
|
class FakePixmap:
|
||||||
|
def tobytes(self, *_args, **_kwargs):
|
||||||
|
return b"jpeg-page"
|
||||||
|
|
||||||
|
|
||||||
|
class FakePage:
|
||||||
|
def get_pixmap(self, **_kwargs):
|
||||||
|
return FakePixmap()
|
||||||
|
|
||||||
|
|
||||||
|
class FakeDocument:
|
||||||
|
page_count = 2
|
||||||
|
|
||||||
|
def __enter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, *_args):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def load_page(self, _index):
|
||||||
|
return FakePage()
|
||||||
|
|
||||||
|
|
||||||
|
def test_scanned_pdf_ocr_preserves_page_order_and_reports_progress(monkeypatch):
|
||||||
|
fake_pymupdf = SimpleNamespace(
|
||||||
|
open=lambda _path: FakeDocument(),
|
||||||
|
Matrix=lambda x, y: (x, y),
|
||||||
|
csRGB="rgb",
|
||||||
|
)
|
||||||
|
monkeypatch.setitem(__import__("sys").modules, "pymupdf", fake_pymupdf)
|
||||||
|
progress = []
|
||||||
|
reservation = SimpleNamespace()
|
||||||
|
config = SimpleNamespace(
|
||||||
|
api_key="configured",
|
||||||
|
ocr_model="qwen-vl-ocr",
|
||||||
|
vision_model="vision",
|
||||||
|
vision_max_tokens=2048,
|
||||||
|
)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("services.pdf_ocr_service.get_chat_model_config", return_value=config),
|
||||||
|
patch("services.pdf_ocr_service.prepare_image", return_value=SimpleNamespace()),
|
||||||
|
patch(
|
||||||
|
"services.pdf_ocr_service.call_vision_model",
|
||||||
|
side_effect=[
|
||||||
|
{"content": "第一页文字", "usage": {"total_tokens": 10}},
|
||||||
|
{"content": "第二页文字", "usage": {"total_tokens": 12}},
|
||||||
|
],
|
||||||
|
),
|
||||||
|
patch("services.pdf_ocr_service.reserve_avatar_tokens", return_value=reservation) as reserve,
|
||||||
|
patch("services.pdf_ocr_service.settle_reservation") as settle,
|
||||||
|
):
|
||||||
|
text = extract_scanned_pdf_text(
|
||||||
|
MagicMock(),
|
||||||
|
SimpleNamespace(id="avatar-1"),
|
||||||
|
"/tmp/scanned.pdf",
|
||||||
|
on_progress=lambda done, total: progress.append((done, total)),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert text == "[第 1 页]\n第一页文字\n\n[第 2 页]\n第二页文字"
|
||||||
|
assert progress == [(1, 2), (2, 2)]
|
||||||
|
assert reserve.call_count == 2
|
||||||
|
assert settle.call_count == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_scanned_pdf_ocr_releases_tokens_after_retries_fail(monkeypatch):
|
||||||
|
fake_document = FakeDocument()
|
||||||
|
fake_document.page_count = 1
|
||||||
|
fake_pymupdf = SimpleNamespace(
|
||||||
|
open=lambda _path: fake_document,
|
||||||
|
Matrix=lambda x, y: (x, y),
|
||||||
|
csRGB="rgb",
|
||||||
|
)
|
||||||
|
monkeypatch.setitem(__import__("sys").modules, "pymupdf", fake_pymupdf)
|
||||||
|
monkeypatch.setenv("KNOWLEDGE_PDF_OCR_ATTEMPTS", "2")
|
||||||
|
reservation = SimpleNamespace()
|
||||||
|
config = SimpleNamespace(
|
||||||
|
api_key="configured",
|
||||||
|
ocr_model="qwen-vl-ocr",
|
||||||
|
vision_model="vision",
|
||||||
|
vision_max_tokens=2048,
|
||||||
|
)
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("services.pdf_ocr_service.get_chat_model_config", return_value=config),
|
||||||
|
patch("services.pdf_ocr_service.prepare_image", return_value=SimpleNamespace()),
|
||||||
|
patch("services.pdf_ocr_service.call_vision_model", side_effect=RuntimeError("timeout")) as call,
|
||||||
|
patch("services.pdf_ocr_service.reserve_avatar_tokens", return_value=reservation),
|
||||||
|
patch("services.pdf_ocr_service.release_reservation") as release,
|
||||||
|
patch("services.pdf_ocr_service.time.sleep"),
|
||||||
|
):
|
||||||
|
with pytest.raises(RuntimeError, match="第 1/1 页识别失败"):
|
||||||
|
extract_scanned_pdf_text(MagicMock(), SimpleNamespace(id="avatar-1"), "/tmp/scanned.pdf")
|
||||||
|
|
||||||
|
assert call.call_count == 2
|
||||||
|
release.assert_called_once()
|
||||||
@@ -1,6 +1,7 @@
|
|||||||
import { createRouter, createWebHashHistory } from 'vue-router'
|
import { createRouter, createWebHashHistory } from 'vue-router'
|
||||||
import type { RouteRecordRaw } from 'vue-router'
|
import type { RouteRecordRaw } from 'vue-router'
|
||||||
import { getAuthToken } from '@/api'
|
import { getAuthToken } from '@/api'
|
||||||
|
import { isInUniWebView } from '@/utils/uniapp-bridge'
|
||||||
|
|
||||||
const routes: RouteRecordRaw[] = [
|
const routes: RouteRecordRaw[] = [
|
||||||
{
|
{
|
||||||
@@ -55,7 +56,7 @@ const routes: RouteRecordRaw[] = [
|
|||||||
path: '/token/charge',
|
path: '/token/charge',
|
||||||
name: 'TokenCharge',
|
name: 'TokenCharge',
|
||||||
component: () => import('@/views/TokenCharge.vue'),
|
component: () => import('@/views/TokenCharge.vue'),
|
||||||
meta: { title: '积分充值', requiresAuth: true }
|
meta: { title: '积分充值', requiresAuth: true, requiresUniWebView: true }
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
path: '/avatar/card',
|
path: '/avatar/card',
|
||||||
@@ -133,6 +134,10 @@ const router = createRouter({
|
|||||||
|
|
||||||
router.beforeEach((to, from, next) => {
|
router.beforeEach((to, from, next) => {
|
||||||
document.title = to.meta.title as string || '会会数字分身'
|
document.title = to.meta.title as string || '会会数字分身'
|
||||||
|
if (to.meta.requiresUniWebView && !isInUniWebView()) {
|
||||||
|
next({ path: '/avatar/manage' })
|
||||||
|
return
|
||||||
|
}
|
||||||
const hasLocalSession = Boolean(localStorage.getItem('hh_app_token'))
|
const hasLocalSession = Boolean(localStorage.getItem('hh_app_token'))
|
||||||
const hasInjectedSession = Boolean(getAuthToken())
|
const hasInjectedSession = Boolean(getAuthToken())
|
||||||
if (to.meta.requiresAuth && !hasLocalSession && !hasInjectedSession) {
|
if (to.meta.requiresAuth && !hasLocalSession && !hasInjectedSession) {
|
||||||
|
|||||||
@@ -20,7 +20,7 @@
|
|||||||
</div>
|
</div>
|
||||||
</section>
|
</section>
|
||||||
|
|
||||||
<!-- 积分余额条:暂时隐藏,保留完整实现便于后续恢复。 -->
|
<!-- 积分余额条:仅在 uni-app 原生壳内开放充值购买。 -->
|
||||||
<section v-if="SHOW_POINTS_BALANCE_CARD" class="token-section">
|
<section v-if="SHOW_POINTS_BALANCE_CARD" class="token-section">
|
||||||
<div class="token-card">
|
<div class="token-card">
|
||||||
<div class="token-info">
|
<div class="token-info">
|
||||||
@@ -28,7 +28,7 @@
|
|||||||
<span class="token-amount">{{ tokenBalance.toLocaleString() }}</span>
|
<span class="token-amount">{{ tokenBalance.toLocaleString() }}</span>
|
||||||
<span class="token-used">累计使用 {{ tokenConsumed.toLocaleString() }}</span>
|
<span class="token-used">累计使用 {{ tokenConsumed.toLocaleString() }}</span>
|
||||||
</div>
|
</div>
|
||||||
<button class="recharge-btn" @click="goToRecharge">充值</button>
|
<button class="recharge-btn" @click="goToRecharge">充值购买</button>
|
||||||
</div>
|
</div>
|
||||||
</section>
|
</section>
|
||||||
|
|
||||||
@@ -89,14 +89,15 @@ import { useAvatarStore } from '@/store/avatar'
|
|||||||
import { useUserStore } from '@/store/user'
|
import { useUserStore } from '@/store/user'
|
||||||
import { createAvatarShareLink } from '@/api'
|
import { createAvatarShareLink } from '@/api'
|
||||||
import { isHuihuiEmbeddedMode } from '@/utils/embed-mode'
|
import { isHuihuiEmbeddedMode } from '@/utils/embed-mode'
|
||||||
|
import { isInUniWebView } from '@/utils/uniapp-bridge'
|
||||||
|
|
||||||
const router = useRouter()
|
const router = useRouter()
|
||||||
const avatarStore = useAvatarStore()
|
const avatarStore = useAvatarStore()
|
||||||
const userStore = useUserStore()
|
const userStore = useUserStore()
|
||||||
const isEmbedded = isHuihuiEmbeddedMode()
|
const isEmbedded = isHuihuiEmbeddedMode()
|
||||||
|
|
||||||
// 临时产品开关:余额卡片代码保留,后续改为 true 即可恢复展示。
|
// 充值购买只在 uni-app 原生壳内提供,避免普通 H5 进入支付链路。
|
||||||
const SHOW_POINTS_BALANCE_CARD = false
|
const SHOW_POINTS_BALANCE_CARD = isInUniWebView()
|
||||||
|
|
||||||
// 当前登录会会用户的资料(头像 / 昵称)
|
// 当前登录会会用户的资料(头像 / 昵称)
|
||||||
const me = computed(() => userStore.user)
|
const me = computed(() => userStore.user)
|
||||||
|
|||||||
@@ -147,7 +147,7 @@ const documentState = (doc: any) => {
|
|||||||
if (['uploaded', 'parsing'].includes(String(doc.status || '').toLowerCase())) {
|
if (['uploaded', 'parsing'].includes(String(doc.status || '').toLowerCase())) {
|
||||||
const stage = String(doc.indexStage || 'queued').toLowerCase()
|
const stage = String(doc.indexStage || 'queued').toLowerCase()
|
||||||
const labels: Record<string, string> = {
|
const labels: Record<string, string> = {
|
||||||
queued: '等待处理', extracting: '解析文档', chunking: '切分文本', embedding: '向量化中'
|
queued: '等待处理', extracting: '解析文档', ocr: '扫描件识别', chunking: '切分文本', embedding: '向量化中'
|
||||||
}
|
}
|
||||||
const progress = Math.max(0, Math.min(99, Number(doc.indexProgress || 0)))
|
const progress = Math.max(0, Math.min(99, Number(doc.indexProgress || 0)))
|
||||||
return { tone: 'pending', label: labels[stage] || '处理中', detail: `${labels[stage] || '正在建立知识索引'} ${progress}%`, progress }
|
return { tone: 'pending', label: labels[stage] || '处理中', detail: `${labels[stage] || '正在建立知识索引'} ${progress}%`, progress }
|
||||||
|
|||||||
Reference in New Issue
Block a user