317 lines
9.8 KiB
Python
317 lines
9.8 KiB
Python
"""
|
|
Knowledge Base management endpoints.
|
|
[AC-ASA-01, AC-ASA-02, AC-ASA-08] Document upload, list, and index job status.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import uuid
|
|
from typing import Annotated, Optional
|
|
|
|
from fastapi import APIRouter, Depends, Query, UploadFile, File, Form
|
|
from fastapi.responses import JSONResponse
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.database import get_session
|
|
from app.core.exceptions import MissingTenantIdException
|
|
from app.core.tenant import get_tenant_id
|
|
from app.models import ErrorResponse
|
|
from app.models.entities import DocumentStatus, IndexJobStatus
|
|
from app.services.kb import KBService
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/admin/kb", tags=["KB Management"])
|
|
|
|
|
|
def get_current_tenant_id() -> str:
|
|
"""Dependency to get current tenant ID or raise exception."""
|
|
tenant_id = get_tenant_id()
|
|
if not tenant_id:
|
|
raise MissingTenantIdException()
|
|
return tenant_id
|
|
|
|
|
|
@router.get(
|
|
"/documents",
|
|
operation_id="listDocuments",
|
|
summary="Query document list",
|
|
description="[AC-ASA-08] Get list of documents with pagination and filtering.",
|
|
responses={
|
|
200: {"description": "Document list with pagination"},
|
|
401: {"description": "Unauthorized", "model": ErrorResponse},
|
|
403: {"description": "Forbidden", "model": ErrorResponse},
|
|
},
|
|
)
|
|
async def list_documents(
|
|
tenant_id: Annotated[str, Depends(get_current_tenant_id)],
|
|
session: Annotated[AsyncSession, Depends(get_session)],
|
|
kb_id: Annotated[Optional[str], Query()] = None,
|
|
status: Annotated[Optional[str], Query()] = None,
|
|
page: int = Query(1, ge=1),
|
|
page_size: int = Query(20, ge=1, le=100),
|
|
) -> JSONResponse:
|
|
"""
|
|
[AC-ASA-08] List documents with filtering and pagination.
|
|
"""
|
|
logger.info(
|
|
f"[AC-ASA-08] Listing documents: tenant={tenant_id}, kb_id={kb_id}, "
|
|
f"status={status}, page={page}, page_size={page_size}"
|
|
)
|
|
|
|
kb_service = KBService(session)
|
|
documents, total = await kb_service.list_documents(
|
|
tenant_id=tenant_id,
|
|
kb_id=kb_id,
|
|
status=status,
|
|
page=page,
|
|
page_size=page_size,
|
|
)
|
|
|
|
total_pages = (total + page_size - 1) // page_size if total > 0 else 0
|
|
|
|
data = [
|
|
{
|
|
"docId": str(doc.id),
|
|
"kbId": doc.kb_id,
|
|
"fileName": doc.file_name,
|
|
"status": doc.status,
|
|
"createdAt": doc.created_at.isoformat() + "Z",
|
|
"updatedAt": doc.updated_at.isoformat() + "Z",
|
|
}
|
|
for doc in documents
|
|
]
|
|
|
|
return JSONResponse(
|
|
content={
|
|
"data": data,
|
|
"pagination": {
|
|
"page": page,
|
|
"pageSize": page_size,
|
|
"total": total,
|
|
"totalPages": total_pages,
|
|
},
|
|
}
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/documents",
|
|
operation_id="uploadDocument",
|
|
summary="Upload/import document",
|
|
description="[AC-ASA-01] Upload document and trigger indexing job.",
|
|
responses={
|
|
202: {"description": "Accepted - async indexing job started"},
|
|
401: {"description": "Unauthorized", "model": ErrorResponse},
|
|
403: {"description": "Forbidden", "model": ErrorResponse},
|
|
},
|
|
)
|
|
async def upload_document(
|
|
tenant_id: Annotated[str, Depends(get_current_tenant_id)],
|
|
session: Annotated[AsyncSession, Depends(get_session)],
|
|
file: UploadFile = File(...),
|
|
kb_id: str = Form(...),
|
|
) -> JSONResponse:
|
|
"""
|
|
[AC-ASA-01] Upload document and create indexing job.
|
|
"""
|
|
logger.info(
|
|
f"[AC-ASA-01] Uploading document: tenant={tenant_id}, "
|
|
f"kb_id={kb_id}, filename={file.filename}"
|
|
)
|
|
|
|
kb_service = KBService(session)
|
|
|
|
kb = await kb_service.get_or_create_kb(tenant_id, kb_id)
|
|
|
|
file_content = await file.read()
|
|
document, job = await kb_service.upload_document(
|
|
tenant_id=tenant_id,
|
|
kb_id=str(kb.id),
|
|
file_name=file.filename or "unknown",
|
|
file_content=file_content,
|
|
file_type=file.content_type,
|
|
)
|
|
|
|
_schedule_indexing(tenant_id, str(job.id), str(document.id), file_content)
|
|
|
|
return JSONResponse(
|
|
status_code=202,
|
|
content={
|
|
"jobId": str(job.id),
|
|
"docId": str(document.id),
|
|
"status": job.status,
|
|
},
|
|
)
|
|
|
|
|
|
def _schedule_indexing(tenant_id: str, job_id: str, doc_id: str, content: bytes):
|
|
"""
|
|
Schedule background indexing task.
|
|
For MVP, we simulate indexing with a simple text extraction.
|
|
In production, this would use a task queue like Celery.
|
|
"""
|
|
import asyncio
|
|
|
|
async def run_indexing():
|
|
from app.core.database import async_session_maker
|
|
from app.services.kb import KBService
|
|
from app.core.qdrant_client import get_qdrant_client
|
|
from qdrant_client.models import PointStruct
|
|
import hashlib
|
|
|
|
await asyncio.sleep(1)
|
|
|
|
async with async_session_maker() as session:
|
|
kb_service = KBService(session)
|
|
try:
|
|
await kb_service.update_job_status(
|
|
tenant_id, job_id, IndexJobStatus.PROCESSING.value, progress=10
|
|
)
|
|
|
|
text = content.decode("utf-8", errors="ignore")
|
|
|
|
chunks = [text[i:i+500] for i in range(0, len(text), 500)]
|
|
|
|
qdrant = await get_qdrant_client()
|
|
await qdrant.ensure_collection_exists(tenant_id)
|
|
|
|
points = []
|
|
for i, chunk in enumerate(chunks):
|
|
hash_obj = hashlib.sha256(chunk.encode())
|
|
hash_bytes = hash_obj.digest()
|
|
embedding = []
|
|
for j in range(0, min(len(hash_bytes) * 8, 1536)):
|
|
byte_idx = j // 8
|
|
bit_idx = j % 8
|
|
if byte_idx < len(hash_bytes):
|
|
val = (hash_bytes[byte_idx] >> bit_idx) & 1
|
|
embedding.append(float(val))
|
|
else:
|
|
embedding.append(0.0)
|
|
while len(embedding) < 1536:
|
|
embedding.append(0.0)
|
|
|
|
points.append(
|
|
PointStruct(
|
|
id=str(uuid.uuid4()),
|
|
vector=embedding[:1536],
|
|
payload={
|
|
"text": chunk,
|
|
"source": doc_id,
|
|
"chunk_index": i,
|
|
},
|
|
)
|
|
)
|
|
|
|
if points:
|
|
await qdrant.upsert_vectors(tenant_id, points)
|
|
|
|
await kb_service.update_job_status(
|
|
tenant_id, job_id, IndexJobStatus.COMPLETED.value, progress=100
|
|
)
|
|
|
|
logger.info(
|
|
f"[AC-ASA-01] Indexing completed: tenant={tenant_id}, "
|
|
f"job_id={job_id}, chunks={len(chunks)}"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"[AC-ASA-01] Indexing failed: {e}")
|
|
await kb_service.update_job_status(
|
|
tenant_id, job_id, IndexJobStatus.FAILED.value,
|
|
progress=0, error_msg=str(e)
|
|
)
|
|
|
|
asyncio.create_task(run_indexing())
|
|
|
|
|
|
@router.get(
|
|
"/index/jobs/{job_id}",
|
|
operation_id="getIndexJob",
|
|
summary="Query index job status",
|
|
description="[AC-ASA-02] Get indexing job status and progress.",
|
|
responses={
|
|
200: {"description": "Job status details"},
|
|
401: {"description": "Unauthorized", "model": ErrorResponse},
|
|
403: {"description": "Forbidden", "model": ErrorResponse},
|
|
},
|
|
)
|
|
async def get_index_job(
|
|
tenant_id: Annotated[str, Depends(get_current_tenant_id)],
|
|
session: Annotated[AsyncSession, Depends(get_session)],
|
|
job_id: str,
|
|
) -> JSONResponse:
|
|
"""
|
|
[AC-ASA-02] Get indexing job status with progress.
|
|
"""
|
|
logger.info(
|
|
f"[AC-ASA-02] Getting job status: tenant={tenant_id}, job_id={job_id}"
|
|
)
|
|
|
|
kb_service = KBService(session)
|
|
job = await kb_service.get_index_job(tenant_id, job_id)
|
|
|
|
if not job:
|
|
return JSONResponse(
|
|
status_code=404,
|
|
content={
|
|
"code": "JOB_NOT_FOUND",
|
|
"message": f"Job {job_id} not found",
|
|
},
|
|
)
|
|
|
|
return JSONResponse(
|
|
content={
|
|
"jobId": str(job.id),
|
|
"docId": str(job.doc_id),
|
|
"status": job.status,
|
|
"progress": job.progress,
|
|
"errorMsg": job.error_msg,
|
|
}
|
|
)
|
|
|
|
|
|
@router.delete(
|
|
"/documents/{doc_id}",
|
|
operation_id="deleteDocument",
|
|
summary="Delete document",
|
|
description="[AC-ASA-08] Delete a document and its associated files.",
|
|
responses={
|
|
200: {"description": "Document deleted"},
|
|
404: {"description": "Document not found"},
|
|
401: {"description": "Unauthorized", "model": ErrorResponse},
|
|
403: {"description": "Forbidden", "model": ErrorResponse},
|
|
},
|
|
)
|
|
async def delete_document(
|
|
tenant_id: Annotated[str, Depends(get_current_tenant_id)],
|
|
session: Annotated[AsyncSession, Depends(get_session)],
|
|
doc_id: str,
|
|
) -> JSONResponse:
|
|
"""
|
|
[AC-ASA-08] Delete a document.
|
|
"""
|
|
logger.info(
|
|
f"[AC-ASA-08] Deleting document: tenant={tenant_id}, doc_id={doc_id}"
|
|
)
|
|
|
|
kb_service = KBService(session)
|
|
deleted = await kb_service.delete_document(tenant_id, doc_id)
|
|
|
|
if not deleted:
|
|
return JSONResponse(
|
|
status_code=404,
|
|
content={
|
|
"code": "DOCUMENT_NOT_FOUND",
|
|
"message": f"Document {doc_id} not found",
|
|
},
|
|
)
|
|
|
|
return JSONResponse(
|
|
content={
|
|
"success": True,
|
|
"message": "Document deleted",
|
|
}
|
|
)
|