diff --git a/backend/.env.example b/backend/.env.example index ad75bce2..fa109b27 100644 --- a/backend/.env.example +++ b/backend/.env.example @@ -27,6 +27,8 @@ LANGSMITH_TRACING=false # 默认建议创建独立数据库，如 ctb；若直接使用 postgres 默认库，也可改为 /postgres。 # APP_DATABASE_URL=postgresql+psycopg://postgres:root@localhost:15432/ctb # APP_DATABASE_ECHO=false +# APP_POSTGRES_VECTOR_DIMENSIONS 必须和所选 embedding 模型维度一致。 +# 例如 text-embedding-3-small 默认是 1536 维。 # APP_POSTGRES_VECTOR_DIMENSIONS=1536 # ============================================================================ diff --git a/backend/core/config.py b/backend/core/config.py index e5718b5e..03702970 100644 --- a/backend/core/config.py +++ b/backend/core/config.py @@ -205,6 +205,7 @@ class Settings(BaseSettings): db_path: Path | None = None database_url: str = "" database_echo: bool = False + postgres_vector_dimensions: int = 1536 # 各类子目录（由 validator 从 runtime_dir 派生，可独立覆盖以便测试） upload_dir: Path | None = None @@ -225,7 +226,7 @@ class Settings(BaseSettings): True # 是否信任系统代理环境变量，Windows 下设为 False 可解 WinError 10054 ) - rag_embedding_model: str = "text-embedding-v3" + rag_embedding_model: str = "text-embedding-3-small" rag_embedding_api_key: str = "" rag_embedding_base_url: str = "" diff --git a/backend/core/model_selection.py b/backend/core/model_selection.py index a2d09cfd..c111de45 100644 --- a/backend/core/model_selection.py +++ b/backend/core/model_selection.py @@ -22,9 +22,11 @@ class LLMSelectionError(Exception): def split_models(model_name: str | None) -> list[str]: - """返回模型名称列表。现在只支持单个模型，但保留列表形式以兼容调用方。""" - name = (model_name or "").strip() - return [name] if name else [] + """返回模型名称列表，兼容逗号分隔的历史存储格式。""" + raw = str(model_name or "").strip() + if not raw: + return [] + return [item.strip() for item in raw.split(",") if item.strip()] def build_managed_provider_context(db): diff --git a/backend/core/rag.py b/backend/core/rag.py index 10a1cf61..bdbd7da4 100644 --- a/backend/core/rag.py +++ b/backend/core/rag.py @@ -317,42 +317,41 @@ def index_question(db: Session, question_id: int) -> bool: vector_json = _serialize_vector(vector) - if existing: - existing.content = chunk_data["content"] - existing.metadata_json = json.dumps(chunk_data["metadata"], ensure_ascii=False) - existing.content_hash = chunk_data["content_hash"] - existing.embedding_model = settings.rag_embedding_model if vector else None - existing.vector_json = vector_json - else: - chunk = RagDocumentChunk( - user_id=question.user_id, - project_id=question.project_id, - source_type="question", - source_id=question_id, - chunk_index=0, - content=chunk_data["content"], - metadata_json=json.dumps(chunk_data["metadata"], ensure_ascii=False), - content_hash=chunk_data["content_hash"], - embedding_model=settings.rag_embedding_model if vector else None, - vector_json=vector_json, - ) - db.add(chunk) - try: - db.flush() - target_chunk = existing or chunk - _write_postgres_vector(db, target_chunk.id, vector) - db.commit() + with db.begin_nested(): + if existing: + existing.content = chunk_data["content"] + existing.metadata_json = json.dumps(chunk_data["metadata"], ensure_ascii=False) + existing.content_hash = chunk_data["content_hash"] + existing.embedding_model = settings.rag_embedding_model if vector else None + existing.vector_json = vector_json + target_chunk = existing + else: + target_chunk = RagDocumentChunk( + user_id=question.user_id, + project_id=question.project_id, + source_type="question", + source_id=question_id, + chunk_index=0, + content=chunk_data["content"], + metadata_json=json.dumps(chunk_data["metadata"], ensure_ascii=False), + content_hash=chunk_data["content_hash"], + embedding_model=settings.rag_embedding_model if vector else None, + vector_json=vector_json, + ) + db.add(target_chunk) + + db.flush() + _write_postgres_vector(db, target_chunk.id, vector) return True except Exception as e: - db.rollback() logger.error("索引题目 %d 失败: %s", question_id, e) return False def delete_question_chunks(db: Session, question_id: int) -> int: """删除错题关联的所有 RAG chunk""" - count = ( + return ( db.query(RagDocumentChunk) .filter( RagDocumentChunk.source_type == "question", @@ -360,13 +359,6 @@ def delete_question_chunks(db: Session, question_id: int) -> int: ) .delete(synchronize_session=False) ) - try: - db.commit() - except Exception as e: - db.rollback() - logger.error("删除题目 %d 的 chunk 失败: %s", question_id, e) - return 0 - return count # --------------------------------------------------------------------------- diff --git a/backend/db/__init__.py b/backend/db/__init__.py index 3397a280..4a53bd5d 100644 --- a/backend/db/__init__.py +++ b/backend/db/__init__.py @@ -2,7 +2,7 @@ 数据库模块：引擎创建、Session 工厂、初始化函数 """ -from sqlalchemy import create_engine, event +from sqlalchemy import create_engine, event, inspect, text from sqlalchemy.orm import sessionmaker import os import sys @@ -154,7 +154,46 @@ def _migrate_schema(): conn.close() +def _prepare_postgresql_extensions(): + """在 PostgreSQL 建表前确保 pgvector 扩展可用。""" + if not is_postgresql_backend(engine): + return + with engine.begin() as conn: + conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + + +def _ensure_postgresql_schema(): + """为 PostgreSQL 现有表补齐 pgvector 列和索引。""" + if not is_postgresql_backend(engine): + return + + inspector = inspect(engine) + table_names = set(inspector.get_table_names()) + if "rag_document_chunks" not in table_names: + return + + columns = {column["name"] for column in inspector.get_columns("rag_document_chunks")} + with engine.begin() as conn: + conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + if "embedding_vector" not in columns: + conn.execute( + text( + "ALTER TABLE rag_document_chunks " + f"ADD COLUMN embedding_vector vector({settings.postgres_vector_dimensions})" + ) + ) + conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_rag_document_chunks_embedding_vector " + "ON rag_document_chunks USING ivfflat (embedding_vector vector_cosine_ops) " + "WITH (lists = 100)" + ) + ) + + def init_db(): """初始化数据库：建表并执行轻量级自动迁移""" + _prepare_postgresql_extensions() Base.metadata.create_all(bind=engine) _migrate_schema() + _ensure_postgresql_schema() diff --git a/backend/db/crud/projects.py b/backend/db/crud/projects.py index 993584ab..cb451301 100644 --- a/backend/db/crud/projects.py +++ b/backend/db/crud/projects.py @@ -6,7 +6,7 @@ from db.models import ( ChatSession, Note, NoteTagMapping, Project, Question, - QuestionEmbedding, QuestionTagMapping, UploadBatch, + QuestionEmbedding, QuestionTagMapping, RagDocumentChunk, UploadBatch, ) @@ -141,6 +141,10 @@ def delete_project(db: Session, project_id: int, user_id=None) -> bool: # 先删除题目和笔记的关联子表，再删除题目/笔记本身 question_ids = [q.id for q in db.query(Question.id).filter(Question.project_id == project.id).all()] if question_ids: + db.query(RagDocumentChunk).filter( + RagDocumentChunk.source_type == "question", + RagDocumentChunk.source_id.in_(question_ids), + ).delete(synchronize_session=False) db.query(QuestionEmbedding).filter(QuestionEmbedding.question_id.in_(question_ids)).delete(synchronize_session=False) db.query(ChatSession).filter(ChatSession.question_id.in_(question_ids)).delete(synchronize_session=False) db.query(QuestionTagMapping).filter(QuestionTagMapping.question_id.in_(question_ids)).delete(synchronize_session=False) diff --git a/backend/db/crud/questions.py b/backend/db/crud/questions.py index 0fbbbda7..7bffb6ca 100644 --- a/backend/db/crud/questions.py +++ b/backend/db/crud/questions.py @@ -358,8 +358,6 @@ def save_questions_to_db( # 科目由编排智能体识别，不再使用关键词匹配 subject = batch_info.get("subject") or "未知" project_id = project_id or batch_info.get("project_id") - if not project_id: - raise ValueError("PROJECT_REQUIRED") # 创建批次记录 batch = UploadBatch( @@ -693,6 +691,10 @@ def delete_question(db: Session, question_id: int, user_id=None) -> bool: try: batch_id = question.batch_id + from core.rag import delete_question_chunks + + delete_question_chunks(db, question_id) + # 删除关联的标签映射 db.query(QuestionTagMapping).filter(QuestionTagMapping.question_id == question_id).delete() diff --git a/backend/db/models.py b/backend/db/models.py index 3b3b9d79..fc2b4ae9 100644 --- a/backend/db/models.py +++ b/backend/db/models.py @@ -4,13 +4,32 @@ import uuid from sqlalchemy import Column, Integer, String, Text, DateTime, Boolean, ForeignKey, UniqueConstraint -from sqlalchemy.orm import relationship -from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.types import TypeDecorator +from sqlalchemy.orm import declarative_base, relationship from datetime import datetime +from core.config import settings + +try: + from pgvector.sqlalchemy import Vector +except ImportError: # pragma: no cover - optional on SQLite-only setups + Vector = None + Base = declarative_base() +class EmbeddingVectorType(TypeDecorator): + """Use pgvector on PostgreSQL and plain text elsewhere.""" + + impl = Text + cache_ok = True + + def load_dialect_impl(self, dialect): + if dialect.name == "postgresql" and Vector is not None: + return dialect.type_descriptor(Vector(settings.postgres_vector_dimensions)) + return dialect.type_descriptor(Text()) + + class User(Base): """用户表""" __tablename__ = "users" @@ -330,5 +349,6 @@ class RagDocumentChunk(Base): content_hash = Column(String(64), default="", index=True) embedding_model = Column(String(100), nullable=True) vector_json = Column(Text, nullable=True) + embedding_vector = Column(EmbeddingVectorType(), nullable=True) created_at = Column(DateTime, default=datetime.utcnow, index=True) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) diff --git a/backend/routes/chat.py b/backend/routes/chat.py index cd19930d..f7acf991 100644 --- a/backend/routes/chat.py +++ b/backend/routes/chat.py @@ -412,21 +412,6 @@ def stream_chat(session_id): user_id = session.get("user_id") with SessionLocal() as db: - try: - selection = resolve_llm_selection( - db, - user_id=user_id, - category=model_provider, - model_name=model_name, - provider_source=provider_source, - provider_id=provider_id, - ) - except LLMSelectionError as e: - return ( - jsonify({"success": False, "code": e.code, "error": e.message}), - e.status_code, - ) - uid = _effective_user_id() cs_query = ( db.query(ChatSessionModel) @@ -446,6 +431,21 @@ def stream_chat(session_id): if not chat_session: return jsonify({"success": False, "error": "对话不存在"}), 404 + try: + selection = resolve_llm_selection( + db, + user_id=user_id, + category=model_provider, + model_name=model_name, + provider_source=provider_source, + provider_id=provider_id, + ) + except LLMSelectionError as e: + return ( + jsonify({"success": False, "code": e.code, "error": e.message}), + e.status_code, + ) + should_consume_quota = bool(user_id) and uses_server_llm_selection( selection["source"], db=db, diff --git a/backend/routes/questions.py b/backend/routes/questions.py index f7bcf035..2caeb494 100644 --- a/backend/routes/questions.py +++ b/backend/routes/questions.py @@ -456,8 +456,10 @@ def rag_reindex(): with SessionLocal() as db: success = index_question(db, qid) if success: + db.commit() indexed += 1 else: + db.rollback() skipped += 1 except Exception as e: logger.warning("索引题目 %d 失败: %s", qid, e) @@ -531,7 +533,10 @@ def update_question(question_id): try: from core.rag import index_question with SessionLocal() as db: - index_question(db, question_id) + if index_question(db, question_id): + db.commit() + else: + db.rollback() except Exception as e: logger.warning(f"更新题目 {question_id} 的 RAG 索引失败: {e}") @@ -560,7 +565,10 @@ def update_question_answer(question_id): try: from core.rag import index_question with SessionLocal() as db: - index_question(db, question_id) + if index_question(db, question_id): + db.commit() + else: + db.rollback() except Exception as e: logger.warning(f"更新题目 {question_id} 的用户作答索引失败: {e}") @@ -624,6 +632,8 @@ def save_to_db(): if not isinstance(selected_uids, list) or not selected_uids: return jsonify({'success': False, 'error': '请选择至少一道题目'}), 400 + if project_id is None: + return jsonify({'success': False, 'error': '请先创建并选择一个错题库'}), 400 run_id = data.get('run_id') record_id = data.get('record_id') @@ -677,15 +687,11 @@ def save_to_db(): with SessionLocal() as db: try: - project_id = ( - crud.require_project_id( - db, - project_id, - user_id=session.get('user_id'), - project_type="question", - ) - if project_id - else None + project_id = crud.resolve_project_id( + db, + project_id, + user_id=session.get('user_id'), + project_type="question", ) except ValueError as exc: if str(exc) == "PROJECT_REQUIRED": @@ -878,7 +884,10 @@ def save_question_answer(question_id): try: from core.rag import index_question with SessionLocal() as db: - index_question(db, question_id) + if index_question(db, question_id): + db.commit() + else: + db.rollback() except Exception as e: logger.warning(f"更新题目 {question_id} 的 RAG 索引失败: {e}") diff --git a/backend/src/utils.py b/backend/src/utils.py index 6997bd8a..5528e720 100644 --- a/backend/src/utils.py +++ b/backend/src/utils.py @@ -85,9 +85,12 @@ def export_wrongbook( os.makedirs(settings.results_dir, exist_ok=True) output_path = os.path.join(settings.results_dir, "wrongbook.md") - # 用 uid 过滤选中的题目，保持原始顺序 + # 优先用 uid 过滤；兼容历史调用中直接传 question_id。 uid_set = set(selected_uids) - selected_questions = [q for q in questions if q.get('uid') in uid_set] + selected_questions = [ + q for q in questions + if q.get('uid') in uid_set or str(q.get('question_id', '')) in uid_set + ] # 构建Markdown内容 md_content = "# 错题本\n\n" @@ -121,23 +124,25 @@ def _fix_html_image_src(html: str) -> str: current_section = _INIT unsorted_started = False serial = 0 # 全局序号 + has_sections = any(q.get('section_title') for q in selected_questions) for q in selected_questions: section = q.get('section_title') - if section: + if has_sections and section: # 有 section_title：正常输出大题分组标题 if section != current_section: current_section = section md_content += f"## {section}\n\n" - else: + elif has_sections: # section=None：首次遇到时输出"（未分类）"节标题 if not unsorted_started: unsorted_started = True md_content += "## （未分类）\n\n" serial += 1 - md_content += f"### {serial}. 题目 {q.get('question_id', '')} ({q.get('question_type', '未知')})\n\n" + heading_level = "###" if has_sections else "##" + md_content += f"{heading_level} {serial}. 题目 {q.get('question_id', '')} ({q.get('question_type', '未知')})\n\n" # 获取图片引用列表，用于填充空的 image block image_refs = q.get('image_refs') or [] @@ -179,7 +184,7 @@ def _fix_html_image_src(html: str) -> str: for image_path in remaining_images: md_content += f"![图片]({_resolve_image_path(image_path)})\n\n" - answer_prefix = "####" if section else "###" + answer_prefix = "####" if has_sections and section else "###" md_content += f"{answer_prefix} 我的答案\n\n" md_content += "_（请在此处填写你的答案）_\n\n" diff --git a/backend/src/workflow.py b/backend/src/workflow.py index b6c7c132..07c96cf9 100644 --- a/backend/src/workflow.py +++ b/backend/src/workflow.py @@ -188,8 +188,8 @@ def _build_overlapping_batches( n_pages = len(ocr_data) if n_pages <= batch_size: - # 只有一批：全部页都是 primary - return [[dict(page, is_primary=True) for page in ocr_data]] + # 只有一批时保留原始结构，默认所有页都视为 primary。 + return [list(ocr_data)] step = batch_size - overlap batches = [] @@ -550,17 +550,15 @@ def _dedup_questions(questions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: # ── 第一轮：按 (section, qid) 复合键去重 ────────────────── groups: Dict[tuple, List[Dict[str, Any]]] = defaultdict(list) - no_id: List[Dict[str, Any]] = [] - for q in questions: qid = q.get("question_id", "") if not qid: - no_id.append(q) + continue else: section = q.get("section_title") or "" groups[(section, qid)].append(q) - after_round1: List[Dict[str, Any]] = list(no_id) + after_round1: List[Dict[str, Any]] = [] for qs in groups.values(): after_round1.append(max(qs, key=_question_richness)) @@ -573,7 +571,7 @@ def _dedup_questions(questions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: by_qid[q.get("question_id", "")].append(q) SIMILARITY_THRESHOLD = 0.75 - final: List[Dict[str, Any]] = list(no_id) + final: List[Dict[str, Any]] = [] round2_removed = 0 for qid, entries in by_qid.items(): diff --git a/backend/tests/README.md b/backend/tests/README.md index 3f9de920..0efa15cc 100644 --- a/backend/tests/README.md +++ b/backend/tests/README.md @@ -56,6 +56,7 @@ backend/tests/ ├── test_web_helpers.py # web_app.py 纯函数测试 ├── test_web_routes.py # Flask 路由集成测试（内存数据库） ├── test_crud.py # 数据库 CRUD 测试 +├── test_db_init.py # PostgreSQL / pgvector 初始化与迁移测试 ├── test_schemas.py # Pydantic schema 校验测试 ├── test_question_tools.py # 题目工具函数测试 ├── test_correct_node.py # 纠错节点合并逻辑测试 @@ -136,6 +137,16 @@ backend/tests/ | `TestGetAllTags` | `get_all_tags` | 2 | 获取全部标签：空库、按科目筛选 | | `TestGetStatistics` | `get_statistics` | 2 | 统计信息：空库、有数据 | +### test_db_init.py + +测试 `backend/db/__init__.py` 中 PostgreSQL 初始化分支的 DDL 逻辑，使用 mock 验证，不依赖真实 PostgreSQL： + +| 测试方法 | 说明 | +|----------|------| +| `test_prepare_postgresql_extensions_creates_vector_extension` | 初始化前会尝试创建 `vector` 扩展 | +| `test_ensure_postgresql_schema_adds_missing_vector_column_and_index` | 缺少 `embedding_vector` 时自动补列并建索引 | +| `test_ensure_postgresql_schema_skips_column_creation_when_present` | 已存在向量列时不重复 `ALTER TABLE` | + ### test_schemas.py 测试 `backend/error_correction_agent/schemas.py` 中 Pydantic 模型的校验逻辑： diff --git a/backend/tests/test_chat_routes.py b/backend/tests/test_chat_routes.py index 74ec4f0d..e32af92d 100644 --- a/backend/tests/test_chat_routes.py +++ b/backend/tests/test_chat_routes.py @@ -20,7 +20,7 @@ from sqlalchemy import create_engine, event from sqlalchemy.orm import sessionmaker -from db.models import Base, ChatMessage, ProviderConfig, User +from db.models import Base, ChatMessage, ProviderConfig, SystemProviderConfig, User from db import crud from tests.conftest import make_question @@ -97,6 +97,21 @@ def _seed_question(test_db): return test_db.query(Question).first() +def _seed_system_openai_provider(test_db, model_name="gpt-4o-mini"): + provider = SystemProviderConfig( + id=str(uuid.uuid4()), + category="openai", + name="平台托管 OpenAI", + is_active=True, + api_key="sk-managed", + base_url="https://example.com", + model_name=model_name, + ) + test_db.add(provider) + test_db.commit() + return provider + + # ═══════════════════════════════════════════════════════════ # PUT /api/question//answer # ═══════════════════════════════════════════════════════════ @@ -288,6 +303,7 @@ def test_server_managed_chat_consumes_quota_after_success(self, client, test_db) user.daily_free_used = 0 user.daily_free_quota_date = datetime.utcnow().date().isoformat() test_db.commit() + _seed_system_openai_provider(test_db) q = _seed_question(test_db) session = crud.create_chat_session(test_db, q.id, user_id=1) @@ -434,6 +450,7 @@ def test_quota_exhausted_chat_returns_429_and_does_not_save_user_message(self, c user.daily_free_used = 5 user.daily_free_quota_date = datetime.utcnow().date().isoformat() test_db.commit() + _seed_system_openai_provider(test_db) q = _seed_question(test_db) session = crud.create_chat_session(test_db, q.id, user_id=1) diff --git a/backend/tests/test_db_init.py b/backend/tests/test_db_init.py new file mode 100644 index 00000000..52faa71c --- /dev/null +++ b/backend/tests/test_db_init.py @@ -0,0 +1,103 @@ +from unittest.mock import Mock, patch + +import db as db_module + + +class _BeginContext: + def __init__(self, conn): + self._conn = conn + + def __enter__(self): + return self._conn + + def __exit__(self, exc_type, exc, tb): + return False + + +class _RecordingConnection: + def __init__(self): + self.statements = [] + + def execute(self, statement): + self.statements.append(" ".join(str(statement).split())) + + +def test_prepare_postgresql_extensions_creates_vector_extension(): + conn = _RecordingConnection() + fake_engine = Mock() + fake_engine.begin.return_value = _BeginContext(conn) + + with ( + patch.object(db_module, "engine", fake_engine), + patch.object(db_module, "is_postgresql_backend", return_value=True), + ): + db_module._prepare_postgresql_extensions() + + assert any( + "CREATE EXTENSION IF NOT EXISTS vector" in statement + for statement in conn.statements + ) + + +def test_ensure_postgresql_schema_adds_missing_vector_column_and_index(): + conn = _RecordingConnection() + fake_engine = Mock() + fake_engine.begin.return_value = _BeginContext(conn) + fake_inspector = Mock() + fake_inspector.get_table_names.return_value = ["rag_document_chunks"] + fake_inspector.get_columns.return_value = [ + {"name": "id"}, + {"name": "content"}, + {"name": "vector_json"}, + ] + + with ( + patch.object(db_module, "engine", fake_engine), + patch.object(db_module, "inspect", return_value=fake_inspector), + patch.object(db_module, "is_postgresql_backend", return_value=True), + ): + db_module._ensure_postgresql_schema() + + assert any( + "CREATE EXTENSION IF NOT EXISTS vector" in statement + for statement in conn.statements + ) + assert any( + "ALTER TABLE rag_document_chunks" in statement + and "embedding_vector" in statement + for statement in conn.statements + ) + assert any( + "CREATE INDEX IF NOT EXISTS idx_rag_document_chunks_embedding_vector" + in statement + for statement in conn.statements + ) + + +def test_ensure_postgresql_schema_skips_column_creation_when_present(): + conn = _RecordingConnection() + fake_engine = Mock() + fake_engine.begin.return_value = _BeginContext(conn) + fake_inspector = Mock() + fake_inspector.get_table_names.return_value = ["rag_document_chunks"] + fake_inspector.get_columns.return_value = [ + {"name": "id"}, + {"name": "embedding_vector"}, + ] + + with ( + patch.object(db_module, "engine", fake_engine), + patch.object(db_module, "inspect", return_value=fake_inspector), + patch.object(db_module, "is_postgresql_backend", return_value=True), + ): + db_module._ensure_postgresql_schema() + + assert not any( + "ALTER TABLE rag_document_chunks ADD COLUMN embedding_vector" in statement + for statement in conn.statements + ) + assert any( + "CREATE INDEX IF NOT EXISTS idx_rag_document_chunks_embedding_vector" + in statement + for statement in conn.statements + ) diff --git a/backend/tests/test_migrate_and_delete_question.py b/backend/tests/test_migrate_and_delete_question.py index 20e4697c..687ab72c 100644 --- a/backend/tests/test_migrate_and_delete_question.py +++ b/backend/tests/test_migrate_and_delete_question.py @@ -1,6 +1,6 @@ from db.crud.questions import delete_question from db.migrate import _ensure_default_question_project -from db.models import Project, Question, UploadBatch, User +from db.models import Project, Question, RagDocumentChunk, UploadBatch, User def test_ensure_default_question_project_creates_immutable_default_project(db): @@ -68,6 +68,19 @@ def test_delete_question_removes_empty_batch_in_same_operation(db): db.add(question) db.commit() + chunk = RagDocumentChunk( + user_id=user.id, + project_id=None, + source_type="question", + source_id=question.id, + chunk_index=0, + content="题目", + vector_json="[0.1,0.2]", + ) + db.add(chunk) + db.commit() + assert delete_question(db, question.id, user_id=user.id) is True assert db.query(Question).filter(Question.id == question.id).first() is None + assert db.query(RagDocumentChunk).filter(RagDocumentChunk.source_id == question.id).first() is None assert db.query(UploadBatch).filter(UploadBatch.id == batch_id).first() is None diff --git a/backend/tests/test_web_routes.py b/backend/tests/test_web_routes.py index 435c0482..2ef9dc8a 100644 --- a/backend/tests/test_web_routes.py +++ b/backend/tests/test_web_routes.py @@ -20,7 +20,7 @@ from sqlalchemy import create_engine, event from sqlalchemy.orm import sessionmaker -from db.models import Base, User, ProviderConfig, SystemProviderConfig, Question +from db.models import Base, User, ProviderConfig, SystemProviderConfig, Project, Question from db import crud # 测试用户 ID，与 client fixture 中的 session['user_id'] 一致 @@ -62,6 +62,13 @@ def _seed_split_run(db, tmp_path, *, user_id, run_id, question_text): return questions +def _create_question_project(db, *, user_id=TEST_USER_ID, name="默认错题库"): + project = Project(user_id=user_id, name=name, project_type="question") + db.add(project) + db.commit() + return project + + @pytest.fixture def test_db(): """内存数据库 + 建表 + 创建测试用户""" @@ -572,7 +579,27 @@ def test_save_to_db_rejects_other_users_run(self, client, test_db, tmp_path): assert resp.status_code == 400 assert test_db.query(Question).count() == 0 + def test_save_to_db_requires_project_id(self, client, test_db, tmp_path): + _seed_split_run( + test_db, + tmp_path, + user_id=TEST_USER_ID, + run_id="own-run-missing-project", + question_text="missing project", + ) + + resp = client.post( + "/api/save-to-db", + json={"run_id": "own-run-missing-project", "selected_ids": ["0"]}, + ) + assert resp.status_code == 400 + data = resp.get_json() + assert data["success"] is False + assert "错题库" in data["error"] + assert test_db.query(Question).count() == 0 + def test_save_to_db_imports_current_users_run(self, client, test_db, tmp_path): + project = _create_question_project(test_db) _seed_split_run( test_db, tmp_path, @@ -583,7 +610,7 @@ def test_save_to_db_imports_current_users_run(self, client, test_db, tmp_path): resp = client.post( "/api/save-to-db", - json={"run_id": "own-run-to-import", "selected_ids": ["0"]}, + json={"run_id": "own-run-to-import", "selected_ids": ["0"], "project_id": project.id}, ) assert resp.status_code == 200 data = resp.get_json() @@ -593,6 +620,7 @@ def test_save_to_db_imports_current_users_run(self, client, test_db, tmp_path): saved = test_db.query(Question).all() assert len(saved) == 1 assert saved[0].user_id == TEST_USER_ID + assert saved[0].project_id == project.id assert "import me" in saved[0].content_json diff --git a/backend/tests/test_workflow_helpers.py b/backend/tests/test_workflow_helpers.py index aa615828..c089d091 100644 --- a/backend/tests/test_workflow_helpers.py +++ b/backend/tests/test_workflow_helpers.py @@ -868,6 +868,7 @@ def _make_ocr_result(n_blocks=2): }] } + @patch.object(PaddleOCRClient, "parse_images_async", new=MagicMock(return_value=[_make_ocr_result.__func__()])) @patch("src.workflow.run_async") @patch.object(PaddleOCRClient, "parse_pdf") @patch.object(PaddleOCRClient, "__init__", return_value=None) @@ -879,9 +880,11 @@ def test_only_images(self, mock_init, mock_pdf, mock_run_async): result = _run_ocr_and_simplify(["a.png", "b.jpg"]) mock_pdf.assert_not_called() + PaddleOCRClient.parse_images_async.assert_called_once() mock_run_async.assert_called_once() assert len(result) >= 1 + @patch.object(PaddleOCRClient, "parse_images_async", new=MagicMock()) @patch("src.workflow.run_async") @patch.object(PaddleOCRClient, "parse_pdf") @patch.object(PaddleOCRClient, "__init__", return_value=None) @@ -892,9 +895,11 @@ def test_only_pdfs(self, mock_init, mock_pdf, mock_run_async): result = _run_ocr_and_simplify(["a.pdf", "b.pdf"]) assert mock_pdf.call_count == 2 + PaddleOCRClient.parse_images_async.assert_not_called() mock_run_async.assert_not_called() assert len(result) >= 1 + @patch.object(PaddleOCRClient, "parse_images_async", new=MagicMock(return_value=[_make_ocr_result.__func__()])) @patch("src.workflow.run_async") @patch.object(PaddleOCRClient, "parse_pdf") @patch.object(PaddleOCRClient, "__init__", return_value=None) @@ -906,10 +911,12 @@ def test_mixed_files(self, mock_init, mock_pdf, mock_run_async): result = _run_ocr_and_simplify(["doc.pdf", "img.png"]) mock_pdf.assert_called_once() + PaddleOCRClient.parse_images_async.assert_called_once() mock_run_async.assert_called_once() # PDF 3 blocks + 图片 2 blocks = 2 pages assert len(result) == 2 + @patch.object(PaddleOCRClient, "parse_images_async", new=MagicMock()) @patch("src.workflow.run_async") @patch.object(PaddleOCRClient, "parse_pdf") @patch.object(PaddleOCRClient, "__init__", return_value=None) @@ -920,6 +927,7 @@ def test_pdf_case_insensitive(self, mock_init, mock_pdf, mock_run_async): _run_ocr_and_simplify(["DOC.PDF", "test.Pdf"]) assert mock_pdf.call_count == 2 + PaddleOCRClient.parse_images_async.assert_not_called() mock_run_async.assert_not_called() @patch("src.workflow.run_async")