Files
amb_rag/backend/app/services/kb_public_service.py
T

188 lines
6.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""公共知识库服务(Phase 10/11/12 统一数据来源)。
HTML/MD/TXT/JSON/搜索全部通过此 Service 获取数据,不各自写查询逻辑。
支持目录树结构和分类过滤。
"""
from datetime import datetime
from app.core.errors import LinkExpiredError, NotFoundError
from app.core.security import hash_token
from app.models.document import Document
from app.models.document_category import DocumentCategory
from app.models.knowledge_base import KnowledgeBase
from app.repositories.doc_repo import DocumentRepository
from app.repositories.kb_repo import KnowledgeBaseRepository
from sqlalchemy import func, select
from sqlalchemy.orm import Session
_TIME_FMT = "%Y-%m-%d %H:%M:%S"
class KbPublicService:
def __init__(self, session: Session) -> None:
self._session = session
self._kb_repo = KnowledgeBaseRepository(session)
self._doc_repo = DocumentRepository(session)
def get_kb_by_token(self, token: str) -> KnowledgeBase:
"""通过 token 获取知识库。
不存在/禁用/删除 → 404;存在但已过期 → 410LinkExpiredError)。
"""
token_hash = hash_token(token)
kb = self._kb_repo.get_by_token_hash(token_hash)
if kb is None or not kb.enabled or kb.status == "DELETED":
raise NotFoundError("知识库不存在。")
if kb.token_expires_at:
try:
expires = datetime.strptime(kb.token_expires_at, _TIME_FMT)
except ValueError:
expires = None
if expires is not None and datetime.now() > expires:
raise LinkExpiredError()
return kb
@staticmethod
def is_expired(kb: KnowledgeBase) -> bool:
"""判断知识库链接是否已过期(管理端展示用)。"""
if not kb.token_expires_at:
return False
try:
expires = datetime.strptime(kb.token_expires_at, _TIME_FMT)
except ValueError:
return False
return datetime.now() > expires
def get_category_tree(self, kb: KnowledgeBase) -> list[dict]:
"""获取目录树(含文档数量),排除已删除目录。"""
stmt = (
select(DocumentCategory)
.where(
DocumentCategory.knowledge_base_id == kb.id,
DocumentCategory.deleted_at.is_(None),
)
.order_by(DocumentCategory.sort_order, DocumentCategory.name)
)
all_cats = list(self._session.scalars(stmt).all())
# 获取每个分类的文档数量
count_stmt = (
select(Document.category_id, func.count())
.where(
Document.knowledge_base_id == kb.id,
Document.status != "DELETED",
)
.group_by(Document.category_id)
)
doc_count_map = dict(self._session.execute(count_stmt).all())
# 构建树
cat_map = {}
for cat in all_cats:
cat_map[cat.id] = {
"id": cat.id,
"name": cat.name,
"path": cat.path or "/",
"is_folder": getattr(cat, 'is_folder', True),
"sort_order": cat.sort_order,
"doc_count": doc_count_map.get(cat.id, 0),
"children": [],
}
roots = []
for cat in all_cats:
node = cat_map[cat.id]
if cat.parent_id and cat.parent_id in cat_map:
cat_map[cat.parent_id]["children"].append(node)
else:
roots.append(node)
return roots
def list_documents(
self, kb: KnowledgeBase, *, category_id: str | None = None, page: int = 1, page_size: int = 50
) -> tuple[list[Document], int]:
"""获取知识库的文档列表(排除 DELETED)。支持按分类过滤。"""
return self._doc_repo.list_by_knowledge_base(
kb.id, category_id=category_id, page=page, page_size=page_size
)
def get_document_by_token(self, kb: KnowledgeBase, doc_token: str) -> Document:
"""通过 token 获取单个文档。"""
doc_token_hash = hash_token(doc_token)
doc = self._doc_repo.get_by_doc_token_hash(doc_token_hash)
if doc is None or doc.knowledge_base_id != kb.id or doc.status != "READY":
raise NotFoundError("文档不存在。")
return doc
def get_document_content(self, doc: Document) -> str:
"""获取文档内容(优先直接内容,其次 Markdown 文件)。"""
# 直接文本内容
if doc.content:
return doc.content
# 从文件读取 Markdown
if doc.markdown_path:
from app.storage.local_storage import get_storage
storage = get_storage()
content = storage.read(doc.markdown_path)
return content.decode("utf-8")
return ""
def search_documents(
self, kb: KnowledgeBase, query: str, *, page: int = 1, page_size: int = 20
) -> tuple[list[dict], int]:
"""关键词搜索文档(Phase 12SQLite LIKE / FTS5)。"""
conditions = [
Document.knowledge_base_id == kb.id,
Document.status == "READY",
]
like_pattern = f"%{query}%"
search_condition = (
Document.title.like(like_pattern)
| Document.description.like(like_pattern)
| Document.keywords.like(like_pattern)
| Document.content_summary.like(like_pattern)
| Document.content.like(like_pattern)
)
conditions.append(search_condition)
count_stmt = select(func.count()).select_from(Document).where(*conditions)
total = self._session.scalar(count_stmt) or 0
stmt = (
select(Document)
.where(*conditions)
.order_by(Document.created_at.desc())
.offset((page - 1) * page_size)
.limit(page_size)
)
docs = list(self._session.scalars(stmt).all())
results = []
for doc in docs:
# 获取分类路径
category_path = ""
if doc.category_id:
cat = self._session.get(DocumentCategory, doc.category_id)
if cat:
category_path = cat.path or ""
results.append({
"id": doc.id,
"title": doc.title or doc.original_filename,
"description": doc.description,
"keywords": doc.keywords,
"file_type": doc.file_ext,
"updated_at": doc.updated_at,
"url_hint": doc.doc_token_hint,
"category_path": category_path,
})
return results, total