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)