feat: v1.2.0 - 图片去重与管理、微信机器人优化、搜索设置可配置
主要功能: - 图片上传时 OCR 内容去重(3个上传端点统一使用公共函数 _check_ocr_duplicate) - 图片管理 Tab:展示所有图片、手动删除、一键去重 - 搜索结果详情弹窗增加删除按钮(带确认弹窗) - 图片管理卡片点击查看详情(复用 showOcrDetailModal) - 搜索限制和 LLM 批量判断数量可通过网站设置 - MiniMax API 调用添加 reasoning_split=True - 企业微信机器人:WebSocket 长连接、图片搜索、配置化搜索数量 - 版本号升级至 1.2.0
This commit is contained in:
84
.env.example
Normal file
84
.env.example
Normal file
@@ -0,0 +1,84 @@
|
||||
# ============================================================
|
||||
# EduBrain - 中文直播教育知识库系统 环境变量配置
|
||||
# ============================================================
|
||||
# 复制此文件为 .env 并根据实际情况修改配置值
|
||||
# cp .env.example .env
|
||||
# ============================================================
|
||||
|
||||
# ── 数据库 ──
|
||||
DATABASE_URL=postgresql+asyncpg://postgres:postgres@db:5432/edu_brain
|
||||
POSTGRES_PASSWORD=postgres
|
||||
|
||||
# ── 嵌入模型配置 ──
|
||||
# 可选值: openai / zhipu / local_bge / dashscope / minimax
|
||||
EMBEDDING_PROVIDER=minimax
|
||||
EMBEDDING_MODEL=MiniMax-Text-Embedding
|
||||
EMBEDDING_DIMENSIONS=1024
|
||||
|
||||
# ── OpenAI 配置 ──
|
||||
OPENAI_API_KEY=sk-your-openai-api-key
|
||||
OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
|
||||
# ── 智谱 AI 配置 ──
|
||||
ZHIPU_API_KEY=your-zhipu-api-key
|
||||
|
||||
# ── 阿里云 DashScope 配置 ──
|
||||
DASHSCOPE_API_KEY=your-dashscope-api-key
|
||||
|
||||
# ── 本地 BGE 模型路径(为空则自动从 HuggingFace 下载) ──
|
||||
LOCAL_BGE_MODEL_PATH=
|
||||
|
||||
# ── OCR 配置 ──
|
||||
# 可选值: paddleocr / aliyun / tencent / deepseek / auto
|
||||
# auto 模式: 本地 PaddleOCR 优先,失败后 fallback 到 DeepSeek / 云端 OCR
|
||||
OCR_PROVIDER=deepseek
|
||||
|
||||
# ── 阿里云 OCR 配置 ──
|
||||
ALIYUN_OCR_ACCESS_KEY=your-aliyun-access-key
|
||||
ALIYUN_OCR_SECRET=your-aliyun-secret
|
||||
|
||||
# ── 腾讯云 OCR 配置 ──
|
||||
TENCENT_OCR_SECRET_ID=your-tencent-secret-id
|
||||
TENCENT_OCR_SECRET_KEY=your-tencent-secret-key
|
||||
|
||||
# ── MiniMax 配置 ──
|
||||
MINIMAX_API_KEY=your-minimax-api-key
|
||||
MINIMAX_BASE_URL=https://api.minimaxi.com/v1
|
||||
MINIMAX_EMBEDDING_MODEL=MiniMax-Text-Embedding
|
||||
MINIMAX_CHAT_MODEL=MiniMax-M2
|
||||
|
||||
# ── DeepSeek 配置 ──
|
||||
DEEPSEEK_API_KEY=your-deepseek-api-key
|
||||
DEEPSEEK_BASE_URL=https://api.deepseek.com
|
||||
DEEPSEEK_OCR_MODEL=deepseek-chat
|
||||
|
||||
# ── LLM 配置(用于查询扩展、实体检测) ──
|
||||
# 可选值: minimax / deepseek / openai
|
||||
LLM_PROVIDER=minimax
|
||||
|
||||
# ── 数据目录 ──
|
||||
DATA_DIR=/app/data
|
||||
|
||||
# ── 文本分块参数 ──
|
||||
CHUNK_SIZE=300
|
||||
CHUNK_OVERLAP=50
|
||||
|
||||
# ── 服务配置 ──
|
||||
HOST=0.0.0.0
|
||||
PORT=8000
|
||||
DEBUG=false
|
||||
|
||||
# ── CORS 配置(多个来源用逗号分隔) ──
|
||||
CORS_ORIGINS=["*"]
|
||||
|
||||
# ── Docker 端口映射 ──
|
||||
APP_PORT=8000
|
||||
DB_PORT=5432
|
||||
|
||||
# ── 企业微信智能机器人(WebSocket 长连接) ──
|
||||
# 作为独立进程运行:python -m app.wework_bot
|
||||
WEWORK_BOT_ENABLED=false
|
||||
WEWORK_BOT_ID=your-bot-id
|
||||
WEWORK_BOT_SECRET=your-bot-secret
|
||||
WEWORK_BOT_SEARCH_LIMIT=3
|
||||
EDUBRAIN_BASE_URL=http://localhost:8765
|
||||
41
.gitignore
vendored
Normal file
41
.gitignore
vendored
Normal file
@@ -0,0 +1,41 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
*.egg
|
||||
|
||||
# Virtual environments
|
||||
venv/
|
||||
.venv/
|
||||
env/
|
||||
|
||||
# Environment variables
|
||||
.env
|
||||
.env.local
|
||||
.env.*.local
|
||||
|
||||
# Database
|
||||
*.db
|
||||
*.db-shm
|
||||
*.db-wal
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Data
|
||||
data/images/
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
36
Dockerfile
Normal file
36
Dockerfile
Normal file
@@ -0,0 +1,36 @@
|
||||
FROM python:3.11-slim
|
||||
|
||||
# ── 系统依赖(PaddleOCR 需要的图形库) ──
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
libgl1-mesa-glx \
|
||||
libglib2.0-0 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
libgomp1 \
|
||||
wget \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# ── 设置工作目录 ──
|
||||
WORKDIR /app
|
||||
|
||||
# ── 安装 Python 依赖(先复制 requirements.txt 利用 Docker 缓存) ──
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt \
|
||||
&& pip install --no-cache-dir paddlepaddle==3.0.0b1 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/
|
||||
|
||||
# ── 复制项目代码 ──
|
||||
COPY . .
|
||||
|
||||
# ── 创建数据目录 ──
|
||||
RUN mkdir -p /app/data/transcripts /app/data/images /app/data/exports
|
||||
|
||||
# ── 暴露端口 ──
|
||||
EXPOSE 8000
|
||||
|
||||
# ── 健康检查 ──
|
||||
HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \
|
||||
CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
|
||||
|
||||
# ── 启动命令 ──
|
||||
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]
|
||||
176
README.md
Normal file
176
README.md
Normal file
@@ -0,0 +1,176 @@
|
||||
# EduBrain - 中文直播教育知识库系统
|
||||
|
||||
基于 FastAPI 的中文直播教育知识库系统,支持直播文字稿导入、OCR 截图识别、AI 语义搜索等功能。
|
||||
|
||||
## 功能特性
|
||||
|
||||
### 知识库管理
|
||||
- **多格式导入**: Markdown / 纯文本 / Word 文档(.docx),自动按标题分页
|
||||
- **批量导入**: 支持多文件批量导入,实时进度显示
|
||||
- **智能分块**: 按段落边界切分文本,保留重叠区域
|
||||
- **向量嵌入**: 支持 MiniMax / OpenAI / 智谱 / DashScope / 本地 BGE 五种嵌入模型
|
||||
|
||||
### 图片 OCR
|
||||
- **OCR 识别**: DeepSeek Vision 模型识别图片文字
|
||||
- **关键词提取**: MiniMax M2.7 自动提取内容标签
|
||||
- **批量处理**: 支持拖拽上传、选择文件、选择文件夹,按文件名去重
|
||||
- **BM25 索引**: 内置倒排索引搜索引擎,jieba 中文分词
|
||||
|
||||
### AI 搜索
|
||||
- **知识库搜索**: 向量语义搜索 + 全文检索混合排序
|
||||
- **图片搜索**: MiniMax M2.7 语义匹配,SSE 流式实时返回结果
|
||||
- **并发池**: 最多 10 个 LLM 请求并发,每批 10 张图片判断
|
||||
|
||||
### 其他
|
||||
- **MCP 协议**: 支持 Model Context Protocol,可被 AI Agent 调用
|
||||
- **Docker 部署**: 完整的 Docker Compose 编排(应用 + PostgreSQL + pgvector)
|
||||
|
||||
## 技术栈
|
||||
|
||||
| 层级 | 技术 |
|
||||
|------|------|
|
||||
| 后端框架 | FastAPI + Uvicorn |
|
||||
| ORM | SQLAlchemy 2.0 (async) |
|
||||
| 数据库 | PostgreSQL (pgvector) / SQLite |
|
||||
| 前端 | 纯 HTML + CSS + JS(苹果风格 SPA) |
|
||||
| OCR | DeepSeek Vision / PaddleOCR / 阿里云 / 腾讯云 |
|
||||
| LLM | MiniMax M2.7 / DeepSeek / Qwen3-8B |
|
||||
| 搜索 | BM25 倒排索引 + 向量语义搜索 |
|
||||
| 部署 | Docker + Docker Compose |
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
edu-brain/
|
||||
├── app/
|
||||
│ ├── main.py # FastAPI 入口,路由注册,生命周期管理
|
||||
│ ├── config.py # pydantic-settings 配置管理
|
||||
│ ├── database.py # 数据库连接(PG/SQLite 双后端)
|
||||
│ ├── models/base.py # ORM 模型(KnowledgePage, OCRImage 等)
|
||||
│ ├── schemas/ # Pydantic 请求/响应 Schema
|
||||
│ ├── api/v1/ # API 路由
|
||||
│ │ ├── pages.py # 知识页面 CRUD
|
||||
│ │ ├── search.py # 语义搜索
|
||||
│ │ ├── images.py # 图片 OCR + AI 搜索(SSE)
|
||||
│ │ ├── import_export.py # 文件导入导出
|
||||
│ │ └── settings.py # 系统设置
|
||||
│ ├── services/
|
||||
│ │ ├── ocr_service.py # OCR 识别服务(多 Provider)
|
||||
│ │ ├── llm_service.py # LLM 服务(标签提取/搜索匹配)
|
||||
│ │ ├── search_engine.py # BM25 倒排索引引擎
|
||||
│ │ ├── search_service.py # 混合搜索服务
|
||||
│ │ ├── embedding_service.py # 嵌入模型服务
|
||||
│ │ ├── import_service.py # 文件导入服务
|
||||
│ │ └── page_service.py # 页面 CRUD 服务
|
||||
│ └── mcp/server.py # MCP Server 实现
|
||||
├── static/index.html # 前端单页应用
|
||||
├── data/images/ # 图片存储目录
|
||||
├── .env.example # 环境变量模板
|
||||
├── requirements.txt # Python 依赖
|
||||
├── Dockerfile # 应用镜像
|
||||
├── Dockerfile.db # 数据库镜像(PG + pgvector)
|
||||
└── docker-compose.yml # Docker 编排
|
||||
```
|
||||
|
||||
## 快速开始
|
||||
|
||||
### 1. 环境准备
|
||||
|
||||
```bash
|
||||
# 克隆项目
|
||||
git clone <repo-url>
|
||||
cd edu-brain
|
||||
|
||||
# 复制配置文件
|
||||
cp .env.example .env
|
||||
|
||||
# 安装依赖
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 2. 配置 .env
|
||||
|
||||
编辑 `.env` 文件,至少配置以下项:
|
||||
|
||||
```env
|
||||
# 数据库(默认 SQLite,无需额外配置)
|
||||
DATABASE_URL=sqlite+aiosqlite:///edu_brain.db
|
||||
|
||||
# OCR(使用 DeepSeek Vision)
|
||||
OCR_PROVIDER=deepseek
|
||||
DEEPSEEK_API_KEY=your-deepseek-api-key
|
||||
DEEPSEEK_BASE_URL=http://your-ollama-host:11434/v1
|
||||
DEEPSEEK_OCR_MODEL=deepseek-ocr:latest
|
||||
|
||||
# LLM(用于标签提取和搜索匹配)
|
||||
MINIMAX_API_KEY=your-minimax-api-key
|
||||
MINIMAX_BASE_URL=https://api.minimaxi.com/v1
|
||||
MINIMAX_CHAT_MODEL=MiniMax-M2.7
|
||||
```
|
||||
|
||||
### 3. 启动服务
|
||||
|
||||
```bash
|
||||
uvicorn app.main:app --host 0.0.0.0 --port 8765
|
||||
```
|
||||
|
||||
访问 http://localhost:8765 即可使用。
|
||||
|
||||
### 4. Docker 部署
|
||||
|
||||
```bash
|
||||
docker-compose up -d
|
||||
```
|
||||
|
||||
## 图片搜索工作流
|
||||
|
||||
```
|
||||
用户上传图片
|
||||
→ DeepSeek Vision OCR 识别文字
|
||||
→ MiniMax M2.7 提取关键词标签
|
||||
→ 存入 SQLite + BM25 索引
|
||||
|
||||
用户搜索 "孩子不想上学"
|
||||
→ 从 DB 倒序取图片(每批 10 张)
|
||||
→ 并发池(最多 10 个请求)发给 MiniMax M2.7 判断
|
||||
→ 匹配结果通过 SSE 实时推送到前端
|
||||
→ 找够指定数量或遍历完 DB 后结束
|
||||
```
|
||||
|
||||
## API 概览
|
||||
|
||||
| 方法 | 路径 | 说明 |
|
||||
|------|------|------|
|
||||
| POST | `/api/v1/images/recognize-direct` | 上传并识别图片 |
|
||||
| POST | `/api/v1/images/batch-recognize` | 批量上传并识别 |
|
||||
| GET | `/api/v1/images/search` | AI 搜索图片(SSE 流式) |
|
||||
| GET | `/api/v1/images/{id}` | 获取图片识别结果 |
|
||||
| GET | `/api/v1/pages` | 知识页面列表 |
|
||||
| POST | `/api/v1/pages` | 创建知识页面 |
|
||||
| POST | `/api/v1/search` | 语义搜索知识库 |
|
||||
| POST | `/api/v1/import/file` | 导入文件 |
|
||||
| GET | `/api/v1/settings` | 获取系统设置 |
|
||||
|
||||
## 配置说明
|
||||
|
||||
### OCR Provider
|
||||
|
||||
| Provider | 说明 | 额外配置 |
|
||||
|----------|------|----------|
|
||||
| `deepseek` | DeepSeek Vision(推荐) | `DEEPSEEK_API_KEY`, `DEEPSEEK_BASE_URL` |
|
||||
| `paddleocr` | 本地 PaddleOCR | 无 |
|
||||
| `aliyun` | 阿里云 OCR | `ALIYUN_OCR_ACCESS_KEY`, `ALIYUN_OCR_SECRET` |
|
||||
| `tencent` | 腾讯云 OCR | `TENCENT_OCR_SECRET_ID`, `TENCENT_OCR_SECRET_KEY` |
|
||||
| `auto` | 自动 fallback | 配置多个 Provider |
|
||||
|
||||
### 嵌入模型
|
||||
|
||||
| Provider | 说明 | 额外配置 |
|
||||
|----------|------|----------|
|
||||
| `minimax` | MiniMax 嵌入 | `MINIMAX_API_KEY` |
|
||||
| `openai` | OpenAI 嵌入 | `OPENAI_API_KEY` |
|
||||
| `zhipu` | 智谱 AI 嵌入 | `ZHIPU_API_KEY` |
|
||||
| `dashscope` | 阿里云 DashScope | `DASHSCOPE_API_KEY` |
|
||||
| `local_bge` | 本地 BGE 模型 | `LOCAL_BGE_MODEL_PATH` |
|
||||
|
||||
|
||||
63
alembic.ini
Normal file
63
alembic.ini
Normal file
@@ -0,0 +1,63 @@
|
||||
# Alembic 数据库迁移配置
|
||||
|
||||
[alembic]
|
||||
# 迁移脚本目录
|
||||
script_location = alembic
|
||||
|
||||
# 数据库连接 URL(通过环境变量覆盖)
|
||||
# 注意:Alembic 使用同步驱动,需要将 asyncpg 替换为 psycopg2
|
||||
sqlalchemy.url = postgresql://postgres:postgres@localhost:5432/edu_brain
|
||||
|
||||
# 模板文件
|
||||
file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
|
||||
|
||||
# 是否自动应用迁移
|
||||
# prepend_sys_path = .
|
||||
|
||||
# 时区设置
|
||||
timezone = Asia/Shanghai
|
||||
|
||||
# 截断长迁移消息
|
||||
truncate_slug_length = 40
|
||||
|
||||
# 修订 ID 格式
|
||||
# revision_environment = false
|
||||
|
||||
# 输出编码
|
||||
# output_encoding = utf-8
|
||||
|
||||
[post_write_hooks]
|
||||
|
||||
[loggers]
|
||||
keys = root,sqlalchemy,alembic
|
||||
|
||||
[handlers]
|
||||
keys = console
|
||||
|
||||
[formatters]
|
||||
keys = generic
|
||||
|
||||
[logger_root]
|
||||
level = WARN
|
||||
handlers = console
|
||||
qualname =
|
||||
|
||||
[logger_sqlalchemy]
|
||||
level = WARN
|
||||
handlers =
|
||||
qualname = sqlalchemy.engine
|
||||
|
||||
[logger_alembic]
|
||||
level = INFO
|
||||
handlers =
|
||||
qualname = alembic
|
||||
|
||||
[handler_console]
|
||||
class = StreamHandler
|
||||
args = (sys.stderr,)
|
||||
level = NOTSET
|
||||
formatter = generic
|
||||
|
||||
[formatter_generic]
|
||||
format = %(levelname)-5.5s [%(name)s] %(message)s
|
||||
datefmt = %H:%M:%S
|
||||
67
alembic/env.py
Normal file
67
alembic/env.py
Normal file
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
Alembic 迁移环境配置
|
||||
支持异步数据库连接
|
||||
"""
|
||||
|
||||
from logging.config import fileConfig
|
||||
|
||||
from alembic import context
|
||||
from sqlalchemy import engine_from_config, pool
|
||||
|
||||
# 导入所有 ORM 模型以确保 Base.metadata 包含完整的表定义
|
||||
from app.models.base import Base
|
||||
|
||||
# Alembic Config 对象
|
||||
config = context.config
|
||||
|
||||
# 设置日志
|
||||
if config.config_file_name is not None:
|
||||
fileConfig(config.config_file_name)
|
||||
|
||||
# 元数据目标,用于自动生成迁移
|
||||
target_metadata = Base.metadata
|
||||
|
||||
|
||||
def run_migrations_offline() -> None:
|
||||
"""
|
||||
以 'offline' 模式运行迁移。
|
||||
只需要 URL,不需要 Engine。调用 context.execute() 将迁移
|
||||
直接发送到数据库。
|
||||
"""
|
||||
url = config.get_main_option("sqlalchemy.url")
|
||||
context.configure(
|
||||
url=url,
|
||||
target_metadata=target_metadata,
|
||||
literal_binds=True,
|
||||
dialect_opts={"paramstyle": "named"},
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
def run_migrations_online() -> None:
|
||||
"""
|
||||
以 'online' 模式运行迁移。
|
||||
创建 Engine 并关联 connection 到 context。
|
||||
"""
|
||||
connectable = engine_from_config(
|
||||
config.get_section(config.config_ini_section, {}),
|
||||
prefix="sqlalchemy.",
|
||||
poolclass=pool.NullPool,
|
||||
)
|
||||
|
||||
with connectable.connect() as connection:
|
||||
context.configure(
|
||||
connection=connection,
|
||||
target_metadata=target_metadata,
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
if context.is_offline_mode():
|
||||
run_migrations_offline()
|
||||
else:
|
||||
run_migrations_online()
|
||||
26
alembic/script.py.mako
Normal file
26
alembic/script.py.mako
Normal file
@@ -0,0 +1,26 @@
|
||||
"""${message}
|
||||
|
||||
Revision ID: ${up_revision}
|
||||
Revises: ${down_revision | comma,n}
|
||||
Create Date: ${create_date}
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
${imports if imports else ""}
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = ${repr(up_revision)}
|
||||
down_revision: Union[str, None] = ${repr(down_revision)}
|
||||
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
${downgrades if downgrades else "pass"}
|
||||
0
alembic/versions/.gitkeep
Normal file
0
alembic/versions/.gitkeep
Normal file
5
app/__init__.py
Normal file
5
app/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""
|
||||
中文直播教育知识库系统 - edu-brain
|
||||
"""
|
||||
|
||||
__version__ = "1.2.0"
|
||||
3
app/api/__init__.py
Normal file
3
app/api/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
API 路由包
|
||||
"""
|
||||
3
app/api/v1/__init__.py
Normal file
3
app/api/v1/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
API v1 路由包
|
||||
"""
|
||||
565
app/api/v1/images.py
Normal file
565
app/api/v1/images.py
Normal file
@@ -0,0 +1,565 @@
|
||||
"""
|
||||
图片 OCR API 路由
|
||||
OCR 识别 + BM25 倒排索引搜索
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, Form, Query
|
||||
from fastapi.responses import StreamingResponse, Response
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import select, func, text, delete
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.config import settings
|
||||
from app.database import get_db
|
||||
from app.models.base import OCRImage
|
||||
from app.services.ocr_service import OCRService
|
||||
from app.services.search_engine import BM25Index, tokenize
|
||||
from app.services.llm_service import LLMService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter()
|
||||
|
||||
ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
|
||||
|
||||
# 全局 BM25 索引实例(在 main.py 中初始化)
|
||||
bm25_index: Optional[BM25Index] = None
|
||||
|
||||
|
||||
def _check_image_ext(filename: str) -> str:
|
||||
suffix = os.path.splitext(filename)[1].lower()
|
||||
if suffix not in ALLOWED_EXTENSIONS:
|
||||
raise HTTPException(status_code=400, detail=f"不支持的图片格式: {suffix}")
|
||||
return suffix
|
||||
|
||||
|
||||
def _save_ocr_result(record: OCRImage, ocr_result, db: AsyncSession):
|
||||
"""保存 OCR 识别结果到数据库"""
|
||||
record.ocr_text = ocr_result.text
|
||||
record.confidence = ocr_result.confidence
|
||||
record.provider = ocr_result.provider
|
||||
record.status = "completed"
|
||||
# 保留 keywords(如果有)
|
||||
if hasattr(ocr_result, 'keywords') and ocr_result.keywords:
|
||||
record.tags = json.dumps(ocr_result.keywords, ensure_ascii=False)
|
||||
|
||||
|
||||
def _add_to_index(record: OCRImage):
|
||||
"""将 OCR 结果添加到 BM25 索引"""
|
||||
if bm25_index is None or not record.ocr_text:
|
||||
return
|
||||
# 索引内容 = OCR 文本(权重最高)
|
||||
index_text = record.ocr_text
|
||||
# 如果有 tags/story_summary 也加入索引
|
||||
if record.tags:
|
||||
try:
|
||||
tags = json.loads(record.tags)
|
||||
if tags:
|
||||
index_text += "\n" + " ".join(tags)
|
||||
except Exception:
|
||||
pass
|
||||
if record.story_summary:
|
||||
index_text += "\n" + record.story_summary
|
||||
|
||||
bm25_index.add_document(
|
||||
doc_id=record.id,
|
||||
text=index_text,
|
||||
metadata={
|
||||
"file_path": record.file_path,
|
||||
"tags": json.loads(record.tags) if record.tags else [],
|
||||
"story_summary": record.story_summary or "",
|
||||
"confidence": record.confidence,
|
||||
"provider": record.provider,
|
||||
"created_at": str(record.created_at),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _check_ocr_duplicate(ocr_text: str, db: AsyncSession) -> Optional[int]:
|
||||
"""检查 OCR 文本是否已存在,返回已存在记录的 ID 或 None"""
|
||||
if not ocr_text:
|
||||
return None
|
||||
dup_query = select(OCRImage.id).where(
|
||||
OCRImage.ocr_text == ocr_text,
|
||||
OCRImage.status == "completed",
|
||||
).limit(1)
|
||||
dup_result = await db.execute(dup_query)
|
||||
return dup_result.scalar_one_or_none()
|
||||
|
||||
|
||||
# ─── 上传 ───
|
||||
|
||||
@router.post("/upload", summary="上传图片")
|
||||
async def upload_image(file: UploadFile, db: AsyncSession = Depends(get_db)):
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="文件名不能为空")
|
||||
_check_image_ext(file.filename)
|
||||
image_record = OCRImage(file_path=file.filename, status="pending")
|
||||
db.add(image_record)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
return {"id": image_record.id, "file_path": image_record.file_path, "original_filename": file.filename, "status": image_record.status, "created_at": str(image_record.created_at)}
|
||||
|
||||
|
||||
# ─── 单文件识别 ───
|
||||
|
||||
@router.post("/{image_id}/recognize", summary="识别图片文字")
|
||||
async def recognize_image(image_id: int, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(OCRImage).where(OCRImage.id == image_id))
|
||||
image_record = result.scalar_one_or_none()
|
||||
if image_record is None:
|
||||
raise HTTPException(status_code=404, detail="图片记录不存在")
|
||||
if not os.path.exists(image_record.file_path):
|
||||
raise HTTPException(status_code=404, detail=f"图片文件不存在: {image_record.file_path}")
|
||||
|
||||
image_record.status = "processing"
|
||||
await db.flush()
|
||||
|
||||
try:
|
||||
ocr_result = await OCRService.recognize(image_record.file_path)
|
||||
_save_ocr_result(image_record, ocr_result, db)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
_add_to_index(image_record)
|
||||
|
||||
return {
|
||||
"id": image_record.id,
|
||||
"file_path": image_record.file_path,
|
||||
"ocr_text": image_record.ocr_text,
|
||||
"tags": json.loads(image_record.tags) if image_record.tags else [],
|
||||
"confidence": image_record.confidence,
|
||||
"provider": image_record.provider,
|
||||
"status": image_record.status,
|
||||
"block_count": len(ocr_result.blocks),
|
||||
}
|
||||
except Exception as e:
|
||||
image_record.status = "failed"
|
||||
image_record.error_message = str(e)
|
||||
await db.flush()
|
||||
raise HTTPException(status_code=500, detail=f"OCR 识别失败: {str(e)}")
|
||||
|
||||
|
||||
# ─── 批量识别 ───
|
||||
|
||||
@router.post("/batch-recognize", summary="批量上传并识别图片")
|
||||
async def batch_recognize(files: List[UploadFile], db: AsyncSession = Depends(get_db)):
|
||||
results = []
|
||||
for file in files:
|
||||
if not file.filename:
|
||||
results.append({"filename": "unknown", "status": "error", "message": "文件名为空"})
|
||||
continue
|
||||
suffix = _check_image_ext(file.filename)
|
||||
dest_path = settings.images_dir / file.filename
|
||||
counter = 1
|
||||
stem = os.path.splitext(file.filename)[0]
|
||||
while dest_path.exists():
|
||||
dest_path = settings.images_dir / f"{stem}_{counter}{suffix}"
|
||||
counter += 1
|
||||
content = await file.read()
|
||||
dest_path.write_bytes(content)
|
||||
image_record = OCRImage(file_path=str(dest_path), status="processing")
|
||||
db.add(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
ocr_result = await OCRService.recognize(str(dest_path))
|
||||
# OCR 内容去重
|
||||
existing_id = await _check_ocr_duplicate(ocr_result.text, db)
|
||||
if existing_id is not None:
|
||||
# 已存在相同内容,回滚当前记录,删除图片文件
|
||||
await db.delete(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
os.unlink(dest_path)
|
||||
except OSError:
|
||||
pass
|
||||
results.append({
|
||||
"filename": file.filename,
|
||||
"file_path": str(dest_path),
|
||||
"status": "duplicate",
|
||||
"duplicate_of": existing_id,
|
||||
"ocr_text": ocr_result.text,
|
||||
})
|
||||
continue
|
||||
_save_ocr_result(image_record, ocr_result, db)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
_add_to_index(image_record)
|
||||
results.append({
|
||||
"id": image_record.id, "filename": file.filename,
|
||||
"file_path": str(dest_path), "status": "completed",
|
||||
"ocr_text": ocr_result.text,
|
||||
"tags": json.loads(image_record.tags) if image_record.tags else [],
|
||||
"confidence": ocr_result.confidence, "provider": ocr_result.provider,
|
||||
"block_count": len(ocr_result.blocks),
|
||||
})
|
||||
except Exception as e:
|
||||
image_record.status = "failed"
|
||||
image_record.error_message = str(e)
|
||||
await db.flush()
|
||||
results.append({"id": image_record.id, "filename": file.filename, "file_path": str(dest_path), "status": "failed", "message": str(e)})
|
||||
return {"total": len(results), "success": sum(1 for r in results if r["status"] == "completed"), "failed": sum(1 for r in results if r["status"] == "failed"), "duplicates": sum(1 for r in results if r["status"] == "duplicate"), "results": results}
|
||||
|
||||
|
||||
# ─── 服务器路径导入 ───
|
||||
|
||||
class PathImportRequest(BaseModel):
|
||||
paths: List[str]
|
||||
recursive: bool = False
|
||||
|
||||
|
||||
@router.post("/import-paths", summary="从服务器路径批量导入图片")
|
||||
async def import_from_paths(data: PathImportRequest, db: AsyncSession = Depends(get_db)):
|
||||
all_paths = []
|
||||
for p in data.paths:
|
||||
if os.path.isfile(p):
|
||||
if os.path.splitext(p)[1].lower() in ALLOWED_EXTENSIONS:
|
||||
all_paths.append(p)
|
||||
elif os.path.isdir(p):
|
||||
walk = os.walk(p) if data.recursive else [(p, [], os.listdir(p))]
|
||||
for root, dirs, files in walk:
|
||||
for f in sorted(files):
|
||||
fp = os.path.join(root, f)
|
||||
if os.path.isfile(fp) and os.path.splitext(f)[1].lower() in ALLOWED_EXTENSIONS:
|
||||
all_paths.append(fp)
|
||||
if not all_paths:
|
||||
return {"total": 0, "success": 0, "failed": 0, "message": "未找到图片文件", "results": []}
|
||||
results = []
|
||||
for file_path in all_paths:
|
||||
image_record = OCRImage(file_path=file_path, status="processing")
|
||||
db.add(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
ocr_result = await OCRService.recognize(file_path)
|
||||
# OCR 内容去重
|
||||
existing_id = await _check_ocr_duplicate(ocr_result.text, db)
|
||||
if existing_id is not None:
|
||||
await db.delete(image_record)
|
||||
await db.flush()
|
||||
results.append({"id": image_record.id, "filename": os.path.basename(file_path), "file_path": file_path, "status": "duplicate", "duplicate_of": existing_id, "ocr_text": ocr_result.text})
|
||||
continue
|
||||
_save_ocr_result(image_record, ocr_result, db)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
_add_to_index(image_record)
|
||||
results.append({"id": image_record.id, "filename": os.path.basename(file_path), "file_path": file_path, "status": "completed", "ocr_text": ocr_result.text, "tags": json.loads(image_record.tags) if image_record.tags else [], "confidence": ocr_result.confidence, "provider": ocr_result.provider})
|
||||
except Exception as e:
|
||||
image_record.status = "failed"
|
||||
image_record.error_message = str(e)
|
||||
await db.flush()
|
||||
results.append({"id": image_record.id, "filename": os.path.basename(file_path), "file_path": file_path, "status": "failed", "message": str(e)})
|
||||
return {"total": len(results), "success": sum(1 for r in results if r["status"] == "completed"), "failed": sum(1 for r in results if r["status"] == "failed"), "results": results}
|
||||
|
||||
|
||||
# ─── 一步上传+识别 ───
|
||||
|
||||
@router.post("/recognize-direct", summary="直接上传并识别图片")
|
||||
async def recognize_direct(file: UploadFile, db: AsyncSession = Depends(get_db)):
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="文件名不能为空")
|
||||
suffix = _check_image_ext(file.filename)
|
||||
dest_path = settings.images_dir / file.filename
|
||||
counter = 1
|
||||
stem = os.path.splitext(file.filename)[0]
|
||||
while dest_path.exists():
|
||||
dest_path = settings.images_dir / f"{stem}_{counter}{suffix}"
|
||||
counter += 1
|
||||
content = await file.read()
|
||||
dest_path.write_bytes(content)
|
||||
image_record = OCRImage(file_path=str(dest_path), status="processing")
|
||||
db.add(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
ocr_result = await OCRService.recognize(str(dest_path))
|
||||
# OCR 内容去重
|
||||
existing_id = await _check_ocr_duplicate(ocr_result.text, db)
|
||||
if existing_id is not None:
|
||||
await db.delete(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
os.unlink(dest_path)
|
||||
except OSError:
|
||||
pass
|
||||
return {
|
||||
"id": image_record.id, "file_path": str(dest_path),
|
||||
"original_filename": file.filename,
|
||||
"status": "duplicate", "duplicate_of": existing_id,
|
||||
"ocr_text": ocr_result.text,
|
||||
}
|
||||
_save_ocr_result(image_record, ocr_result, db)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
_add_to_index(image_record)
|
||||
return {
|
||||
"id": image_record.id, "file_path": image_record.file_path,
|
||||
"original_filename": file.filename,
|
||||
"ocr_text": image_record.ocr_text,
|
||||
"tags": json.loads(image_record.tags) if image_record.tags else [],
|
||||
"confidence": image_record.confidence, "provider": image_record.provider,
|
||||
"status": image_record.status,
|
||||
"blocks": [{"text": b.text, "bbox": b.bbox, "confidence": b.confidence} for b in ocr_result.blocks],
|
||||
}
|
||||
except Exception as e:
|
||||
image_record.status = "failed"
|
||||
image_record.error_message = str(e)
|
||||
await db.flush()
|
||||
raise HTTPException(status_code=500, detail=f"OCR 识别失败: {str(e)}")
|
||||
|
||||
|
||||
# ─── LLM 精排搜索(SSE 流式) ───
|
||||
|
||||
@router.get("/search", summary="AI 搜索图片(流式)")
|
||||
async def search_images(
|
||||
keyword: str,
|
||||
limit: int = Query(5, ge=1, le=100),
|
||||
sort: str = Query("time_desc"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
AI 搜索图片(SSE 流式返回)。
|
||||
1. 从 DB 取所有图片,按指定排序
|
||||
2. 每批 5 条发给 Qwen3-8B 判断
|
||||
3. 并发池最多 10 个请求
|
||||
4. 匹配结果实时 SSE 推送
|
||||
5. 找够 limit 个或遍历完 DB 后结束
|
||||
"""
|
||||
|
||||
async def event_stream():
|
||||
found = 0
|
||||
offset = 0
|
||||
batch_size = settings.JUDGE_BATCH_SIZE
|
||||
max_concurrent = 10
|
||||
semaphore = asyncio.Semaphore(max_concurrent)
|
||||
|
||||
# 发送开始事件
|
||||
yield f"data: {json.dumps({'type': 'start', 'keyword': keyword, 'limit': limit}, ensure_ascii=False)}\n\n"
|
||||
|
||||
# 先获取总记录数
|
||||
count_result = await db.execute(
|
||||
select(func.count(OCRImage.id)).where(OCRImage.status == "completed")
|
||||
)
|
||||
total_records = count_result.scalar() or 0
|
||||
|
||||
if total_records == 0:
|
||||
yield f"data: {json.dumps({'type': 'done', 'total_found': 0, 'total_checked': 0}, ensure_ascii=False)}\n\n"
|
||||
return
|
||||
|
||||
# 排序
|
||||
order_clause = OCRImage.created_at.desc()
|
||||
if sort == "time_asc":
|
||||
order_clause = OCRImage.created_at.asc()
|
||||
elif sort == "random":
|
||||
order_clause = text("RANDOM()")
|
||||
|
||||
# 用于收集结果的队列
|
||||
result_queue = asyncio.Queue()
|
||||
checked_count = 0
|
||||
|
||||
async def process_batch(batch_items: list, batch_offset: int):
|
||||
"""处理一批图片:发给 LLM 判断"""
|
||||
nonlocal checked_count
|
||||
async with semaphore:
|
||||
articles = [{"id": item.id, "text": item.ocr_text or ""} for item in batch_items]
|
||||
matched_ids = await LLMService.judge_batch(keyword, articles)
|
||||
|
||||
checked_count += len(batch_items)
|
||||
|
||||
# 发送进度
|
||||
progress_data = {
|
||||
"type": "progress",
|
||||
"checked": checked_count,
|
||||
"total": total_records,
|
||||
"found": found,
|
||||
}
|
||||
await result_queue.put(progress_data)
|
||||
|
||||
# 如果有匹配且未达上限,发送结果
|
||||
for match_id in matched_ids:
|
||||
if found >= limit:
|
||||
break
|
||||
item = next((i for i in batch_items if i.id == match_id), None)
|
||||
if item:
|
||||
file_name = item.file_path.split('/')[-1] if item.file_path else ''
|
||||
# 读取图片并转为 base64(用于发送到企业微信等)
|
||||
image_base64 = ""
|
||||
if file_name:
|
||||
img_path = settings.images_dir / file_name
|
||||
if img_path.exists():
|
||||
image_base64 = base64.b64encode(img_path.read_bytes()).decode()
|
||||
result_data = {
|
||||
"type": "result",
|
||||
"id": item.id,
|
||||
"file_path": item.file_path,
|
||||
"image_url": file_name,
|
||||
"image_base64": image_base64,
|
||||
"ocr_text": item.ocr_text,
|
||||
"ocr_text_preview": (item.ocr_text[:150] + "...") if item.ocr_text and len(item.ocr_text) > 150 else item.ocr_text,
|
||||
"tags": json.loads(item.tags) if item.tags else [],
|
||||
"confidence": item.confidence,
|
||||
"provider": item.provider,
|
||||
"created_at": str(item.created_at),
|
||||
}
|
||||
await result_queue.put(result_data)
|
||||
|
||||
# 启动所有批次的异步任务
|
||||
tasks = []
|
||||
while offset < total_records:
|
||||
query = select(OCRImage).where(OCRImage.status == "completed").order_by(order_clause).offset(offset).limit(batch_size)
|
||||
result = await db.execute(query)
|
||||
batch_items = result.scalars().all()
|
||||
|
||||
if not batch_items:
|
||||
break
|
||||
|
||||
tasks.append(asyncio.create_task(process_batch(batch_items, offset)))
|
||||
offset += batch_size
|
||||
|
||||
# 等待所有任务完成
|
||||
async def wait_all():
|
||||
await asyncio.gather(*tasks)
|
||||
await result_queue.put({"type": "all_done"})
|
||||
|
||||
asyncio.create_task(wait_all())
|
||||
|
||||
# 从队列读取结果并 yield
|
||||
while True:
|
||||
try:
|
||||
data = await asyncio.wait_for(result_queue.get(), timeout=300)
|
||||
except asyncio.TimeoutError:
|
||||
break
|
||||
|
||||
if data.get("type") == "all_done":
|
||||
break
|
||||
|
||||
if data.get("type") == "result":
|
||||
found += 1
|
||||
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
||||
if found >= limit:
|
||||
# 找够了,发送结束事件
|
||||
yield f"data: {json.dumps({'type': 'done', 'total_found': found, 'total_checked': checked_count}, ensure_ascii=False)}\n\n"
|
||||
return
|
||||
elif data.get("type") == "progress":
|
||||
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
||||
|
||||
# 发送结束事件
|
||||
yield f"data: {json.dumps({'type': 'done', 'total_found': found, 'total_checked': checked_count}, ensure_ascii=False)}\n\n"
|
||||
|
||||
return StreamingResponse(
|
||||
event_stream(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ─── 删除单张图片 ───
|
||||
|
||||
@router.delete("/{image_id}", summary="删除单张图片", status_code=204)
|
||||
async def delete_image(image_id: int, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(OCRImage).where(OCRImage.id == image_id))
|
||||
image_record = result.scalar_one_or_none()
|
||||
if image_record is None:
|
||||
raise HTTPException(status_code=404, detail="图片记录不存在")
|
||||
# 从 BM25 索引移除
|
||||
if bm25_index is not None:
|
||||
bm25_index.remove_document(image_record.id)
|
||||
# 删除磁盘上的图片文件
|
||||
try:
|
||||
os.unlink(image_record.file_path)
|
||||
except OSError:
|
||||
pass
|
||||
# 删除数据库记录
|
||||
await db.delete(image_record)
|
||||
await db.flush()
|
||||
return Response(status_code=204)
|
||||
|
||||
|
||||
# ─── 其他接口 ───
|
||||
|
||||
@router.get("/{image_id}", summary="获取图片识别结果")
|
||||
async def get_image_result(image_id: int, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(OCRImage).where(OCRImage.id == image_id))
|
||||
image_record = result.scalar_one_or_none()
|
||||
if image_record is None:
|
||||
raise HTTPException(status_code=404, detail="图片记录不存在")
|
||||
return image_record.to_dict()
|
||||
|
||||
|
||||
@router.get("", summary="获取图片列表")
|
||||
async def list_images(page: int = 1, page_size: int = 20, status: Optional[str] = None, db: AsyncSession = Depends(get_db)):
|
||||
query = select(OCRImage)
|
||||
count_query = select(func.count(OCRImage.id))
|
||||
if status:
|
||||
query = query.where(OCRImage.status == status)
|
||||
count_query = count_query.where(OCRImage.status == status)
|
||||
total_result = await db.execute(count_query)
|
||||
total = total_result.scalar() or 0
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(OCRImage.created_at.desc()).offset(offset).limit(page_size)
|
||||
result = await db.execute(query)
|
||||
items = result.scalars().all()
|
||||
return {"total": total, "page": page, "page_size": page_size, "items": [item.to_dict() for item in items]}
|
||||
|
||||
|
||||
# ─── 一键去重 ───
|
||||
|
||||
@router.post("/dedup", summary="一键去重:删除 OCR 内容重复的图片")
|
||||
async def dedup_images(db: AsyncSession = Depends(get_db)):
|
||||
# 查询所有 status='completed' 且 ocr_text 不为空的记录
|
||||
query = select(OCRImage).where(
|
||||
OCRImage.status == "completed",
|
||||
OCRImage.ocr_text.isnot(None),
|
||||
OCRImage.ocr_text != "",
|
||||
).order_by(OCRImage.created_at.asc())
|
||||
result = await db.execute(query)
|
||||
all_records = result.scalars().all()
|
||||
|
||||
total_checked = len(all_records)
|
||||
# 按 ocr_text 分组
|
||||
groups: dict[str, list[OCRImage]] = defaultdict(list)
|
||||
for record in all_records:
|
||||
groups[record.ocr_text].append(record)
|
||||
|
||||
duplicates_found = 0
|
||||
deleted = 0
|
||||
kept = 0
|
||||
|
||||
for ocr_text, group in groups.items():
|
||||
if len(group) <= 1:
|
||||
kept += 1
|
||||
continue
|
||||
# 保留 created_at 最早的一条(已按 asc 排序,第一条即最早)
|
||||
duplicates_found += len(group) - 1
|
||||
kept += 1
|
||||
for record in group[1:]:
|
||||
# 从 BM25 索引移除
|
||||
if bm25_index is not None:
|
||||
bm25_index.remove_document(record.id)
|
||||
# 删除磁盘上的图片文件
|
||||
try:
|
||||
os.unlink(record.file_path)
|
||||
except OSError:
|
||||
pass
|
||||
# 删除数据库记录
|
||||
await db.delete(record)
|
||||
deleted += 1
|
||||
|
||||
await db.flush()
|
||||
return {
|
||||
"total_checked": total_checked,
|
||||
"duplicates_found": duplicates_found,
|
||||
"deleted": deleted,
|
||||
"kept": kept,
|
||||
}
|
||||
215
app/api/v1/import_export.py
Normal file
215
app/api/v1/import_export.py
Normal file
@@ -0,0 +1,215 @@
|
||||
"""
|
||||
导入导出 API 路由
|
||||
提供文件导入和知识库导出接口
|
||||
支持 .md / .txt / .docx 格式
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.config import settings
|
||||
from app.database import get_db
|
||||
from app.services.import_service import ImportService
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# 支持的导入文件格式
|
||||
SUPPORTED_EXTENSIONS = (".md", ".txt", ".docx")
|
||||
|
||||
|
||||
@router.post("/file", summary="从文件导入知识")
|
||||
async def import_from_file(
|
||||
file: UploadFile = File(..., description="上传文件(.md / .txt / .docx)"),
|
||||
course_name: Optional[str] = Form(None, description="课程名称"),
|
||||
teacher_name: Optional[str] = Form(None, description="讲师名称"),
|
||||
live_date: Optional[str] = Form(None, description="直播日期 (YYYY-MM-DD)"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
上传文件并导入到知识库。
|
||||
|
||||
- Markdown (.md):解析 frontmatter 元数据,按 ## 标题分页
|
||||
- 纯文本 (.txt):整体作为一个页面
|
||||
- Word 文档 (.docx):提取文本,按标题分页
|
||||
- 导入后自动分块、生成嵌入向量、建立全文索引
|
||||
"""
|
||||
# 验证文件类型
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="文件名不能为空")
|
||||
|
||||
suffix = os.path.splitext(file.filename)[1].lower()
|
||||
if suffix not in SUPPORTED_EXTENSIONS:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"仅支持 {', '.join(SUPPORTED_EXTENSIONS)} 文件",
|
||||
)
|
||||
|
||||
# 保存上传文件到 transcripts 目录
|
||||
dest_dir = settings.transcripts_dir
|
||||
dest_path = dest_dir / file.filename
|
||||
|
||||
# 如果文件已存在,添加序号
|
||||
counter = 1
|
||||
original_stem = os.path.splitext(file.filename)[0]
|
||||
while dest_path.exists():
|
||||
dest_path = dest_dir / f"{original_stem}_{counter}{suffix}"
|
||||
counter += 1
|
||||
|
||||
# 写入文件
|
||||
content = await file.read()
|
||||
dest_path.write_bytes(content)
|
||||
|
||||
# 执行导入
|
||||
service = ImportService(db)
|
||||
try:
|
||||
result = await service.import_file(
|
||||
file_path=str(dest_path),
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
live_date=live_date,
|
||||
)
|
||||
return {
|
||||
"message": "导入成功",
|
||||
"file": str(dest_path),
|
||||
**result,
|
||||
}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"导入失败: {str(e)}")
|
||||
|
||||
|
||||
@router.post("/directory", summary="从目录批量导入")
|
||||
async def import_from_directory(
|
||||
directory: str = Form(..., description="数据目录中的子目录名(相对于 transcripts 目录)"),
|
||||
course_name: Optional[str] = Form(None, description="课程名称"),
|
||||
teacher_name: Optional[str] = Form(None, description="讲师名称"),
|
||||
live_date: Optional[str] = Form(None, description="直播日期"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
批量导入指定目录下的所有文件(.md / .txt / .docx)。
|
||||
"""
|
||||
target_dir = settings.transcripts_dir / directory
|
||||
if not target_dir.exists():
|
||||
raise HTTPException(status_code=404, detail=f"目录不存在: {directory}")
|
||||
|
||||
service = ImportService(db)
|
||||
total_pages = 0
|
||||
total_chunks = 0
|
||||
files_processed = 0
|
||||
errors: list[str] = []
|
||||
|
||||
for filepath in sorted(target_dir.iterdir()):
|
||||
if filepath.suffix.lower() not in SUPPORTED_EXTENSIONS:
|
||||
continue
|
||||
|
||||
try:
|
||||
result = await service.import_file(
|
||||
file_path=str(filepath),
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
live_date=live_date,
|
||||
)
|
||||
total_pages += result["pages"]
|
||||
total_chunks += result["chunks"]
|
||||
files_processed += 1
|
||||
except Exception as e:
|
||||
errors.append(f"{filepath.name}: {str(e)}")
|
||||
|
||||
return {
|
||||
"message": "批量导入完成",
|
||||
"files_processed": files_processed,
|
||||
"total_pages": total_pages,
|
||||
"total_chunks": total_chunks,
|
||||
"errors": errors if errors else None,
|
||||
}
|
||||
|
||||
|
||||
@router.get("/export", summary="导出知识库")
|
||||
async def export_knowledge_base(
|
||||
format: str = "json",
|
||||
course_name: Optional[str] = None,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
导出知识库内容。
|
||||
|
||||
Args:
|
||||
format: 导出格式,支持 json / markdown
|
||||
course_name: 按课程名称过滤导出
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
query = text("""
|
||||
SELECT id, title, content, source_file, course_name, teacher_name,
|
||||
live_date, page_number, metadata_json, created_at
|
||||
FROM knowledge_pages
|
||||
WHERE (:course_name IS NULL OR course_name = :course_name)
|
||||
ORDER BY created_at DESC
|
||||
""")
|
||||
result = await db.execute(query, {"course_name": course_name})
|
||||
pages = result.fetchall()
|
||||
|
||||
if format == "markdown":
|
||||
# 导出为 Markdown 文件
|
||||
lines = []
|
||||
for row in pages:
|
||||
lines.append(f"# {row.title}")
|
||||
lines.append("")
|
||||
meta_parts = []
|
||||
if row.course_name:
|
||||
meta_parts.append(f"课程: {row.course_name}")
|
||||
if row.teacher_name:
|
||||
meta_parts.append(f"讲师: {row.teacher_name}")
|
||||
if row.live_date:
|
||||
meta_parts.append(f"日期: {row.live_date}")
|
||||
if meta_parts:
|
||||
lines.append(" | ".join(meta_parts))
|
||||
lines.append("")
|
||||
lines.append(row.content)
|
||||
lines.append("")
|
||||
lines.append("---")
|
||||
lines.append("")
|
||||
|
||||
export_content = "\n".join(lines)
|
||||
ext = "md"
|
||||
else:
|
||||
# 导出为 JSON
|
||||
pages_data = []
|
||||
for row in pages:
|
||||
pages_data.append({
|
||||
"id": row.id,
|
||||
"title": row.title,
|
||||
"content": row.content,
|
||||
"source_file": row.source_file,
|
||||
"course_name": row.course_name,
|
||||
"teacher_name": row.teacher_name,
|
||||
"live_date": row.live_date,
|
||||
"page_number": row.page_number,
|
||||
"metadata_json": row.metadata_json,
|
||||
"created_at": str(row.created_at),
|
||||
})
|
||||
export_content = json.dumps(pages_data, ensure_ascii=False, indent=2)
|
||||
ext = "json"
|
||||
|
||||
# 保存到 exports 目录
|
||||
export_dir = settings.exports_dir
|
||||
export_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"export_{timestamp}.{ext}"
|
||||
export_path = export_dir / filename
|
||||
export_path.write_text(export_content, encoding="utf-8")
|
||||
|
||||
return {
|
||||
"message": "导出成功",
|
||||
"file": filename,
|
||||
"pages": len(pages),
|
||||
"format": format,
|
||||
}
|
||||
112
app/api/v1/pages.py
Normal file
112
app/api/v1/pages.py
Normal file
@@ -0,0 +1,112 @@
|
||||
"""
|
||||
知识页面 API 路由
|
||||
提供知识页面的 CRUD 接口
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.database import get_db
|
||||
from app.schemas.page import PageCreate, PageDetailResponse, PageListResponse, PageResponse, PageUpdate
|
||||
from app.services.page_service import PageService
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("", response_model=PageListResponse, summary="获取知识页面列表")
|
||||
async def list_pages(
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=100, description="每页数量"),
|
||||
course_name: Optional[str] = Query(None, description="按课程名称过滤"),
|
||||
teacher_name: Optional[str] = Query(None, description="按讲师名称过滤"),
|
||||
keyword: Optional[str] = Query(None, description="关键词搜索"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""分页查询知识页面列表,支持按课程/讲师/关键词过滤"""
|
||||
service = PageService(db)
|
||||
result = await service.list_pages(
|
||||
page=page,
|
||||
page_size=page_size,
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
keyword=keyword,
|
||||
)
|
||||
return PageListResponse(
|
||||
total=result["total"],
|
||||
page=result["page"],
|
||||
page_size=result["page_size"],
|
||||
items=[PageResponse.model_validate(p) for p in result["items"]],
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{page_id}", response_model=PageDetailResponse, summary="获取知识页面详情")
|
||||
async def get_page(
|
||||
page_id: int,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""根据 ID 获取知识页面详情,包含关联的分块信息"""
|
||||
service = PageService(db)
|
||||
page = await service.get_page(page_id)
|
||||
if page is None:
|
||||
raise HTTPException(status_code=404, detail="知识页面不存在")
|
||||
|
||||
chunks = await service.get_page_chunks(page_id)
|
||||
return PageDetailResponse(
|
||||
**page.to_dict(),
|
||||
chunks=[c.to_dict() for c in chunks],
|
||||
)
|
||||
|
||||
|
||||
@router.post("", response_model=PageResponse, status_code=201, summary="创建知识页面")
|
||||
async def create_page(
|
||||
data: PageCreate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""创建新的知识页面"""
|
||||
service = PageService(db)
|
||||
page = await service.create_page(data)
|
||||
return PageResponse.model_validate(page)
|
||||
|
||||
|
||||
@router.put("/{page_id}", response_model=PageResponse, summary="更新知识页面")
|
||||
async def update_page(
|
||||
page_id: int,
|
||||
data: PageUpdate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""更新知识页面信息"""
|
||||
service = PageService(db)
|
||||
page = await service.update_page(page_id, data)
|
||||
if page is None:
|
||||
raise HTTPException(status_code=404, detail="知识页面不存在")
|
||||
return PageResponse.model_validate(page)
|
||||
|
||||
|
||||
@router.delete("/{page_id}", status_code=204, summary="删除知识页面")
|
||||
async def delete_page(
|
||||
page_id: int,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""删除知识页面及其关联的所有分块"""
|
||||
service = PageService(db)
|
||||
success = await service.delete_page(page_id)
|
||||
if not success:
|
||||
raise HTTPException(status_code=404, detail="知识页面不存在")
|
||||
|
||||
|
||||
@router.post("/{page_id}/reindex", summary="重新索引知识页面")
|
||||
async def reindex_page(
|
||||
page_id: int,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""重新对知识页面进行分块和向量化(页面内容更新后使用)"""
|
||||
service = PageService(db)
|
||||
try:
|
||||
result = await service.reindex_page(page_id)
|
||||
return result
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
30
app/api/v1/search.py
Normal file
30
app/api/v1/search.py
Normal file
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
语义搜索 API 路由
|
||||
提供向量搜索 + 全文搜索混合查询接口
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.database import get_db
|
||||
from app.schemas.search import SearchRequest, SearchResponse
|
||||
from app.services.search_service import SearchService
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("", response_model=SearchResponse, summary="语义搜索")
|
||||
async def semantic_search(
|
||||
request: SearchRequest,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
对知识库进行语义搜索。
|
||||
|
||||
支持向量相似度搜索和中文全文搜索的混合排序。
|
||||
可按课程名称、讲师名称、直播日期范围过滤结果。
|
||||
"""
|
||||
service = SearchService(db)
|
||||
return await service.search(request)
|
||||
223
app/api/v1/settings.py
Normal file
223
app/api/v1/settings.py
Normal file
@@ -0,0 +1,223 @@
|
||||
"""
|
||||
系统设置 API 路由
|
||||
提供运行时配置的读取和修改接口
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class SettingsResponse(BaseModel):
|
||||
"""设置响应(隐藏敏感信息的中间位)"""
|
||||
embedding_provider: str
|
||||
embedding_model: str
|
||||
embedding_dimensions: int
|
||||
ocr_provider: str
|
||||
llm_provider: str
|
||||
minimax_base_url: str
|
||||
minimax_embedding_model: str
|
||||
minimax_chat_model: str
|
||||
deepseek_base_url: str
|
||||
deepseek_ocr_model: str
|
||||
openai_base_url: str | None
|
||||
chunk_size: int
|
||||
chunk_overlap: int
|
||||
search_limit: int
|
||||
judge_batch_size: int
|
||||
# API Key 只返回是否已配置(不返回实际值)
|
||||
has_minimax_key: bool
|
||||
has_deepseek_key: bool
|
||||
has_openai_key: bool
|
||||
has_zhipu_key: bool
|
||||
has_dashscope_key: bool
|
||||
has_aliyun_ocr: bool
|
||||
has_tencent_ocr: bool
|
||||
|
||||
|
||||
@router.get("", response_model=SettingsResponse, summary="获取当前设置")
|
||||
async def get_settings():
|
||||
"""获取当前系统设置(API Key 脱敏)"""
|
||||
return SettingsResponse(
|
||||
embedding_provider=settings.EMBEDDING_PROVIDER.value,
|
||||
embedding_model=settings.EMBEDDING_MODEL,
|
||||
embedding_dimensions=settings.EMBEDDING_DIMENSIONS,
|
||||
ocr_provider=settings.OCR_PROVIDER.value,
|
||||
llm_provider=settings.LLM_PROVIDER,
|
||||
minimax_base_url=settings.MINIMAX_BASE_URL,
|
||||
minimax_embedding_model=settings.MINIMAX_EMBEDDING_MODEL,
|
||||
minimax_chat_model=settings.MINIMAX_CHAT_MODEL,
|
||||
deepseek_base_url=settings.DEEPSEEK_BASE_URL,
|
||||
deepseek_ocr_model=settings.DEEPSEEK_OCR_MODEL,
|
||||
openai_base_url=settings.OPENAI_BASE_URL,
|
||||
chunk_size=settings.CHUNK_SIZE,
|
||||
chunk_overlap=settings.CHUNK_OVERLAP,
|
||||
search_limit=settings.SEARCH_LIMIT,
|
||||
judge_batch_size=settings.JUDGE_BATCH_SIZE,
|
||||
has_minimax_key=bool(settings.MINIMAX_API_KEY),
|
||||
has_deepseek_key=bool(settings.DEEPSEEK_API_KEY),
|
||||
has_openai_key=bool(settings.OPENAI_API_KEY),
|
||||
has_zhipu_key=bool(settings.ZHIPU_API_KEY),
|
||||
has_dashscope_key=bool(settings.DASHSCOPE_API_KEY),
|
||||
has_aliyun_ocr=bool(settings.ALIYUN_OCR_ACCESS_KEY),
|
||||
has_tencent_ocr=bool(settings.TENCENT_OCR_SECRET_ID),
|
||||
)
|
||||
|
||||
|
||||
class SettingsUpdate(BaseModel):
|
||||
"""设置更新请求(所有字段可选)"""
|
||||
embedding_provider: str | None = None
|
||||
embedding_model: str | None = None
|
||||
embedding_dimensions: int | None = None
|
||||
ocr_provider: str | None = None
|
||||
llm_provider: str | None = None
|
||||
minimax_api_key: str | None = None
|
||||
minimax_base_url: str | None = None
|
||||
minimax_embedding_model: str | None = None
|
||||
minimax_chat_model: str | None = None
|
||||
deepseek_api_key: str | None = None
|
||||
deepseek_base_url: str | None = None
|
||||
deepseek_ocr_model: str | None = None
|
||||
openai_api_key: str | None = None
|
||||
openai_base_url: str | None = None
|
||||
zhipu_api_key: str | None = None
|
||||
dashscope_api_key: str | None = None
|
||||
aliyun_ocr_access_key: str | None = None
|
||||
aliyun_ocr_secret: str | None = None
|
||||
tencent_ocr_secret_id: str | None = None
|
||||
tencent_ocr_secret_key: str | None = None
|
||||
chunk_size: int | None = None
|
||||
chunk_overlap: int | None = None
|
||||
search_limit: int | None = None
|
||||
judge_batch_size: int | None = None
|
||||
|
||||
|
||||
@router.put("", summary="更新设置")
|
||||
async def update_settings(data: SettingsUpdate):
|
||||
"""
|
||||
更新系统设置(运行时生效)。
|
||||
修改嵌入模型或 OCR 提供商后,对应的单例服务会自动重置。
|
||||
"""
|
||||
from app.config import EmbeddingProvider, OCRProvider
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
from app.services.ocr_service import OCRService
|
||||
from app.services.llm_service import LLMService
|
||||
|
||||
update_data = data.model_dump(exclude_none=True)
|
||||
if not update_data:
|
||||
raise HTTPException(status_code=400, detail="没有需要更新的字段")
|
||||
|
||||
provider_changed = False
|
||||
|
||||
# 映射字段名(API 用 snake_case,Settings 用 UPPER_CASE)
|
||||
field_map = {
|
||||
"embedding_provider": "EMBEDDING_PROVIDER",
|
||||
"embedding_model": "EMBEDDING_MODEL",
|
||||
"embedding_dimensions": "EMBEDDING_DIMENSIONS",
|
||||
"ocr_provider": "OCR_PROVIDER",
|
||||
"llm_provider": "LLM_PROVIDER",
|
||||
"minimax_api_key": "MINIMAX_API_KEY",
|
||||
"minimax_base_url": "MINIMAX_BASE_URL",
|
||||
"minimax_embedding_model": "MINIMAX_EMBEDDING_MODEL",
|
||||
"minimax_chat_model": "MINIMAX_CHAT_MODEL",
|
||||
"deepseek_api_key": "DEEPSEEK_API_KEY",
|
||||
"deepseek_base_url": "DEEPSEEK_BASE_URL",
|
||||
"deepseek_ocr_model": "DEEPSEEK_OCR_MODEL",
|
||||
"openai_api_key": "OPENAI_API_KEY",
|
||||
"openai_base_url": "OPENAI_BASE_URL",
|
||||
"zhipu_api_key": "ZHIPU_API_KEY",
|
||||
"dashscope_api_key": "DASHSCOPE_API_KEY",
|
||||
"aliyun_ocr_access_key": "ALIYUN_OCR_ACCESS_KEY",
|
||||
"aliyun_ocr_secret": "ALIYUN_OCR_SECRET",
|
||||
"tencent_ocr_secret_id": "TENCENT_OCR_SECRET_ID",
|
||||
"tencent_ocr_secret_key": "TENCENT_OCR_SECRET_KEY",
|
||||
"chunk_size": "CHUNK_SIZE",
|
||||
"chunk_overlap": "CHUNK_OVERLAP",
|
||||
"search_limit": "SEARCH_LIMIT",
|
||||
"judge_batch_size": "JUDGE_BATCH_SIZE",
|
||||
}
|
||||
|
||||
for api_field, config_field in field_map.items():
|
||||
if api_field in update_data:
|
||||
value = update_data[api_field]
|
||||
# 枚举类型需要转换
|
||||
if config_field == "EMBEDDING_PROVIDER":
|
||||
value = EmbeddingProvider(value)
|
||||
provider_changed = True
|
||||
elif config_field == "OCR_PROVIDER":
|
||||
value = OCRProvider(value)
|
||||
provider_changed = True
|
||||
setattr(settings, config_field, value)
|
||||
|
||||
# 如果提供商变了,重置单例
|
||||
if provider_changed:
|
||||
EmbeddingService.reset()
|
||||
OCRService.reset()
|
||||
LLMService.reset()
|
||||
|
||||
return {"message": "设置已更新", "provider_changed": provider_changed}
|
||||
|
||||
|
||||
@router.get("/stats", summary="获取知识库统计")
|
||||
async def get_stats():
|
||||
"""获取知识库统计信息"""
|
||||
from sqlalchemy import text
|
||||
from app.database import async_session_factory
|
||||
|
||||
async with async_session_factory() as session:
|
||||
result = await session.execute(text("""
|
||||
SELECT
|
||||
(SELECT COUNT(*) FROM knowledge_pages) AS page_count,
|
||||
(SELECT COUNT(*) FROM knowledge_chunks) AS chunk_count,
|
||||
(SELECT COUNT(*) FROM knowledge_chunks WHERE embedding IS NOT NULL) AS embedded_count,
|
||||
(SELECT COUNT(*) FROM ocr_images) AS image_count,
|
||||
(SELECT COUNT(*) FROM ocr_images WHERE status = 'completed') AS ocr_completed_count
|
||||
"""))
|
||||
row = result.fetchone()
|
||||
|
||||
return {
|
||||
"page_count": row.page_count,
|
||||
"chunk_count": row.chunk_count,
|
||||
"embedded_count": row.embedded_count,
|
||||
"image_count": row.image_count,
|
||||
"ocr_completed_count": row.ocr_completed_count,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/test/embedding", summary="测试嵌入模型连接")
|
||||
async def test_embedding():
|
||||
"""测试当前嵌入模型是否可用"""
|
||||
try:
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
result = await EmbeddingService.embed_single("测试")
|
||||
return {"success": True, "dimension": len(result), "message": f"嵌入模型正常,维度: {len(result)}"}
|
||||
except Exception as e:
|
||||
return {"success": False, "message": f"嵌入模型测试失败: {str(e)}"}
|
||||
|
||||
|
||||
@router.post("/test/llm", summary="测试 LLM 连接")
|
||||
async def test_llm():
|
||||
"""测试当前 LLM 是否可用"""
|
||||
try:
|
||||
from app.services.llm_service import LLMService
|
||||
result = await LLMService.chat([
|
||||
{"role": "user", "content": "请回复\"连接正常\""}
|
||||
], max_tokens=20)
|
||||
return {"success": True, "message": f"LLM 正常,回复: {result[:100]}"}
|
||||
except Exception as e:
|
||||
return {"success": False, "message": f"LLM 测试失败: {str(e)}"}
|
||||
|
||||
|
||||
@router.post("/test/ocr", summary="测试 OCR(使用示例文本)")
|
||||
async def test_ocr():
|
||||
"""测试当前 OCR 提供商是否可用(不实际上传图片,只检查配置)"""
|
||||
try:
|
||||
from app.services.ocr_service import OCRService
|
||||
provider = OCRService.get_instance()
|
||||
return {"success": True, "message": f"OCR 提供商 {provider.name} 已就绪"}
|
||||
except Exception as e:
|
||||
return {"success": False, "message": f"OCR 测试失败: {str(e)}"}
|
||||
216
app/config.py
Normal file
216
app/config.py
Normal file
@@ -0,0 +1,216 @@
|
||||
"""
|
||||
配置管理模块
|
||||
使用 pydantic-settings 从环境变量 / .env 文件加载配置
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class EmbeddingProvider(str, Enum):
|
||||
"""嵌入模型提供商"""
|
||||
OPENAI = "openai"
|
||||
ZHIPU = "zhipu"
|
||||
LOCAL_BGE = "local_bge"
|
||||
DASHSCOPE = "dashscope"
|
||||
MINIMAX = "minimax"
|
||||
|
||||
|
||||
class OCRProvider(str, Enum):
|
||||
"""OCR 识别提供商"""
|
||||
PADDLEOCR = "paddleocr"
|
||||
ALIYUN = "aliyun"
|
||||
TENCENT = "tencent"
|
||||
AUTO = "auto"
|
||||
DEEPSEEK = "deepseek"
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
"""
|
||||
全局配置类,所有配置项均可通过环境变量覆盖。
|
||||
环境变量优先级:系统环境变量 > .env 文件 > 默认值
|
||||
"""
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=".env",
|
||||
env_file_encoding="utf-8",
|
||||
case_sensitive=False,
|
||||
extra="ignore",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 数据库 ────────────────────────────
|
||||
DATABASE_URL: str = Field(
|
||||
default="postgresql+asyncpg://postgres:postgres@db:5432/edu_brain",
|
||||
description="PostgreSQL 异步连接字符串",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 嵌入模型 ────────────────────────────
|
||||
EMBEDDING_PROVIDER: EmbeddingProvider = Field(
|
||||
default=EmbeddingProvider.LOCAL_BGE,
|
||||
description="嵌入模型提供商: openai / zhipu / local_bge / dashscope",
|
||||
)
|
||||
EMBEDDING_MODEL: str = Field(
|
||||
default="BAAI/bge-large-zh-v1.5",
|
||||
description="嵌入模型名称",
|
||||
)
|
||||
EMBEDDING_DIMENSIONS: int = Field(
|
||||
default=1024,
|
||||
description="嵌入向量维度",
|
||||
)
|
||||
|
||||
# ── OpenAI ──
|
||||
OPENAI_API_KEY: Optional[str] = Field(default=None, description="OpenAI API Key")
|
||||
OPENAI_BASE_URL: Optional[str] = Field(
|
||||
default=None,
|
||||
description="OpenAI 兼容接口地址(可用于代理或第三方兼容服务)",
|
||||
)
|
||||
|
||||
# ── 智谱 AI ──
|
||||
ZHIPU_API_KEY: Optional[str] = Field(default=None, description="智谱 AI API Key")
|
||||
|
||||
# ── 阿里云 DashScope ──
|
||||
DASHSCOPE_API_KEY: Optional[str] = Field(
|
||||
default=None, description="阿里云 DashScope API Key"
|
||||
)
|
||||
|
||||
# ── 本地 BGE 模型 ──
|
||||
LOCAL_BGE_MODEL_PATH: Optional[str] = Field(
|
||||
default=None,
|
||||
description="本地 BGE 模型路径,为空则自动从 HuggingFace 下载",
|
||||
)
|
||||
|
||||
# ──────────────────────────── OCR ────────────────────────────
|
||||
OCR_PROVIDER: OCRProvider = Field(
|
||||
default=OCRProvider.AUTO,
|
||||
description="OCR 提供商: paddleocr / aliyun / tencent / auto",
|
||||
)
|
||||
|
||||
# ── 阿里云 OCR ──
|
||||
ALIYUN_OCR_ACCESS_KEY: Optional[str] = Field(
|
||||
default=None, description="阿里云 OCR AccessKey ID"
|
||||
)
|
||||
ALIYUN_OCR_SECRET: Optional[str] = Field(
|
||||
default=None, description="阿里云 OCR AccessKey Secret"
|
||||
)
|
||||
|
||||
# ── 腾讯云 OCR ──
|
||||
TENCENT_OCR_SECRET_ID: Optional[str] = Field(
|
||||
default=None, description="腾讯云 OCR SecretId"
|
||||
)
|
||||
TENCENT_OCR_SECRET_KEY: Optional[str] = Field(
|
||||
default=None, description="腾讯云 OCR SecretKey"
|
||||
)
|
||||
|
||||
# ── MiniMax ──
|
||||
MINIMAX_API_KEY: Optional[str] = Field(default=None, description="MiniMax API Key")
|
||||
MINIMAX_BASE_URL: str = Field(
|
||||
default="https://api.minimaxi.com/v1",
|
||||
description="MiniMax API 地址",
|
||||
)
|
||||
MINIMAX_EMBEDDING_MODEL: str = Field(
|
||||
default="MiniMax-Text-Embedding",
|
||||
description="MiniMax 嵌入模型名称",
|
||||
)
|
||||
MINIMAX_CHAT_MODEL: str = Field(
|
||||
default="MiniMax-M2",
|
||||
description="MiniMax Chat 模型名称",
|
||||
)
|
||||
|
||||
# ── DeepSeek ──
|
||||
DEEPSEEK_API_KEY: Optional[str] = Field(default=None, description="DeepSeek API Key")
|
||||
DEEPSEEK_BASE_URL: str = Field(
|
||||
default="https://api.deepseek.com",
|
||||
description="DeepSeek API 地址(自部署时改为你的服务器地址)",
|
||||
)
|
||||
DEEPSEEK_OCR_MODEL: str = Field(
|
||||
default="deepseek-chat",
|
||||
description="DeepSeek OCR 使用的模型名称",
|
||||
)
|
||||
|
||||
# ── LLM(用于查询扩展、实体检测等) ──
|
||||
LLM_PROVIDER: str = Field(
|
||||
default="minimax",
|
||||
description="LLM 提供商: minimax / deepseek / openai",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 数据目录 ────────────────────────────
|
||||
DATA_DIR: str = Field(
|
||||
default="/app/data",
|
||||
description="数据根目录( transcripts / images / exports 的父目录)",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 文本分块 ────────────────────────────
|
||||
CHUNK_SIZE: int = Field(default=300, description="文本分块大小(字符数)")
|
||||
CHUNK_OVERLAP: int = Field(default=50, description="分块重叠字符数")
|
||||
|
||||
# ──────────────────────────── 图片搜索 ────────────────────────────
|
||||
SEARCH_LIMIT: int = Field(default=3, description="图片搜索默认返回数量(1-10)")
|
||||
JUDGE_BATCH_SIZE: int = Field(default=10, description="每次LLM判断的文章数量(1-50)")
|
||||
|
||||
# ──────────────────────────── 服务端口 ────────────────────────────
|
||||
HOST: str = Field(default="0.0.0.0", description="服务监听地址")
|
||||
PORT: int = Field(default=8000, description="服务监听端口")
|
||||
DEBUG: bool = Field(default=False, description="调试模式")
|
||||
|
||||
# ──────────────────────────── CORS ────────────────────────────
|
||||
CORS_ORIGINS: list[str] = Field(
|
||||
default=["*"],
|
||||
description="允许的跨域来源列表",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 企业微信智能机器人(WebSocket 长连接) ────────────────────────────
|
||||
WEWORK_BOT_ENABLED: bool = Field(
|
||||
default=False,
|
||||
description="是否启用企业微信智能机器人(独立进程运行)",
|
||||
)
|
||||
WEWORK_BOT_ID: Optional[str] = Field(
|
||||
default=None,
|
||||
description="企业微信智能机器人 ID",
|
||||
)
|
||||
WEWORK_BOT_SECRET: Optional[str] = Field(
|
||||
default=None,
|
||||
description="企业微信智能机器人 Secret",
|
||||
)
|
||||
WEWORK_BOT_SEARCH_LIMIT: int = Field(
|
||||
default=3,
|
||||
ge=1,
|
||||
le=10,
|
||||
description="企业微信机器人搜索返回图片数量上限",
|
||||
)
|
||||
EDUBRAIN_BASE_URL: str = Field(
|
||||
default="http://localhost:8765",
|
||||
description="EduBrain 服务地址(供企业微信机器人调用搜索接口)",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 属性方法 ────────────────────────────
|
||||
|
||||
@property
|
||||
def transcripts_dir(self) -> Path:
|
||||
"""直播文字稿目录"""
|
||||
return Path(self.DATA_DIR) / "transcripts"
|
||||
|
||||
@property
|
||||
def images_dir(self) -> Path:
|
||||
"""聊天截图目录"""
|
||||
return Path(self.DATA_DIR) / "images"
|
||||
|
||||
@property
|
||||
def exports_dir(self) -> Path:
|
||||
"""导出文件目录"""
|
||||
return Path(self.DATA_DIR) / "exports"
|
||||
|
||||
def ensure_dirs(self) -> None:
|
||||
"""确保所有数据目录存在"""
|
||||
for d in (self.transcripts_dir, self.images_dir, self.exports_dir):
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
# ──────────────────────────── 全局单例 ────────────────────────────
|
||||
|
||||
settings = Settings()
|
||||
107
app/database.py
Normal file
107
app/database.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""
|
||||
数据库连接模块
|
||||
支持 PostgreSQL(生产)和 SQLite(本地开发)两种后端
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from pathlib import Path
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import (
|
||||
AsyncSession,
|
||||
async_sessionmaker,
|
||||
create_async_engine,
|
||||
)
|
||||
from sqlalchemy.pool import NullPool
|
||||
|
||||
from app.config import settings
|
||||
|
||||
# ──────────────────────────── 判断数据库类型 ────────────────────────────
|
||||
|
||||
DB_URL = settings.DATABASE_URL
|
||||
IS_SQLITE = DB_URL.startswith("sqlite")
|
||||
|
||||
# ──────────────────────────── 异步引擎 ────────────────────────────
|
||||
|
||||
if IS_SQLITE:
|
||||
# SQLite 配置
|
||||
sqlite_path = DB_URL.replace("sqlite+aiosqlite:///", "")
|
||||
if not sqlite_path.startswith("/"):
|
||||
sqlite_path = str(Path(__file__).parent.parent / sqlite_path)
|
||||
|
||||
engine = create_async_engine(
|
||||
f"sqlite+aiosqlite:///{sqlite_path}",
|
||||
echo=settings.DEBUG,
|
||||
)
|
||||
else:
|
||||
# PostgreSQL 配置
|
||||
engine = create_async_engine(
|
||||
DB_URL,
|
||||
echo=settings.DEBUG,
|
||||
poolclass=NullPool,
|
||||
pool_pre_ping=True,
|
||||
)
|
||||
|
||||
# ──────────────────────────── 会话工厂 ────────────────────────────
|
||||
|
||||
async_session_factory = async_sessionmaker(
|
||||
bind=engine,
|
||||
class_=AsyncSession,
|
||||
expire_on_commit=False,
|
||||
autoflush=False,
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 依赖注入 ────────────────────────────
|
||||
|
||||
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
||||
"""
|
||||
FastAPI 依赖项:获取异步数据库会话。
|
||||
"""
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
yield session
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
|
||||
# ──────────────────────────── 生命周期辅助 ────────────────────────────
|
||||
|
||||
async def init_db() -> None:
|
||||
"""初始化数据库:建表"""
|
||||
from app.models.base import Base
|
||||
|
||||
async with engine.begin() as conn:
|
||||
if IS_SQLITE:
|
||||
# SQLite 启用 WAL 模式和外键约束
|
||||
await conn.execute(text("PRAGMA journal_mode=WAL"))
|
||||
await conn.execute(text("PRAGMA foreign_keys=ON"))
|
||||
# 创建所有表
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
if not IS_SQLITE:
|
||||
# PostgreSQL 需要创建扩展
|
||||
async with engine.begin() as conn:
|
||||
try:
|
||||
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS zhparser"))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
await conn.execute(text(
|
||||
"CREATE TEXT SEARCH CONFIGURATION IF NOT EXISTS chinese_zh (PARSER = zhparser)"
|
||||
))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def close_db() -> None:
|
||||
"""关闭数据库引擎"""
|
||||
await engine.dispose()
|
||||
171
app/main.py
Normal file
171
app/main.py
Normal file
@@ -0,0 +1,171 @@
|
||||
"""
|
||||
FastAPI 应用入口
|
||||
挂载 v1 API 路由、CORS 中间件、生命周期管理
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
|
||||
|
||||
from app.config import settings
|
||||
from app.database import close_db, init_db
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
"""
|
||||
应用生命周期管理:
|
||||
- 启动时:初始化数据目录、验证数据库连接、预热嵌入服务、初始化 BM25 索引
|
||||
- 关闭时:释放数据库连接
|
||||
"""
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
# ── 启动 ──
|
||||
settings.ensure_dirs()
|
||||
await init_db()
|
||||
|
||||
# 预热嵌入服务(懒加载,首次调用时初始化)
|
||||
try:
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
EmbeddingService.get_instance()
|
||||
except Exception as exc:
|
||||
_logger.warning("嵌入服务预热失败(将在首次使用时重试): %s", exc)
|
||||
|
||||
# ── 初始化 BM25 索引 ──
|
||||
try:
|
||||
from app.services.search_engine import BM25Index
|
||||
import app.api.v1.images as images_module
|
||||
|
||||
_images_bm25 = BM25Index()
|
||||
|
||||
# 加载数据库中已有的 OCR 记录到索引
|
||||
from app.database import async_session_factory
|
||||
from app.models.base import OCRImage
|
||||
from sqlalchemy import select
|
||||
|
||||
async with async_session_factory() as session:
|
||||
result = await session.execute(
|
||||
select(OCRImage).where(OCRImage.status == "completed")
|
||||
)
|
||||
records = result.scalars().all()
|
||||
for record in records:
|
||||
if record.ocr_text:
|
||||
index_text = record.ocr_text
|
||||
if record.tags:
|
||||
try:
|
||||
tags = json.loads(record.tags)
|
||||
if tags:
|
||||
index_text += "\n" + " ".join(tags)
|
||||
except Exception:
|
||||
pass
|
||||
if record.story_summary:
|
||||
index_text += "\n" + record.story_summary
|
||||
_images_bm25.add_document(
|
||||
doc_id=record.id,
|
||||
text=index_text,
|
||||
metadata={
|
||||
"file_path": record.file_path,
|
||||
"tags": json.loads(record.tags) if record.tags else [],
|
||||
"story_summary": record.story_summary or "",
|
||||
"confidence": record.confidence,
|
||||
"provider": record.provider,
|
||||
"created_at": str(record.created_at),
|
||||
},
|
||||
)
|
||||
|
||||
images_module.bm25_index = _images_bm25
|
||||
_logger.info("BM25 索引加载完成: %d 条记录", _images_bm25.size)
|
||||
except Exception as exc:
|
||||
_logger.warning("BM25 索引初始化失败: %s", exc)
|
||||
|
||||
yield
|
||||
|
||||
# ── 关闭 ──
|
||||
await close_db()
|
||||
|
||||
|
||||
# ──────────────────────────── 创建应用 ────────────────────────────
|
||||
|
||||
app = FastAPI(
|
||||
title="EduBrain - 中文直播教育知识库",
|
||||
description="基于向量检索的中文直播教育知识库系统,支持直播文字稿导入、OCR 截图识别、语义搜索等功能。",
|
||||
version="1.2.0",
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
# ──────────────────────────── CORS 中间件 ────────────────────────────
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.CORS_ORIGINS,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 全局异常处理 ────────────────────────────
|
||||
|
||||
@app.exception_handler(Exception)
|
||||
async def global_exception_handler(request: Request, exc: Exception) -> JSONResponse:
|
||||
"""全局未捕获异常处理"""
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error("未处理异常: %s", exc, exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content={"detail": "服务器内部错误", "message": str(exc) if settings.DEBUG else "请稍后重试"},
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 注册路由 ────────────────────────────
|
||||
|
||||
from app.api.v1.pages import router as pages_router
|
||||
from app.api.v1.search import router as search_router
|
||||
from app.api.v1.import_export import router as import_export_router
|
||||
from app.api.v1.images import router as images_router
|
||||
from app.api.v1.settings import router as settings_router
|
||||
|
||||
app.include_router(pages_router, prefix="/api/v1/pages", tags=["知识页面"])
|
||||
app.include_router(search_router, prefix="/api/v1/search", tags=["语义搜索"])
|
||||
app.include_router(import_export_router, prefix="/api/v1/import", tags=["导入导出"])
|
||||
app.include_router(images_router, prefix="/api/v1/images", tags=["图片OCR"])
|
||||
app.include_router(settings_router, prefix="/api/v1/settings", tags=["系统设置"])
|
||||
|
||||
|
||||
# ──────────────────────────── 健康检查 ────────────────────────────
|
||||
|
||||
@app.get("/health", tags=["系统"])
|
||||
async def health_check():
|
||||
"""健康检查接口"""
|
||||
return {"status": "ok", "version": "1.2.0"}
|
||||
|
||||
|
||||
# ──────────────────────────── 图片文件访问 ────────────────────────────
|
||||
|
||||
@app.get("/data/images/{image_name}", tags=["前端"])
|
||||
async def serve_image(image_name: str):
|
||||
"""访问 images 目录下的图片文件(用于前端预览)"""
|
||||
img_path = settings.images_dir / image_name
|
||||
if not img_path.exists():
|
||||
return JSONResponse(status_code=404, content={"detail": "图片不存在"})
|
||||
return FileResponse(str(img_path))
|
||||
|
||||
|
||||
# ──────────────────────────── 前端页面 ────────────────────────────
|
||||
|
||||
@app.get("/", tags=["前端"])
|
||||
async def index():
|
||||
"""返回前端管理页面"""
|
||||
static_dir = Path(__file__).parent.parent / "static"
|
||||
index_path = static_dir / "index.html"
|
||||
if index_path.exists():
|
||||
return FileResponse(str(index_path))
|
||||
return RedirectResponse(url="/docs")
|
||||
4
app/mcp/__init__.py
Normal file
4
app/mcp/__init__.py
Normal file
@@ -0,0 +1,4 @@
|
||||
"""
|
||||
MCP (Model Context Protocol) Server 包
|
||||
提供 MCP 协议接口,允许 AI Agent 直接查询知识库
|
||||
"""
|
||||
297
app/mcp/server.py
Normal file
297
app/mcp/server.py
Normal file
@@ -0,0 +1,297 @@
|
||||
"""
|
||||
MCP Server 实现
|
||||
基于 MCP SDK 提供知识库查询工具,供 AI Agent 调用
|
||||
|
||||
工具列表:
|
||||
1. search_knowledge - 语义搜索知识库
|
||||
2. get_page - 获取知识页面详情
|
||||
3. list_courses - 列出所有课程
|
||||
4. get_page_chunks - 获取页面的分块内容
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ──────────────────────────── 创建 MCP Server ────────────────────────────
|
||||
|
||||
mcp = FastMCP(
|
||||
name="EduBrain-Knowledge",
|
||||
instructions=(
|
||||
"中文直播教育知识库助手。你可以使用以下工具来搜索和查询知识库中的内容:\n"
|
||||
"- search_knowledge: 语义搜索知识库,支持按课程、讲师、日期过滤\n"
|
||||
"- get_page: 获取指定知识页面的完整内容\n"
|
||||
"- list_courses: 列出知识库中所有课程\n"
|
||||
"- get_page_chunks: 获取页面的分块内容(用于精确定位)\n"
|
||||
"- search_images: 根据关键字搜索已 OCR 识别的聊天截图,返回图片路径"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 工具定义 ────────────────────────────
|
||||
|
||||
@mcp.tool()
|
||||
async def search_knowledge(
|
||||
query: str,
|
||||
top_k: int = 5,
|
||||
course_name: str = "",
|
||||
teacher_name: str = "",
|
||||
threshold: float = 0.5,
|
||||
) -> str:
|
||||
"""
|
||||
语义搜索知识库。
|
||||
|
||||
Args:
|
||||
query: 搜索查询文本(自然语言描述你要查找的知识点)
|
||||
top_k: 返回结果数量,默认 5
|
||||
course_name: 按课程名称过滤(可选)
|
||||
teacher_name: 按讲师名称过滤(可选)
|
||||
threshold: 相似度阈值(0-1),默认 0.5
|
||||
|
||||
Returns:
|
||||
搜索结果列表,包含标题、内容摘要和相似度分数
|
||||
"""
|
||||
from app.database import async_session_factory
|
||||
from app.schemas.search import SearchRequest
|
||||
from app.services.search_service import SearchService
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
request = SearchRequest(
|
||||
query=query,
|
||||
top_k=top_k,
|
||||
course_name=course_name or None,
|
||||
teacher_name=teacher_name or None,
|
||||
threshold=threshold,
|
||||
use_fulltext=True,
|
||||
)
|
||||
service = SearchService(session)
|
||||
response = await service.search(request)
|
||||
|
||||
if not response.results:
|
||||
return f"未找到与「{query}」相关的知识内容。"
|
||||
|
||||
# 格式化搜索结果
|
||||
output_parts = [
|
||||
f"找到 {response.total} 条相关结果(耗时 {response.elapsed_ms:.0f}ms):\n"
|
||||
]
|
||||
for i, result in enumerate(response.results, 1):
|
||||
output_parts.append(
|
||||
f"【{i}】{result.page_title} "
|
||||
f"(相似度: {result.score:.2%})"
|
||||
)
|
||||
if result.course_name:
|
||||
output_parts.append(f" 课程: {result.course_name}")
|
||||
if result.teacher_name:
|
||||
output_parts.append(f" 讲师: {result.teacher_name}")
|
||||
# 截断内容避免过长
|
||||
content = result.content[:200]
|
||||
if len(result.content) > 200:
|
||||
content += "..."
|
||||
output_parts.append(f" 内容: {content}")
|
||||
output_parts.append("")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 搜索失败: %s", e, exc_info=True)
|
||||
return f"搜索出错: {str(e)}"
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_page(page_id: int) -> str:
|
||||
"""
|
||||
获取指定知识页面的完整内容。
|
||||
|
||||
Args:
|
||||
page_id: 知识页面 ID
|
||||
|
||||
Returns:
|
||||
页面的完整内容,包括标题、正文、元数据
|
||||
"""
|
||||
from app.database import async_session_factory
|
||||
from app.services.page_service import PageService
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
service = PageService(session)
|
||||
page = await service.get_page(page_id)
|
||||
|
||||
if page is None:
|
||||
return f"未找到 ID 为 {page_id} 的知识页面。"
|
||||
|
||||
parts = [
|
||||
f"标题: {page.title}",
|
||||
]
|
||||
if page.course_name:
|
||||
parts.append(f"课程: {page.course_name}")
|
||||
if page.teacher_name:
|
||||
parts.append(f"讲师: {page.teacher_name}")
|
||||
if page.live_date:
|
||||
parts.append(f"直播日期: {page.live_date}")
|
||||
parts.append(f"\n{page.content}")
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 获取页面失败: %s", e, exc_info=True)
|
||||
return f"获取页面出错: {str(e)}"
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def list_courses() -> str:
|
||||
"""
|
||||
列出知识库中所有课程。
|
||||
|
||||
Returns:
|
||||
课程列表,包含课程名称、讲师、页面数量等信息
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
from app.database import async_session_factory
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
result = await session.execute(text("""
|
||||
SELECT
|
||||
course_name,
|
||||
teacher_name,
|
||||
COUNT(*) as page_count,
|
||||
MIN(live_date) as earliest_date,
|
||||
MAX(live_date) as latest_date
|
||||
FROM knowledge_pages
|
||||
WHERE course_name IS NOT NULL
|
||||
GROUP BY course_name, teacher_name
|
||||
ORDER BY page_count DESC
|
||||
"""))
|
||||
rows = result.fetchall()
|
||||
|
||||
if not rows:
|
||||
return "知识库中暂无课程数据。"
|
||||
|
||||
parts = [f"共 {len(rows)} 个课程:\n"]
|
||||
for i, row in enumerate(rows, 1):
|
||||
parts.append(
|
||||
f"{i}. {row.course_name}"
|
||||
f"(讲师: {row.teacher_name or '未知'},"
|
||||
f"页面: {row.page_count},"
|
||||
f"日期: {row.earliest_date or '未知'} ~ {row.latest_date or '未知'})"
|
||||
)
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 列出课程失败: %s", e, exc_info=True)
|
||||
return f"列出课程出错: {str(e)}"
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def search_images(
|
||||
keyword: str,
|
||||
top_k: int = 10,
|
||||
) -> str:
|
||||
"""
|
||||
根据关键字搜索已 OCR 识别的聊天截图。
|
||||
|
||||
Args:
|
||||
keyword: 搜索关键字(在 OCR 识别文本中查找)
|
||||
top_k: 返回结果数量,默认 10
|
||||
|
||||
Returns:
|
||||
匹配的图片列表,包含图片路径和 OCR 文本预览
|
||||
"""
|
||||
from sqlalchemy import select
|
||||
from app.database import async_session_factory
|
||||
from app.models.base import OCRImage
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
like_pattern = f"%{keyword}%"
|
||||
result = await session.execute(
|
||||
select(OCRImage)
|
||||
.where(
|
||||
OCRImage.status == "completed",
|
||||
OCRImage.ocr_text.ilike(like_pattern),
|
||||
)
|
||||
.order_by(OCRImage.created_at.desc())
|
||||
.limit(top_k)
|
||||
)
|
||||
images = result.scalars().all()
|
||||
|
||||
if not images:
|
||||
return f"未找到包含「{keyword}」的图片。"
|
||||
|
||||
parts = [f"找到 {len(images)} 张包含「{keyword}」的图片:\n"]
|
||||
for i, img in enumerate(images, 1):
|
||||
preview = (img.ocr_text[:100] + "...") if img.ocr_text and len(img.ocr_text) > 100 else (img.ocr_text or "")
|
||||
parts.append(
|
||||
f"【{i}】图片路径: {img.file_path}\n"
|
||||
f" OCR 预览: {preview}\n"
|
||||
f" 置信度: {img.confidence:.0%} | 识别方式: {img.provider}\n"
|
||||
)
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 图片搜索失败: %s", e, exc_info=True)
|
||||
return f"图片搜索出错: {str(e)}"
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_page_chunks(page_id: int) -> str:
|
||||
"""
|
||||
获取知识页面的分块内容(用于精确定位知识点)。
|
||||
|
||||
Args:
|
||||
page_id: 知识页面 ID
|
||||
|
||||
Returns:
|
||||
该页面的所有分块内容列表
|
||||
"""
|
||||
from app.database import async_session_factory
|
||||
from app.services.page_service import PageService
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
service = PageService(session)
|
||||
chunks = await service.get_page_chunks(page_id)
|
||||
|
||||
if not chunks:
|
||||
return f"页面 {page_id} 没有分块数据。"
|
||||
|
||||
parts = [f"页面 {page_id} 共 {len(chunks)} 个分块:\n"]
|
||||
for i, chunk in enumerate(chunks, 1):
|
||||
content = chunk.content[:300]
|
||||
if len(chunk.content) > 300:
|
||||
content += "..."
|
||||
parts.append(f"【分块 {i}】(ID: {chunk.id})\n{content}\n")
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 获取分块失败: %s", e, exc_info=True)
|
||||
return f"获取分块出错: {str(e)}"
|
||||
|
||||
|
||||
# ──────────────────────────── 启动入口 ────────────────────────────
|
||||
|
||||
def run_mcp_server(transport: str = "stdio", port: int = 8001):
|
||||
"""
|
||||
启动 MCP Server。
|
||||
|
||||
Args:
|
||||
transport: 传输协议,支持 "stdio"(默认)或 "sse"
|
||||
port: SSE 模式下的端口号
|
||||
"""
|
||||
if transport == "sse":
|
||||
mcp.run(transport="sse", port=port)
|
||||
else:
|
||||
mcp.run(transport="stdio")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_mcp_server()
|
||||
7
app/models/__init__.py
Normal file
7
app/models/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
SQLAlchemy ORM 模型包
|
||||
"""
|
||||
|
||||
from app.models.base import Base, TimestampMixin
|
||||
|
||||
__all__ = ["Base", "TimestampMixin"]
|
||||
123
app/models/base.py
Normal file
123
app/models/base.py
Normal file
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
SQLAlchemy 基础模型 + 通用 Mixin
|
||||
提供 DeclarativeBase、时间戳 Mixin 和公共序列化方法
|
||||
兼容 PostgreSQL(pgvector)和 SQLite(本地开发)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict
|
||||
|
||||
from sqlalchemy import DateTime, Float, Integer, String, Text, func
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||
|
||||
from app.database import IS_SQLITE
|
||||
|
||||
# SQLite 下不使用 pgvector,用 Text 存储向量(JSON 格式)
|
||||
if IS_SQLITE:
|
||||
VectorType = Text
|
||||
else:
|
||||
from pgvector.sqlalchemy import Vector
|
||||
VectorType = Vector
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
"""所有 ORM 模型的基类"""
|
||||
pass
|
||||
|
||||
|
||||
class TimestampMixin:
|
||||
"""
|
||||
时间戳 Mixin,为模型自动添加 created_at / updated_at 字段。
|
||||
"""
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
comment="创建时间",
|
||||
)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
onupdate=func.now(),
|
||||
comment="更新时间",
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""将模型实例转换为字典"""
|
||||
result: Dict[str, Any] = {}
|
||||
for column in self.__table__.columns: # type: ignore[attr-defined]
|
||||
value = getattr(self, column.name, None)
|
||||
if isinstance(value, datetime):
|
||||
value = value.isoformat()
|
||||
elif hasattr(value, "tolist"):
|
||||
value = value.tolist()
|
||||
result[column.name] = value
|
||||
return result
|
||||
|
||||
|
||||
# ──────────────────────────── 业务模型 ────────────────────────────
|
||||
|
||||
class KnowledgePage(Base, TimestampMixin):
|
||||
"""知识页面模型"""
|
||||
__tablename__ = "knowledge_pages"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
title: Mapped[str] = mapped_column(String(500), nullable=False, comment="页面标题")
|
||||
content: Mapped[str] = mapped_column(Text, nullable=False, comment="页面正文内容")
|
||||
source_file: Mapped[str | None] = mapped_column(String(500), nullable=True, comment="来源文件名")
|
||||
course_name: Mapped[str | None] = mapped_column(String(200), nullable=True, comment="课程名称")
|
||||
teacher_name: Mapped[str | None] = mapped_column(String(100), nullable=True, comment="讲师名称")
|
||||
live_date: Mapped[str | None] = mapped_column(String(20), nullable=True, comment="直播日期")
|
||||
page_number: Mapped[int | None] = mapped_column(Integer, nullable=True, comment="原始页码")
|
||||
metadata_json: Mapped[str | None] = mapped_column(Text, nullable=True, comment="额外元数据")
|
||||
|
||||
|
||||
class KnowledgeChunk(Base, TimestampMixin):
|
||||
"""知识分块模型"""
|
||||
__tablename__ = "knowledge_chunks"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
page_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True, comment="关联页面 ID")
|
||||
chunk_index: Mapped[int] = mapped_column(Integer, nullable=False, comment="分块序号")
|
||||
content: Mapped[str] = mapped_column(Text, nullable=False, comment="分块文本内容")
|
||||
embedding: Mapped[list | None] = mapped_column(
|
||||
VectorType(768) if not IS_SQLITE else Text,
|
||||
nullable=True,
|
||||
comment="嵌入向量",
|
||||
)
|
||||
search_vector: Mapped[Any | None] = mapped_column(
|
||||
Text, nullable=True, comment="全文搜索向量(仅 PG)",
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result = super().to_dict()
|
||||
result.pop("embedding", None)
|
||||
result.pop("search_vector", None)
|
||||
return result
|
||||
|
||||
|
||||
class OCRImage(Base, TimestampMixin):
|
||||
"""OCR 图片记录模型"""
|
||||
__tablename__ = "ocr_images"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
file_path: Mapped[str] = mapped_column(String(500), nullable=False, comment="图片文件路径")
|
||||
ocr_text: Mapped[str | None] = mapped_column(Text, nullable=True, comment="OCR 识别全文")
|
||||
confidence: Mapped[float | None] = mapped_column(Float, nullable=True, comment="平均置信度")
|
||||
provider: Mapped[str | None] = mapped_column(String(50), nullable=True, comment="OCR 提供商")
|
||||
status: Mapped[str] = mapped_column(String(20), default="pending", comment="处理状态")
|
||||
error_message: Mapped[str | None] = mapped_column(Text, nullable=True, comment="错误信息")
|
||||
tags: Mapped[str | None] = mapped_column(Text, nullable=True, comment="LLM 提取的标签,JSON 数组格式")
|
||||
story_summary: Mapped[str | None] = mapped_column(Text, nullable=True, comment="Qwen3-8B 提炼的故事摘要")
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result = super().to_dict()
|
||||
# 将 tags 从 JSON 字符串解析为列表返回
|
||||
if result.get("tags") and isinstance(result["tags"], str):
|
||||
try:
|
||||
import json
|
||||
result["tags"] = json.loads(result["tags"])
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
return result
|
||||
32
app/schemas/__init__.py
Normal file
32
app/schemas/__init__.py
Normal file
@@ -0,0 +1,32 @@
|
||||
"""
|
||||
Pydantic Schemas 包
|
||||
用于 API 请求/响应的数据校验与序列化
|
||||
"""
|
||||
|
||||
from app.schemas.page import (
|
||||
PageCreate,
|
||||
PageResponse,
|
||||
PageListResponse,
|
||||
)
|
||||
from app.schemas.search import (
|
||||
SearchRequest,
|
||||
SearchResponse,
|
||||
SearchResult,
|
||||
)
|
||||
from app.schemas.ocr import (
|
||||
OCRResultSchema,
|
||||
OCRBlockSchema,
|
||||
ImageUploadResponse,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"PageCreate",
|
||||
"PageResponse",
|
||||
"PageListResponse",
|
||||
"SearchRequest",
|
||||
"SearchResponse",
|
||||
"SearchResult",
|
||||
"OCRResultSchema",
|
||||
"OCRBlockSchema",
|
||||
"ImageUploadResponse",
|
||||
]
|
||||
47
app/schemas/ocr.py
Normal file
47
app/schemas/ocr.py
Normal file
@@ -0,0 +1,47 @@
|
||||
"""
|
||||
OCR 相关的 Pydantic Schema
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class OCRBlockSchema(BaseModel):
|
||||
"""OCR 识别的文本块"""
|
||||
text: str = Field(..., description="文本块内容")
|
||||
bbox: Optional[list[float]] = Field(None, description="边界框 [x1, y1, x2, y2]")
|
||||
confidence: float = Field(..., description="置信度 (0-1)")
|
||||
|
||||
|
||||
class OCRResultSchema(BaseModel):
|
||||
"""OCR 识别结果"""
|
||||
text: str = Field(..., description="识别全文")
|
||||
blocks: list[OCRBlockSchema] = Field(default_factory=list, description="文本块列表")
|
||||
confidence: float = Field(..., description="平均置信度 (0-1)")
|
||||
provider: str = Field(..., description="使用的 OCR 提供商")
|
||||
|
||||
|
||||
class ImageUploadResponse(BaseModel):
|
||||
"""图片上传响应"""
|
||||
id: int
|
||||
file_path: str
|
||||
status: str
|
||||
created_at: datetime
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class ImageOCRResponse(BaseModel):
|
||||
"""图片 OCR 结果响应"""
|
||||
id: int
|
||||
file_path: str
|
||||
ocr_text: Optional[str] = None
|
||||
confidence: Optional[float] = None
|
||||
provider: Optional[str] = None
|
||||
status: str
|
||||
error_message: Optional[str] = None
|
||||
created_at: datetime
|
||||
model_config = {"from_attributes": True}
|
||||
63
app/schemas/page.py
Normal file
63
app/schemas/page.py
Normal file
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
知识页面相关的 Pydantic Schema
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class PageCreate(BaseModel):
|
||||
"""创建知识页面请求"""
|
||||
title: str = Field(..., max_length=500, description="页面标题")
|
||||
content: str = Field(..., description="页面正文内容")
|
||||
source_file: Optional[str] = Field(None, max_length=500, description="来源文件名")
|
||||
course_name: Optional[str] = Field(None, max_length=200, description="课程名称")
|
||||
teacher_name: Optional[str] = Field(None, max_length=100, description="讲师名称")
|
||||
live_date: Optional[str] = Field(None, max_length=20, description="直播日期 (YYYY-MM-DD)")
|
||||
page_number: Optional[int] = Field(None, description="原始页码")
|
||||
metadata_json: Optional[str] = Field(None, description="额外元数据 (JSON)")
|
||||
|
||||
|
||||
class PageUpdate(BaseModel):
|
||||
"""更新知识页面请求(所有字段可选)"""
|
||||
title: Optional[str] = Field(None, max_length=500, description="页面标题")
|
||||
content: Optional[str] = Field(None, description="页面正文内容")
|
||||
course_name: Optional[str] = Field(None, max_length=200, description="课程名称")
|
||||
teacher_name: Optional[str] = Field(None, max_length=100, description="讲师名称")
|
||||
live_date: Optional[str] = Field(None, max_length=20, description="直播日期")
|
||||
page_number: Optional[int] = Field(None, description="原始页码")
|
||||
metadata_json: Optional[str] = Field(None, description="额外元数据")
|
||||
|
||||
|
||||
class PageResponse(BaseModel):
|
||||
"""知识页面响应"""
|
||||
id: int
|
||||
title: str
|
||||
content: str
|
||||
source_file: Optional[str] = None
|
||||
course_name: Optional[str] = None
|
||||
teacher_name: Optional[str] = None
|
||||
live_date: Optional[str] = None
|
||||
page_number: Optional[int] = None
|
||||
metadata_json: Optional[str] = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class PageListResponse(BaseModel):
|
||||
"""知识页面列表响应(带分页)"""
|
||||
total: int = Field(..., description="总记录数")
|
||||
page: int = Field(..., description="当前页码")
|
||||
page_size: int = Field(..., description="每页数量")
|
||||
items: list[PageResponse] = Field(default_factory=list, description="页面列表")
|
||||
|
||||
|
||||
class PageDetailResponse(PageResponse):
|
||||
"""知识页面详情响应(包含分块信息)"""
|
||||
chunks: list[dict[str, Any]] = Field(default_factory=list, description="关联的分块列表")
|
||||
42
app/schemas/search.py
Normal file
42
app/schemas/search.py
Normal file
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
语义搜索相关的 Pydantic Schema
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
"""语义搜索请求"""
|
||||
query: str = Field(..., min_length=1, max_length=1000, description="搜索查询文本")
|
||||
top_k: int = Field(default=10, ge=1, le=100, description="返回结果数量")
|
||||
course_name: Optional[str] = Field(None, description="按课程名称过滤")
|
||||
teacher_name: Optional[str] = Field(None, description="按讲师名称过滤")
|
||||
live_date_from: Optional[str] = Field(None, description="直播日期起始 (YYYY-MM-DD)")
|
||||
live_date_to: Optional[str] = Field(None, description="直播日期结束 (YYYY-MM-DD)")
|
||||
threshold: float = Field(default=0.5, ge=0.0, le=1.0, description="相似度阈值")
|
||||
use_fulltext: bool = Field(default=True, description="是否同时使用全文搜索混合排序")
|
||||
|
||||
|
||||
class SearchResult(BaseModel):
|
||||
"""单条搜索结果"""
|
||||
chunk_id: int = Field(..., description="分块 ID")
|
||||
page_id: int = Field(..., description="知识页面 ID")
|
||||
page_title: str = Field(..., description="页面标题")
|
||||
content: str = Field(..., description="分块文本内容")
|
||||
score: float = Field(..., description="相似度分数")
|
||||
course_name: Optional[str] = None
|
||||
teacher_name: Optional[str] = None
|
||||
live_date: Optional[str] = None
|
||||
highlight: Optional[str] = Field(None, description="高亮摘要")
|
||||
|
||||
|
||||
class SearchResponse(BaseModel):
|
||||
"""搜索结果响应"""
|
||||
query: str = Field(..., description="原始查询文本")
|
||||
total: int = Field(..., description="结果总数")
|
||||
results: list[SearchResult] = Field(default_factory=list, description="搜索结果列表")
|
||||
elapsed_ms: float = Field(..., description="搜索耗时(毫秒)")
|
||||
17
app/services/__init__.py
Normal file
17
app/services/__init__.py
Normal file
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
业务逻辑服务层包
|
||||
"""
|
||||
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
from app.services.ocr_service import OCRService
|
||||
from app.services.search_service import SearchService
|
||||
from app.services.import_service import ImportService
|
||||
from app.services.page_service import PageService
|
||||
|
||||
__all__ = [
|
||||
"EmbeddingService",
|
||||
"OCRService",
|
||||
"SearchService",
|
||||
"ImportService",
|
||||
"PageService",
|
||||
]
|
||||
328
app/services/embedding_service.py
Normal file
328
app/services/embedding_service.py
Normal file
@@ -0,0 +1,328 @@
|
||||
"""
|
||||
嵌入服务模块
|
||||
支持多种嵌入模型提供商:OpenAI / 智谱 / 阿里云 DashScope / 本地 BGE
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from app.config import EmbeddingProvider, settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ──────────────────────────── 抽象基类 ────────────────────────────
|
||||
|
||||
class EmbeddingProviderBase(abc.ABC):
|
||||
"""嵌入模型提供商抽象基类"""
|
||||
|
||||
@abc.abstractmethod
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
批量生成文本嵌入向量。
|
||||
|
||||
Args:
|
||||
texts: 待嵌入的文本列表
|
||||
|
||||
Returns:
|
||||
嵌入向量列表,每个向量是 float 的列表
|
||||
|
||||
Raises:
|
||||
Exception: 嵌入生成失败时抛出异常
|
||||
"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
"""生成单条文本的嵌入向量"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def dimension(self) -> int:
|
||||
"""嵌入向量维度"""
|
||||
...
|
||||
|
||||
|
||||
# ──────────────────────────── OpenAI 实现 ────────────────────────────
|
||||
|
||||
class OpenAIEmbedding(EmbeddingProviderBase):
|
||||
"""OpenAI 嵌入模型(兼容 OpenAI API 格式的服务)"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
self._client = AsyncOpenAI(
|
||||
api_key=settings.OPENAI_API_KEY or "EMPTY",
|
||||
base_url=settings.OPENAI_BASE_URL,
|
||||
)
|
||||
self._model = settings.EMBEDDING_MODEL or "text-embedding-3-small"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用 OpenAI 兼容接口批量生成嵌入"""
|
||||
# OpenAI 接口单次最多 2048 条,自动分批
|
||||
batch_size = 2048
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
response = await self._client.embeddings.create(
|
||||
input=batch,
|
||||
model=self._model,
|
||||
dimensions=self.dimension if self.dimension <= 3072 else None,
|
||||
)
|
||||
batch_embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 智谱 AI 实现 ────────────────────────────
|
||||
|
||||
class ZhipuEmbedding(EmbeddingProviderBase):
|
||||
"""智谱 AI 嵌入模型"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from zhipuai import ZhipuAI
|
||||
|
||||
self._client = ZhipuAI(api_key=settings.ZHIPU_API_KEY)
|
||||
self._model = settings.EMBEDDING_MODEL or "embedding-3"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用智谱 AI 接口批量生成嵌入"""
|
||||
import asyncio
|
||||
|
||||
# 智谱 API 单次最多 16 条
|
||||
batch_size = 16
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
# 智谱 SDK 是同步的,用线程池包装
|
||||
loop = asyncio.get_event_loop()
|
||||
response = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._client.embeddings.create(
|
||||
input=batch,
|
||||
model=self._model,
|
||||
),
|
||||
)
|
||||
batch_embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 阿里云 DashScope 实现 ────────────────────────────
|
||||
|
||||
class DashscopeEmbedding(EmbeddingProviderBase):
|
||||
"""阿里云 DashScope 嵌入模型"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
import dashscope
|
||||
dashscope.api_key = settings.DASHSCOPE_API_KEY
|
||||
self._model = settings.EMBEDDING_MODEL or "text-embedding-v3"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用 DashScope 接口批量生成嵌入"""
|
||||
import asyncio
|
||||
from dashscope import TextEmbedding
|
||||
|
||||
batch_size = 25 # DashScope 单次最多 25 条
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _call() -> list:
|
||||
resp = TextEmbedding.call(
|
||||
model=self._model,
|
||||
input=batch,
|
||||
dimension=self.dimension,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise RuntimeError(f"DashScope 嵌入失败: {resp.code} - {resp.message}")
|
||||
return [item["embedding"] for item in resp.output["embeddings"]]
|
||||
|
||||
batch_embeddings = await loop.run_in_executor(None, _call)
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 本地 BGE 实现 ────────────────────────────
|
||||
|
||||
class LocalBGEEmbedding(EmbeddingProviderBase):
|
||||
"""
|
||||
本地 BGE 嵌入模型
|
||||
使用 sentence-transformers 加载模型,在 CPU/GPU 上本地推理
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._model = None
|
||||
self._dim = settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
def _load_model(self):
|
||||
"""延迟加载模型(首次调用时初始化)"""
|
||||
if self._model is not None:
|
||||
return
|
||||
|
||||
logger.info("正在加载本地 BGE 嵌入模型: %s", settings.EMBEDDING_MODEL)
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
model_path = settings.LOCAL_BGE_MODEL_PATH or settings.EMBEDDING_MODEL
|
||||
self._model = SentenceTransformer(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
# 获取实际维度
|
||||
test_embedding = self._model.encode(["测试"])
|
||||
self._dim = len(test_embedding[0])
|
||||
logger.info("BGE 模型加载完成,维度: %d", self._dim)
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return self._dim
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""使用本地模型批量生成嵌入"""
|
||||
import asyncio
|
||||
|
||||
self._load_model()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _encode() -> List[List[float]]:
|
||||
embeddings = self._model.encode(texts, normalize_embeddings=True)
|
||||
return embeddings.tolist()
|
||||
|
||||
return await loop.run_in_executor(None, _encode)
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── MiniMax 实现 ────────────────────────────
|
||||
|
||||
class MiniMaxEmbedding(EmbeddingProviderBase):
|
||||
"""
|
||||
MiniMax 嵌入模型
|
||||
通过 MiniMax 的 OpenAI 兼容接口调用嵌入 API
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
self._client = AsyncOpenAI(
|
||||
api_key=settings.MINIMAX_API_KEY or "EMPTY",
|
||||
base_url=settings.MINIMAX_BASE_URL,
|
||||
)
|
||||
self._model = settings.MINIMAX_EMBEDDING_MODEL
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用 MiniMax 兼容接口批量生成嵌入"""
|
||||
batch_size = 64
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
response = await self._client.embeddings.create(
|
||||
input=batch,
|
||||
model=self._model,
|
||||
)
|
||||
batch_embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 工厂类 ────────────────────────────
|
||||
|
||||
class EmbeddingService:
|
||||
"""
|
||||
嵌入服务工厂类
|
||||
根据配置自动选择嵌入模型提供商,提供全局单例访问
|
||||
"""
|
||||
|
||||
_instance: Optional[EmbeddingProviderBase] = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> EmbeddingProviderBase:
|
||||
"""获取嵌入服务单例实例"""
|
||||
if cls._instance is not None:
|
||||
return cls._instance
|
||||
|
||||
provider_map = {
|
||||
EmbeddingProvider.OPENAI: OpenAIEmbedding,
|
||||
EmbeddingProvider.ZHIPU: ZhipuEmbedding,
|
||||
EmbeddingProvider.DASHSCOPE: DashscopeEmbedding,
|
||||
EmbeddingProvider.LOCAL_BGE: LocalBGEEmbedding,
|
||||
EmbeddingProvider.MINIMAX: MiniMaxEmbedding,
|
||||
}
|
||||
|
||||
provider_cls = provider_map.get(settings.EMBEDDING_PROVIDER)
|
||||
if provider_cls is None:
|
||||
raise ValueError(f"不支持的嵌入模型提供商: {settings.EMBEDDING_PROVIDER}")
|
||||
|
||||
cls._instance = provider_cls()
|
||||
logger.info(
|
||||
"嵌入服务初始化完成: provider=%s, model=%s",
|
||||
settings.EMBEDDING_PROVIDER.value,
|
||||
settings.EMBEDDING_MODEL,
|
||||
)
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
async def embed_batch(cls, texts: List[str]) -> List[List[float]]:
|
||||
"""便捷方法:批量生成嵌入"""
|
||||
instance = cls.get_instance()
|
||||
return await instance.embed_batch(texts)
|
||||
|
||||
@classmethod
|
||||
async def embed_single(cls, text: str) -> List[float]:
|
||||
"""便捷方法:生成单条嵌入"""
|
||||
instance = cls.get_instance()
|
||||
return await instance.embed_single(text)
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""重置单例(用于测试或切换配置)"""
|
||||
cls._instance = None
|
||||
452
app/services/import_service.py
Normal file
452
app/services/import_service.py
Normal file
@@ -0,0 +1,452 @@
|
||||
"""
|
||||
导入服务
|
||||
负责从 Markdown / 纯文本 / Word 文件导入知识页面,并进行文本分块和向量化
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.config import settings
|
||||
from app.database import IS_SQLITE
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _parse_frontmatter(content: str) -> tuple[dict, str]:
|
||||
"""
|
||||
手动解析 YAML frontmatter。
|
||||
格式:--- 开头和结尾包围的 YAML 块。
|
||||
"""
|
||||
import yaml
|
||||
|
||||
if not content.startswith("---"):
|
||||
return {}, content
|
||||
|
||||
parts = content.split("---", 2)
|
||||
if len(parts) < 3:
|
||||
return {}, content
|
||||
|
||||
try:
|
||||
metadata = yaml.safe_load(parts[1]) or {}
|
||||
except Exception:
|
||||
metadata = {}
|
||||
|
||||
body = parts[2].strip()
|
||||
return metadata, body
|
||||
|
||||
|
||||
class ImportService:
|
||||
"""知识导入服务"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self.db = db
|
||||
|
||||
async def import_file(
|
||||
self,
|
||||
file_path: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
从文件导入知识内容。
|
||||
|
||||
支持的格式:
|
||||
- Markdown (.md): 解析 frontmatter 元数据,按标题分页
|
||||
- 纯文本 (.txt): 整体作为一个页面
|
||||
- Word 文档 (.docx): 提取文本,按标题分页
|
||||
|
||||
Args:
|
||||
file_path: 文件路径
|
||||
course_name: 课程名称(覆盖 frontmatter)
|
||||
teacher_name: 讲师名称(覆盖 frontmatter)
|
||||
live_date: 直播日期(覆盖 frontmatter)
|
||||
|
||||
Returns:
|
||||
导入结果统计
|
||||
"""
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"文件不存在: {file_path}")
|
||||
|
||||
suffix = path.suffix.lower()
|
||||
|
||||
if suffix == ".md":
|
||||
content = path.read_text(encoding="utf-8")
|
||||
return await self._import_markdown(
|
||||
content, path.name, course_name, teacher_name, live_date
|
||||
)
|
||||
elif suffix == ".txt":
|
||||
content = path.read_text(encoding="utf-8")
|
||||
return await self._import_text(
|
||||
content, path.name, course_name, teacher_name, live_date
|
||||
)
|
||||
elif suffix == ".docx":
|
||||
return await self._import_docx(
|
||||
str(path), path.name, course_name, teacher_name, live_date
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"不支持的文件格式: {suffix},仅支持 .md、.txt 和 .docx")
|
||||
|
||||
async def _import_markdown(
|
||||
self,
|
||||
content: str,
|
||||
filename: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""导入 Markdown 文件,解析 frontmatter 并按标题分页"""
|
||||
# 解析 frontmatter
|
||||
fm_data, body = _parse_frontmatter(content)
|
||||
|
||||
# frontmatter 中的元数据可作为默认值
|
||||
fm_course = fm_data.get("course", fm_data.get("course_name"))
|
||||
fm_teacher = fm_data.get("teacher", fm_data.get("teacher_name"))
|
||||
fm_date = fm_data.get("date", fm_data.get("live_date"))
|
||||
|
||||
final_course = course_name or fm_course
|
||||
final_teacher = teacher_name or fm_teacher
|
||||
final_date = live_date or (str(fm_date) if fm_date else None)
|
||||
|
||||
# 按二级标题(##)拆分页面
|
||||
sections = self._split_markdown_sections(body)
|
||||
|
||||
if not sections:
|
||||
# 没有二级标题,整体作为一个页面
|
||||
sections = [{"title": Path(filename).stem, "content": body}]
|
||||
|
||||
imported_pages = 0
|
||||
imported_chunks = 0
|
||||
|
||||
for idx, section in enumerate(sections):
|
||||
page_number = idx + 1
|
||||
|
||||
# 插入知识页面
|
||||
page_id = await self._insert_page(
|
||||
title=section["title"],
|
||||
content=section["content"],
|
||||
source_file=filename,
|
||||
course_name=final_course,
|
||||
teacher_name=final_teacher,
|
||||
live_date=final_date,
|
||||
page_number=page_number,
|
||||
)
|
||||
imported_pages += 1
|
||||
|
||||
# 分块并向量化
|
||||
chunks = self._chunk_text(section["content"])
|
||||
chunk_count = await self._insert_chunks(page_id, chunks)
|
||||
imported_chunks += chunk_count
|
||||
|
||||
logger.info(
|
||||
"Markdown 导入完成: file=%s, pages=%d, chunks=%d",
|
||||
filename, imported_pages, imported_chunks,
|
||||
)
|
||||
|
||||
return {
|
||||
"file": filename,
|
||||
"pages": imported_pages,
|
||||
"chunks": imported_chunks,
|
||||
}
|
||||
|
||||
async def _import_text(
|
||||
self,
|
||||
content: str,
|
||||
filename: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""导入纯文本文件,整体作为一个页面"""
|
||||
title = Path(filename).stem
|
||||
|
||||
page_id = await self._insert_page(
|
||||
title=title,
|
||||
content=content,
|
||||
source_file=filename,
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
live_date=live_date,
|
||||
page_number=1,
|
||||
)
|
||||
|
||||
chunks = self._chunk_text(content)
|
||||
chunk_count = await self._insert_chunks(page_id, chunks)
|
||||
|
||||
logger.info(
|
||||
"文本导入完成: file=%s, chunks=%d",
|
||||
filename, chunk_count,
|
||||
)
|
||||
|
||||
return {
|
||||
"file": filename,
|
||||
"pages": 1,
|
||||
"chunks": chunk_count,
|
||||
}
|
||||
|
||||
async def _import_docx(
|
||||
self,
|
||||
file_path: str,
|
||||
filename: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
导入 Word 文档(.docx)。
|
||||
提取所有段落文本,按标题(Heading 1/2)拆分为多个知识页面。
|
||||
"""
|
||||
from docx import Document
|
||||
|
||||
doc = Document(file_path)
|
||||
|
||||
# 提取所有段落,保留标题层级信息
|
||||
sections: List[dict] = [] # [{"title": "...", "content": "...", "level": int}]
|
||||
current_title = Path(filename).stem
|
||||
current_content: List[str] = []
|
||||
current_level = 0
|
||||
|
||||
for para in doc.paragraphs:
|
||||
style_name = (para.style.name or "").lower()
|
||||
|
||||
# 判断是否为标题段落
|
||||
if style_name.startswith("heading"):
|
||||
try:
|
||||
level = int(style_name.replace("heading", "").strip())
|
||||
except ValueError:
|
||||
level = 1
|
||||
|
||||
# 保存上一个段落
|
||||
if current_content:
|
||||
text = "\n\n".join(current_content).strip()
|
||||
if text:
|
||||
sections.append({
|
||||
"title": current_title,
|
||||
"content": text,
|
||||
"level": current_level,
|
||||
})
|
||||
|
||||
# 开始新段落
|
||||
current_title = para.text.strip() or f"第 {len(sections) + 1} 节"
|
||||
current_content = []
|
||||
current_level = level
|
||||
else:
|
||||
text = para.text.strip()
|
||||
if text:
|
||||
current_content.append(text)
|
||||
|
||||
# 保存最后一个段落
|
||||
if current_content:
|
||||
text = "\n\n".join(current_content).strip()
|
||||
if text:
|
||||
sections.append({
|
||||
"title": current_title,
|
||||
"content": text,
|
||||
"level": current_level,
|
||||
})
|
||||
|
||||
if not sections:
|
||||
raise ValueError(f"Word 文档内容为空: {filename}")
|
||||
|
||||
# 如果只有一个段落且没有标题,整体作为一个页面
|
||||
if len(sections) == 1 and sections[0]["level"] == 0:
|
||||
sections[0]["title"] = Path(filename).stem
|
||||
|
||||
imported_pages = 0
|
||||
imported_chunks = 0
|
||||
|
||||
for idx, section in enumerate(sections):
|
||||
page_id = await self._insert_page(
|
||||
title=section["title"],
|
||||
content=section["content"],
|
||||
source_file=filename,
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
live_date=live_date,
|
||||
page_number=idx + 1,
|
||||
)
|
||||
imported_pages += 1
|
||||
|
||||
chunks = self._chunk_text(section["content"])
|
||||
chunk_count = await self._insert_chunks(page_id, chunks)
|
||||
imported_chunks += chunk_count
|
||||
|
||||
logger.info(
|
||||
"Word 文档导入完成: file=%s, pages=%d, chunks=%d",
|
||||
filename, imported_pages, imported_chunks,
|
||||
)
|
||||
|
||||
return {
|
||||
"file": filename,
|
||||
"pages": imported_pages,
|
||||
"chunks": imported_chunks,
|
||||
}
|
||||
|
||||
async def _insert_page(
|
||||
self,
|
||||
title: str,
|
||||
content: str,
|
||||
source_file: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
page_number: Optional[int] = None,
|
||||
) -> int:
|
||||
"""插入知识页面记录,返回页面 ID"""
|
||||
sql = text("""
|
||||
INSERT INTO knowledge_pages (title, content, source_file, course_name, teacher_name, live_date, page_number)
|
||||
VALUES (:title, :content, :source_file, :course_name, :teacher_name, :live_date, :page_number)
|
||||
RETURNING id
|
||||
""")
|
||||
result = await self.db.execute(sql, {
|
||||
"title": title,
|
||||
"content": content,
|
||||
"source_file": source_file,
|
||||
"course_name": course_name,
|
||||
"teacher_name": teacher_name,
|
||||
"live_date": live_date,
|
||||
"page_number": page_number,
|
||||
})
|
||||
row = result.fetchone()
|
||||
await self.db.flush()
|
||||
return row.id
|
||||
|
||||
async def _insert_chunks(self, page_id: int, chunks: List[str]) -> int:
|
||||
"""
|
||||
批量插入分块记录并生成嵌入向量。
|
||||
|
||||
Args:
|
||||
page_id: 关联的知识页面 ID
|
||||
chunks: 分块文本列表
|
||||
|
||||
Returns:
|
||||
插入的分块数量
|
||||
"""
|
||||
if not chunks:
|
||||
return 0
|
||||
|
||||
# 批量生成嵌入向量
|
||||
try:
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
embeddings = await EmbeddingService.embed_batch(chunks)
|
||||
except Exception as exc:
|
||||
logger.error("嵌入向量生成失败,分块将不包含向量: %s", exc)
|
||||
embeddings = [None] * len(chunks)
|
||||
|
||||
# 批量插入
|
||||
for idx, (chunk_text, embedding) in enumerate(zip(chunks, embeddings)):
|
||||
if IS_SQLITE:
|
||||
# SQLite: 向量存为 JSON 文本
|
||||
import json as _json
|
||||
embedding_str = _json.dumps(embedding) if embedding else None
|
||||
sql = text("""
|
||||
INSERT INTO knowledge_chunks (page_id, chunk_index, content, embedding)
|
||||
VALUES (:page_id, :chunk_index, :content, :embedding)
|
||||
""")
|
||||
else:
|
||||
# PostgreSQL: 使用 pgvector 类型
|
||||
embedding_str = "[" + ",".join(str(x) for x in embedding) + "]" if embedding else None
|
||||
sql = text("""
|
||||
INSERT INTO knowledge_chunks (page_id, chunk_index, content, embedding)
|
||||
VALUES (:page_id, :chunk_index, :content, :embedding::vector)
|
||||
""")
|
||||
await self.db.execute(sql, {
|
||||
"page_id": page_id,
|
||||
"chunk_index": idx,
|
||||
"content": chunk_text,
|
||||
"embedding": embedding_str,
|
||||
})
|
||||
|
||||
# 更新全文搜索向量(仅 PostgreSQL)
|
||||
if not IS_SQLITE:
|
||||
await self.db.execute(text("""
|
||||
UPDATE knowledge_chunks
|
||||
SET search_vector = to_tsvector('chinese_zh', content)
|
||||
WHERE page_id = :page_id
|
||||
"""), {"page_id": page_id})
|
||||
|
||||
await self.db.flush()
|
||||
return len(chunks)
|
||||
|
||||
@staticmethod
|
||||
def _split_markdown_sections(body: str) -> List[dict]:
|
||||
"""
|
||||
按二级标题(##)拆分 Markdown 内容为多个段落。
|
||||
|
||||
Returns:
|
||||
[{"title": "...", "content": "..."}, ...]
|
||||
"""
|
||||
sections: List[dict] = []
|
||||
# 匹配 ## 开头的标题
|
||||
pattern = re.compile(r"^##\s+(.+)$", re.MULTILINE)
|
||||
matches = list(pattern.finditer(body))
|
||||
|
||||
if not matches:
|
||||
return []
|
||||
|
||||
for i, match in enumerate(matches):
|
||||
title = match.group(1).strip()
|
||||
start = match.end()
|
||||
end = matches[i + 1].start() if i + 1 < len(matches) else len(body)
|
||||
content = body[start:end].strip()
|
||||
if content:
|
||||
sections.append({"title": title, "content": content})
|
||||
|
||||
return sections
|
||||
|
||||
@staticmethod
|
||||
def _chunk_text(text: str) -> List[str]:
|
||||
"""
|
||||
将文本按固定大小分块,保留重叠部分。
|
||||
|
||||
使用简单的字符数分块策略,按段落边界切分以保持语义完整性。
|
||||
"""
|
||||
chunk_size = settings.CHUNK_SIZE
|
||||
overlap = settings.CHUNK_OVERLAP
|
||||
|
||||
if not text or len(text) <= chunk_size:
|
||||
return [text] if text.strip() else []
|
||||
|
||||
# 按段落分割
|
||||
paragraphs = re.split(r"\n{2,}", text)
|
||||
chunks: List[str] = []
|
||||
current_chunk = ""
|
||||
|
||||
for para in paragraphs:
|
||||
para = para.strip()
|
||||
if not para:
|
||||
continue
|
||||
|
||||
if len(current_chunk) + len(para) + 2 <= chunk_size:
|
||||
# 当前块还能容纳这个段落
|
||||
if current_chunk:
|
||||
current_chunk += "\n\n" + para
|
||||
else:
|
||||
current_chunk = para
|
||||
else:
|
||||
# 当前块已满,保存并开始新块
|
||||
if current_chunk:
|
||||
chunks.append(current_chunk)
|
||||
|
||||
if len(para) > chunk_size:
|
||||
# 单个段落超过 chunk_size,强制切分
|
||||
for i in range(0, len(para), chunk_size - overlap):
|
||||
piece = para[i : i + chunk_size]
|
||||
if piece.strip():
|
||||
chunks.append(piece)
|
||||
current_chunk = ""
|
||||
else:
|
||||
current_chunk = para
|
||||
|
||||
if current_chunk.strip():
|
||||
chunks.append(current_chunk)
|
||||
|
||||
return chunks
|
||||
479
app/services/llm_service.py
Normal file
479
app/services/llm_service.py
Normal file
@@ -0,0 +1,479 @@
|
||||
"""
|
||||
LLM 服务模块
|
||||
提供统一的 LLM 调用接口,用于查询扩展、标签提取等任务
|
||||
支持 MiniMax / DeepSeek / OpenAI(均为 OpenAI 兼容格式)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Optional
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMService:
|
||||
"""
|
||||
LLM 服务类
|
||||
根据配置的 LLM_PROVIDER 自动选择对应的 API 端点。
|
||||
所有支持的提供商均兼容 OpenAI Chat Completions 格式。
|
||||
"""
|
||||
|
||||
_client = None
|
||||
|
||||
@classmethod
|
||||
def _get_client(cls):
|
||||
"""延迟初始化 OpenAI 兼容客户端"""
|
||||
if cls._client is not None:
|
||||
return cls._client
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
provider = settings.LLM_PROVIDER.lower()
|
||||
|
||||
if provider == "minimax":
|
||||
cls._client = AsyncOpenAI(
|
||||
api_key=settings.MINIMAX_API_KEY or "EMPTY",
|
||||
base_url=settings.MINIMAX_BASE_URL,
|
||||
)
|
||||
cls._model = settings.MINIMAX_CHAT_MODEL
|
||||
elif provider == "deepseek":
|
||||
cls._client = AsyncOpenAI(
|
||||
api_key=settings.DEEPSEEK_API_KEY or "EMPTY",
|
||||
base_url=settings.DEEPSEEK_BASE_URL,
|
||||
)
|
||||
cls._model = settings.DEEPSEEK_OCR_MODEL
|
||||
elif provider == "openai":
|
||||
cls._client = AsyncOpenAI(
|
||||
api_key=settings.OPENAI_API_KEY or "EMPTY",
|
||||
base_url=settings.OPENAI_BASE_URL,
|
||||
)
|
||||
cls._model = settings.EMBEDDING_MODEL
|
||||
else:
|
||||
raise ValueError(f"不支持的 LLM 提供商: {provider}")
|
||||
|
||||
logger.info("LLM 服务初始化完成: provider=%s, model=%s", provider, cls._model)
|
||||
return cls._client
|
||||
|
||||
@staticmethod
|
||||
def _clean_reasoning_content(content: str) -> str:
|
||||
"""清理推理模型的思考过程,只保留最终回答"""
|
||||
import re
|
||||
|
||||
# 方式1: 清理 🧠...💠 标签(MiniMax M2.7 reasoning 模式)
|
||||
content = re.sub(r'🧠[\s\S]*?💠', '', content)
|
||||
|
||||
# 方式2: 清理 💭(U+1F4AD) 包裹的思考块
|
||||
if '💭' in content:
|
||||
parts = content.split('💭')
|
||||
answer_parts = [parts[i] for i in range(len(parts)) if i % 2 == 1]
|
||||
if answer_parts:
|
||||
content = ''.join(answer_parts)
|
||||
|
||||
# 方式3: 清理 </think 标签(M2.7 有时使用)
|
||||
think_end = content.find('</think')
|
||||
if think_end >= 0:
|
||||
after = content[think_end + len('</think'):]
|
||||
after = re.sub(r'^[>\n\s]*', '', after)
|
||||
content = after
|
||||
|
||||
# 方式4: 清理 \x0f 分隔的思考块
|
||||
if '\x0f' in content:
|
||||
parts = content.split('\x0f')
|
||||
content = ''.join(parts[1::2]) if len(parts) > 1 else parts[-1] if parts else content
|
||||
|
||||
# 方式5: 清理 0x00 等特殊字符
|
||||
content = content.replace('\x00', '').strip()
|
||||
|
||||
# 方式6: 如果内容仍然很长(说明还有思考残留),取最后一个短段落作为答案
|
||||
if len(content) > 100 and '\n\n' in content:
|
||||
segments = [s.strip() for s in content.split('\n\n')]
|
||||
for seg in reversed(segments):
|
||||
if seg and len(seg) <= 100:
|
||||
content = seg
|
||||
break
|
||||
|
||||
return content.strip()
|
||||
|
||||
@classmethod
|
||||
async def chat(cls, messages: List[dict], temperature: float = 0.3, max_tokens: int = 1024) -> str:
|
||||
"""
|
||||
发送聊天请求。
|
||||
|
||||
Args:
|
||||
messages: 消息列表,格式 [{"role": "system/user/assistant", "content": "..."}]
|
||||
temperature: 采样温度
|
||||
max_tokens: 最大生成 token 数
|
||||
|
||||
Returns:
|
||||
模型回复文本(已清理思考过程)
|
||||
"""
|
||||
client = cls._get_client()
|
||||
|
||||
# MiniMax 模型默认开启思考模式,使用 reasoning_split 将思考内容
|
||||
# 分离到 reasoning_details 字段,content 只保留最终回答
|
||||
extra_body = {}
|
||||
if settings.LLM_PROVIDER.lower() == "minimax":
|
||||
extra_body["reasoning_split"] = True
|
||||
|
||||
response = await client.chat.completions.create(
|
||||
model=cls._model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
**({"extra_body": extra_body} if extra_body else {}),
|
||||
)
|
||||
content = response.choices[0].message.content or ""
|
||||
|
||||
# 兜底清理:以防 reasoning_split 未生效或使用了其他模型
|
||||
content = cls._clean_reasoning_content(content)
|
||||
|
||||
return content
|
||||
|
||||
@classmethod
|
||||
async def expand_query(cls, query: str) -> List[str]:
|
||||
"""
|
||||
查询扩展:将原始查询改写为多个语义相近的表述。
|
||||
用于混合搜索时提高召回率。
|
||||
|
||||
Args:
|
||||
query: 原始查询文本
|
||||
|
||||
Returns:
|
||||
扩展后的查询列表(包含原始查询)
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个搜索查询优化助手。用户会给你一个搜索查询,"
|
||||
"你需要生成 2 个语义相近但表述不同的替代查询。"
|
||||
"只输出 JSON 数组格式,不要添加任何其他文字。"
|
||||
'例如:["原始查询", "替代查询1", "替代查询2"]'
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"请为以下查询生成 2 个替代表述:{query}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
result = await cls.chat(messages, temperature=0.3, max_tokens=256)
|
||||
# 尝试解析 JSON
|
||||
result = result.strip()
|
||||
if result.startswith("```"):
|
||||
result = result.split("\n", 1)[-1]
|
||||
if result.endswith("```"):
|
||||
result = result[:-3]
|
||||
result = result.strip()
|
||||
queries = json.loads(result)
|
||||
if isinstance(queries, list):
|
||||
# 去重并限制数量
|
||||
seen = set()
|
||||
unique = []
|
||||
for q in queries:
|
||||
q = str(q).strip()
|
||||
if q and q not in seen:
|
||||
seen.add(q)
|
||||
unique.append(q)
|
||||
return unique[:3]
|
||||
except Exception as exc:
|
||||
logger.warning("查询扩展失败: %s,使用原始查询", exc)
|
||||
|
||||
return [query]
|
||||
|
||||
@classmethod
|
||||
async def detect_entities(cls, text: str) -> List[dict]:
|
||||
"""
|
||||
实体检测:从文本中提取知识点、人物、概念等实体。
|
||||
|
||||
Args:
|
||||
text: 待分析的文本
|
||||
|
||||
Returns:
|
||||
实体列表,格式 [{"name": "...", "type": "knowledge_point|person|concept|technique", "description": "..."}]
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个教育内容分析助手。用户会给你一段直播课的文字稿,"
|
||||
"你需要从中提取关键实体(知识点、技巧、人物、概念等)。"
|
||||
"只输出 JSON 数组格式,不要添加任何其他文字。\n"
|
||||
'格式:[{"name": "实体名称", "type": "knowledge_point|technique|person|concept", "description": "简要描述"}]\n'
|
||||
"type 说明:\n"
|
||||
"- knowledge_point: 知识点(如公式、定理、定义)\n"
|
||||
"- technique: 技巧/方法(如解题技巧、记忆方法)\n"
|
||||
"- person: 人物\n"
|
||||
"- concept: 概念/术语"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"请从以下文字中提取关键实体:\n\n{text[:3000]}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
result = await cls.chat(messages, temperature=0.1, max_tokens=2048)
|
||||
result = result.strip()
|
||||
if result.startswith("```"):
|
||||
result = result.split("\n", 1)[-1]
|
||||
if result.endswith("```"):
|
||||
result = result[:-3]
|
||||
result = result.strip()
|
||||
entities = json.loads(result)
|
||||
if isinstance(entities, list):
|
||||
return entities
|
||||
except Exception as exc:
|
||||
logger.warning("实体检测失败: %s", exc)
|
||||
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
async def extract_tags(cls, text: str) -> List[str]:
|
||||
"""
|
||||
从文本中提取语义标签。
|
||||
用于图片 OCR 内容的标签化和搜索查询的标签化。
|
||||
|
||||
Args:
|
||||
text: 待提取标签的文本
|
||||
|
||||
Returns:
|
||||
标签列表,如 ["教育", "高考数学", "二次方程", "解题技巧"]
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个内容标签提取助手。用户会给你一段文字内容,"
|
||||
"你需要从中提取 3-10 个能概括内容主题的标签。\n"
|
||||
"标签要求:\n"
|
||||
"1. 用简短的词语(2-6个字)\n"
|
||||
"2. 涵盖主题、场景、情感、关键实体等维度\n"
|
||||
"3. 考虑同义词和近义词(如'不想工作'也打上'躺平'标签)\n"
|
||||
"4. 只输出 JSON 数组,不要其他文字\n"
|
||||
'例如:["高考数学", "二次方程", "韦达定理", "解题技巧"]'
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"请从以下内容中提取标签:\n\n{text[:2000]}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
result = await cls.chat(messages, temperature=0.3, max_tokens=256)
|
||||
result = result.strip()
|
||||
if result.startswith("```"):
|
||||
result = result.split("\n", 1)[-1]
|
||||
if result.endswith("```"):
|
||||
result = result[:-3]
|
||||
result = result.strip()
|
||||
tags = json.loads(result)
|
||||
if isinstance(tags, list):
|
||||
# 去重、清洗
|
||||
seen = set()
|
||||
clean = []
|
||||
for t in tags:
|
||||
t = str(t).strip()
|
||||
if 1 <= len(t) <= 20 and t not in seen:
|
||||
seen.add(t)
|
||||
clean.append(t)
|
||||
return clean[:15]
|
||||
except Exception as exc:
|
||||
logger.warning("标签提取失败: %s", exc)
|
||||
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
async def extract_search_tags(cls, query: str) -> List[str]:
|
||||
"""
|
||||
从用户搜索查询中提取标签,用于匹配数据库中存储的标签。
|
||||
和 extract_tags 类似,但更侧重于提取搜索意图相关的标签。
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个搜索意图分析助手。用户会给你一个搜索查询,"
|
||||
"你需要提取 3-8 个可能匹配的标签词。\n"
|
||||
"标签要求:\n"
|
||||
"1. 用简短的词语(2-6个字)\n"
|
||||
"2. 包含同义词和近义词\n"
|
||||
"3. 包含上义词和下义词\n"
|
||||
"4. 只输出 JSON 数组\n"
|
||||
'例如:用户搜"孩子躺平",输出 ["躺平", "不上班", "年轻人", "就业", "摆烂", "啃老"]'
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"请从以下搜索查询中提取匹配标签:{query}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
result = await cls.chat(messages, temperature=0.3, max_tokens=256)
|
||||
result = result.strip()
|
||||
if result.startswith("```"):
|
||||
result = result.split("\n", 1)[-1]
|
||||
if result.endswith("```"):
|
||||
result = result[:-3]
|
||||
result = result.strip()
|
||||
tags = json.loads(result)
|
||||
if isinstance(tags, list):
|
||||
seen = set()
|
||||
clean = []
|
||||
for t in tags:
|
||||
t = str(t).strip()
|
||||
if 1 <= len(t) <= 20 and t not in seen:
|
||||
seen.add(t)
|
||||
clean.append(t)
|
||||
return clean[:10]
|
||||
except Exception as exc:
|
||||
logger.warning("搜索标签提取失败: %s", exc)
|
||||
|
||||
# 降级:直接用原始查询词作为标签
|
||||
return [query.strip()] if query.strip() else []
|
||||
|
||||
@classmethod
|
||||
async def summarize_story(cls, text: str) -> str:
|
||||
"""
|
||||
用 Qwen3-8B 提炼图片 OCR 文本的故事内容。
|
||||
非思考模式,直接输出。
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个内容提炼助手。用户会给你一段从图片中识别出的文字内容,"
|
||||
"你需要提炼出其中讲述的故事或事件。\n"
|
||||
"要求:\n"
|
||||
"1. 用简洁的自然语言描述图片内容讲述的故事或事件\n"
|
||||
"2. 保留关键人物、事件、情感、场景等核心信息\n"
|
||||
"3. 50-200字\n"
|
||||
"4. 直接输出提炼内容,不要加前缀"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"/no_think\n请提炼以下内容的故事:\n\n{text[:3000]}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
from openai import AsyncOpenAI
|
||||
client = AsyncOpenAI(
|
||||
api_key=settings.OPENAI_API_KEY,
|
||||
base_url=settings.OPENAI_BASE_URL,
|
||||
)
|
||||
response = await client.chat.completions.create(
|
||||
model="Qwen/Qwen3-8B",
|
||||
messages=messages,
|
||||
temperature=0.3,
|
||||
max_tokens=512,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||
)
|
||||
content = response.choices[0].message.content or ""
|
||||
# 清理可能的 /no_think 回显
|
||||
content = content.replace("/no_think", "").strip()
|
||||
return content
|
||||
except Exception as exc:
|
||||
logger.warning("故事提炼失败: %s", exc)
|
||||
return ""
|
||||
|
||||
@classmethod
|
||||
async def judge_batch(cls, query: str, articles: List[dict]) -> List[int]:
|
||||
"""
|
||||
用 LLM 批量判断哪些文章符合查询。
|
||||
|
||||
Args:
|
||||
query: 用户搜索查询
|
||||
articles: [{"id": int, "text": str}, ...],最多10个
|
||||
|
||||
Returns:
|
||||
匹配的文章 id 列表(空列表表示都不匹配)
|
||||
"""
|
||||
if not articles:
|
||||
return []
|
||||
|
||||
# 构建文章列表
|
||||
articles_text = ""
|
||||
for i, article in enumerate(articles, 1):
|
||||
text = (article.get("text") or "")[:800]
|
||||
articles_text += f"{i}. {text}\n"
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个严格的语义搜索匹配引擎。用户会给你一个搜索查询和多篇文章片段,"
|
||||
"你需要判断哪些文章与查询在**核心语义上高度相关**。\n\n"
|
||||
"判断原则:\n"
|
||||
"- 文章的**核心主题**必须与查询直接相关,而不是仅仅提及了查询中的某个词\n"
|
||||
"- 例如:查询\"爸爸出轨\",文章只提到\"爸爸带孩子玩\"是不相关的,"
|
||||
"只有文章讨论出轨、婚姻危机、夫妻矛盾等才算相关\n"
|
||||
"- 例如:查询\"孩子不想上学\",文章提到\"孩子打游戏不去学校\"是相关的,"
|
||||
"但只提到\"孩子\"或\"学校\"其他方面则不相关\n\n"
|
||||
"不相关的典型情况(必须排除):\n"
|
||||
"- 仅包含查询中的某个词,但讨论的是完全不同的话题\n"
|
||||
"- 主题相近但具体内容无关(如查询出轨,文章讲的是正常的夫妻相处)\n"
|
||||
"- 只有一个宽泛的关联词,没有实质性的语义联系\n\n"
|
||||
"只输出匹配的文章编号,用逗号分隔。\n"
|
||||
"如果没有文章与查询相关,只输出 0。\n"
|
||||
"不要输出任何解释。宁可漏判,不可误判。"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"搜索查询: {query}\n\n{articles_text}\n请输出匹配的文章编号:",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
from openai import AsyncOpenAI
|
||||
client = AsyncOpenAI(
|
||||
api_key=settings.MINIMAX_API_KEY,
|
||||
base_url=settings.MINIMAX_BASE_URL,
|
||||
http_client=httpx.AsyncClient(timeout=60.0),
|
||||
)
|
||||
response = await client.chat.completions.create(
|
||||
model=settings.MINIMAX_CHAT_MODEL,
|
||||
messages=messages,
|
||||
temperature=0.0,
|
||||
max_tokens=512,
|
||||
stream=False,
|
||||
extra_body={"reasoning_split": True},
|
||||
)
|
||||
content = response.choices[0].message.content or ""
|
||||
content = cls._clean_reasoning_content(content)
|
||||
|
||||
# 解析返回的编号
|
||||
numbers = re.findall(r'\d+', content)
|
||||
|
||||
if not numbers:
|
||||
return []
|
||||
|
||||
matched_ids = []
|
||||
for num_str in numbers:
|
||||
num = int(num_str)
|
||||
if num == 0:
|
||||
continue # 0 表示都不匹配
|
||||
if 1 <= num <= len(articles):
|
||||
matched_ids.append(articles[num - 1]["id"])
|
||||
|
||||
return matched_ids
|
||||
except Exception as exc:
|
||||
logger.warning("LLM 批量判断失败: %s", exc)
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""重置客户端(用于测试或切换配置)"""
|
||||
cls._client = None
|
||||
540
app/services/ocr_service.py
Normal file
540
app/services/ocr_service.py
Normal file
@@ -0,0 +1,540 @@
|
||||
"""
|
||||
OCR 识别服务模块
|
||||
支持多种 OCR 提供商:PaddleOCR(本地)/ 阿里云 / 腾讯云
|
||||
auto 模式下本地优先 + 云端 fallback
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import asyncio
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from app.config import OCRProvider, settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ──────────────────────────── 数据结构 ────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class OCRBlock:
|
||||
"""OCR 识别的文本块"""
|
||||
text: str
|
||||
bbox: Optional[List[float]] = None # [x1, y1, x2, y2]
|
||||
confidence: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class OCRResult:
|
||||
"""OCR 识别结果"""
|
||||
text: str
|
||||
blocks: List[OCRBlock] = field(default_factory=list)
|
||||
confidence: float = 0.0
|
||||
provider: str = ""
|
||||
keywords: List[str] = field(default_factory=list) # OCR 时提取的关键词标签
|
||||
|
||||
|
||||
# ──────────────────────────── 抽象基类 ────────────────────────────
|
||||
|
||||
class OCRProviderBase(abc.ABC):
|
||||
"""OCR 提供商抽象基类"""
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def name(self) -> str:
|
||||
"""提供商名称"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""
|
||||
识别图片中的文字。
|
||||
|
||||
Args:
|
||||
image_path: 图片文件路径
|
||||
|
||||
Returns:
|
||||
OCRResult 包含识别文本、文本块和置信度
|
||||
|
||||
Raises:
|
||||
Exception: 识别失败时抛出异常
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
# ──────────────────────────── PaddleOCR 实现(本地) ────────────────────────────
|
||||
|
||||
class PaddleOCRProvider(OCRProviderBase):
|
||||
"""
|
||||
PaddleOCR 本地识别
|
||||
优点:免费、离线可用、中文识别效果好
|
||||
缺点:需要较多内存、首次加载模型较慢
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._ocr = None
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "paddleocr"
|
||||
|
||||
def _init_ocr(self):
|
||||
"""延迟初始化 PaddleOCR 引擎"""
|
||||
if self._ocr is not None:
|
||||
return
|
||||
|
||||
logger.info("正在初始化 PaddleOCR 引擎...")
|
||||
from paddleocr import PaddleOCR
|
||||
|
||||
self._ocr = PaddleOCR(
|
||||
use_angle_cls=True,
|
||||
lang="ch",
|
||||
use_gpu=False, # 默认 CPU,生产环境可通过环境变量控制
|
||||
show_log=False,
|
||||
enable_mkldnn=True, # 启用 MKLDNN 加速
|
||||
)
|
||||
logger.info("PaddleOCR 引擎初始化完成")
|
||||
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""使用 PaddleOCR 识别图片"""
|
||||
self._init_ocr()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _run_ocr():
|
||||
result = self._ocr.ocr(image_path, cls=True)
|
||||
return result
|
||||
|
||||
raw_result = await loop.run_in_executor(None, _run_ocr)
|
||||
|
||||
# 解析 PaddleOCR 返回结果
|
||||
blocks: List[OCRBlock] = []
|
||||
all_text_parts: List[str] = []
|
||||
total_confidence = 0.0
|
||||
|
||||
if raw_result and raw_result[0]:
|
||||
for line in raw_result[0]:
|
||||
bbox_points = line[0]
|
||||
text = line[1][0]
|
||||
confidence = float(line[1][1])
|
||||
|
||||
# 将四个角点转为 [x1, y1, x2, y2]
|
||||
xs = [p[0] for p in bbox_points]
|
||||
ys = [p[1] for p in bbox_points]
|
||||
bbox = [min(xs), min(ys), max(xs), max(ys)]
|
||||
|
||||
blocks.append(OCRBlock(text=text, bbox=bbox, confidence=confidence))
|
||||
all_text_parts.append(text)
|
||||
total_confidence += confidence
|
||||
|
||||
full_text = "\n".join(all_text_parts)
|
||||
avg_confidence = total_confidence / len(blocks) if blocks else 0.0
|
||||
|
||||
return OCRResult(
|
||||
text=full_text,
|
||||
blocks=blocks,
|
||||
confidence=avg_confidence,
|
||||
provider=self.name,
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 阿里云 OCR 实现 ────────────────────────────
|
||||
|
||||
class AliyunOCRProvider(OCRProviderBase):
|
||||
"""
|
||||
阿里云 OCR 识别
|
||||
使用阿里云文字识别 API(通用文字识别)
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
if not settings.ALIYUN_OCR_ACCESS_KEY or not settings.ALIYUN_OCR_SECRET:
|
||||
raise ValueError("阿里云 OCR 需要配置 ALIYUN_OCR_ACCESS_KEY 和 ALIYUN_OCR_SECRET")
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "aliyun"
|
||||
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""调用阿里云 OCR API 识别图片"""
|
||||
import base64
|
||||
import json
|
||||
|
||||
import asyncio
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _call_api():
|
||||
from alibabacloud_tea_openapi.models import Config
|
||||
from alibabacloud_ocr_api20210707.client import Client
|
||||
from alibabacloud_ocr_api20210707.models import RecognizeGeneralRequest
|
||||
from alibabacloud_tea_util.models import RuntimeOptions
|
||||
|
||||
config = Config(
|
||||
access_key_id=settings.ALIYUN_OCR_ACCESS_KEY,
|
||||
access_key_secret=settings.ALIYUN_OCR_SECRET,
|
||||
endpoint="ocr-api.cn-hangzhou.aliyuncs.com",
|
||||
)
|
||||
client = Client(config)
|
||||
|
||||
# 读取图片并 Base64 编码
|
||||
with open(image_path, "rb") as f:
|
||||
image_bytes = f.read()
|
||||
image_b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||||
|
||||
request = RecognizeGeneralRequest(
|
||||
body=json.dumps({"url": f"data:image/png;base64,{image_b64}"}),
|
||||
)
|
||||
runtime = RuntimeOptions()
|
||||
|
||||
response = client.recognize_general_with_options(request, runtime)
|
||||
return response
|
||||
|
||||
response = await loop.run_in_executor(None, _call_api)
|
||||
|
||||
if response.body and response.body.data:
|
||||
data = json.loads(response.body.data)
|
||||
blocks: List[OCRBlock] = []
|
||||
all_text_parts: List[str] = []
|
||||
total_confidence = 0.0
|
||||
|
||||
# 阿里云返回格式可能包含多个文本块
|
||||
if isinstance(data, dict) and "content" in data:
|
||||
text = data["content"]
|
||||
all_text_parts.append(text)
|
||||
blocks.append(OCRBlock(text=text, confidence=1.0))
|
||||
total_confidence = 1.0
|
||||
elif isinstance(data, list):
|
||||
for item in data:
|
||||
text = item.get("text", item.get("word", ""))
|
||||
score = float(item.get("score", item.get("confidence", 1.0)))
|
||||
blocks.append(OCRBlock(text=text, confidence=score))
|
||||
all_text_parts.append(text)
|
||||
total_confidence += score
|
||||
|
||||
full_text = "\n".join(all_text_parts)
|
||||
avg_confidence = total_confidence / len(blocks) if blocks else 0.0
|
||||
|
||||
return OCRResult(
|
||||
text=full_text,
|
||||
blocks=blocks,
|
||||
confidence=avg_confidence,
|
||||
provider=self.name,
|
||||
)
|
||||
|
||||
return OCRResult(text="", blocks=[], confidence=0.0, provider=self.name)
|
||||
|
||||
|
||||
# ──────────────────────────── 腾讯云 OCR 实现 ────────────────────────────
|
||||
|
||||
class TencentOCRProvider(OCRProviderBase):
|
||||
"""
|
||||
腾讯云 OCR 识别
|
||||
使用腾讯云文字识别 API(通用印刷体识别)
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
if not settings.TENCENT_OCR_SECRET_ID or not settings.TENCENT_OCR_SECRET_KEY:
|
||||
raise ValueError("腾讯云 OCR 需要配置 TENCENT_OCR_SECRET_ID 和 TENCENT_OCR_SECRET_KEY")
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "tencent"
|
||||
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""调用腾讯云 OCR API 识别图片"""
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _call_api():
|
||||
from tencentcloud.common import credential
|
||||
from tencentcloud.common.profile.client_profile import ClientProfile
|
||||
from tencentcloud.common.profile.http_profile import HttpProfile
|
||||
from tencentcloud.ocr.v20181119 import ocr_client, models
|
||||
|
||||
cred = credential.Credential(
|
||||
settings.TENCENT_OCR_SECRET_ID,
|
||||
settings.TENCENT_OCR_SECRET_KEY,
|
||||
)
|
||||
http_profile = HttpProfile()
|
||||
http_profile.endpoint = "ocr.tencentcloudapi.com"
|
||||
|
||||
client_profile = ClientProfile()
|
||||
client_profile.httpProfile = http_profile
|
||||
|
||||
client = ocr_client.OcrClient(cred, "", client_profile)
|
||||
|
||||
# 读取图片并 Base64 编码
|
||||
with open(image_path, "rb") as f:
|
||||
image_bytes = f.read()
|
||||
image_b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||||
|
||||
req = models.GeneralBasicOCRRequest()
|
||||
req.ImageBase64 = image_b64
|
||||
|
||||
resp = client.GeneralBasicOCR(req)
|
||||
return resp
|
||||
|
||||
resp = await loop.run_in_executor(None, _call_api)
|
||||
|
||||
blocks: List[OCRBlock] = []
|
||||
all_text_parts: List[str] = []
|
||||
total_confidence = 0.0
|
||||
|
||||
if resp.TextDetections:
|
||||
for item in resp.TextDetections:
|
||||
text = item.DetectedText or ""
|
||||
confidence = item.Confidence / 100.0 if item.Confidence else 0.0
|
||||
|
||||
# 解析多边形顶点
|
||||
bbox = None
|
||||
if item.Polygon:
|
||||
xs = [p.X for p in item.Polygon]
|
||||
ys = [p.Y for p in item.Polygon]
|
||||
bbox = [min(xs), min(ys), max(xs), max(ys)]
|
||||
|
||||
blocks.append(OCRBlock(text=text, bbox=bbox, confidence=confidence))
|
||||
all_text_parts.append(text)
|
||||
total_confidence += confidence
|
||||
|
||||
full_text = "\n".join(all_text_parts)
|
||||
avg_confidence = total_confidence / len(blocks) if blocks else 0.0
|
||||
|
||||
return OCRResult(
|
||||
text=full_text,
|
||||
blocks=blocks,
|
||||
confidence=avg_confidence,
|
||||
provider=self.name,
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── DeepSeek OCR 实现 ────────────────────────────
|
||||
|
||||
class DeepSeekOCRProvider(OCRProviderBase):
|
||||
"""
|
||||
DeepSeek Vision OCR
|
||||
通过 OpenAI 兼容接口发送图片 base64,让视觉模型识别文字。
|
||||
支持 DeepSeek 官方 API 和 Ollama 自部署(如 deepseek-ocr 模型)。
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
if not settings.DEEPSEEK_API_KEY:
|
||||
raise ValueError("DeepSeek OCR 需要配置 DEEPSEEK_API_KEY")
|
||||
import httpx
|
||||
from openai import AsyncOpenAI
|
||||
self._client = AsyncOpenAI(
|
||||
api_key=settings.DEEPSEEK_API_KEY,
|
||||
base_url=settings.DEEPSEEK_BASE_URL,
|
||||
http_client=httpx.AsyncClient(timeout=120.0), # OCR 较慢,超时 120 秒
|
||||
)
|
||||
self._model = settings.DEEPSEEK_OCR_MODEL
|
||||
# 判断是否为 Ollama 部署
|
||||
self._is_ollama = "ollama" in settings.DEEPSEEK_API_KEY.lower() or \
|
||||
"11434" in settings.DEEPSEEK_BASE_URL
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "deepseek"
|
||||
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""使用 DeepSeek Vision 模型识别图片中的文字"""
|
||||
import base64
|
||||
|
||||
# 读取图片并编码为 base64
|
||||
with open(image_path, "rb") as f:
|
||||
image_bytes = f.read()
|
||||
base64_image = base64.b64encode(image_bytes).decode("utf-8")
|
||||
|
||||
# 根据图片格式确定 MIME 类型
|
||||
suffix = Path(image_path).suffix.lower()
|
||||
mime_map = {
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".bmp": "image/bmp",
|
||||
".webp": "image/webp",
|
||||
}
|
||||
mime_type = mime_map.get(suffix, "image/png")
|
||||
|
||||
system_prompt = (
|
||||
"你是一个 OCR 文字识别工具。"
|
||||
"请识别图片中的所有文字,逐行输出,不要添加任何解释。"
|
||||
)
|
||||
|
||||
response = await self._client.chat.completions.create(
|
||||
model=self._model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": system_prompt,
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:{mime_type};base64,{base64_image}"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "请识别这张图片中的所有文字内容。",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
max_tokens=4096,
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
text = response.choices[0].message.content or ""
|
||||
|
||||
# 清理特殊 token(deepseek-ocr 模型可能输出 <|begin of sentence|> 等)
|
||||
import re
|
||||
text = re.sub(r'<|[^>]*|>', '', text)
|
||||
text = re.sub(r'<\|[^>]*\|>', '', text)
|
||||
|
||||
# 清理可能的 markdown 包裹
|
||||
text = text.strip()
|
||||
if text.startswith("```"):
|
||||
text = text.split("\n", 1)[-1]
|
||||
if text.endswith("```"):
|
||||
text = text[:-3]
|
||||
text = text.strip()
|
||||
|
||||
lines = [line for line in text.split("\n") if line.strip()]
|
||||
blocks = [OCRBlock(text=line, confidence=1.0) for line in lines]
|
||||
|
||||
return OCRResult(
|
||||
text=text,
|
||||
blocks=blocks,
|
||||
confidence=1.0,
|
||||
provider=self.name,
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 工厂类 ────────────────────────────
|
||||
|
||||
class OCRService:
|
||||
"""
|
||||
OCR 服务工厂类
|
||||
根据配置自动选择 OCR 提供商。
|
||||
auto 模式下:本地 PaddleOCR 优先,失败后 fallback 到云端 OCR。
|
||||
"""
|
||||
|
||||
_instance: Optional[OCRProviderBase] = None
|
||||
|
||||
@classmethod
|
||||
def _create_provider(cls, provider: OCRProvider) -> OCRProviderBase:
|
||||
"""根据提供商枚举创建对应的 OCR 实例"""
|
||||
provider_map = {
|
||||
OCRProvider.PADDLEOCR: PaddleOCRProvider,
|
||||
OCRProvider.ALIYUN: AliyunOCRProvider,
|
||||
OCRProvider.TENCENT: TencentOCRProvider,
|
||||
OCRProvider.DEEPSEEK: DeepSeekOCRProvider,
|
||||
}
|
||||
provider_cls = provider_map.get(provider)
|
||||
if provider_cls is None:
|
||||
raise ValueError(f"不支持的 OCR 提供商: {provider}")
|
||||
return provider_cls()
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> OCRProviderBase:
|
||||
"""获取 OCR 服务单例实例"""
|
||||
if cls._instance is not None:
|
||||
return cls._instance
|
||||
|
||||
if settings.OCR_PROVIDER == OCRProvider.AUTO:
|
||||
# auto 模式返回 PaddleOCR,fallback 逻辑在 recognize_auto 中处理
|
||||
cls._instance = PaddleOCRProvider()
|
||||
else:
|
||||
cls._instance = cls._create_provider(settings.OCR_PROVIDER)
|
||||
|
||||
logger.info("OCR 服务初始化完成: provider=%s", cls._instance.name)
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
async def recognize(cls, image_path: str) -> OCRResult:
|
||||
"""
|
||||
识别图片文字。
|
||||
auto 模式下:PaddleOCR 优先,失败后依次尝试阿里云、腾讯云。
|
||||
"""
|
||||
if settings.OCR_PROVIDER == OCRProvider.AUTO:
|
||||
return await cls._recognize_auto(image_path)
|
||||
|
||||
instance = cls.get_instance()
|
||||
return await instance.recognize(image_path)
|
||||
|
||||
@classmethod
|
||||
async def _recognize_auto(cls, image_path: str) -> OCRResult:
|
||||
"""
|
||||
auto 模式识别逻辑:
|
||||
1. 先尝试本地 PaddleOCR
|
||||
2. 如果失败且配置了阿里云,尝试阿里云 OCR
|
||||
3. 如果仍失败且配置了腾讯云,尝试腾讯云 OCR
|
||||
"""
|
||||
# 第一步:本地 PaddleOCR
|
||||
try:
|
||||
provider = PaddleOCRProvider()
|
||||
result = await provider.recognize(image_path)
|
||||
if result.text.strip():
|
||||
logger.info("PaddleOCR 识别成功,文本长度: %d", len(result.text))
|
||||
return result
|
||||
logger.warning("PaddleOCR 返回空文本,尝试 fallback")
|
||||
except Exception as exc:
|
||||
logger.warning("PaddleOCR 识别失败: %s,尝试 fallback", exc)
|
||||
|
||||
# 第二步:DeepSeek OCR(如果配置了)
|
||||
if settings.DEEPSEEK_API_KEY:
|
||||
try:
|
||||
provider = DeepSeekOCRProvider()
|
||||
result = await provider.recognize(image_path)
|
||||
if result.text.strip():
|
||||
logger.info("DeepSeek OCR 识别成功,文本长度: %d", len(result.text))
|
||||
return result
|
||||
logger.warning("DeepSeek OCR 返回空文本,尝试 fallback")
|
||||
except Exception as exc:
|
||||
logger.warning("DeepSeek OCR 识别失败: %s,尝试 fallback", exc)
|
||||
|
||||
# 第三步:阿里云 OCR
|
||||
if settings.ALIYUN_OCR_ACCESS_KEY and settings.ALIYUN_OCR_SECRET:
|
||||
try:
|
||||
provider = AliyunOCRProvider()
|
||||
result = await provider.recognize(image_path)
|
||||
if result.text.strip():
|
||||
logger.info("阿里云 OCR 识别成功,文本长度: %d", len(result.text))
|
||||
return result
|
||||
logger.warning("阿里云 OCR 返回空文本,继续 fallback")
|
||||
except Exception as exc:
|
||||
logger.warning("阿里云 OCR 识别失败: %s,继续 fallback", exc)
|
||||
|
||||
# 第四步:腾讯云 OCR
|
||||
if settings.TENCENT_OCR_SECRET_ID and settings.TENCENT_OCR_SECRET_KEY:
|
||||
try:
|
||||
provider = TencentOCRProvider()
|
||||
result = await provider.recognize(image_path)
|
||||
if result.text.strip():
|
||||
logger.info("腾讯云 OCR 识别成功,文本长度: %d", len(result.text))
|
||||
return result
|
||||
logger.warning("腾讯云 OCR 返回空文本")
|
||||
except Exception as exc:
|
||||
logger.warning("腾讯云 OCR 识别失败: %s", exc)
|
||||
|
||||
# 所有 OCR 都失败
|
||||
return OCRResult(
|
||||
text="",
|
||||
blocks=[],
|
||||
confidence=0.0,
|
||||
provider="none",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""重置单例(用于测试或切换配置)"""
|
||||
cls._instance = None
|
||||
202
app/services/page_service.py
Normal file
202
app/services/page_service.py
Normal file
@@ -0,0 +1,202 @@
|
||||
"""
|
||||
知识页面服务
|
||||
提供知识页面的 CRUD 操作
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import func, select, text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models.base import KnowledgeChunk, KnowledgePage
|
||||
from app.schemas.page import PageCreate, PageResponse, PageUpdate
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PageService:
|
||||
"""知识页面服务"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self.db = db
|
||||
|
||||
async def create_page(self, data: PageCreate) -> KnowledgePage:
|
||||
"""创建知识页面"""
|
||||
page = KnowledgePage(
|
||||
title=data.title,
|
||||
content=data.content,
|
||||
source_file=data.source_file,
|
||||
course_name=data.course_name,
|
||||
teacher_name=data.teacher_name,
|
||||
live_date=data.live_date,
|
||||
page_number=data.page_number,
|
||||
metadata_json=data.metadata_json,
|
||||
)
|
||||
self.db.add(page)
|
||||
await self.db.flush()
|
||||
await self.db.refresh(page)
|
||||
return page
|
||||
|
||||
async def get_page(self, page_id: int) -> Optional[KnowledgePage]:
|
||||
"""根据 ID 获取知识页面"""
|
||||
result = await self.db.execute(
|
||||
select(KnowledgePage).where(KnowledgePage.id == page_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def list_pages(
|
||||
self,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
keyword: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
分页查询知识页面列表。
|
||||
|
||||
Args:
|
||||
page: 页码(从 1 开始)
|
||||
page_size: 每页数量
|
||||
course_name: 按课程名称过滤
|
||||
teacher_name: 按讲师名称过滤
|
||||
keyword: 关键词搜索(标题或内容)
|
||||
|
||||
Returns:
|
||||
{"total": int, "page": int, "page_size": int, "items": list}
|
||||
"""
|
||||
query = select(KnowledgePage)
|
||||
count_query = select(func.count(KnowledgePage.id))
|
||||
|
||||
# 过滤条件
|
||||
if course_name:
|
||||
query = query.where(KnowledgePage.course_name == course_name)
|
||||
count_query = count_query.where(KnowledgePage.course_name == course_name)
|
||||
if teacher_name:
|
||||
query = query.where(KnowledgePage.teacher_name == teacher_name)
|
||||
count_query = count_query.where(KnowledgePage.teacher_name == teacher_name)
|
||||
if keyword:
|
||||
like_pattern = f"%{keyword}%"
|
||||
query = query.where(
|
||||
KnowledgePage.title.ilike(like_pattern)
|
||||
| KnowledgePage.content.ilike(like_pattern)
|
||||
)
|
||||
count_query = count_query.where(
|
||||
KnowledgePage.title.ilike(like_pattern)
|
||||
| KnowledgePage.content.ilike(like_pattern)
|
||||
)
|
||||
|
||||
# 总数
|
||||
total_result = await self.db.execute(count_query)
|
||||
total = total_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(KnowledgePage.created_at.desc()).offset(offset).limit(page_size)
|
||||
result = await self.db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
"items": items,
|
||||
}
|
||||
|
||||
async def update_page(self, page_id: int, data: PageUpdate) -> Optional[KnowledgePage]:
|
||||
"""更新知识页面"""
|
||||
page = await self.get_page(page_id)
|
||||
if page is None:
|
||||
return None
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
for field, value in update_data.items():
|
||||
setattr(page, field, value)
|
||||
|
||||
await self.db.flush()
|
||||
await self.db.refresh(page)
|
||||
return page
|
||||
|
||||
async def delete_page(self, page_id: int) -> bool:
|
||||
"""删除知识页面及其关联分块"""
|
||||
page = await self.get_page(page_id)
|
||||
if page is None:
|
||||
return False
|
||||
|
||||
# 先删除关联的分块
|
||||
await self.db.execute(
|
||||
text("DELETE FROM knowledge_chunks WHERE page_id = :page_id"),
|
||||
{"page_id": page_id},
|
||||
)
|
||||
# 再删除页面
|
||||
await self.db.delete(page)
|
||||
await self.db.flush()
|
||||
return True
|
||||
|
||||
async def get_page_chunks(self, page_id: int) -> List[KnowledgeChunk]:
|
||||
"""获取知识页面的所有分块"""
|
||||
result = await self.db.execute(
|
||||
select(KnowledgeChunk)
|
||||
.where(KnowledgeChunk.page_id == page_id)
|
||||
.order_by(KnowledgeChunk.chunk_index)
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def reindex_page(self, page_id: int) -> dict:
|
||||
"""
|
||||
重新索引知识页面:重新分块并生成嵌入向量。
|
||||
用于页面内容更新后刷新向量索引。
|
||||
"""
|
||||
page = await self.get_page(page_id)
|
||||
if page is None:
|
||||
raise ValueError(f"知识页面不存在: {page_id}")
|
||||
|
||||
# 删除旧分块
|
||||
await self.db.execute(
|
||||
text("DELETE FROM knowledge_chunks WHERE page_id = :page_id"),
|
||||
{"page_id": page_id},
|
||||
)
|
||||
|
||||
# 重新分块
|
||||
from app.services.import_service import ImportService
|
||||
chunks = ImportService._chunk_text(page.content)
|
||||
|
||||
if not chunks:
|
||||
await self.db.flush()
|
||||
return {"page_id": page_id, "chunks": 0}
|
||||
|
||||
# 生成嵌入并插入
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
try:
|
||||
embeddings = await EmbeddingService.embed_batch(chunks)
|
||||
except Exception as exc:
|
||||
logger.error("重新索引嵌入失败: %s", exc)
|
||||
embeddings = [None] * len(chunks)
|
||||
|
||||
for idx, (chunk_text, embedding) in enumerate(zip(chunks, embeddings)):
|
||||
embedding_str = None
|
||||
if embedding:
|
||||
embedding_str = "[" + ",".join(str(x) for x in embedding) + "]"
|
||||
|
||||
await self.db.execute(text("""
|
||||
INSERT INTO knowledge_chunks (page_id, chunk_index, content, embedding)
|
||||
VALUES (:page_id, :chunk_index, :content, :embedding::vector)
|
||||
"""), {
|
||||
"page_id": page_id,
|
||||
"chunk_index": idx,
|
||||
"content": chunk_text,
|
||||
"embedding": embedding_str,
|
||||
})
|
||||
|
||||
# 更新全文搜索向量
|
||||
await self.db.execute(text("""
|
||||
UPDATE knowledge_chunks
|
||||
SET search_vector = to_tsvector('chinese_zh', content)
|
||||
WHERE page_id = :page_id
|
||||
"""), {"page_id": page_id})
|
||||
|
||||
await self.db.flush()
|
||||
return {"page_id": page_id, "chunks": len(chunks)}
|
||||
185
app/services/search_engine.py
Normal file
185
app/services/search_engine.py
Normal file
@@ -0,0 +1,185 @@
|
||||
"""
|
||||
BM25 倒排索引搜索引擎
|
||||
用于图片 OCR 内容的全文搜索
|
||||
参考 Meilisearch / Whoosh 设计
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import math
|
||||
import logging
|
||||
from collections import Counter, defaultdict
|
||||
from typing import List, Dict, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 尝试导入 jieba,失败则用简单的字符分割
|
||||
try:
|
||||
import jieba
|
||||
jieba.setLogLevel(logging.WARNING)
|
||||
def tokenize(text: str) -> List[str]:
|
||||
"""中文分词"""
|
||||
return [w for w in jieba.cut_for_search(text) if len(w.strip()) > 1]
|
||||
except ImportError:
|
||||
logger.warning("jieba 未安装,使用简单分词(pip install jieba 获得更好效果)")
|
||||
def tokenize(text: str) -> List[str]:
|
||||
"""简单分词:按空格和标点分割"""
|
||||
return [w for w in re.split(r'[\s,.\-;:!?"\'/\\|()\[\]{}<>,。;:!?""''、()【】]+', text) if len(w.strip()) > 1]
|
||||
|
||||
|
||||
class BM25Index:
|
||||
"""
|
||||
BM25 倒排索引
|
||||
- 支持增量添加文档
|
||||
- 支持删除文档
|
||||
- 支持中文分词(jieba)
|
||||
- BM25 评分排序
|
||||
"""
|
||||
|
||||
def __init__(self, k1: float = 1.5, b: float = 0.75):
|
||||
"""
|
||||
Args:
|
||||
k1: 词频饱和参数(默认 1.5)
|
||||
b: 文档长度归一化参数(默认 0.75)
|
||||
"""
|
||||
self.k1 = k1
|
||||
self.b = b
|
||||
|
||||
# 倒排索引:term -> {doc_id: tf}
|
||||
self.inverted_index: Dict[str, Dict[int, int]] = defaultdict(lambda: defaultdict(int))
|
||||
|
||||
# 文档信息
|
||||
self.doc_lengths: Dict[int, int] = {} # doc_id -> 文档长度(词数)
|
||||
self.doc_count: int = 0
|
||||
self.avg_doc_length: float = 0.0
|
||||
|
||||
# 文档存储(存完整数据)
|
||||
self.documents: Dict[int, dict] = {}
|
||||
|
||||
# IDF 缓存
|
||||
self._idf_cache: Dict[str, float] = {}
|
||||
|
||||
def add_document(self, doc_id: int, text: str, metadata: Optional[dict] = None):
|
||||
"""
|
||||
添加文档到索引
|
||||
|
||||
Args:
|
||||
doc_id: 文档 ID
|
||||
text: 待索引的文本内容
|
||||
metadata: 附加元数据(如 file_path, tags 等)
|
||||
"""
|
||||
# 如果已存在,先删除
|
||||
if doc_id in self.documents:
|
||||
self.remove_document(doc_id)
|
||||
|
||||
tokens = tokenize(text)
|
||||
self.doc_lengths[doc_id] = len(tokens)
|
||||
self.documents[doc_id] = {
|
||||
"text": text,
|
||||
"tokens": tokens,
|
||||
"metadata": metadata or {},
|
||||
}
|
||||
|
||||
# 构建倒排索引
|
||||
tf = Counter(tokens)
|
||||
for term, count in tf.items():
|
||||
self.inverted_index[term][doc_id] = count
|
||||
|
||||
self.doc_count += 1
|
||||
self._update_avg_doc_length()
|
||||
self._idf_cache.clear() # 文档数变了,IDF 需要重算
|
||||
|
||||
def remove_document(self, doc_id: int):
|
||||
"""从索引中删除文档"""
|
||||
if doc_id not in self.documents:
|
||||
return
|
||||
|
||||
# 从倒排索引中移除
|
||||
tokens = self.documents[doc_id].get("tokens", [])
|
||||
for term in set(tokens):
|
||||
if term in self.inverted_index and doc_id in self.inverted_index[term]:
|
||||
del self.inverted_index[term][doc_id]
|
||||
if not self.inverted_index[term]:
|
||||
del self.inverted_index[term]
|
||||
|
||||
del self.documents[doc_id]
|
||||
del self.doc_lengths[doc_id]
|
||||
self.doc_count -= 1
|
||||
self._update_avg_doc_length()
|
||||
self._idf_cache.clear()
|
||||
|
||||
def _update_avg_doc_length(self):
|
||||
"""更新平均文档长度"""
|
||||
if self.doc_count > 0:
|
||||
self.avg_doc_length = sum(self.doc_lengths.values()) / self.doc_count
|
||||
else:
|
||||
self.avg_doc_length = 0.0
|
||||
|
||||
def _idf(self, term: str) -> float:
|
||||
"""
|
||||
计算逆文档频率 IDF
|
||||
IDF(t) = ln((N - df(t) + 0.5) / (df(t) + 0.5) + 1)
|
||||
其中 N 是文档总数,df(t) 是包含 term 的文档数
|
||||
"""
|
||||
if term in self._idf_cache:
|
||||
return self._idf_cache[term]
|
||||
|
||||
df = len(self.inverted_index.get(term, {}))
|
||||
idf = math.log((self.doc_count - df + 0.5) / (df + 0.5) + 1)
|
||||
|
||||
self._idf_cache[term] = idf
|
||||
return idf
|
||||
|
||||
def search(self, query: str, top_k: int = 10) -> List[Tuple[int, float]]:
|
||||
"""
|
||||
BM25 搜索
|
||||
|
||||
Args:
|
||||
query: 搜索查询(会自动分词)
|
||||
top_k: 返回前 K 个结果
|
||||
|
||||
Returns:
|
||||
[(doc_id, score), ...] 按 BM25 分数降序排列
|
||||
"""
|
||||
if self.doc_count == 0 or self.avg_doc_length == 0:
|
||||
return []
|
||||
|
||||
query_tokens = tokenize(query)
|
||||
if not query_tokens:
|
||||
return []
|
||||
|
||||
scores: Dict[int, float] = defaultdict(float)
|
||||
|
||||
for term in query_tokens:
|
||||
idf = self._idf(term)
|
||||
postings = self.inverted_index.get(term, {})
|
||||
|
||||
for doc_id, tf in postings.items():
|
||||
dl = self.doc_lengths[doc_id]
|
||||
# BM25 公式
|
||||
numerator = tf * (self.k1 + 1)
|
||||
denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_doc_length)
|
||||
scores[doc_id] += idf * numerator / denominator
|
||||
|
||||
# 按分数降序排序
|
||||
ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
|
||||
return ranked[:top_k]
|
||||
|
||||
def get_document(self, doc_id: int) -> Optional[dict]:
|
||||
"""获取文档完整数据"""
|
||||
return self.documents.get(doc_id)
|
||||
|
||||
def clear(self):
|
||||
"""清空索引"""
|
||||
self.inverted_index.clear()
|
||||
self.doc_lengths.clear()
|
||||
self.documents.clear()
|
||||
self._idf_cache.clear()
|
||||
self.doc_count = 0
|
||||
self.avg_doc_length = 0.0
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
"""索引中的文档数量"""
|
||||
return self.doc_count
|
||||
286
app/services/search_service.py
Normal file
286
app/services/search_service.py
Normal file
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
语义搜索服务
|
||||
支持 PostgreSQL(向量+全文混合)和 SQLite(LIKE 关键词搜索)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.config import settings
|
||||
from app.database import IS_SQLITE
|
||||
from app.schemas.search import SearchRequest, SearchResult, SearchResponse
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SearchService:
|
||||
"""语义搜索服务"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self.db = db
|
||||
|
||||
async def search(self, request: SearchRequest) -> SearchResponse:
|
||||
"""
|
||||
执行搜索。
|
||||
- PostgreSQL: 向量搜索 + 中文全文搜索混合排序
|
||||
- SQLite: LIKE 关键词搜索(无向量能力)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
if IS_SQLITE:
|
||||
return await self._search_sqlite(request, start_time)
|
||||
else:
|
||||
return await self._search_postgres(request, start_time)
|
||||
|
||||
# ──────────────────────────── SQLite 搜索 ────────────────────────────
|
||||
|
||||
async def _search_sqlite(self, request: SearchRequest, start_time: float) -> SearchResponse:
|
||||
"""SQLite 模式:使用 LIKE 关键词搜索"""
|
||||
# LLM 查询扩展(可选)
|
||||
expanded_queries = [request.query]
|
||||
try:
|
||||
from app.services.llm_service import LLMService
|
||||
expanded_queries = await LLMService.expand_query(request.query)
|
||||
except Exception as exc:
|
||||
logger.debug("查询扩展跳过: %s", exc)
|
||||
|
||||
results: list[SearchResult] = []
|
||||
seen_chunk_ids = set()
|
||||
|
||||
for q in expanded_queries:
|
||||
like_pattern = f"%{q}%"
|
||||
sql = text("""
|
||||
SELECT
|
||||
kc.id AS chunk_id,
|
||||
kp.id AS page_id,
|
||||
kp.title AS page_title,
|
||||
kc.content,
|
||||
kp.course_name,
|
||||
kp.teacher_name,
|
||||
kp.live_date
|
||||
FROM knowledge_chunks kc
|
||||
JOIN knowledge_pages kp ON kc.page_id = kp.id
|
||||
WHERE kc.content LIKE :pattern
|
||||
AND (:course IS NULL OR kp.course_name = :course)
|
||||
AND (:teacher IS NULL OR kp.teacher_name = :teacher)
|
||||
ORDER BY kc.id
|
||||
LIMIT :limit
|
||||
""")
|
||||
params = {
|
||||
"pattern": like_pattern,
|
||||
"course": request.course_name,
|
||||
"teacher": request.teacher_name,
|
||||
"limit": request.top_k,
|
||||
}
|
||||
result = await self.db.execute(sql, params)
|
||||
rows = result.fetchall()
|
||||
|
||||
for row in rows:
|
||||
if row.chunk_id not in seen_chunk_ids:
|
||||
seen_chunk_ids.add(row.chunk_id)
|
||||
# 简单的相关性评分:关键词出现次数
|
||||
score = min(1.0, row.content.lower().count(q.lower()) * 0.2 + 0.5)
|
||||
results.append(SearchResult(
|
||||
chunk_id=row.chunk_id,
|
||||
page_id=row.page_id,
|
||||
page_title=row.page_title,
|
||||
content=row.content,
|
||||
score=round(score, 4),
|
||||
course_name=row.course_name,
|
||||
teacher_name=row.teacher_name,
|
||||
live_date=row.live_date,
|
||||
highlight=None,
|
||||
))
|
||||
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
results = results[:request.top_k]
|
||||
|
||||
elapsed_ms = (time.time() - start_time) * 1000
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
total=len(results),
|
||||
results=results,
|
||||
elapsed_ms=round(elapsed_ms, 2),
|
||||
)
|
||||
|
||||
# ──────────────────────────── PostgreSQL 搜索 ────────────────────────────
|
||||
|
||||
async def _search_postgres(self, request: SearchRequest, start_time: float) -> SearchResponse:
|
||||
"""PostgreSQL 模式:向量搜索 + 中文全文搜索混合"""
|
||||
# LLM 查询扩展
|
||||
expanded_queries = [request.query]
|
||||
if request.use_fulltext:
|
||||
try:
|
||||
from app.services.llm_service import LLMService
|
||||
expanded_queries = await LLMService.expand_query(request.query)
|
||||
except Exception as exc:
|
||||
logger.debug("查询扩展跳过: %s", exc)
|
||||
|
||||
# 生成查询向量
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
all_query_embeddings = await EmbeddingService.embed_batch(expanded_queries)
|
||||
|
||||
# 构建过滤条件
|
||||
where_clauses = []
|
||||
params: dict = {}
|
||||
|
||||
if request.course_name:
|
||||
where_clauses.append("kp.course_name = :course_name")
|
||||
params["course_name"] = request.course_name
|
||||
if request.teacher_name:
|
||||
where_clauses.append("kp.teacher_name = :teacher_name")
|
||||
params["teacher_name"] = request.teacher_name
|
||||
if request.live_date_from:
|
||||
where_clauses.append("kp.live_date >= :live_date_from")
|
||||
params["live_date_from"] = request.live_date_from
|
||||
if request.live_date_to:
|
||||
where_clauses.append("kp.live_date <= :live_date_to")
|
||||
params["live_date_to"] = request.live_date_to
|
||||
|
||||
where_sql = ""
|
||||
if where_clauses:
|
||||
where_sql = "AND " + " AND ".join(where_clauses)
|
||||
|
||||
# 向量搜索
|
||||
vector_results = []
|
||||
params["limit"] = request.top_k * 2
|
||||
for q_embedding in all_query_embeddings:
|
||||
embedding_str = "[" + ",".join(str(x) for x in q_embedding) + "]"
|
||||
vector_sql = f"""
|
||||
SELECT
|
||||
kc.id AS chunk_id,
|
||||
kp.id AS page_id,
|
||||
kp.title AS page_title,
|
||||
kc.content,
|
||||
1 - (kc.embedding <=> :embedding::vector) AS vector_score,
|
||||
0.0 AS text_score,
|
||||
kp.course_name,
|
||||
kp.teacher_name,
|
||||
kp.live_date
|
||||
FROM knowledge_chunks kc
|
||||
JOIN knowledge_pages kp ON kc.page_id = kp.id
|
||||
WHERE kc.embedding IS NOT NULL
|
||||
{where_sql}
|
||||
ORDER BY kc.embedding <=> :embedding::vector
|
||||
LIMIT :limit
|
||||
"""
|
||||
v_params = {**params, "embedding": embedding_str}
|
||||
result = await self.db.execute(text(vector_sql), v_params)
|
||||
vector_results.extend(result.fetchall())
|
||||
|
||||
# 全文搜索
|
||||
text_results = []
|
||||
if request.use_fulltext:
|
||||
for q in expanded_queries:
|
||||
clean_query = q.replace("'", "''")
|
||||
fulltext_sql = f"""
|
||||
SELECT
|
||||
kc.id AS chunk_id,
|
||||
kp.id AS page_id,
|
||||
kp.title AS page_title,
|
||||
kc.content,
|
||||
0.0 AS vector_score,
|
||||
ts_rank_cd(kc.search_vector, to_tsquery('chinese_zh', :query)) AS text_score,
|
||||
kp.course_name,
|
||||
kp.teacher_name,
|
||||
kp.live_date,
|
||||
ts_headline('chinese_zh', kc.content, to_tsquery('chinese_zh', :query),
|
||||
'MaxWords=50, MinWords=20, ShortWord=2') AS highlight
|
||||
FROM knowledge_chunks kc
|
||||
JOIN knowledge_pages kp ON kc.page_id = kp.id
|
||||
WHERE kc.search_vector @@ to_tsquery('chinese_zh', :query)
|
||||
{where_sql}
|
||||
ORDER BY text_score DESC
|
||||
LIMIT :limit
|
||||
"""
|
||||
ft_params = {**params, "query": clean_query}
|
||||
ft_result = await self.db.execute(text(fulltext_sql), ft_params)
|
||||
text_results.extend(ft_result.fetchall())
|
||||
|
||||
# 合并结果
|
||||
merged: dict[int, dict] = {}
|
||||
|
||||
for row in vector_results:
|
||||
chunk_id = row.chunk_id
|
||||
if chunk_id not in merged:
|
||||
merged[chunk_id] = {
|
||||
"chunk_id": chunk_id,
|
||||
"page_id": row.page_id,
|
||||
"page_title": row.page_title,
|
||||
"content": row.content,
|
||||
"vector_score": row.vector_score,
|
||||
"text_score": 0.0,
|
||||
"course_name": row.course_name,
|
||||
"teacher_name": row.teacher_name,
|
||||
"live_date": row.live_date,
|
||||
"highlight": None,
|
||||
}
|
||||
else:
|
||||
merged[chunk_id]["vector_score"] = max(
|
||||
merged[chunk_id]["vector_score"], row.vector_score
|
||||
)
|
||||
|
||||
for row in text_results:
|
||||
chunk_id = row.chunk_id
|
||||
if chunk_id not in merged:
|
||||
merged[chunk_id] = {
|
||||
"chunk_id": chunk_id,
|
||||
"page_id": row.page_id,
|
||||
"page_title": row.page_title,
|
||||
"content": row.content,
|
||||
"vector_score": 0.0,
|
||||
"text_score": row.text_score or 0.0,
|
||||
"course_name": row.course_name,
|
||||
"teacher_name": row.teacher_name,
|
||||
"live_date": row.live_date,
|
||||
"highlight": row.highlight,
|
||||
}
|
||||
else:
|
||||
merged[chunk_id]["text_score"] = max(
|
||||
merged[chunk_id]["text_score"], row.text_score or 0.0
|
||||
)
|
||||
if row.highlight:
|
||||
merged[chunk_id]["highlight"] = row.highlight
|
||||
|
||||
# 计算综合得分
|
||||
max_text_score = max(
|
||||
(r["text_score"] for r in merged.values()), default=1.0
|
||||
) or 1.0
|
||||
|
||||
scored_results = []
|
||||
for item in merged.values():
|
||||
normalized_text = item["text_score"] / max_text_score if max_text_score > 0 else 0
|
||||
combined_score = 0.7 * item["vector_score"] + 0.3 * normalized_text
|
||||
|
||||
if combined_score >= request.threshold:
|
||||
scored_results.append(
|
||||
SearchResult(
|
||||
chunk_id=item["chunk_id"],
|
||||
page_id=item["page_id"],
|
||||
page_title=item["page_title"],
|
||||
content=item["content"],
|
||||
score=round(combined_score, 4),
|
||||
course_name=item["course_name"],
|
||||
teacher_name=item["teacher_name"],
|
||||
live_date=item["live_date"],
|
||||
highlight=item["highlight"],
|
||||
)
|
||||
)
|
||||
|
||||
scored_results.sort(key=lambda x: x.score, reverse=True)
|
||||
scored_results = scored_results[:request.top_k]
|
||||
|
||||
elapsed_ms = (time.time() - start_time) * 1000
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
total=len(scored_results),
|
||||
results=scored_results,
|
||||
elapsed_ms=round(elapsed_ms, 2),
|
||||
)
|
||||
362
app/wework_bot.py
Normal file
362
app/wework_bot.py
Normal file
@@ -0,0 +1,362 @@
|
||||
"""
|
||||
企业微信智能机器人服务(WebSocket 长连接模式)
|
||||
|
||||
使用官方 SDK (wecom-aibot-python-sdk) 的 WSClient 类,
|
||||
封装了连接管理、心跳、断线重连等能力。
|
||||
|
||||
作为独立进程运行:python -m app.wework_bot
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
|
||||
import httpx
|
||||
from aibot import WSClient, WSClientOptions, generate_req_id
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# 加载 .env 文件(独立运行时需要)
|
||||
load_dotenv()
|
||||
|
||||
logger = logging.getLogger("wework_bot")
|
||||
|
||||
# 分片上传参数(参考官方文档)
|
||||
_CHUNK_SIZE = 512 * 1024 # 单个分片最大 512KB(Base64 编码前)
|
||||
|
||||
|
||||
class WeComBotService:
|
||||
"""企业微信智能机器人服务(WebSocket 长连接模式)"""
|
||||
|
||||
def __init__(self):
|
||||
self.bot_id = os.getenv("WEWORK_BOT_ID", "")
|
||||
self.bot_secret = os.getenv("WEWORK_BOT_SECRET", "")
|
||||
self.edubrain_base_url = os.getenv(
|
||||
"EDUBRAIN_BASE_URL", "http://localhost:8765"
|
||||
)
|
||||
self.search_limit = int(os.getenv("WEWORK_BOT_SEARCH_LIMIT", "3"))
|
||||
self.enabled = os.getenv("WEWORK_BOT_ENABLED", "false").lower() == "true"
|
||||
self.ws_client: WSClient | None = None
|
||||
|
||||
def start(self):
|
||||
"""启动机器人服务"""
|
||||
if not self.enabled:
|
||||
logger.info("企业微信机器人未启用 (WEWORK_BOT_ENABLED=false)")
|
||||
return
|
||||
|
||||
if not self.bot_id or not self.bot_secret:
|
||||
logger.error("WEWORK_BOT_ID 或 WEWORK_BOT_SECRET 未配置")
|
||||
return
|
||||
|
||||
self.ws_client = WSClient(
|
||||
WSClientOptions(
|
||||
bot_id=self.bot_id,
|
||||
secret=self.bot_secret,
|
||||
)
|
||||
)
|
||||
|
||||
# 注册事件
|
||||
self.ws_client.on("authenticated", self._on_authenticated)
|
||||
self.ws_client.on("message.text", self._on_text_message)
|
||||
self.ws_client.on("event.enter_chat", self._on_enter_chat)
|
||||
self.ws_client.on("disconnected", self._on_disconnected)
|
||||
self.ws_client.on("error", self._on_error)
|
||||
|
||||
logger.info("启动企业微信智能机器人...")
|
||||
self.ws_client.run()
|
||||
|
||||
# ──────────────────────────── 事件处理 ────────────────────────────
|
||||
|
||||
def _on_authenticated(self):
|
||||
"""认证成功回调"""
|
||||
logger.info("企业微信机器人认证成功")
|
||||
|
||||
async def _on_enter_chat(self, frame: dict):
|
||||
"""进入会话事件 -> 回复欢迎语"""
|
||||
await self.ws_client.reply_welcome(
|
||||
frame,
|
||||
{
|
||||
"msgtype": "text",
|
||||
"text": {
|
||||
"content": "你好!我是知识库图片搜索助手,发送关键词即可搜索相关图片。"
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
async def _on_text_message(self, frame: dict):
|
||||
"""收到文本消息 -> 搜索图片 -> 流式回复文本 + 逐张发送图片"""
|
||||
body = frame.get("body", {})
|
||||
text_content = body.get("text", {}).get("content", "")
|
||||
|
||||
# 去掉 @机器人 的前缀
|
||||
text_content = re.sub(r"@.*?\s*", "", text_content).strip()
|
||||
|
||||
if not text_content:
|
||||
stream_id = generate_req_id("stream")
|
||||
await self.ws_client.reply_stream(
|
||||
frame, stream_id, "请输入搜索关键词", True
|
||||
)
|
||||
return
|
||||
|
||||
# 清理口语化前缀,提取关键词
|
||||
keyword = re.sub(
|
||||
r"(帮我|请|找|搜|搜索|查|查一下|来|给我|要|想要|的?图|图片|相关|关于|一下)",
|
||||
"", text_content,
|
||||
).strip()
|
||||
if not keyword:
|
||||
keyword = text_content # 兜底
|
||||
|
||||
stream_id = generate_req_id("stream")
|
||||
|
||||
# 立即回复:开始搜索
|
||||
await self.ws_client.reply_stream(
|
||||
frame,
|
||||
stream_id,
|
||||
f"\U0001f50d 正在搜索「{keyword}」...",
|
||||
False,
|
||||
)
|
||||
|
||||
# 调用 EduBrain 搜索
|
||||
try:
|
||||
results = await self._search_images(keyword)
|
||||
|
||||
if not results:
|
||||
await self.ws_client.reply_stream(
|
||||
frame,
|
||||
stream_id,
|
||||
f"未找到与「{text_content}」相关的图片",
|
||||
True,
|
||||
)
|
||||
return
|
||||
|
||||
# 构建文本回复内容
|
||||
content_parts = [f"找到 {len(results)} 张相关图片:"]
|
||||
for i, r in enumerate(results, 1):
|
||||
preview = r.get("ocr_text_preview", "")[:80]
|
||||
content_parts.append(f"\n{i}. {preview}...")
|
||||
|
||||
content = "\n".join(content_parts)
|
||||
|
||||
# 完成流式消息(仅文本,不使用 msg_item)
|
||||
await self.ws_client.reply_stream(
|
||||
frame,
|
||||
stream_id,
|
||||
content,
|
||||
True,
|
||||
)
|
||||
|
||||
# 逐张上传临时素材并发送图片消息
|
||||
for i, r in enumerate(results, 1):
|
||||
b64 = r.get("image_base64", "")
|
||||
if not b64:
|
||||
continue
|
||||
|
||||
try:
|
||||
# 解码 base64 获取原始图片数据
|
||||
image_data = base64.b64decode(b64)
|
||||
file_name = r.get("image_url", f"image_{i}.jpg")
|
||||
|
||||
# 上传临时素材(分片上传)
|
||||
media_id = await self._upload_temp_media(
|
||||
image_data, file_name
|
||||
)
|
||||
if not media_id:
|
||||
logger.warning("第 %d 张图片上传失败,跳过", i)
|
||||
continue
|
||||
|
||||
# 发送图片消息(通过 aibot_respond_msg)
|
||||
await self.ws_client.reply(
|
||||
frame,
|
||||
{
|
||||
"msgtype": "image",
|
||||
"image": {"media_id": media_id},
|
||||
},
|
||||
)
|
||||
logger.info("已发送第 %d 张图片 (media_id: %s)", i, media_id)
|
||||
except Exception as e:
|
||||
logger.error("发送第 %d 张图片失败: %s", i, e, exc_info=True)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("搜索失败: %s", e, exc_info=True)
|
||||
await self.ws_client.reply_stream(
|
||||
frame,
|
||||
stream_id,
|
||||
f"搜索出错:{e}",
|
||||
True,
|
||||
)
|
||||
|
||||
def _on_disconnected(self, reason: str):
|
||||
"""连接断开回调"""
|
||||
logger.warning("企业微信连接断开: %s", reason)
|
||||
|
||||
def _on_error(self, error):
|
||||
"""错误回调"""
|
||||
logger.error("企业微信机器人错误: %s", error)
|
||||
|
||||
# ──────────────────────── 上传临时素材 ────────────────────────
|
||||
|
||||
async def _upload_temp_media(
|
||||
self, file_data: bytes, filename: str
|
||||
) -> str | None:
|
||||
"""
|
||||
通过 WebSocket 通道上传临时素材(三步分片上传)。
|
||||
|
||||
流程:aibot_upload_media_init → aibot_upload_media_chunk × N → aibot_upload_media_finish
|
||||
|
||||
Args:
|
||||
file_data: 文件原始二进制数据
|
||||
filename: 文件名
|
||||
|
||||
Returns:
|
||||
成功返回 media_id,失败返回 None
|
||||
"""
|
||||
total_size = len(file_data)
|
||||
total_chunks = math.ceil(total_size / _CHUNK_SIZE)
|
||||
md5 = hashlib.md5(file_data).hexdigest()
|
||||
ws = self.ws_client._ws_manager
|
||||
|
||||
# 第一步:上传初始化
|
||||
init_req_id = generate_req_id("upload")
|
||||
try:
|
||||
init_result = await ws.send_reply(
|
||||
init_req_id,
|
||||
{
|
||||
"type": "image",
|
||||
"filename": filename,
|
||||
"total_size": total_size,
|
||||
"total_chunks": total_chunks,
|
||||
"md5": md5,
|
||||
},
|
||||
cmd="aibot_upload_media_init",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("上传初始化失败: %s", e)
|
||||
return None
|
||||
|
||||
upload_id = init_result.get("body", {}).get("upload_id", "")
|
||||
if not upload_id:
|
||||
logger.error("上传初始化未返回 upload_id: %s", init_result)
|
||||
return None
|
||||
|
||||
logger.info(
|
||||
"上传初始化成功: upload_id=%s, total_size=%d, chunks=%d",
|
||||
upload_id, total_size, total_chunks,
|
||||
)
|
||||
|
||||
# 第二步:分片上传
|
||||
for chunk_idx in range(total_chunks):
|
||||
start = chunk_idx * _CHUNK_SIZE
|
||||
end = min(start + _CHUNK_SIZE, total_size)
|
||||
chunk_data = file_data[start:end]
|
||||
chunk_b64 = base64.b64encode(chunk_data).decode("utf-8")
|
||||
|
||||
chunk_req_id = generate_req_id("upload_chunk")
|
||||
try:
|
||||
await ws.send_reply(
|
||||
chunk_req_id,
|
||||
{
|
||||
"upload_id": upload_id,
|
||||
"chunk_index": chunk_idx, # 分片索引从 0 开始
|
||||
"base64_data": chunk_b64,
|
||||
},
|
||||
cmd="aibot_upload_media_chunk",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"分片上传失败 (chunk %d/%d): %s",
|
||||
chunk_idx + 1, total_chunks, e,
|
||||
)
|
||||
return None
|
||||
|
||||
logger.info("所有分片上传完成: upload_id=%s", upload_id)
|
||||
|
||||
# 第三步:完成上传
|
||||
finish_req_id = generate_req_id("upload_finish")
|
||||
try:
|
||||
finish_result = await ws.send_reply(
|
||||
finish_req_id,
|
||||
{"upload_id": upload_id},
|
||||
cmd="aibot_upload_media_finish",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("完成上传失败: %s", e)
|
||||
return None
|
||||
|
||||
media_id = finish_result.get("body", {}).get("media_id", "")
|
||||
if not media_id:
|
||||
logger.error("完成上传未返回 media_id: %s", finish_result)
|
||||
return None
|
||||
|
||||
logger.info("上传临时素材成功: media_id=%s", media_id)
|
||||
return media_id
|
||||
|
||||
# ──────────────────────────── 搜索调用 ────────────────────────────
|
||||
|
||||
async def _get_search_limit(self) -> int:
|
||||
"""从后端 API 实时获取当前 search_limit 配置"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
resp = await client.get(f"{self.edubrain_base_url}/api/v1/settings")
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
return data.get("search_limit", self.search_limit)
|
||||
except Exception as e:
|
||||
logger.warning("获取 search_limit 失败,使用默认值: %s", e)
|
||||
return self.search_limit
|
||||
|
||||
async def _search_images(self, keyword: str) -> list:
|
||||
"""
|
||||
调用 EduBrain 搜索接口(SSE 流式),收集所有匹配结果。
|
||||
|
||||
Args:
|
||||
keyword: 搜索关键词
|
||||
|
||||
Returns:
|
||||
匹配结果列表
|
||||
"""
|
||||
results = []
|
||||
url = f"{self.edubrain_base_url}/api/v1/images/search"
|
||||
limit = await self._get_search_limit()
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
async with client.stream(
|
||||
"GET",
|
||||
url,
|
||||
params={
|
||||
"keyword": keyword,
|
||||
"limit": limit,
|
||||
"sort": "time_desc",
|
||||
},
|
||||
) as resp:
|
||||
async for line in resp.aiter_lines():
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
try:
|
||||
data = json.loads(line[6:])
|
||||
if data.get("type") == "result":
|
||||
results.append(data)
|
||||
elif data.get("type") == "done":
|
||||
break
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def run_wework_bot():
|
||||
"""入口函数"""
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
|
||||
)
|
||||
service = WeComBotService()
|
||||
service.start()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_wework_bot()
|
||||
BIN
data/transcripts/2026年1月16日家庭关系299转写文字.docx
Normal file
BIN
data/transcripts/2026年1月16日家庭关系299转写文字.docx
Normal file
Binary file not shown.
BIN
data/transcripts/2026年1月21日+第10营训练营.docx
Normal file
BIN
data/transcripts/2026年1月21日+第10营训练营.docx
Normal file
Binary file not shown.
BIN
data/transcripts/2026年2月5日家庭关系299转写原文.docx
Normal file
BIN
data/transcripts/2026年2月5日家庭关系299转写原文.docx
Normal file
Binary file not shown.
BIN
data/transcripts/2026年2月6日家庭关系299转写原文.docx
Normal file
BIN
data/transcripts/2026年2月6日家庭关系299转写原文.docx
Normal file
Binary file not shown.
46
data/transcripts/test-demo.md
Normal file
46
data/transcripts/test-demo.md
Normal file
@@ -0,0 +1,46 @@
|
||||
---
|
||||
course: 高考数学冲刺班
|
||||
teacher: 张老师
|
||||
date: 2026-04-10
|
||||
---
|
||||
|
||||
## 二次方程的解法技巧
|
||||
|
||||
二次方程 ax² + bx + c = 0 是高考数学中的高频考点。解法主要有三种:
|
||||
|
||||
### 因式分解法
|
||||
|
||||
当方程可以分解为两个一次因式的乘积时,直接因式分解是最快的方法。例如 x² - 5x + 6 = 0 可以分解为 (x-2)(x-3) = 0,所以 x=2 或 x=3。
|
||||
|
||||
**技巧提示**:先看常数项 c 的因数对,再验证是否满足 b。
|
||||
|
||||
### 配方法
|
||||
|
||||
配方法适用于任何二次方程,步骤如下:
|
||||
1. 将方程化为 x² + px + q = 0 的形式
|
||||
2. 移项:x² + px = -q
|
||||
3. 两边加上 (p/2)²
|
||||
4. 化为完全平方形式
|
||||
|
||||
### 公式法
|
||||
|
||||
万能公式 x = (-b ± √(b²-4ac)) / 2a 适用于所有情况。但要注意判别式 Δ = b²-4ac 的符号:
|
||||
- Δ > 0:两个不相等的实数根
|
||||
- Δ = 0:两个相等的实数根(重根)
|
||||
- Δ < 0:无实数根
|
||||
|
||||
## 韦达定理的应用
|
||||
|
||||
韦达定理:若 x₁、x₂ 是方程 ax² + bx + c = 0 的两个根,则:
|
||||
- x₁ + x₂ = -b/a
|
||||
- x₁ · x₂ = c/a
|
||||
|
||||
这个定理在选择题和填空题中非常有用,可以快速求出两根之和与积,而不需要实际求解方程。
|
||||
|
||||
**常见考法**:已知一根求另一根、求参数范围、证明不等式等。
|
||||
|
||||
## 真题演练
|
||||
|
||||
2025年全国卷第12题:已知方程 x² - 3x + m = 0 的两个根都是正整数,求 m 的所有可能值。
|
||||
|
||||
解题思路:设两根为 x₁、x₂,由韦达定理 x₁+x₂=3,x₁·x₂=m。因为都是正整数,所以只能是 1 和 2,因此 m=2。
|
||||
46
data/transcripts/test-demo_1.md
Normal file
46
data/transcripts/test-demo_1.md
Normal file
@@ -0,0 +1,46 @@
|
||||
---
|
||||
course: 高考数学冲刺班
|
||||
teacher: 张老师
|
||||
date: 2026-04-10
|
||||
---
|
||||
|
||||
## 二次方程的解法技巧
|
||||
|
||||
二次方程 ax² + bx + c = 0 是高考数学中的高频考点。解法主要有三种:
|
||||
|
||||
### 因式分解法
|
||||
|
||||
当方程可以分解为两个一次因式的乘积时,直接因式分解是最快的方法。例如 x² - 5x + 6 = 0 可以分解为 (x-2)(x-3) = 0,所以 x=2 或 x=3。
|
||||
|
||||
**技巧提示**:先看常数项 c 的因数对,再验证是否满足 b。
|
||||
|
||||
### 配方法
|
||||
|
||||
配方法适用于任何二次方程,步骤如下:
|
||||
1. 将方程化为 x² + px + q = 0 的形式
|
||||
2. 移项:x² + px = -q
|
||||
3. 两边加上 (p/2)²
|
||||
4. 化为完全平方形式
|
||||
|
||||
### 公式法
|
||||
|
||||
万能公式 x = (-b ± √(b²-4ac)) / 2a 适用于所有情况。但要注意判别式 Δ = b²-4ac 的符号:
|
||||
- Δ > 0:两个不相等的实数根
|
||||
- Δ = 0:两个相等的实数根(重根)
|
||||
- Δ < 0:无实数根
|
||||
|
||||
## 韦达定理的应用
|
||||
|
||||
韦达定理:若 x₁、x₂ 是方程 ax² + bx + c = 0 的两个根,则:
|
||||
- x₁ + x₂ = -b/a
|
||||
- x₁ · x₂ = c/a
|
||||
|
||||
这个定理在选择题和填空题中非常有用,可以快速求出两根之和与积,而不需要实际求解方程。
|
||||
|
||||
**常见考法**:已知一根求另一根、求参数范围、证明不等式等。
|
||||
|
||||
## 真题演练
|
||||
|
||||
2025年全国卷第12题:已知方程 x² - 3x + m = 0 的两个根都是正整数,求 m 的所有可能值。
|
||||
|
||||
解题思路:设两根为 x₁、x₂,由韦达定理 x₁+x₂=3,x₁·x₂=m。因为都是正整数,所以只能是 1 和 2,因此 m=2。
|
||||
46
data/transcripts/test-demo_2.md
Normal file
46
data/transcripts/test-demo_2.md
Normal file
@@ -0,0 +1,46 @@
|
||||
---
|
||||
course: 高考数学冲刺班
|
||||
teacher: 张老师
|
||||
date: 2026-04-10
|
||||
---
|
||||
|
||||
## 二次方程的解法技巧
|
||||
|
||||
二次方程 ax² + bx + c = 0 是高考数学中的高频考点。解法主要有三种:
|
||||
|
||||
### 因式分解法
|
||||
|
||||
当方程可以分解为两个一次因式的乘积时,直接因式分解是最快的方法。例如 x² - 5x + 6 = 0 可以分解为 (x-2)(x-3) = 0,所以 x=2 或 x=3。
|
||||
|
||||
**技巧提示**:先看常数项 c 的因数对,再验证是否满足 b。
|
||||
|
||||
### 配方法
|
||||
|
||||
配方法适用于任何二次方程,步骤如下:
|
||||
1. 将方程化为 x² + px + q = 0 的形式
|
||||
2. 移项:x² + px = -q
|
||||
3. 两边加上 (p/2)²
|
||||
4. 化为完全平方形式
|
||||
|
||||
### 公式法
|
||||
|
||||
万能公式 x = (-b ± √(b²-4ac)) / 2a 适用于所有情况。但要注意判别式 Δ = b²-4ac 的符号:
|
||||
- Δ > 0:两个不相等的实数根
|
||||
- Δ = 0:两个相等的实数根(重根)
|
||||
- Δ < 0:无实数根
|
||||
|
||||
## 韦达定理的应用
|
||||
|
||||
韦达定理:若 x₁、x₂ 是方程 ax² + bx + c = 0 的两个根,则:
|
||||
- x₁ + x₂ = -b/a
|
||||
- x₁ · x₂ = c/a
|
||||
|
||||
这个定理在选择题和填空题中非常有用,可以快速求出两根之和与积,而不需要实际求解方程。
|
||||
|
||||
**常见考法**:已知一根求另一根、求参数范围、证明不等式等。
|
||||
|
||||
## 真题演练
|
||||
|
||||
2025年全国卷第12题:已知方程 x² - 3x + m = 0 的两个根都是正整数,求 m 的所有可能值。
|
||||
|
||||
解题思路:设两根为 x₁、x₂,由韦达定理 x₁+x₂=3,x₁·x₂=m。因为都是正整数,所以只能是 1 和 2,因此 m=2。
|
||||
46
data/transcripts/test-demo_3.md
Normal file
46
data/transcripts/test-demo_3.md
Normal file
@@ -0,0 +1,46 @@
|
||||
---
|
||||
course: 高考数学冲刺班
|
||||
teacher: 张老师
|
||||
date: 2026-04-10
|
||||
---
|
||||
|
||||
## 二次方程的解法技巧
|
||||
|
||||
二次方程 ax² + bx + c = 0 是高考数学中的高频考点。解法主要有三种:
|
||||
|
||||
### 因式分解法
|
||||
|
||||
当方程可以分解为两个一次因式的乘积时,直接因式分解是最快的方法。例如 x² - 5x + 6 = 0 可以分解为 (x-2)(x-3) = 0,所以 x=2 或 x=3。
|
||||
|
||||
**技巧提示**:先看常数项 c 的因数对,再验证是否满足 b。
|
||||
|
||||
### 配方法
|
||||
|
||||
配方法适用于任何二次方程,步骤如下:
|
||||
1. 将方程化为 x² + px + q = 0 的形式
|
||||
2. 移项:x² + px = -q
|
||||
3. 两边加上 (p/2)²
|
||||
4. 化为完全平方形式
|
||||
|
||||
### 公式法
|
||||
|
||||
万能公式 x = (-b ± √(b²-4ac)) / 2a 适用于所有情况。但要注意判别式 Δ = b²-4ac 的符号:
|
||||
- Δ > 0:两个不相等的实数根
|
||||
- Δ = 0:两个相等的实数根(重根)
|
||||
- Δ < 0:无实数根
|
||||
|
||||
## 韦达定理的应用
|
||||
|
||||
韦达定理:若 x₁、x₂ 是方程 ax² + bx + c = 0 的两个根,则:
|
||||
- x₁ + x₂ = -b/a
|
||||
- x₁ · x₂ = c/a
|
||||
|
||||
这个定理在选择题和填空题中非常有用,可以快速求出两根之和与积,而不需要实际求解方程。
|
||||
|
||||
**常见考法**:已知一根求另一根、求参数范围、证明不等式等。
|
||||
|
||||
## 真题演练
|
||||
|
||||
2025年全国卷第12题:已知方程 x² - 3x + m = 0 的两个根都是正整数,求 m 的所有可能值。
|
||||
|
||||
解题思路:设两根为 x₁、x₂,由韦达定理 x₁+x₂=3,x₁·x₂=m。因为都是正整数,所以只能是 1 和 2,因此 m=2。
|
||||
BIN
data/transcripts/test-english.docx
Normal file
BIN
data/transcripts/test-english.docx
Normal file
Binary file not shown.
BIN
data/transcripts/test-english_1.docx
Normal file
BIN
data/transcripts/test-english_1.docx
Normal file
Binary file not shown.
55
docker-compose.yml
Normal file
55
docker-compose.yml
Normal file
@@ -0,0 +1,55 @@
|
||||
version: "3.8"
|
||||
|
||||
services:
|
||||
# ── 应用服务 ──
|
||||
app:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: edu-brain-app
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "${APP_PORT:-8000}:8000"
|
||||
volumes:
|
||||
- ./data:/app/data # 数据目录挂载
|
||||
env_file:
|
||||
- .env
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
environment:
|
||||
- DATABASE_URL=postgresql+asyncpg://postgres:${POSTGRES_PASSWORD:-postgres}@db:5432/edu_brain
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"]
|
||||
interval: 30s
|
||||
timeout: 10s
|
||||
retries: 3
|
||||
start_period: 60s
|
||||
|
||||
# ── PostgreSQL 数据库(pgvector + zhparser 自定义镜像) ──
|
||||
db:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile.db
|
||||
container_name: edu-brain-db
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "${DB_PORT:-5432}:5432"
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
- ./sql/init.sql:/docker-entrypoint-initdb.d/01-init.sql:ro
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-postgres}
|
||||
POSTGRES_DB: edu_brain
|
||||
POSTGRES_HOST_AUTH_METHOD: scram-sha-256
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U postgres -d edu_brain"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
start_period: 30s
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
driver: local
|
||||
438
docs/IMAGE_API_GUIDE.md
Normal file
438
docs/IMAGE_API_GUIDE.md
Normal file
@@ -0,0 +1,438 @@
|
||||
# EduBrain 图片功能集成指南
|
||||
|
||||
> 本文档介绍如何在 ADK Agent 项目中集成 EduBrain 的图片 OCR 和搜索功能。
|
||||
|
||||
## 服务地址
|
||||
|
||||
```
|
||||
BASE_URL = http://localhost:8765
|
||||
```
|
||||
|
||||
所有接口前缀:`/api/v1/images`
|
||||
|
||||
---
|
||||
|
||||
## 一、接口总览
|
||||
|
||||
| 方法 | 路径 | 说明 | 请求类型 |
|
||||
|------|------|------|----------|
|
||||
| POST | `/recognize-direct` | 上传并识别单张图片 | multipart/form-data |
|
||||
| POST | `/batch-recognize` | 批量上传并识别 | multipart/form-data |
|
||||
| POST | `/import-paths` | 从服务器路径导入 | JSON |
|
||||
| GET | `/search` | AI 搜索图片(SSE 流式) | query params |
|
||||
| GET | `/{image_id}` | 获取单张图片详情 | path param |
|
||||
| GET | `` | 获取图片列表(分页) | query params |
|
||||
|
||||
---
|
||||
|
||||
## 二、上传并识别图片
|
||||
|
||||
### 接口
|
||||
|
||||
```
|
||||
POST /api/v1/images/recognize-direct
|
||||
Content-Type: multipart/form-data
|
||||
```
|
||||
|
||||
### 参数
|
||||
|
||||
| 字段 | 类型 | 必填 | 说明 |
|
||||
|------|------|------|------|
|
||||
| file | File | 是 | 图片文件,支持 .png .jpg .jpeg .bmp .webp |
|
||||
|
||||
### 返回
|
||||
|
||||
```json
|
||||
{
|
||||
"id": 386,
|
||||
"file_path": "data/images/test.png",
|
||||
"original_filename": "test.png",
|
||||
"ocr_text": "识别出的文字内容...",
|
||||
"tags": [],
|
||||
"confidence": 1.0,
|
||||
"provider": "deepseek",
|
||||
"status": "completed",
|
||||
"blocks": [
|
||||
{"text": "第一段文字", "bbox": [x1,y1,x2,y2], "confidence": 0.98}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### Python 调用示例
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
async def recognize_image(file_path: str) -> dict:
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
with open(file_path, "rb") as f:
|
||||
resp = await client.post(
|
||||
f"{BASE_URL}/api/v1/images/recognize-direct",
|
||||
files={"file": (file_path, f)}
|
||||
)
|
||||
return resp.json()
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 三、批量上传并识别
|
||||
|
||||
### 接口
|
||||
|
||||
```
|
||||
POST /api/v1/images/batch-recognize
|
||||
Content-Type: multipart/form-data
|
||||
```
|
||||
|
||||
### 参数
|
||||
|
||||
| 字段 | 类型 | 必填 | 说明 |
|
||||
|------|------|------|------|
|
||||
| files | File[] | 是 | 多个图片文件 |
|
||||
|
||||
### 返回
|
||||
|
||||
```json
|
||||
{
|
||||
"total": 3,
|
||||
"success": 2,
|
||||
"failed": 1,
|
||||
"results": [
|
||||
{
|
||||
"id": 387,
|
||||
"filename": "img1.png",
|
||||
"file_path": "data/images/img1.png",
|
||||
"status": "completed",
|
||||
"ocr_text": "识别内容...",
|
||||
"tags": [],
|
||||
"confidence": 1.0,
|
||||
"provider": "deepseek"
|
||||
},
|
||||
{
|
||||
"id": 388,
|
||||
"filename": "img2.jpg",
|
||||
"file_path": "data/images/img2.jpg",
|
||||
"status": "failed",
|
||||
"message": "OCR 识别失败: ..."
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 四、从服务器路径导入
|
||||
|
||||
### 接口
|
||||
|
||||
```
|
||||
POST /api/v1/images/import-paths
|
||||
Content-Type: application/json
|
||||
```
|
||||
|
||||
### 请求体
|
||||
|
||||
```json
|
||||
{
|
||||
"paths": ["/data/images/001.png", "/data/images/002.jpg"],
|
||||
"recursive": false
|
||||
}
|
||||
```
|
||||
|
||||
### 返回
|
||||
|
||||
与批量上传相同格式。
|
||||
|
||||
---
|
||||
|
||||
## 五、AI 搜索图片(SSE 流式)⚠️ 重点
|
||||
|
||||
### 接口
|
||||
|
||||
```
|
||||
GET /api/v1/images/search?keyword=搜索词&limit=5&sort=time_desc
|
||||
```
|
||||
|
||||
**这不是普通 HTTP 请求,而是 SSE(Server-Sent Events)流式接口。**
|
||||
|
||||
### 参数
|
||||
|
||||
| 参数 | 类型 | 必填 | 默认值 | 说明 |
|
||||
|------|------|------|--------|------|
|
||||
| keyword | string | 是 | - | 搜索关键词 |
|
||||
| limit | int | 否 | 5 | 最多返回多少条匹配结果 |
|
||||
| sort | string | 否 | time_desc | 排序方式:`time_desc`(时间倒序)、`time_asc`(时间正序)、`random`(随机) |
|
||||
|
||||
### 搜索原理
|
||||
|
||||
```
|
||||
用户搜索 "孩子不想上学"
|
||||
→ 从数据库取所有已识别图片(按 sort 排序)
|
||||
→ 每批 10 张发给 MiniMax M2.7 判断是否匹配
|
||||
→ 并发池最多 10 个 LLM 请求同时执行
|
||||
→ 匹配结果通过 SSE 实时推送(找到就立刻返回)
|
||||
→ 找够 limit 条或遍历完所有图片后结束
|
||||
```
|
||||
|
||||
### SSE 事件格式
|
||||
|
||||
连接后,服务端会持续推送 JSON 事件,每条格式为:
|
||||
|
||||
```
|
||||
data: {"type": "事件类型", ...其他字段}\n\n
|
||||
```
|
||||
|
||||
### 事件类型
|
||||
|
||||
#### 1. start - 搜索开始
|
||||
|
||||
```json
|
||||
{"type": "start", "keyword": "孩子不想上学", "limit": 5}
|
||||
```
|
||||
|
||||
#### 2. progress - 进度更新
|
||||
|
||||
```json
|
||||
{"type": "progress", "checked": 170, "total": 385, "found": 3}
|
||||
```
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
|------|------|------|
|
||||
| checked | int | 已检查的图片数量 |
|
||||
| total | int | 数据库中图片总数 |
|
||||
| found | int | 已找到的匹配数量 |
|
||||
|
||||
#### 3. result - 匹配结果(实时推送)
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "result",
|
||||
"id": 42,
|
||||
"file_path": "data/images/001.png",
|
||||
"image_url": "001.png",
|
||||
"image_base64": "/9j/4QE3RXhpZgAATU0AKgAAAAgABgEAAAQ...",
|
||||
"ocr_text": "完整的 OCR 识别文本...",
|
||||
"ocr_text_preview": "前150个字符的预览...",
|
||||
"tags": ["亲子关系", "沟通"],
|
||||
"confidence": 1.0,
|
||||
"provider": "deepseek",
|
||||
"created_at": "2026-04-11T21:59:56"
|
||||
}
|
||||
```
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
|------|------|------|
|
||||
| id | int | 图片记录 ID |
|
||||
| file_path | string | 图片文件路径(服务器本地路径) |
|
||||
| image_url | string | 图片文件名(访问地址为 `BASE_URL/data/images/{image_url}`) |
|
||||
| image_base64 | string | 图片的 Base64 编码,可直接用于发送到企业微信等 |
|
||||
| ocr_text | string | 完整 OCR 文本 |
|
||||
| ocr_text_preview | string | OCR 文本预览(前150字符) |
|
||||
| tags | string[] | 关键词标签 |
|
||||
| confidence | float | OCR 置信度(0-1) |
|
||||
| provider | string | OCR 提供商 |
|
||||
| created_at | string | 创建时间(ISO 8601) |
|
||||
|
||||
#### 4. done - 搜索结束
|
||||
|
||||
```json
|
||||
{"type": "done", "total_found": 5, "total_checked": 385}
|
||||
```
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
|------|------|------|
|
||||
| total_found | int | 总共找到的匹配数量 |
|
||||
| total_checked | int | 总共检查的图片数量 |
|
||||
|
||||
### Python 调用示例(推荐)
|
||||
|
||||
```python
|
||||
import httpx
|
||||
import json
|
||||
from typing import AsyncGenerator
|
||||
|
||||
BASE_URL = "http://localhost:8765"
|
||||
|
||||
async def search_images(
|
||||
keyword: str,
|
||||
limit: int = 5,
|
||||
sort: str = "time_desc",
|
||||
) -> AsyncGenerator[dict, None]:
|
||||
"""
|
||||
AI 搜索图片(SSE 流式)
|
||||
|
||||
Yields:
|
||||
dict: 每个 SSE 事件解析后的字典
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
async with client.stream(
|
||||
"GET",
|
||||
f"{BASE_URL}/api/v1/images/search",
|
||||
params={"keyword": keyword, "limit": limit, "sort": sort},
|
||||
) as resp:
|
||||
async for line in resp.aiter_lines():
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
data = json.loads(line[6:])
|
||||
yield data
|
||||
|
||||
|
||||
# ── 使用示例 ──
|
||||
|
||||
async def demo_search():
|
||||
results = []
|
||||
|
||||
async for event in search_images("孩子不想上学", limit=5):
|
||||
event_type = event.get("type")
|
||||
|
||||
if event_type == "start":
|
||||
print(f"开始搜索: {event['keyword']}")
|
||||
|
||||
elif event_type == "progress":
|
||||
print(f"进度: {event['checked']}/{event['total']},已找到 {event['found']} 条")
|
||||
|
||||
elif event_type == "result":
|
||||
results.append(event)
|
||||
print(f"找到匹配: [{event['id']}] {event['ocr_text_preview'][:50]}")
|
||||
|
||||
elif event_type == "done":
|
||||
print(f"搜索完成: 共找到 {event['total_found']} 条,检查了 {event['total_checked']} 张")
|
||||
|
||||
return results
|
||||
```
|
||||
|
||||
### ADK Agent 集成示例
|
||||
|
||||
在 ADK Agent 中,可以将搜索功能封装为一个 Tool:
|
||||
|
||||
```python
|
||||
from google.adk import Tool
|
||||
|
||||
class ImageSearchTool(Tool):
|
||||
"""搜索图片知识库"""
|
||||
|
||||
name = "search_images"
|
||||
description = "搜索已识别的图片库,返回与查询语义相关的图片及其 Base64 数据。支持自然语言查询。"
|
||||
|
||||
async def execute(self, keyword: str, limit: int = 5) -> str:
|
||||
results = []
|
||||
async for event in search_images(keyword, limit=limit):
|
||||
if event["type"] == "result":
|
||||
results.append({
|
||||
"id": event["id"],
|
||||
"image_base64": event["image_base64"],
|
||||
"image_url": event["image_url"],
|
||||
"ocr_text_preview": event["ocr_text_preview"],
|
||||
"tags": event.get("tags", []),
|
||||
})
|
||||
elif event["type"] == "done":
|
||||
pass # 搜索结束
|
||||
|
||||
if not results:
|
||||
return f"未找到与 '{keyword}' 相关的图片"
|
||||
|
||||
# 返回 JSON,方便 Agent 后续处理(如发送到企业微信)
|
||||
return json.dumps({
|
||||
"keyword": keyword,
|
||||
"total": len(results),
|
||||
"images": results,
|
||||
}, ensure_ascii=False)
|
||||
```
|
||||
|
||||
> **提示**:`image_base64` 字段包含图片的完整 Base64 编码,可直接用于企业微信发送图片接口。
|
||||
> 企业微信发送图片需要先通过「上传临时素材」接口上传 base64 获取 `media_id`,再发送图片消息。
|
||||
|
||||
---
|
||||
|
||||
## 六、获取图片详情
|
||||
|
||||
### 接口
|
||||
|
||||
```
|
||||
GET /api/v1/images/{image_id}
|
||||
```
|
||||
|
||||
### 返回
|
||||
|
||||
```json
|
||||
{
|
||||
"id": 1,
|
||||
"file_path": "data/images/001.png",
|
||||
"ocr_text": "完整 OCR 文本...",
|
||||
"confidence": 1.0,
|
||||
"provider": "deepseek",
|
||||
"status": "completed",
|
||||
"tags": null,
|
||||
"story_summary": null,
|
||||
"created_at": "2026-04-11T20:00:00",
|
||||
"updated_at": "2026-04-11T20:00:05"
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 七、获取图片列表
|
||||
|
||||
### 接口
|
||||
|
||||
```
|
||||
GET /api/v1/images?page=1&page_size=20&status=completed
|
||||
```
|
||||
|
||||
### 参数
|
||||
|
||||
| 参数 | 类型 | 必填 | 默认值 | 说明 |
|
||||
|------|------|------|--------|------|
|
||||
| page | int | 否 | 1 | 页码 |
|
||||
| page_size | int | 否 | 20 | 每页数量 |
|
||||
| status | string | 否 | - | 筛选状态:`completed`、`pending`、`processing`、`failed` |
|
||||
|
||||
### 返回
|
||||
|
||||
```json
|
||||
{
|
||||
"total": 385,
|
||||
"page": 1,
|
||||
"page_size": 20,
|
||||
"items": [
|
||||
{
|
||||
"id": 385,
|
||||
"file_path": "data/images/001.png",
|
||||
"ocr_text": "...",
|
||||
"confidence": 1.0,
|
||||
"provider": "deepseek",
|
||||
"status": "completed",
|
||||
"tags": null,
|
||||
"story_summary": null,
|
||||
"created_at": "2026-04-11T21:59:56",
|
||||
"updated_at": "2026-04-11T21:59:59"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 八、图片文件访问
|
||||
|
||||
识别后的图片可通过静态文件路径访问:
|
||||
|
||||
```
|
||||
GET /data/images/{文件名}
|
||||
```
|
||||
|
||||
例如:
|
||||
```
|
||||
http://localhost:8765/data/images/001.png
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 九、注意事项
|
||||
|
||||
1. **搜索接口是 SSE 流式**,不能用普通的 `requests.get()` 或 `httpx.get()` 直接获取 JSON,必须用流式读取(`stream()` / `iter_lines()`)
|
||||
2. **搜索耗时较长**:385 张图片全部遍历约需 2-5 分钟,建议设置较长的超时(300 秒)
|
||||
3. **搜索结果实时推送**:不需要等所有图片检查完毕,匹配结果会逐条推送
|
||||
4. **limit 参数控制返回数量**:找到 `limit` 条匹配后自动停止搜索
|
||||
5. **OCR 识别耗时**:单张图片约 5-15 秒,建议设置 120 秒超时
|
||||
6. **文件格式**:仅支持 .png .jpg .jpeg .bmp .webp
|
||||
BIN
docs/软著材料/源代码_前30页.docx
Normal file
BIN
docs/软著材料/源代码_前30页.docx
Normal file
Binary file not shown.
BIN
docs/软著材料/源代码_后30页.docx
Normal file
BIN
docs/软著材料/源代码_后30页.docx
Normal file
Binary file not shown.
BIN
docs/软著材料/用户操作手册.docx
Normal file
BIN
docs/软著材料/用户操作手册.docx
Normal file
Binary file not shown.
BIN
docs/软著材料/申请表.docx
Normal file
BIN
docs/软著材料/申请表.docx
Normal file
Binary file not shown.
62
pyproject.toml
Normal file
62
pyproject.toml
Normal file
@@ -0,0 +1,62 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=75.0", "wheel"]
|
||||
build-backend = "setuptools.backends._legacy:_Backend"
|
||||
|
||||
[project]
|
||||
name = "edu-brain"
|
||||
version = "1.2.0"
|
||||
description = "中文直播教育知识库系统 - 基于向量检索的知识管理平台"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
license = {text = "MIT"}
|
||||
authors = [
|
||||
{name = "EduBrain Team"},
|
||||
]
|
||||
classifiers = [
|
||||
"Development Status :: 3 - Alpha",
|
||||
"Framework :: FastAPI",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Topic :: Education",
|
||||
]
|
||||
|
||||
dependencies = [
|
||||
"fastapi>=0.115.0",
|
||||
"uvicorn[standard]>=0.32.0",
|
||||
"pydantic-settings>=2.6.0",
|
||||
"sqlalchemy[asyncio]>=2.0.36",
|
||||
"asyncpg>=0.30.0",
|
||||
"alembic>=1.14.0",
|
||||
"pgvector>=0.3.6",
|
||||
"openai>=1.58.0",
|
||||
"zhipuai>=2.4.0",
|
||||
"dashscope>=1.20.0",
|
||||
"sentence-transformers>=3.3.0",
|
||||
"paddleocr>=2.9.0",
|
||||
"mcp[cli]>=1.2.0",
|
||||
"frontmatter>=3.2.0",
|
||||
"python-multipart>=0.0.18",
|
||||
"numpy>=1.26.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=8.0.0",
|
||||
"pytest-asyncio>=0.24.0",
|
||||
"httpx>=0.28.0",
|
||||
"ruff>=0.8.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
edu-brain = "app.main:app"
|
||||
edu-brain-mcp = "app.mcp.server:run_mcp_server"
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py311"
|
||||
line-length = 120
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "N", "W", "UP"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests"]
|
||||
41
requirements.txt
Normal file
41
requirements.txt
Normal file
@@ -0,0 +1,41 @@
|
||||
# ── Web 框架 ──
|
||||
fastapi>=0.115.0,<1.0.0
|
||||
uvicorn[standard]>=0.32.0,<1.0.0
|
||||
pydantic-settings>=2.6.0,<3.0.0
|
||||
python-multipart>=0.0.18
|
||||
|
||||
# ── 数据库 ──
|
||||
sqlalchemy[asyncio]>=2.0.36,<3.0.0
|
||||
aiosqlite>=0.20.0,<1.0.0
|
||||
asyncpg>=0.30.0,<1.0.0
|
||||
alembic>=1.14.0,<2.0.0
|
||||
pgvector>=0.3.6,<1.0.0
|
||||
|
||||
# ── 嵌入模型 SDK ──
|
||||
openai>=1.58.0,<2.0.0
|
||||
zhipuai>=2.1.0,<3.0.0
|
||||
dashscope>=1.20.0,<2.0.0
|
||||
sentence-transformers>=3.3.0,<4.0.0
|
||||
httpx>=0.27.0,<1.0.0
|
||||
|
||||
# ── OCR(可选,按需安装) ──
|
||||
# paddleocr>=2.9.0,<3.0.0
|
||||
# paddlepaddle 单独安装
|
||||
|
||||
# ── MCP SDK ──
|
||||
mcp[cli]>=1.2.0,<2.0.0
|
||||
|
||||
# ── Markdown 解析 ──
|
||||
python-frontmatter>=1.0.0,<2.0.0
|
||||
|
||||
# ── Word 文档解析 ──
|
||||
python-docx>=1.1.0,<2.0.0
|
||||
|
||||
# ── 中文分词(BM25 搜索引擎) ──
|
||||
jieba>=0.42.1
|
||||
|
||||
# ── 工具库 ──
|
||||
numpy>=1.26.0,<3.0.0
|
||||
|
||||
# ── 企业微信智能机器人 SDK(WebSocket 长连接) ──
|
||||
wecom-aibot-python-sdk>=1.0.0
|
||||
129
sql/init.sql
Normal file
129
sql/init.sql
Normal file
@@ -0,0 +1,129 @@
|
||||
-- ============================================================
|
||||
-- EduBrain 数据库初始化脚本
|
||||
-- 包含 pgvector 扩展、zhparser 中文分词扩展
|
||||
-- ============================================================
|
||||
|
||||
-- 创建 pgvector 扩展(向量相似度搜索)
|
||||
CREATE EXTENSION IF NOT EXISTS vector;
|
||||
|
||||
-- 创建 zhparser 中文分词扩展(全文搜索)
|
||||
CREATE EXTENSION IF NOT EXISTS zhparser;
|
||||
|
||||
-- 创建基于 zhparser 的中文全文搜索配置
|
||||
CREATE TEXT SEARCH CONFIGURATION chinese_zh (PARSER = zhparser);
|
||||
|
||||
-- 配置中文全文搜索映射:名词、动词、形容词、成语、叹词、习语使用 simple 字典
|
||||
ALTER TEXT SEARCH CONFIGURATION chinese_zh ADD MAPPING FOR n,v,a,i,e,l WITH simple;
|
||||
|
||||
-- ============================================================
|
||||
-- 知识页面表
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS knowledge_pages (
|
||||
id SERIAL PRIMARY KEY,
|
||||
title VARCHAR(500) NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
source_file VARCHAR(500),
|
||||
course_name VARCHAR(200),
|
||||
teacher_name VARCHAR(100),
|
||||
live_date VARCHAR(20),
|
||||
page_number INTEGER,
|
||||
metadata_json TEXT,
|
||||
created_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ DEFAULT NOW()
|
||||
);
|
||||
|
||||
-- 索引
|
||||
CREATE INDEX IF NOT EXISTS idx_pages_course ON knowledge_pages(course_name);
|
||||
CREATE INDEX IF NOT EXISTS idx_pages_teacher ON knowledge_pages(teacher_name);
|
||||
CREATE INDEX IF NOT EXISTS idx_pages_date ON knowledge_pages(live_date);
|
||||
CREATE INDEX IF NOT EXISTS idx_pages_created ON knowledge_pages(created_at DESC);
|
||||
|
||||
-- ============================================================
|
||||
-- 知识分块表(向量检索单元)
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS knowledge_chunks (
|
||||
id SERIAL PRIMARY KEY,
|
||||
page_id INTEGER NOT NULL REFERENCES knowledge_pages(id) ON DELETE CASCADE,
|
||||
chunk_index INTEGER NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
embedding vector(1024), -- 嵌入向量(默认 1024 维)
|
||||
search_vector TSVECTOR, -- 中文全文搜索向量
|
||||
created_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ DEFAULT NOW()
|
||||
);
|
||||
|
||||
-- 索引
|
||||
CREATE INDEX IF NOT EXISTS idx_chunks_page ON knowledge_chunks(page_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_chunks_page_index ON knowledge_chunks(page_id, chunk_index);
|
||||
|
||||
-- HNSW 向量索引(适合高维向量,查询速度快)
|
||||
CREATE INDEX IF NOT EXISTS idx_chunks_embedding
|
||||
ON knowledge_chunks
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
WITH (m = 16, ef_construction = 200);
|
||||
|
||||
-- GIN 全文搜索索引
|
||||
CREATE INDEX IF NOT EXISTS idx_chunks_search
|
||||
ON knowledge_chunks
|
||||
USING gin(search_vector);
|
||||
|
||||
-- ============================================================
|
||||
-- OCR 图片记录表
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS ocr_images (
|
||||
id SERIAL PRIMARY KEY,
|
||||
file_path VARCHAR(500) NOT NULL,
|
||||
ocr_text TEXT,
|
||||
confidence FLOAT,
|
||||
provider VARCHAR(50),
|
||||
status VARCHAR(20) DEFAULT 'pending',
|
||||
error_message TEXT,
|
||||
created_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ DEFAULT NOW()
|
||||
);
|
||||
|
||||
-- 索引
|
||||
CREATE INDEX IF NOT EXISTS idx_ocr_status ON ocr_images(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_ocr_created ON ocr_images(created_at DESC);
|
||||
|
||||
-- ============================================================
|
||||
-- updated_at 自动更新触发器
|
||||
-- ============================================================
|
||||
CREATE OR REPLACE FUNCTION update_updated_at_column()
|
||||
RETURNS TRIGGER AS $$
|
||||
BEGIN
|
||||
NEW.updated_at = NOW();
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
-- 为所有表添加 updated_at 触发器
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_trigger
|
||||
WHERE tgname = 'tr_pages_updated_at'
|
||||
) THEN
|
||||
CREATE TRIGGER tr_pages_updated_at
|
||||
BEFORE UPDATE ON knowledge_pages
|
||||
FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();
|
||||
END IF;
|
||||
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_trigger
|
||||
WHERE tgname = 'tr_chunks_updated_at'
|
||||
) THEN
|
||||
CREATE TRIGGER tr_chunks_updated_at
|
||||
BEFORE UPDATE ON knowledge_chunks
|
||||
FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();
|
||||
END IF;
|
||||
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_trigger
|
||||
WHERE tgname = 'tr_ocr_updated_at'
|
||||
) THEN
|
||||
CREATE TRIGGER tr_ocr_updated_at
|
||||
BEFORE UPDATE ON ocr_images
|
||||
FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();
|
||||
END IF;
|
||||
END $$;
|
||||
2782
static/index.html
Normal file
2782
static/index.html
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user