# src/implementations/document_service.py
from fastapi import HTTPException
from pathlib import Path
import shutil
import os
from uuid import uuid4
from typing import List, Tuple, Dict
from fastapi import UploadFile, BackgroundTasks
from src.vectorstores.chroma_vectorstore import ChromaVectorStore
from src.utils.document_processor import DocumentProcessor
from src.models import DocumentResponse, DocumentInfo, BatchUploadResponse
from src.utils.logger import logger
from src.db.mongodb_store import MongoDBStore
from config.config import settings
class DocumentService:
def __init__(
self,
doc_processor: DocumentProcessor,
mongodb: MongoDBStore
):
self.doc_processor = doc_processor
self.mongodb = mongodb
self.permanent_dir = Path(settings.UPLOADS_DIR)
self.permanent_dir.mkdir(parents=True, exist_ok=True)
async def check_duplicate_filename(self, filename: str) -> bool:
"""
Check if a file with the same name exists
Args:
filename (str): Original filename to check
Returns:
bool: True if duplicate exists, False otherwise
"""
documents = await self.mongodb.get_all_documents()
return any(doc.get('filename') == filename for doc in documents)
async def process_documents(
self,
files: List[UploadFile],
vector_store: ChromaVectorStore,
background_tasks: BackgroundTasks
) -> BatchUploadResponse:
"""Process multiple document uploads"""
processed_files, failed_files = await self._handle_file_uploads(
files,
vector_store,
background_tasks
)
return BatchUploadResponse(
message=f"Processed {len(processed_files)} documents with {len(failed_files)} failures",
processed_files=processed_files,
failed_files=failed_files
)
async def _handle_file_uploads(
self,
files: List[UploadFile],
vector_store: ChromaVectorStore,
background_tasks: BackgroundTasks
) -> Tuple[List[DocumentResponse], List[dict]]:
"""Handle individual file uploads and processing"""
processed_files = []
failed_files = []
for file in files:
try:
# Check for duplicate filename
if await self.check_duplicate_filename(file.filename):
failed_files.append(self._create_failed_file_entry(
file.filename,
"A document with this name already exists. Please upload another document."
))
continue
if not self._is_supported_format(file.filename):
failed_files.append(self._create_failed_file_entry(
file.filename,
"Unsupported file format"
))
continue
document_response = await self._process_single_file(
file,
vector_store,
background_tasks
)
processed_files.append(document_response)
except Exception as e:
logger.error(
f"Error processing file {file.filename}: {str(e)}")
failed_files.append(self._create_failed_file_entry(
file.filename,
str(e)
))
return processed_files, failed_files
async def _process_single_file(
self,
file: UploadFile,
vector_store: ChromaVectorStore,
background_tasks: BackgroundTasks
) -> DocumentResponse:
"""Process a single file upload with proper handle closure"""
document_id = str(uuid4())
filename = f"{document_id}_{file.filename}"
file_path = self.permanent_dir / filename
url_path = f"/docs/{filename}"
try:
# Save file to permanent location using a context manager
with open(file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
# Close the uploaded file explicitly
await file.close()
# Process document with proper cleanup for Excel files
try:
processed_doc = await self.doc_processor.process_document(file_path)
# For Excel files, ensure pandas closes the file
if file_path.suffix.lower() in ['.xlsx', '.xls']:
import gc
gc.collect() # Help cleanup any lingering file handles
except Exception as proc_error:
logger.error(f"Error processing document: {str(proc_error)}")
raise
# Store in MongoDB with url_path
await self.mongodb.store_document(
document_id=document_id,
filename=file.filename,
content_type=file.content_type,
file_size=os.path.getsize(file_path),
url_path=url_path,
source="user_upload"
)
# Process for vector store in background
background_tasks.add_task(
self._process_for_vector_store,
processed_doc['chunks'],
vector_store,
document_id,
file.filename
)
return DocumentResponse(
message="Document uploaded successfully",
document_id=document_id,
status="processing",
document_info=DocumentInfo(
original_filename=file.filename,
size=os.path.getsize(file_path),
content_type=file.content_type,
url_path=url_path
)
)
except Exception as e:
# Clean up file if it was created
if file_path.exists():
try:
file_path.unlink()
except Exception as cleanup_error:
logger.error(
f"Error cleaning up file {file_path}: {str(cleanup_error)}")
# Clean up from MongoDB if document was created
try:
await self.mongodb.delete_document(document_id)
except Exception as db_cleanup_error:
logger.error(
f"Error cleaning up MongoDB document {document_id}: {str(db_cleanup_error)}")
logger.error(f"Error processing file {file.filename}: {str(e)}")
raise
async def _process_for_vector_store(
self,
chunks: List[str],
vector_store: ChromaVectorStore,
document_id: str,
filename: str
):
"""Process document content for vector store"""
try:
# Generate chunk IDs using document_id
chunk_ids = [
f"{document_id}-chunk-{i}" for i in range(len(chunks))]
# Get embeddings
embeddings = vector_store.embedding_function(chunks)
# Prepare metadata for each chunk
metadatas = [{
'document_id': document_id,
'source_file': filename,
'chunk_index': i,
'total_chunks': len(chunks)
} for i in range(len(chunks))]
# Store in vector store
vector_store.add_documents(
documents=chunks,
embeddings=embeddings,
metadatas=metadatas,
ids=chunk_ids
)
logger.info(
f"Successfully processed document {filename} (ID: {document_id}) into {len(chunks)} chunks")
except Exception as e:
logger.error(
f"Error processing document {filename} (ID: {document_id}) for vector store: {str(e)}")
raise
def _is_supported_format(self, filename: str) -> bool:
"""Check if file format is supported"""
return any(filename.lower().endswith(ext)
for ext in self.doc_processor.supported_formats)
def _create_failed_file_entry(self, filename: str, error: str) -> dict:
"""Create a failed file entry"""
return {
"filename": filename,
"error": error
}
async def delete_document(self, document_id: str) -> bool:
"""Delete document from storage and MongoDB"""
try:
# Get document details from MongoDB
doc = await self.mongodb.get_document(document_id)
if doc:
# Get filename from url_path
filename = doc['url_path'].split('/')[-1]
file_path = self.permanent_dir / filename
# Delete physical file if it exists
if file_path.exists():
file_path.unlink()
# Delete from MongoDB
return await self.mongodb.delete_document(document_id)
return False
except Exception as e:
logger.error(f"Error deleting document: {str(e)}")
raise
def cleanup(self):
"""Clean up permanent directory if empty"""
if self.permanent_dir.exists() and not any(self.permanent_dir.iterdir()):
self.permanent_dir.rmdir()