102 lines
4.9 KiB
Python
102 lines
4.9 KiB
Python
"""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
|