vent_ai_platform/app/api/v1/endpoints/knowledge.py

from __future__ import annotations

import json

from fastapi import APIRouter, Depends, Query, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from app.db.repositories import knowledge as knowledge_repo
from app.db.repositories.base import create_model
from app.db.session import get_db
from app.schemas.common import Page
from app.schemas.knowledge import (
    KnowledgeAskRequest,
    KnowledgeAskResponse,
    KnowledgeCatalogResponse,
    KnowledgeGraphResponse,
    KnowledgeItemCreate,
    KnowledgeItemRead,
    KnowledgeNoteDetail,
    KnowledgeOverviewResponse,
    KnowledgeSearchHit,
    KnowledgeSearchRequest,
    KnowledgeSearchResponse,
)
from app.services.knowledge_hub_service import (
    build_answer,
    build_catalog,
    build_graph,
    build_note_detail,
    build_overview,
    stream_answer,
)
from app.services.knowledge_service import search_knowledge

router = APIRouter(prefix="/knowledge", tags=["知识中心"])


@router.get("", response_model=Page[KnowledgeItemRead], summary="知识条目列表")
def list_items(
    mine_id: str | None = Query(None),
    knowledge_type: str | None = Query(None),
    category: str | None = Query(None),
    keyword: str | None = Query(None),
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=200),
    db: Session = Depends(get_db),
):
    items, total = knowledge_repo.list_knowledge(
        db,
        mine_id=mine_id,
        knowledge_type=knowledge_type,
        category=category,
        keyword=keyword,
        page=page,
        page_size=page_size,
    )
    return Page(items=items, total=total, page=page, page_size=page_size)


@router.post("", response_model=KnowledgeItemRead, status_code=201, summary="新增知识")
def create_knowledge(payload: KnowledgeItemCreate, db: Session = Depends(get_db)):
    data = payload.model_dump(exclude={"metadata"})
    data["metadata_json"] = payload.metadata
    return create_model(db, knowledge_repo.KnowledgeItem, data)


@router.post(
    "/search",
    response_model=KnowledgeSearchResponse,
    summary="知识检索",
)
def search(payload: KnowledgeSearchRequest, db: Session = Depends(get_db)):
    hits = search_knowledge(db, payload)
    return KnowledgeSearchResponse(
        query=payload.query,
        total=len(hits),
        items=[
            KnowledgeSearchHit(
                item=KnowledgeItemRead.model_validate(hit["item"]),
                score=hit["score"],
                highlight=hit["highlight"],
            )
            for hit in hits
        ],
    )


@router.get(
    "/overview",
    response_model=KnowledgeOverviewResponse,
    summary="知识中心总览",
)
async def overview() -> KnowledgeOverviewResponse:
    return await build_overview()


@router.get(
    "/catalog",
    response_model=KnowledgeCatalogResponse,
    summary="知识中心目录",
)
async def catalog(
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=200),
    category: str | None = Query(None),
    keyword: str | None = Query(None),
    tag: str | None = Query(None),
) -> KnowledgeCatalogResponse:
    return await build_catalog(
        page=page,
        page_size=page_size,
        category=category,
        keyword=keyword,
        tag=tag,
    )


@router.get(
    "/catalog/{note_id:path}",
    response_model=KnowledgeNoteDetail,
    summary="知识条目详情",
)
async def note_detail(note_id: str) -> KnowledgeNoteDetail:
    return await build_note_detail(note_id)


@router.post(
    "/ask",
    response_model=KnowledgeAskResponse,
    summary="知识检索问答",
)
async def ask(payload: KnowledgeAskRequest, db: Session = Depends(get_db)) -> KnowledgeAskResponse:
    return await build_answer(db, payload)


@router.post(
    "/rag",
    response_model=KnowledgeAskResponse,
    summary="RAG 知识问答",
)
async def rag(payload: KnowledgeAskRequest, db: Session = Depends(get_db)) -> KnowledgeAskResponse:
    rag_payload = payload.model_copy(update={"generate_answer": True})
    return await build_answer(db, rag_payload)


@router.post(
    "/rag/stream",
    summary="RAG 知识问答流式接口",
)
async def rag_stream(payload: KnowledgeAskRequest, db: Session = Depends(get_db)):
    return _rag_stream_response(payload, db)


@router.get(
    "/rag/stream",
    summary="RAG 知识问答流式接口（EventSource）",
)
async def rag_stream_get(
    question: str = Query(min_length=1),
    top_k: int = Query(5, ge=1, le=20),
    db: Session = Depends(get_db),
):
    payload = KnowledgeAskRequest(
        question=question,
        top_k=top_k,
        generate_answer=True,
    )
    return _rag_stream_response(payload, db)


def _rag_stream_response(payload: KnowledgeAskRequest, db: Session):
    async def event_stream():
        async for item in stream_answer(db, payload):
            data = json.dumps(item["data"], ensure_ascii=False)
            yield f"event: {item['event']}\ndata: {data}\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@router.websocket("/rag/ws")
async def rag_ws(websocket: WebSocket, db: Session = Depends(get_db)):
    await websocket.accept()
    try:
        payload = KnowledgeAskRequest.model_validate(await websocket.receive_json())
        async for item in stream_answer(db, payload):
            await websocket.send_json(item)
    except WebSocketDisconnect:
        return
    except Exception as exc:
        await websocket.send_json(
            {
                "event": "error",
                "data": {
                    "message": f"流式问答失败：{type(exc).__name__}",
                },
            }
        )
    finally:
        try:
            await websocket.close()
        except RuntimeError:
            pass


@router.get(
    "/graph",
    response_model=KnowledgeGraphResponse,
    summary="风险关联知识图谱",
)
async def graph(topic: str = Query("瓦斯超限风险事件")) -> KnowledgeGraphResponse:
    return await build_graph(topic)