Aigenthix/Graph_RAG

1"""Document processing and management service"""2 3import os4from pathlib import Path5from typing import Optional, List, Dict, Any6from datetime import datetime7from uuid import uuid48import logging9 10logger = logging.getLogger(__name__)11 12 13class DocumentService:14    """Service for managing documents"""15 16    def __init__(self, upload_dir: str):17        self.upload_dir = Path(upload_dir)18        self.upload_dir.mkdir(parents=True, exist_ok=True)19        self.documents: Dict[str, Dict[str, Any]] = {}20 21    async def save_document(22        self,23        file_content: bytes,24        filename: str,25        document_type: str,26        metadata: Optional[Dict[str, Any]] = None,27    ) -> Dict[str, Any]:28        """Save uploaded document"""29        doc_id = str(uuid4())30        file_path = self.upload_dir / f"{doc_id}_{filename}"31 32        # Save file33        file_path.write_bytes(file_content)34 35        # Store metadata36        doc_info = {37            "id": doc_id,38            "name": filename,39            "document_type": document_type,40            "upload_date": datetime.now(),41            "file_size": len(file_content),42            "file_path": str(file_path),43            "chunks_count": 0,44            "metadata": metadata or {},45        }46        self.documents[doc_id] = doc_info47 48        logger.info(f"Document saved: {doc_id} ({filename})")49        return doc_info50 51    def get_document(self, doc_id: str) -> Optional[Dict[str, Any]]:52        """Get document info"""53        return self.documents.get(doc_id)54 55    def list_documents(self) -> List[Dict[str, Any]]:56        """List all documents"""57        return list(self.documents.values())58 59    def delete_document(self, doc_id: str) -> bool:60        """Delete document"""61        if doc_id in self.documents:62            doc = self.documents[doc_id]63            file_path = Path(doc["file_path"])64            if file_path.exists():65                file_path.unlink()66            del self.documents[doc_id]67            logger.info(f"Document deleted: {doc_id}")68            return True69        return False70 71    def update_chunks_count(self, doc_id: str, count: int) -> None:72        """Update chunk count for document"""73        if doc_id in self.documents:74            self.documents[doc_id]["chunks_count"] = count75