import logging from datetime import datetime, timezone from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile from sqlalchemy.ext.asyncio import AsyncSession from db.session import get_session from models.requests import DocumentIntentsUpdateRequest, DocumentRawUpdateRequest from models.responses import ( ChunkDetail, ChunkPreview, DocumentChunksResponse, DocumentDeleteResponse, DocumentInfo, DocumentIntentsResponse, DocumentListResponse, DocumentUploadResponse, ) from services import document_service, intent_document_service from services.document_processor import process_document, rechunk_raw_text logger = logging.getLogger(__name__) router = APIRouter(prefix="/documents", tags=["documents"]) ALLOWED_EXTENSIONS = {".pdf", ".docx", ".doc", ".txt", ".md"} MAX_FILE_SIZE = 50 * 1024 * 1024 # 50 MB @router.post("/upload", response_model=DocumentUploadResponse) async def upload_document( file: UploadFile = File(...), document_name: str | None = Form(None), session: AsyncSession = Depends(get_session), ): from main import vectorstore_service if vectorstore_service is None: raise HTTPException(status_code=503, detail="Service not ready") filename = file.filename or "unknown" ext = "." + filename.rsplit(".", 1)[-1].lower() if "." in filename else "" if ext not in ALLOWED_EXTENSIONS: raise HTTPException( status_code=400, detail=f"Unsupported file format: {ext}. Allowed: {', '.join(ALLOWED_EXTENSIONS)}", ) file_bytes = await file.read() if len(file_bytes) > MAX_FILE_SIZE: raise HTTPException(status_code=400, detail="File too large (max 50 MB)") if len(file_bytes) == 0: raise HTTPException(status_code=400, detail="Empty file") display_name = document_name or filename try: document_id, raw_text, sections, chunks = process_document(file_bytes, filename) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) except Exception as e: logger.exception("Failed to process document: %s", filename) raise HTTPException(status_code=500, detail=f"Error processing document: {e}") if not chunks: raise HTTPException(status_code=400, detail="No content could be extracted from the document") file_type = ext.lstrip(".") await document_service.save_document_raw( session=session, document_id=document_id, name=display_name, file_type=file_type, raw_text=raw_text, ) chunks_count = vectorstore_service.add_document( document_id=document_id, document_name=display_name, file_type=file_type, chunks=[ { "text": c.text, "section": c.section, "page_number": c.page_number, "chunk_index": c.chunk_index, } for c in chunks ], ) chunks_prev = [ ChunkPreview( index=c.chunk_index, section=c.section, page_number=c.page_number, text_preview=c.text[:300], char_length=len(c.text), ) for c in chunks[:3] ] return DocumentUploadResponse( document_id=document_id, name=display_name, chunks_count=chunks_count, status="indexed", created_at=datetime.now(timezone.utc).isoformat(), chunks_preview=chunks_prev, ) @router.get("", response_model=DocumentListResponse) async def list_documents(): from main import vectorstore_service if vectorstore_service is None: raise HTTPException(status_code=503, detail="Service not ready") docs = vectorstore_service.list_documents() return DocumentListResponse( documents=[DocumentInfo(**d) for d in docs], total=len(docs), ) @router.get("/{document_id}/chunks", response_model=DocumentChunksResponse) async def get_document_chunks(document_id: str): from main import vectorstore_service if vectorstore_service is None: raise HTTPException(status_code=503, detail="Service not ready") raw_chunks = vectorstore_service.get_document_chunks(document_id) if not raw_chunks: raise HTTPException(status_code=404, detail="Document not found") meta0 = raw_chunks[0]["metadata"] chunks = [ ChunkDetail( index=c["metadata"].get("chunk_index", 0), section=c["metadata"].get("section", ""), page_number=c["metadata"].get("page_number", 0), text=c["text"], char_length=len(c["text"]), embedding=c.get("embedding", []), embedding_dim=len(c.get("embedding", [])), ) for c in raw_chunks ] return DocumentChunksResponse( document_id=document_id, name=meta0.get("document_name", ""), file_type=meta0.get("file_type", ""), chunks_count=len(chunks), chunks=chunks, ) @router.delete("/{document_id}", response_model=DocumentDeleteResponse) async def delete_document(document_id: str, session: AsyncSession = Depends(get_session)): from main import vectorstore_service if vectorstore_service is None: raise HTTPException(status_code=503, detail="Service not ready") deleted = vectorstore_service.delete_document(document_id) await document_service.delete_document_raw(session, document_id) if deleted == 0: raise HTTPException(status_code=404, detail="Document not found") return DocumentDeleteResponse(ok=True, deleted_chunks=deleted) @router.post("/{document_id}/reindex", response_model=DocumentUploadResponse) async def reindex_document(document_id: str, session: AsyncSession = Depends(get_session)): """Переразметить документ с актуальными правилами чанкера на основе сохранённого raw_text.""" from main import vectorstore_service if vectorstore_service is None: raise HTTPException(status_code=503, detail="Service not ready") doc = await document_service.get_document_raw(session, document_id) if doc is None: raise HTTPException(status_code=404, detail="Document raw_text not found — reindex невозможен") chunks = rechunk_raw_text(doc.raw_text) if not chunks: raise HTTPException(status_code=400, detail="После переразметки не осталось чанков") vectorstore_service.delete_document(document_id) chunks_count = vectorstore_service.add_document( document_id=document_id, document_name=doc.name, file_type=doc.file_type, chunks=[ { "text": c.text, "section": c.section, "page_number": c.page_number, "chunk_index": c.chunk_index, } for c in chunks ], ) chunks_prev = [ ChunkPreview( index=c.chunk_index, section=c.section, page_number=c.page_number, text_preview=c.text[:300], char_length=len(c.text), ) for c in chunks[:3] ] return DocumentUploadResponse( document_id=document_id, name=doc.name, chunks_count=chunks_count, status="reindexed", created_at=datetime.now(timezone.utc).isoformat(), chunks_preview=chunks_prev, ) @router.get("/{document_id}/raw") async def get_document_raw_text(document_id: str, session: AsyncSession = Depends(get_session)): """Отдать исходный текст документа — для редактирования в UI.""" doc = await document_service.get_document_raw(session, document_id) if doc is None: raise HTTPException(status_code=404, detail="Document not found") return { "document_id": doc.id, "name": doc.name, "file_type": doc.file_type, "raw_text": doc.raw_text, } @router.put("/{document_id}/raw", response_model=DocumentUploadResponse) async def update_document_raw_text( document_id: str, req: DocumentRawUpdateRequest, session: AsyncSession = Depends(get_session), ): """Обновить raw_text + сразу переразметить и переиндексировать в Chroma. Деструктивная операция: старые чанки удаляются, новые берут их место. Имя и file_type не меняются. """ from main import vectorstore_service if vectorstore_service is None: raise HTTPException(status_code=503, detail="Service not ready") doc = await document_service.get_document_raw(session, document_id) if doc is None: raise HTTPException(status_code=404, detail="Document not found") new_raw = req.raw_text.strip() if not new_raw: raise HTTPException(status_code=400, detail="raw_text не может быть пустым") chunks = rechunk_raw_text(new_raw) if not chunks: raise HTTPException(status_code=400, detail="После переразметки не осталось чанков") # 1) обновляем raw_text в SQLite await document_service.save_document_raw( session=session, document_id=document_id, name=doc.name, file_type=doc.file_type, raw_text=new_raw, ) # 2) переиндексация в Chroma — те же шаги, что в /reindex vectorstore_service.delete_document(document_id) chunks_count = vectorstore_service.add_document( document_id=document_id, document_name=doc.name, file_type=doc.file_type, chunks=[ { "text": c.text, "section": c.section, "page_number": c.page_number, "chunk_index": c.chunk_index, } for c in chunks ], ) chunks_prev = [ ChunkPreview( index=c.chunk_index, section=c.section, page_number=c.page_number, text_preview=c.text[:300], char_length=len(c.text), ) for c in chunks[:3] ] return DocumentUploadResponse( document_id=document_id, name=doc.name, chunks_count=chunks_count, status="updated", created_at=datetime.now(timezone.utc).isoformat(), chunks_preview=chunks_prev, ) @router.post("/reindex-all") async def reindex_all(session: AsyncSession = Depends(get_session)): """Переразметить все документы, у которых есть raw_text в SQLite.""" from main import vectorstore_service if vectorstore_service is None: raise HTTPException(status_code=503, detail="Service not ready") docs = await document_service.list_documents_raw(session) results = [] for doc in docs: chunks = rechunk_raw_text(doc.raw_text) if not chunks: results.append({"document_id": doc.id, "name": doc.name, "status": "empty"}) continue vectorstore_service.delete_document(doc.id) n = vectorstore_service.add_document( document_id=doc.id, document_name=doc.name, file_type=doc.file_type, chunks=[ { "text": c.text, "section": c.section, "page_number": c.page_number, "chunk_index": c.chunk_index, } for c in chunks ], ) results.append({"document_id": doc.id, "name": doc.name, "status": "reindexed", "chunks_count": n}) return {"total": len(results), "results": results} @router.get("/{document_id}/intents", response_model=DocumentIntentsResponse) async def list_document_intents(document_id: str, session: AsyncSession = Depends(get_session)): doc = await document_service.get_document_raw(session, document_id) if doc is None: raise HTTPException(status_code=404, detail="Document not found") intent_codes = await intent_document_service.list_intents_for_document(session, document_id) return DocumentIntentsResponse(document_id=document_id, intent_codes=intent_codes) @router.put("/{document_id}/intents", response_model=DocumentIntentsResponse) async def set_document_intents( document_id: str, req: DocumentIntentsUpdateRequest, session: AsyncSession = Depends(get_session), ): doc = await document_service.get_document_raw(session, document_id) if doc is None: raise HTTPException(status_code=404, detail="Document not found") intent_codes = await intent_document_service.set_intents_for_document( session, document_id, req.intent_codes, ) return DocumentIntentsResponse(document_id=document_id, intent_codes=intent_codes)