feat(server): add workbench API, DSH planner worker and RAG pipeline

server/app/training_scope.py

@@ -0,0 +1,101 @@
+"""Training-map library/category scope, shared by retrieval, full reading and estimates."""
+from sqlmodel import select
+from fastapi import HTTPException
+from .models import KnowledgeBase, KnowledgeCategory, KnowledgeFile
+from .rag import store
+from .rag.quarantine import file_exclusions
+
+
+def normalize_scope(session, kb_ids, mapping):
+    if not isinstance(mapping, dict):
+        raise HTTPException(422, '知识库分类选择格式无效')
+    selected = set(kb_ids)
+    result = {}
+    for key, values in mapping.items():
+        try:
+            kid = int(key)
+        except (ValueError, TypeError):
+            raise HTTPException(422, '知识库编号无效') from None
+        if str(kid) != str(key) or kid not in selected:
+            raise HTTPException(422, '分类选择包含未选中的知识库')
+        if not isinstance(values, list) or not values:
+            raise HTTPException(422, '请至少选择一个分类，或切换为全部分类')
+        categories = []
+        for value in values:
+            if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+                raise HTTPException(422, '分类编号无效')
+            cat = session.get(KnowledgeCategory, value) if value else None
+            if value and (not cat or cat.kb_id != kid):
+                raise HTTPException(422, '所选分类已删除或不属于对应知识库，请重新选择')
+            categories.append(value)
+        result[str(kid)] = sorted(set(categories))
+    return {key: result[key] for key in sorted(result, key=int)}
+
+
+def file_scope(session, kb_id, mapping):
+    """None means whole library; an empty set means no matching files (never whole library)."""
+    excluded = file_exclusions(session, kb_id)
+    values = (mapping or {}).get(str(kb_id))
+    if values is None:
+        if not excluded:
+            return None
+        registered = session.exec(select(KnowledgeFile.file_name).where(KnowledgeFile.kb_id == kb_id)).all()
+        return (set(store.count_by_file(kb_id)) | set(registered)) - excluded
+    normalize_scope(session, [kb_id], {str(kb_id): values})
+    cats = session.exec(select(KnowledgeCategory).where(KnowledgeCategory.kb_id == kb_id)).all()
+    included = set(values)
+    changed = True
+    while changed:
+        changed = False
+        for cat in cats:
+            if cat.parent_id in included and cat.id not in included:
+                included.add(cat.id); changed = True
+    files = session.exec(select(KnowledgeFile).where(KnowledgeFile.kb_id == kb_id)).all()
+    return {file.file_name for file in files
+            if (file.category_id if file.category_id is not None else file.ai_category_id or 0) in included} - excluded
+
+
+def category_options(session, kb_id):
+    counts = store.count_by_file(kb_id)
+    files = session.exec(select(KnowledgeFile).where(KnowledgeFile.kb_id == kb_id)).all()
+    cats = session.exec(select(KnowledgeCategory).where(KnowledgeCategory.kb_id == kb_id)
+                        .order_by(KnowledgeCategory.sort_order, KnowledgeCategory.id)).all()
+    # Read the tree and files once; repeated per-category queries make large
+    # libraries noticeably slow when their category drawer first opens.
+    parents = {cat.id: cat.parent_id for cat in cats}
+    names_by_category = {cat.id: set() for cat in cats}
+    unclassified = set()
+    for file in files:
+        current = file.category_id if file.category_id is not None else file.ai_category_id or 0
+        if not current:
+            unclassified.add(file.file_name)
+        visited = set()
+        while current in parents and current not in visited:
+            visited.add(current)
+            names_by_category[current].add(file.file_name)
+            current = parents[current]
+    categories = []
+    for cat in cats:
+        names = names_by_category[cat.id]
+        categories.append({'id': cat.id, 'name': cat.name, 'parent_id': cat.parent_id,
+                           'file_count': len(names), 'chunk_count': sum(counts.get(name, 0) for name in names)})
+    registered = {file.file_name for file in files}
+    return {'categories': categories, 'unclassified': {'file_count': len(unclassified),
+            'chunk_count': sum(counts.get(name, 0) for name in unclassified)},
+            'unregistered_index_files': len(set(counts) - registered)}
+
+
+def scope_counts(session, kb_ids, mapping):
+    normalized = normalize_scope(session, kb_ids, mapping or {})
+    result = []
+    for kid in dict.fromkeys(kb_ids):
+        kb = session.get(KnowledgeBase, kid)
+        allowed = file_scope(session, kid, normalized)
+        if allowed is None:
+            result.append({'kb_id': kid, 'file_count': max(0, kb.file_count or 0),
+                           'chunk_count': max(0, kb.chunk_count or 0), 'chunk_target': kb.chunk_target})
+        else:
+            counts = store.count_by_file(kid)
+            result.append({'kb_id': kid, 'file_count': len(allowed),
+                           'chunk_count': sum(counts.get(name, 0) for name in allowed), 'chunk_target': kb.chunk_target})
+    return result