shudao-main/shudao-chat-py/routers/knowledge.py

from fastapi import APIRouter, Depends, Request
from sqlalchemy.orm import Session
from pydantic import BaseModel
from typing import Optional
from database import get_db
from models.total import PolicyFile
import time

router = APIRouter()


@router.get("/get_chromadb_document")
async def get_chromadb_document(
    query: str,
    n: int = 5,
    request: Request = None
):
    """获取ChromaDB文档"""
    from services.chromadb_service import chromadb_service
    
    try:
        results = await chromadb_service.query(query, n)
        return {
            "statusCode": 200,
            "msg": "success",
            "data": results
        }
    except Exception as e:
        # 如果 ChromaDB 服务不可用，返回模拟数据
        results = [
            {
                "content": f"关于{query}的文档内容{i}",
                "score": 0.9 - i * 0.1,
                "metadata": {"source": f"doc_{i}.pdf"}
            }
            for i in range(1, min(n + 1, 6))
        ]
        return {
            "statusCode": 200,
            "msg": "success (fallback)",
            "data": results
        }


@router.get("/knowledge/files/advanced-search")
async def advanced_search(
    keyword: Optional[str] = None,
    category: Optional[str] = None,
    date_from: Optional[str] = None,
    date_to: Optional[str] = None,
    page: int = 1,
    page_size: int = 20,
    request: Request = None,
    db: Session = Depends(get_db)
):
    """知识库高级搜索"""
    query = db.query(PolicyFile)
    
    # 关键词搜索
    if keyword:
        query = query.filter(PolicyFile.policy_name.like(f"%{keyword}%"))
    
    # 分类筛选
    if category:
        category_map = {
            "国家规范": 1,
            "行业规范": 2,
            "地方规范": 3,
            "内部条例": 4
        }
        if category in category_map:
            query = query.filter(PolicyFile.policy_type == category_map[category])
    
    # 日期筛选
    if date_from:
        query = query.filter(PolicyFile.created_at >= int(time.mktime(time.strptime(date_from, "%Y-%m-%d"))))
    if date_to:
        query = query.filter(PolicyFile.created_at <= int(time.mktime(time.strptime(date_to, "%Y-%m-%d"))))
    
    # 分页
    total = query.count()
    files = query.offset((page - 1) * page_size).limit(page_size).all()
    
    return {
        "statusCode": 200,
        "msg": "success",
        "data": {
            "total": total,
            "page": page,
            "page_size": page_size,
            "items": [
                {
                    "id": f.id,
                    "policy_name": f.policy_name,
                    "policy_file_url": f.policy_file_url,
                    "policy_type": f.policy_type,
                    "view_count": f.view_count,
                    "file_type": f.file_type,
                    "created_at": f.created_at
                }
                for f in files
            ]
        }
    }