diff --git a/backend/alembic/env.py b/backend/alembic/env.py index 9fc807d..c41635f 100644 --- a/backend/alembic/env.py +++ b/backend/alembic/env.py @@ -2,12 +2,12 @@ from logging.config import fileConfig -from alembic import context from sqlalchemy import engine_from_config, pool +from alembic import context from chatmaster.config import get_settings -from chatmaster.db.base import Base from chatmaster.db import models # noqa: F401 +from chatmaster.db.base import Base config = context.config diff --git a/backend/alembic/versions/0001_baseline.py b/backend/alembic/versions/0001_baseline.py index 0bd657e..0ba24ec 100644 --- a/backend/alembic/versions/0001_baseline.py +++ b/backend/alembic/versions/0001_baseline.py @@ -7,9 +7,10 @@ from __future__ import annotations -from alembic import op import sqlalchemy as sa +from alembic import op + revision = "0001_baseline" down_revision = None branch_labels = None diff --git a/backend/alembic/versions/0002_reliability_fields.py b/backend/alembic/versions/0002_reliability_fields.py index 8d20a71..d27ccfd 100644 --- a/backend/alembic/versions/0002_reliability_fields.py +++ b/backend/alembic/versions/0002_reliability_fields.py @@ -7,9 +7,9 @@ from __future__ import annotations -from alembic import op import sqlalchemy as sa +from alembic import op revision = "0002_reliability_fields" down_revision = "0001_baseline" diff --git a/backend/alembic/versions/0003_persona_management.py b/backend/alembic/versions/0003_persona_management.py index 7a66d67..9c2dbac 100644 --- a/backend/alembic/versions/0003_persona_management.py +++ b/backend/alembic/versions/0003_persona_management.py @@ -7,9 +7,9 @@ from __future__ import annotations -from alembic import op import sqlalchemy as sa +from alembic import op revision = "0003_persona_management" down_revision = "0002_reliability_fields" diff --git a/backend/chatmaster/ai/models.py b/backend/chatmaster/ai/models.py index 6b2af05..b66bfdb 100644 --- a/backend/chatmaster/ai/models.py +++ b/backend/chatmaster/ai/models.py @@ -115,6 +115,10 @@ def build_embeddings(identity: IdentityConfig | None = None) -> Embeddings: provider = cfg.provider.lower() if provider == "huggingface": + validate_provider_url( + cfg.huggingface_endpoint, + allow_private_network=settings.allow_private_provider_urls, + ) return _build_hf_embeddings(model, cfg.huggingface_endpoint) if provider in _OPENAI_COMPATIBLE: api_key = cfg.api_key or settings.openai_api_key diff --git a/backend/chatmaster/ai/providers.py b/backend/chatmaster/ai/providers.py index 121af89..f88bd95 100644 --- a/backend/chatmaster/ai/providers.py +++ b/backend/chatmaster/ai/providers.py @@ -13,12 +13,11 @@ import threading import uuid from functools import lru_cache - -from sqlalchemy import select -from sqlalchemy.exc import SQLAlchemyError from typing import Literal from pydantic import BaseModel, Field +from sqlalchemy import select +from sqlalchemy.exc import SQLAlchemyError from chatmaster.config import get_settings diff --git a/backend/chatmaster/chat/graph.py b/backend/chatmaster/chat/graph.py index d0da503..eb48296 100644 --- a/backend/chatmaster/chat/graph.py +++ b/backend/chatmaster/chat/graph.py @@ -16,8 +16,8 @@ from chatmaster.conversations.service import load_history as load_history_from_db from chatmaster.db.models import Conversation, Message, utc_now from chatmaster.db.session import SessionLocal -from chatmaster.identities.service import get_identity_config from chatmaster.identities.schema import IdentityConfig +from chatmaster.identities.service import get_identity_config from chatmaster.retrieval.retriever import retrieve from chatmaster.retrieval.schemas import RetrievedChunk diff --git a/backend/chatmaster/chat/service.py b/backend/chatmaster/chat/service.py index 81cb02a..534c440 100644 --- a/backend/chatmaster/chat/service.py +++ b/backend/chatmaster/chat/service.py @@ -78,9 +78,7 @@ def _begin_turn( conversation = db.get(Conversation, conversation_id) if conversation is not None and conversation.title == "新对话": has_messages = db.scalars( - select(Message.id) - .where(Message.conversation_id == conversation_id) - .limit(1) + select(Message.id).where(Message.conversation_id == conversation_id).limit(1) ).first() if has_messages is None: conversation.title = title_from_message(message) diff --git a/backend/chatmaster/db/init_db.py b/backend/chatmaster/db/init_db.py index df5e0cf..11a6142 100644 --- a/backend/chatmaster/db/init_db.py +++ b/backend/chatmaster/db/init_db.py @@ -15,13 +15,13 @@ def init_db() -> None: def migrate_db() -> None: """Upgrade the runtime database, adopting legacy create_all databases once.""" import shutil - from datetime import datetime + from datetime import datetime, timezone from pathlib import Path - from alembic import command from alembic.config import Config from sqlalchemy import inspect + from alembic import command from chatmaster.config import get_settings settings = get_settings() @@ -34,7 +34,7 @@ def migrate_db() -> None: if settings.database_url.startswith("sqlite:///"): db_path = Path(settings.database_url.removeprefix("sqlite:///")) if db_path.exists(): - suffix = datetime.now().strftime("%Y%m%d%H%M%S") + suffix = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S") shutil.copy2(db_path, db_path.with_suffix(db_path.suffix + f".{suffix}.bak")) message_columns = ( {column["name"] for column in inspector.get_columns("messages")} diff --git a/backend/chatmaster/db/models.py b/backend/chatmaster/db/models.py index ee19f64..dd5adef 100644 --- a/backend/chatmaster/db/models.py +++ b/backend/chatmaster/db/models.py @@ -9,7 +9,7 @@ from datetime import datetime, timezone from typing import Any -from sqlalchemy import Boolean, ForeignKey, Integer, JSON, String, Text, UniqueConstraint +from sqlalchemy import JSON, Boolean, ForeignKey, Integer, String, Text, UniqueConstraint from sqlalchemy.orm import Mapped, mapped_column, relationship from chatmaster.db.base import Base diff --git a/backend/chatmaster/documents/jobs.py b/backend/chatmaster/documents/jobs.py index 990c910..05a6055 100644 --- a/backend/chatmaster/documents/jobs.py +++ b/backend/chatmaster/documents/jobs.py @@ -6,11 +6,13 @@ import threading from concurrent.futures import ThreadPoolExecutor -from sqlalchemy import select +from sqlalchemy import select, update +from sqlalchemy.orm import aliased from chatmaster.config import get_settings from chatmaster.db.models import Document, IngestJob from chatmaster.db.session import SessionLocal +from chatmaster.documents.service import document_operation_lock from chatmaster.services.ingest_service import ingest logger = logging.getLogger(__name__) @@ -29,46 +31,80 @@ def _executor() -> ThreadPoolExecutor: return _EXECUTOR +def _claim_job(db, *, job_id: str, document_id: str) -> bool: + """Atomically claim a pending job when no sibling job is already running.""" + other_job = aliased(IngestJob) + claimed = db.execute( + update(IngestJob) + .where( + IngestJob.id == job_id, + IngestJob.status == "pending", + ~select(other_job.id) + .where( + other_job.document_id == document_id, + other_job.status == "running", + other_job.id != job_id, + ) + .exists(), + ) + .values(status="running", error=None) + ).rowcount + return claimed == 1 + + def _run_job(job_id: str) -> None: with SessionLocal() as db: - job = db.get(IngestJob, job_id) - if job is None or job.status not in {"pending", "running"}: + initial_job = db.get(IngestJob, job_id) + if initial_job is None: return - document = db.get(Document, job.document_id) - if document is None: - job.status = "failed" - job.error = "Document not found" + document_id = initial_job.document_id + with document_operation_lock(initial_job.document_id): + if not _claim_job(db, job_id=job_id, document_id=document_id): + db.rollback() + return db.commit() - return - job.status = "running" - document.status = "ingesting" - db.commit() - try: - result = ingest( - document.identity_id, - [__import__("pathlib").Path(document.storage_path)], - target=document.namespace, - workspace_id=document.workspace_id, - db=db, - ) - errors = [item.error for item in result.files if item.error] - if errors or result.total_chunks == 0: - raise RuntimeError("; ".join(errors) or "No indexable content was produced") - job.status = "completed" - job.total_chunks = result.total_chunks - job.error = None - document.status = "indexed" - except Exception as exc: # noqa: BLE001 - logger.exception("Ingest job failed job_id=%s document_id=%s", job.id, document.id) - db.rollback() job = db.get(IngestJob, job_id) - document = db.get(Document, job.document_id) if job is not None else None - if job is not None: + if job is None: + return + document = db.get(Document, job.document_id) + if document is None: job.status = "failed" - job.error = f"{type(exc).__name__}: {exc}" - if document is not None: - document.status = "failed" - db.commit() + job.error = "Document not found" + db.commit() + return + document.status = "ingesting" + db.commit() + try: + result = ingest( + document.identity_id, + [__import__("pathlib").Path(document.storage_path)], + target=document.namespace, + workspace_id=document.workspace_id, + db=db, + ) + db.expire_all() + document = db.get(Document, job.document_id) + job = db.get(IngestJob, job_id) + if document is None or job is None or job.status != "running": + raise RuntimeError("Document was removed while it was being indexed") + errors = [item.error for item in result.files if item.error] + if errors or result.total_chunks == 0: + raise RuntimeError("; ".join(errors) or "No indexable content was produced") + job.status = "completed" + job.total_chunks = result.total_chunks + job.error = None + document.status = "indexed" + except Exception as exc: + logger.exception("Ingest job failed job_id=%s document_id=%s", job_id, document_id) + db.rollback() + job = db.get(IngestJob, job_id) + document = db.get(Document, job.document_id) if job is not None else None + if job is not None: + job.status = "failed" + job.error = f"{type(exc).__name__}: {exc}" + if document is not None: + document.status = "failed" + db.commit() def enqueue_job(job_id: str) -> None: @@ -109,7 +145,7 @@ def run() -> None: identity_id=identity_id, target=target, ) - except Exception: # noqa: BLE001 + except Exception: logger.exception( "Index rebuild failed workspace_id=%s identity_id=%s target=%s", workspace_id, diff --git a/backend/chatmaster/documents/service.py b/backend/chatmaster/documents/service.py index 4fe6703..c30986d 100644 --- a/backend/chatmaster/documents/service.py +++ b/backend/chatmaster/documents/service.py @@ -5,8 +5,10 @@ import hashlib import shutil import tempfile +import threading import uuid from collections.abc import Callable +from contextlib import contextmanager from pathlib import Path from sqlalchemy import select @@ -24,6 +26,22 @@ class UnsupportedDocumentType(ValueError): pass +class DocumentOperationConflict(RuntimeError): + """Raised when a document already has an active operation.""" + + +_DOCUMENT_LOCKS_GUARD = threading.Lock() +_DOCUMENT_LOCKS: dict[str, threading.RLock] = {} + + +@contextmanager +def document_operation_lock(document_id: str): + with _DOCUMENT_LOCKS_GUARD: + lock = _DOCUMENT_LOCKS.setdefault(document_id, threading.RLock()) + with lock: + yield + + def _sha256(data: bytes) -> str: return hashlib.sha256(data).hexdigest() @@ -376,46 +394,70 @@ def retry_ingest_job(db: Session, *, workspace_id: str, job_id: str) -> IngestJo document = db.get(Document, old.document_id) if document is None: raise KeyError(old.document_id) - job = IngestJob( - id=str(uuid.uuid4()), - workspace_id=workspace_id, - document_id=document.id, - status="pending", - error=None, - total_chunks=0, - ) - document.status = "pending" - db.add(job) - db.commit() - db.refresh(job) - return job + with document_operation_lock(document.id): + db.refresh(old) + if old.status != "failed": + raise DocumentOperationConflict("Only failed ingest jobs can be retried") + active = db.scalars( + select(IngestJob).where( + IngestJob.document_id == document.id, + IngestJob.status.in_(("pending", "running")), + ) + ).first() + if active is not None: + raise DocumentOperationConflict("Document already has an active ingest job") + job = IngestJob( + id=str(uuid.uuid4()), + workspace_id=workspace_id, + document_id=document.id, + status="pending", + error=None, + total_chunks=0, + ) + document.status = "pending" + db.add(job) + db.commit() + db.refresh(job) + return job def delete_document(db: Session, *, workspace_id: str, document_id: str) -> None: document = db.get(Document, document_id) if document is None or document.workspace_id != workspace_id: raise KeyError(document_id) - try: - from chatmaster.ai.vectorstore import delete_points + with document_operation_lock(document_id): + db.refresh(document) + active = db.scalars( + select(IngestJob).where( + IngestJob.document_id == document_id, + IngestJob.status.in_(("pending", "running")), + ) + ).first() + if active is not None: + raise DocumentOperationConflict("Document is still being indexed") + try: + from chatmaster.ai.vectorstore import delete_points - chunks = list( - db.scalars(select(DocumentChunk).where(DocumentChunk.document_id == document_id)) - ) - by_collection: dict[str, list[str]] = {} - for chunk in chunks: - version = db.get(IndexVersion, chunk.index_version_id) - if version is not None: - by_collection.setdefault(version.collection_name, []).append(chunk.qdrant_point_id) - for collection, point_ids in by_collection.items(): - delete_points(collection, point_ids) - stored_path = Path(document.storage_path) - stored_path.unlink(missing_ok=True) - db.delete(document) - db.commit() - except Exception: - db.rollback() - document = db.get(Document, document_id) - if document is not None: - document.status = "delete_failed" + chunks = list( + db.scalars(select(DocumentChunk).where(DocumentChunk.document_id == document_id)) + ) + by_collection: dict[str, list[str]] = {} + for chunk in chunks: + version = db.get(IndexVersion, chunk.index_version_id) + if version is not None: + by_collection.setdefault(version.collection_name, []).append( + chunk.qdrant_point_id + ) + for collection, point_ids in by_collection.items(): + delete_points(collection, point_ids) + stored_path = Path(document.storage_path) + stored_path.unlink(missing_ok=True) + db.delete(document) db.commit() - raise + except Exception: + db.rollback() + document = db.get(Document, document_id) + if document is not None: + document.status = "delete_failed" + db.commit() + raise diff --git a/backend/chatmaster/identities/__init__.py b/backend/chatmaster/identities/__init__.py index 6fbbcf9..9903ea1 100644 --- a/backend/chatmaster/identities/__init__.py +++ b/backend/chatmaster/identities/__init__.py @@ -9,8 +9,8 @@ __all__ = [ "IdentityConfig", - "IdentityOut", "IdentityNotFound", + "IdentityOut", "IdentityRegistry", "RetrievalConfig", "get_registry", diff --git a/backend/chatmaster/identities/service.py b/backend/chatmaster/identities/service.py index e1d7bfe..947b46a 100644 --- a/backend/chatmaster/identities/service.py +++ b/backend/chatmaster/identities/service.py @@ -53,11 +53,7 @@ def list_identity_models( stmt = _query(db, workspace_id) if not include_archived: stmt = stmt.where(Identity.is_active.is_(True)) - return list( - db.scalars( - stmt.order_by(Identity.is_system.desc(), Identity.created_at.asc()) - ) - ) + return list(db.scalars(stmt.order_by(Identity.is_system.desc(), Identity.created_at.asc()))) def to_config(identity: Identity) -> IdentityConfig: diff --git a/backend/chatmaster/providers/service.py b/backend/chatmaster/providers/service.py index 8c29692..3438447 100644 --- a/backend/chatmaster/providers/service.py +++ b/backend/chatmaster/providers/service.py @@ -132,6 +132,10 @@ def save_provider_config( payload.embedding.base_url, allow_private_network=getattr(settings, "allow_private_provider_urls", False), ) + validate_provider_url( + payload.embedding.huggingface_endpoint, + allow_private_network=getattr(settings, "allow_private_provider_urls", False), + ) current = get_provider_config(db, workspace_id, settings) row = _find_row(db, workspace_id) if row is None: diff --git a/backend/chatmaster/retrieval/indexes.py b/backend/chatmaster/retrieval/indexes.py index 3e04653..2a12b71 100644 --- a/backend/chatmaster/retrieval/indexes.py +++ b/backend/chatmaster/retrieval/indexes.py @@ -7,7 +7,7 @@ import uuid from pathlib import Path -from sqlalchemy import select +from sqlalchemy import and_, or_, select from sqlalchemy.orm import Session from chatmaster.ai.models import build_embeddings @@ -76,11 +76,25 @@ class IndexRebuildRequired(RuntimeError): """The configured embedding model no longer matches an active index.""" -def assert_indexes_fresh(db: Session, *, workspace_id: str) -> None: +def assert_indexes_fresh( + db: Session, + *, + workspace_id: str, + identity_id: str | None = None, + include_private: bool = False, +) -> None: + scopes = [("common", None)] + if include_private and identity_id: + scopes.append(("private", identity_id)) + scope_filters = [ + and_(IndexVersion.namespace == namespace, IndexVersion.identity_id == scope_identity) + for namespace, scope_identity in scopes + ] stale = db.scalars( select(IndexVersion).where( IndexVersion.workspace_id == workspace_id, IndexVersion.status == "stale", + or_(*scope_filters), ) ).first() if stale is not None: diff --git a/backend/chatmaster/retrieval/retriever.py b/backend/chatmaster/retrieval/retriever.py index eac6eed..4c940bf 100644 --- a/backend/chatmaster/retrieval/retriever.py +++ b/backend/chatmaster/retrieval/retriever.py @@ -12,8 +12,8 @@ from chatmaster.ai.vectorstore import get_store from chatmaster.config import get_settings from chatmaster.identities.schema import IdentityConfig -from chatmaster.retrieval.schemas import RetrievedChunk, SearchHit from chatmaster.retrieval.indexes import active_collection, assert_indexes_fresh +from chatmaster.retrieval.schemas import RetrievedChunk, SearchHit _RRF_K = 60 @@ -114,7 +114,12 @@ async def retrieve( ) -> list[RetrievedChunk]: settings = get_settings() cfg = identity.retrieval - assert_indexes_fresh(db, workspace_id=workspace_id) + assert_indexes_fresh( + db, + workspace_id=workspace_id, + identity_id=identity.id, + include_private=identity.uses_private_knowledge, + ) common_collection = active_collection( db, diff --git a/backend/chatmaster/routers/chat.py b/backend/chatmaster/routers/chat.py index eb2a632..79e83ff 100644 --- a/backend/chatmaster/routers/chat.py +++ b/backend/chatmaster/routers/chat.py @@ -3,15 +3,15 @@ from __future__ import annotations from fastapi import APIRouter, Depends, HTTPException -from sse_starlette.sse import EventSourceResponse from sqlalchemy.orm import Session +from sse_starlette.sse import EventSourceResponse +from chatmaster.chat.service import request_cancel, stream_chat from chatmaster.conversations.service import ConversationNotFound, get_conversation from chatmaster.core.auth import get_current_user_id, get_current_workspace_id from chatmaster.db.session import get_db from chatmaster.identities.service import IdentityNotFound, get_identity_model from chatmaster.schemas.api import ChatRequest -from chatmaster.chat.service import request_cancel, stream_chat router = APIRouter(prefix="/api", tags=["chat"]) diff --git a/backend/chatmaster/routers/conversations.py b/backend/chatmaster/routers/conversations.py index b29d272..0e30297 100644 --- a/backend/chatmaster/routers/conversations.py +++ b/backend/chatmaster/routers/conversations.py @@ -16,7 +16,12 @@ from chatmaster.core.auth import get_current_workspace_id from chatmaster.db.session import get_db from chatmaster.identities.service import IdentityNotFound -from chatmaster.schemas.api import ConversationCreate, ConversationOut, ConversationUpdate, MessageOut +from chatmaster.schemas.api import ( + ConversationCreate, + ConversationOut, + ConversationUpdate, + MessageOut, +) router = APIRouter(tags=["conversations"]) diff --git a/backend/chatmaster/routers/documents.py b/backend/chatmaster/routers/documents.py index 0855ea6..7f7c096 100644 --- a/backend/chatmaster/routers/documents.py +++ b/backend/chatmaster/routers/documents.py @@ -11,6 +11,7 @@ from chatmaster.db.models import Document, IngestJob from chatmaster.db.session import get_db from chatmaster.documents.service import ( + DocumentOperationConflict, UnsupportedDocumentType, delete_document, list_documents, @@ -156,6 +157,8 @@ async def retry_job( job = retry_ingest_job(db, workspace_id=workspace_id, job_id=job_id) except KeyError: raise HTTPException(status_code=404, detail="Ingest job not found") from None + except DocumentOperationConflict as exc: + raise HTTPException(status_code=409, detail=str(exc)) from None from chatmaster.documents.jobs import enqueue_job enqueue_job(job.id) @@ -172,7 +175,9 @@ async def remove_document( delete_document(db, workspace_id=workspace_id, document_id=document_id) except KeyError: raise HTTPException(status_code=404, detail="Document not found") from None - except Exception as exc: # noqa: BLE001 + except DocumentOperationConflict as exc: + raise HTTPException(status_code=409, detail=str(exc)) from None + except Exception as exc: raise HTTPException(status_code=503, detail="Document cleanup failed") from exc diff --git a/backend/chatmaster/routers/identities.py b/backend/chatmaster/routers/identities.py index 7050d14..409c28b 100644 --- a/backend/chatmaster/routers/identities.py +++ b/backend/chatmaster/routers/identities.py @@ -136,8 +136,6 @@ async def clone_identity( db: Session = Depends(get_db), ): try: - return to_detail( - duplicate_identity(db, workspace_id=workspace_id, identity_id=identity_id) - ) + return to_detail(duplicate_identity(db, workspace_id=workspace_id, identity_id=identity_id)) except IdentityNotFound: raise HTTPException(status_code=404, detail="Identity not found") from None diff --git a/backend/chatmaster/routers/providers.py b/backend/chatmaster/routers/providers.py index de139c8..6a1a18e 100644 --- a/backend/chatmaster/routers/providers.py +++ b/backend/chatmaster/routers/providers.py @@ -8,8 +8,8 @@ from __future__ import annotations -import logging import asyncio +import logging from fastapi import APIRouter, Depends from fastapi.concurrency import run_in_threadpool @@ -104,7 +104,7 @@ def _test_chat() -> str: model = build_chat_model(dummy) model.invoke([HumanMessage(content="ping")]) return "ok" - except Exception: # noqa: BLE001 + except Exception: logger.exception("Provider chat test failed") return "failed (see server log)" @@ -113,7 +113,7 @@ def _test_embedding() -> str: emb = build_embeddings() vec = emb.embed_query("dimension probe") return f"ok (dim={len(vec)})" - except Exception: # noqa: BLE001 + except Exception: logger.exception("Provider embedding test failed") return "failed (see server log)" diff --git a/backend/chatmaster/schemas/api.py b/backend/chatmaster/schemas/api.py index 14735dd..72ecbab 100644 --- a/backend/chatmaster/schemas/api.py +++ b/backend/chatmaster/schemas/api.py @@ -3,7 +3,6 @@ from __future__ import annotations from datetime import datetime - from typing import Literal from uuid import UUID diff --git a/backend/chatmaster/services/ingest_service.py b/backend/chatmaster/services/ingest_service.py index 00d2595..902905e 100644 --- a/backend/chatmaster/services/ingest_service.py +++ b/backend/chatmaster/services/ingest_service.py @@ -2,21 +2,24 @@ from __future__ import annotations +import logging import uuid from pathlib import Path -from sqlalchemy import delete +from sqlalchemy import delete, select from sqlalchemy.orm import Session from chatmaster.ai.chunkers import split_documents from chatmaster.ai.loaders import UnsupportedFileType, load_file from chatmaster.ai.models import build_embeddings -from chatmaster.ai.vectorstore import get_store +from chatmaster.ai.vectorstore import delete_points, get_store from chatmaster.config import get_settings -from chatmaster.db.models import Document, DocumentChunk +from chatmaster.db.models import Document, DocumentChunk, IndexVersion from chatmaster.identities.service import get_identity_model, to_config from chatmaster.schemas.api import IngestFileResult, IngestResult +logger = logging.getLogger(__name__) + class IngestFailed(RuntimeError): """Raised when one or more files could not be indexed.""" @@ -65,7 +68,12 @@ def ingest( ) from chatmaster.retrieval.indexes import active_collection, assert_indexes_fresh - assert_indexes_fresh(db, workspace_id=workspace_id) + assert_indexes_fresh( + db, + workspace_id=workspace_id, + identity_id=identity_id, + include_private=target == "private", + ) collection = active_collection( db, workspace_id=workspace_id, @@ -79,19 +87,28 @@ def ingest( if db is not None and workspace_id is not None: from chatmaster.retrieval.indexes import ensure_active_version - version = ensure_active_version( - db, - workspace_id=workspace_id, - logical_name=logical_collection, - identity_id=None if target == "common" else identity_id, - embeddings=embeddings, - ) + if collection_name is not None: + version = db.scalars( + select(IndexVersion).where( + IndexVersion.workspace_id == workspace_id, + IndexVersion.collection_name == collection_name, + ) + ).first() + if version is None: + version = ensure_active_version( + db, + workspace_id=workspace_id, + logical_name=logical_collection, + identity_id=None if target == "common" else identity_id, + embeddings=embeddings, + ) store = get_store(collection, embeddings) file_results: list[IngestFileResult] = [] total_chunks = 0 for path in files: + written_point_ids: list[str] = [] try: docs = load_file(path, identity_id=identity_id or "common") chunks = split_documents(docs) @@ -138,8 +155,23 @@ def ingest( metadata_json=dict(chunk.metadata), ) ) - store.add_documents(chunks, ids=point_ids) - db.commit() + written_point_ids = point_ids + try: + store.add_documents(chunks, ids=point_ids) + db.commit() + except Exception: + db.rollback() + if written_point_ids: + try: + delete_points(collection, written_point_ids) + except Exception as cleanup_exc: # noqa: BLE001 + logger.warning( + "Failed to compensate Qdrant points collection=%s count=%d: %s", + collection, + len(written_point_ids), + cleanup_exc, + ) + raise else: store.add_documents(chunks) file_results.append(IngestFileResult(file=path.name, chunks=len(chunks))) diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 2ec1709..63eeeb8 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -67,3 +67,8 @@ filterwarnings = [ [tool.ruff] line-length = 100 target-version = "py310" + +[tool.ruff.lint] +# FastAPI, Typer, and dependency injection declarations intentionally call +# framework factories in defaults; B008 is not actionable for these APIs. +ignore = ["B008"] diff --git a/backend/scripts/seed_sample_docs.py b/backend/scripts/seed_sample_docs.py index 6fe180f..dd354f3 100644 --- a/backend/scripts/seed_sample_docs.py +++ b/backend/scripts/seed_sample_docs.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """Import all bundled sample documents into SQLite + Qdrant. Run from the backend/ directory inside the chatmaster conda environment: diff --git a/backend/tests/chat/test_graph.py b/backend/tests/chat/test_graph.py index e30879e..185512f 100644 --- a/backend/tests/chat/test_graph.py +++ b/backend/tests/chat/test_graph.py @@ -228,8 +228,8 @@ async def stream_answer(_state) -> AsyncIterator[str]: @pytest.mark.asyncio async def test_durable_chat_turn_is_idempotent_and_returns_persisted_id(monkeypatch) -> None: - from chatmaster.chat.graph import ChatRuntime from chatmaster.chat import service + from chatmaster.chat.graph import ChatRuntime engine = create_engine("sqlite:///:memory:", future=True) test_session = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True) diff --git a/backend/tests/db/test_seed.py b/backend/tests/db/test_seed.py index 1abe92d..f982e29 100644 --- a/backend/tests/db/test_seed.py +++ b/backend/tests/db/test_seed.py @@ -23,7 +23,9 @@ def _session() -> Session: return SessionLocal() -def test_seed_local_data_inserts_workspace_user_yaml_identities_and_fallback(tmp_path: Path) -> None: +def test_seed_local_data_inserts_workspace_user_yaml_identities_and_fallback( + tmp_path: Path, +) -> None: from chatmaster.db.seed import seed_local_data identities_yaml = tmp_path / "identities.yaml" diff --git a/backend/tests/documents/test_router.py b/backend/tests/documents/test_router.py index 6a66820..ac81a07 100644 --- a/backend/tests/documents/test_router.py +++ b/backend/tests/documents/test_router.py @@ -12,7 +12,7 @@ from chatmaster.routers.documents import router -def _client(): +def _client(*, job_status: str = "completed"): engine = create_engine( "sqlite:///:memory:", connect_args={"check_same_thread": False}, @@ -41,7 +41,7 @@ def _client(): id="job-1", workspace_id="local", document_id="document-1", - status="completed", + status=job_status, total_chunks=2, ) ) @@ -78,3 +78,19 @@ def test_list_ingest_jobs_endpoint_returns_persisted_jobs() -> None: payload = response.json() assert payload[0]["id"] == "job-1" assert payload[0]["status"] == "completed" + + +def test_retry_active_ingest_job_returns_conflict() -> None: + client = _client(job_status="pending") + + response = client.post("/api/ingest-jobs/job-1/retry") + + assert response.status_code == 409 + + +def test_delete_document_with_active_job_returns_conflict() -> None: + client = _client(job_status="running") + + response = client.delete("/api/documents/document-1") + + assert response.status_code == 409 diff --git a/backend/tests/documents/test_service.py b/backend/tests/documents/test_service.py index 58d3d95..d28d8f8 100644 --- a/backend/tests/documents/test_service.py +++ b/backend/tests/documents/test_service.py @@ -7,7 +7,14 @@ from sqlalchemy.orm import Session, sessionmaker from chatmaster.db.base import Base -from chatmaster.db.models import Document, Identity, IngestJob, Workspace +from chatmaster.db.models import ( + Document, + DocumentChunk, + Identity, + IndexVersion, + IngestJob, + Workspace, +) def _session() -> Session: @@ -197,3 +204,149 @@ async def read(self, _size: int) -> bytes: assert documents[0].scope_key == "common" assert first.document_id == second.document_id assert second.duplicate is True + + +def test_retry_only_allows_failed_jobs() -> None: + from chatmaster.documents.service import DocumentOperationConflict, retry_ingest_job + + with _session() as db: + _seed_identity(db) + document = Document( + id="document-1", + workspace_id="local", + identity_id="legal_expert", + namespace="private", + scope_key="legal_expert", + filename="note.txt", + storage_path="note.txt", + sha256="b" * 64, + status="pending", + ) + db.add_all( + [ + document, + IngestJob( + id="job-1", + workspace_id="local", + document_id=document.id, + status="pending", + total_chunks=0, + ), + ] + ) + db.commit() + + with pytest.raises(DocumentOperationConflict, match="Only failed"): + retry_ingest_job(db, workspace_id="local", job_id="job-1") + + +def test_worker_claim_allows_only_one_running_job_per_document() -> None: + from chatmaster.documents.jobs import _claim_job + + with _session() as db: + _seed_identity(db) + document = Document( + id="document-1", + workspace_id="local", + identity_id="legal_expert", + namespace="private", + scope_key="legal_expert", + filename="note.txt", + storage_path="note.txt", + sha256="d" * 64, + status="pending", + ) + db.add_all( + [ + document, + IngestJob( + id="job-1", + workspace_id="local", + document_id=document.id, + status="pending", + total_chunks=0, + ), + IngestJob( + id="job-2", + workspace_id="local", + document_id=document.id, + status="pending", + total_chunks=0, + ), + ] + ) + db.commit() + + assert _claim_job(db, job_id="job-1", document_id=document.id) + db.commit() + assert not _claim_job(db, job_id="job-2", document_id=document.id) + + +def test_delete_document_cleans_vectors_and_persisted_records(tmp_path: Path, monkeypatch) -> None: + from chatmaster.documents.service import delete_document + + stored_path = tmp_path / "note.txt" + stored_path.write_text("hello", encoding="utf-8") + deleted: list[tuple[str, list[str]]] = [] + monkeypatch.setattr( + "chatmaster.ai.vectorstore.delete_points", + lambda collection, ids: deleted.append((collection, ids)), + ) + + with _session() as db: + _seed_identity(db) + document = Document( + id="document-1", + workspace_id="local", + identity_id="legal_expert", + namespace="private", + scope_key="legal_expert", + filename="note.txt", + storage_path=str(stored_path), + sha256="c" * 64, + status="indexed", + ) + version = IndexVersion( + id="version-1", + workspace_id="local", + namespace="private", + identity_id="legal_expert", + collection_name="private-v1", + embedding_provider="test", + embedding_model="test", + embedding_dim=3, + status="active", + ) + db.add_all( + [ + document, + version, + DocumentChunk( + id="chunk-1", + workspace_id="local", + document_id=document.id, + index_version_id=version.id, + qdrant_point_id="point-1", + chunk_index=0, + text="hello", + metadata_json={}, + ), + IngestJob( + id="job-1", + workspace_id="local", + document_id=document.id, + status="completed", + total_chunks=1, + ), + ] + ) + db.commit() + + delete_document(db, workspace_id="local", document_id=document.id) + + assert db.get(Document, document.id) is None + assert db.query(DocumentChunk).count() == 0 + assert db.query(IngestJob).count() == 0 + + assert deleted == [("private-v1", ["point-1"])] + assert not stored_path.exists() diff --git a/backend/tests/providers/test_service.py b/backend/tests/providers/test_service.py index ad63b18..506d2a3 100644 --- a/backend/tests/providers/test_service.py +++ b/backend/tests/providers/test_service.py @@ -1,5 +1,6 @@ from __future__ import annotations +import pytest from sqlalchemy import create_engine, select from sqlalchemy.orm import Session, sessionmaker @@ -24,6 +25,14 @@ class DummySettings: allow_private_provider_urls = True +class PrivateBlockedSettings(DummySettings): + allow_private_provider_urls = False + + +def _public_addresses(*items: str): + return [(2, 1, 6, "", (item, 0)) for item in items] + + def _session() -> Session: engine = create_engine("sqlite:///:memory:", future=True) Base.metadata.create_all(bind=engine) @@ -45,9 +54,13 @@ def test_get_provider_config_seeds_from_settings_when_missing() -> None: assert cfg.embedding.provider == "huggingface" -def test_save_provider_config_replaces_new_keys() -> None: +def test_save_provider_config_replaces_new_keys(monkeypatch) -> None: from chatmaster.providers.service import get_provider_config, save_provider_config + monkeypatch.setattr( + "chatmaster.providers.security.socket.getaddrinfo", + lambda *_: _public_addresses("93.184.216.34"), + ) with _session() as db: db.add(Workspace(id="local", name="Local Workspace")) db.commit() @@ -78,9 +91,13 @@ def test_save_provider_config_replaces_new_keys() -> None: assert row.chat_api_key_encrypted.startswith("fernet:") -def test_save_provider_config_keeps_previous_keys_when_payload_is_masked() -> None: +def test_save_provider_config_keeps_previous_keys_when_payload_is_masked(monkeypatch) -> None: from chatmaster.providers.service import get_provider_config, save_provider_config + monkeypatch.setattr( + "chatmaster.providers.security.socket.getaddrinfo", + lambda *_: _public_addresses("93.184.216.34"), + ) with _session() as db: db.add(Workspace(id="local", name="Local Workspace")) db.commit() @@ -123,3 +140,26 @@ def test_save_provider_config_keeps_previous_keys_when_payload_is_masked() -> No assert saved.chat.api_key == "sk-secret-chat" assert saved.chat.base_url == "https://api.changed.com/v1" assert saved.embedding.api_key == "sk-secret-embedding" + + +def test_save_provider_config_validates_huggingface_endpoint(monkeypatch) -> None: + from chatmaster.providers.security import UnsafeProviderUrl + from chatmaster.providers.service import save_provider_config + + monkeypatch.setattr( + "chatmaster.providers.security.socket.getaddrinfo", + lambda *_: _public_addresses("192.168.1.20"), + ) + payload = ProvidersConfig( + chat=ChatProviderConfig(model="gpt-4o-mini"), + embedding=EmbeddingProviderConfig( + provider="huggingface", + model="BAAI/bge-small-zh-v1.5", + huggingface_endpoint="http://provider.test", + ), + ) + with _session() as db: + db.add(Workspace(id="local", name="Local Workspace")) + db.commit() + with pytest.raises(UnsafeProviderUrl): + save_provider_config(db, "local", payload, PrivateBlockedSettings()) diff --git a/backend/tests/retrieval/test_indexes.py b/backend/tests/retrieval/test_indexes.py index bb69c9b..619d0cf 100644 --- a/backend/tests/retrieval/test_indexes.py +++ b/backend/tests/retrieval/test_indexes.py @@ -2,6 +2,7 @@ from pathlib import Path +import pytest from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker @@ -20,6 +21,85 @@ class _Settings: common_collection = "common" +def test_assert_indexes_fresh_checks_only_required_scopes() -> None: + from chatmaster.retrieval.indexes import IndexRebuildRequired, assert_indexes_fresh + + SessionLocal = _session() + with SessionLocal() as db: + db.add(Workspace(id="local", name="Local")) + db.add_all( + [ + IndexVersion( + id="common", + workspace_id="local", + namespace="common", + identity_id=None, + collection_name="common", + embedding_provider="test", + embedding_model="test", + embedding_dim=3, + status="active", + ), + IndexVersion( + id="private-a", + workspace_id="local", + namespace="private", + identity_id="persona-a", + collection_name="private-a", + embedding_provider="test", + embedding_model="test", + embedding_dim=3, + status="stale", + ), + ] + ) + db.commit() + + assert_indexes_fresh( + db, + workspace_id="local", + identity_id="persona-b", + include_private=True, + ) + with pytest.raises(IndexRebuildRequired): + assert_indexes_fresh( + db, + workspace_id="local", + identity_id="persona-a", + include_private=True, + ) + + +def test_assert_indexes_fresh_always_checks_common_scope() -> None: + from chatmaster.retrieval.indexes import IndexRebuildRequired, assert_indexes_fresh + + SessionLocal = _session() + with SessionLocal() as db: + db.add(Workspace(id="local", name="Local")) + db.add( + IndexVersion( + id="common", + workspace_id="local", + namespace="common", + identity_id=None, + collection_name="common", + embedding_provider="test", + embedding_model="test", + embedding_dim=3, + status="stale", + ) + ) + db.commit() + + with pytest.raises(IndexRebuildRequired): + assert_indexes_fresh( + db, + workspace_id="local", + identity_id="persona-b", + include_private=False, + ) + + def test_active_collection_uses_active_version(monkeypatch) -> None: from chatmaster.retrieval import indexes @@ -102,7 +182,7 @@ def test_rebuild_activates_new_version_only_after_success(tmp_path: Path, monkey embedding_provider="old", embedding_model="old", embedding_dim=3, - status="active", + status="stale", ) ) db.commit() diff --git a/frontend/src/features/knowledge/KnowledgePage.test.tsx b/frontend/src/features/knowledge/KnowledgePage.test.tsx index 8694b2a..25f109a 100644 --- a/frontend/src/features/knowledge/KnowledgePage.test.tsx +++ b/frontend/src/features/knowledge/KnowledgePage.test.tsx @@ -58,4 +58,42 @@ describe("KnowledgePage", () => { fireEvent.click(screen.getByRole("button", { name: "确认删除" })); await waitFor(() => expect(mocks.deleteDocument).toHaveBeenCalledWith("doc-1")); }); + + it("rebuilds a private index for the identity recorded on that index", async () => { + mocks.getDocuments.mockResolvedValue([]); + mocks.getIngestJobs.mockResolvedValue([]); + mocks.getIndexes.mockResolvedValue([ + { + id: "index-private-2", + namespace: "private", + identity_id: "persona-2", + collection_name: "persona-2-v1", + logical_name: "persona-2", + embedding_provider: "huggingface", + embedding_model: "BAAI/bge-small-zh-v1.5", + embedding_dim: 384, + config_fingerprint: "test", + status: "stale", + created_at: "2026-07-29T08:00:00Z", + updated_at: "2026-07-29T08:00:00Z", + }, + ]); + mocks.rebuildIndex.mockResolvedValue(undefined); + + const client = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + render( + + undefined} /> + + ); + + const rebuildButton = await screen.findByRole("button", { name: "重建" }); + fireEvent.click(rebuildButton); + + await waitFor(() => + expect(mocks.rebuildIndex).toHaveBeenCalledWith("private", "persona-2") + ); + }); }); diff --git a/frontend/src/features/knowledge/KnowledgePage.tsx b/frontend/src/features/knowledge/KnowledgePage.tsx index 9a64f20..0387c54 100644 --- a/frontend/src/features/knowledge/KnowledgePage.tsx +++ b/frontend/src/features/knowledge/KnowledgePage.tsx @@ -19,6 +19,7 @@ import { retryIngestJob, } from "../../api/client"; import { DocumentUpload } from "../../components/DocumentUpload"; +import type { IndexVersionOut } from "../../types/api"; interface Props { identityId: string | null; @@ -34,6 +35,7 @@ export function KnowledgePage({ identityId, onBack }: Props) { const [deletingId, setDeletingId] = useState(null); const [deleteBusy, setDeleteBusy] = useState(false); const [deleteError, setDeleteError] = useState(null); + const [rebuildingId, setRebuildingId] = useState(null); const documents = useQuery({ queryKey: ["documents", identityId, scope, statusFilter], queryFn: () => @@ -76,9 +78,14 @@ export function KnowledgePage({ identityId, onBack }: Props) { } }; - const rebuild = async (target: "private" | "common") => { - await rebuildIndex(target, target === "private" ? identityId : null); - await indexes.refetch(); + const rebuild = async (index: IndexVersionOut) => { + setRebuildingId(index.id); + try { + await rebuildIndex(index.namespace, index.namespace === "private" ? index.identity_id : null); + await indexes.refetch(); + } finally { + setRebuildingId(null); + } }; return ( @@ -307,10 +314,13 @@ export function KnowledgePage({ identityId, onBack }: Props) { {index.status === "stale" && ( )}