src/main.py

# src/main.py
from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware  # Add this import
from typing import List
import uuid
from datetime import datetime
from pathlib import Path
import os

# Import custom modules1
from src.agents.rag_agent import RAGAgent
from src.models.document import AllDocumentsResponse, StoredDocument
from src.utils.document_processor import DocumentProcessor
from src.utils.conversation_summarizer import ConversationSummarizer
from src.utils.logger import logger
from src.utils.llm_utils import get_llm_instance, get_vector_store
from src.db.mongodb_store import MongoDBStore
from src.implementations.document_service import DocumentService
from src.models import (
    ChatRequest, 
    ChatResponse, 
    DocumentResponse,
    BatchUploadResponse,
    SummarizeRequest,
    SummaryResponse,
    FeedbackRequest
)
from config.config import settings

app = FastAPI(title="Chatbot API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:8080"],  # Add your frontend URL
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Initialize MongoDB
mongodb = MongoDBStore(settings.MONGODB_URI)

# Initialize core components
doc_processor = DocumentProcessor()
summarizer = ConversationSummarizer()
document_service = DocumentService(doc_processor, mongodb)

# Create uploads directory if it doesn't exist
UPLOADS_DIR = Path("uploads")
UPLOADS_DIR.mkdir(exist_ok=True)

# Mount the uploads directory for static file serving
app.mount("/docs", StaticFiles(directory=str(UPLOADS_DIR)), name="documents")

@app.get("/documents")
async def get_all_documents():
    """Get all documents from MongoDB"""
    try:
        documents = await mongodb.get_all_documents()
        
        formatted_documents = []
        for doc in documents:
            try:
                formatted_doc = {
                    "document_id": doc.get("document_id"),
                    "filename": doc.get("filename"),
                    "content_type": doc.get("content_type"),
                    "file_size": doc.get("file_size"),
                    "url_path": doc.get("url_path"),
                    "upload_timestamp": doc.get("upload_timestamp")
                }
                formatted_documents.append(formatted_doc)
            except Exception as e:
                logger.error(f"Error formatting document {doc.get('document_id', 'unknown')}: {str(e)}")
                continue
        
        return {
            "total_documents": len(formatted_documents),
            "documents": formatted_documents
        }
    except Exception as e:
        logger.error(f"Error retrieving documents: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/documents/{document_id}/download")
async def get_document_file(document_id: str):
    """Serve a document file by its ID"""
    try:
        # Get document info from MongoDB
        doc = await mongodb.get_document(document_id)
        if not doc:
            raise HTTPException(status_code=404, detail="Document not found")
        
        # Extract filename from url_path
        filename = doc["url_path"].split("/")[-1]
        file_path = UPLOADS_DIR / filename
        
        if not file_path.exists():
            raise HTTPException(
                status_code=404, 
                detail=f"File not found on server: {filename}"
            )
            
        return FileResponse(
            path=str(file_path),
            filename=doc["filename"],
            media_type=doc["content_type"]
        )
            
    except Exception as e:
        logger.error(f"Error serving document file: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/documents/upload", response_model=BatchUploadResponse)
async def upload_documents(
    files: List[UploadFile] = File(...),
    background_tasks: BackgroundTasks = BackgroundTasks()
):
    """Upload and process multiple documents"""
    try:
        vector_store, _ = await get_vector_store()
        response = await document_service.process_documents(
            files, 
            vector_store, 
            background_tasks
        )
        return response
    except Exception as e:
        logger.error(f"Error in document upload: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))    

@app.get("/documentchunks/{document_id}")
async def get_document_chunks(document_id: str):
    """Get all chunks for a specific document"""
    try:
        vector_store, _ = await get_vector_store()
        chunks = vector_store.get_document_chunks(document_id)
        
        if not chunks:
            raise HTTPException(status_code=404, detail="Document not found")
            
        return {
            "document_id": document_id,
            "total_chunks": len(chunks),
            "chunks": chunks
        }
    except Exception as e:
        logger.error(f"Error retrieving document chunks: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/documents/{document_id}")
async def delete_document(document_id: str):
    """Delete document from MongoDB, ChromaDB, and physical storage"""
    try:
        # First get document details from MongoDB to get file path
        document = await mongodb.get_document(document_id)
        if not document:
            raise HTTPException(status_code=404, detail="Document not found")
        
        # Get vector store instance
        vector_store, _ = await get_vector_store()
        
        # Delete physical file using document service
        deletion_success = await document_service.delete_document(document_id)
        if not deletion_success:
            logger.warning(f"Failed to delete physical file for document {document_id}")
            
        # Delete from vector store
        try:
            vector_store.delete_document(document_id)
        except Exception as e:
            logger.error(f"Error deleting document from vector store: {str(e)}")
            raise HTTPException(
                status_code=500, 
                detail=f"Failed to delete document from vector store: {str(e)}"
            )
            
        # Delete from MongoDB - don't check return value since document might already be deleted
        await mongodb.delete_document(document_id)
        
        return {
            "status": "success",
            "message": f"Document {document_id} successfully deleted from all stores"
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in delete_document endpoint: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))
    
@app.post("/chat", response_model=ChatResponse)
async def chat_endpoint(
    request: ChatRequest,
    background_tasks: BackgroundTasks
):
    """Chat endpoint with RAG support"""
    try:
        vector_store, embedding_model = await get_vector_store()
        llm = get_llm_instance(request.llm_provider)
        
        # Initialize RAG agent with required MongoDB
        rag_agent = RAGAgent(
            llm=llm,
            embedding=embedding_model,
            vector_store=vector_store,
            mongodb=mongodb
        )
        
        # Use provided conversation ID or create new one
        conversation_id = request.conversation_id or str(uuid.uuid4())
        query = request.query + ". The response should be short and to the point. make sure, to not add any irrelevant information. Stick to the point is very very important."
        # Generate response
        response = await rag_agent.generate_response(
            query=query,
            conversation_id=conversation_id,
            temperature=request.temperature
        )
        
        # Store message in chat history
        await mongodb.store_message(
            conversation_id=conversation_id,
            query=request.query,
            response=response.response,
            context=response.context_docs,
            sources=response.sources,
            llm_provider=request.llm_provider
        )
        
        return ChatResponse(
            response=response.response,
            context=response.context_docs,
            sources=response.sources,
            conversation_id=conversation_id,
            timestamp=datetime.now(),
            relevant_doc_scores=response.scores if hasattr(response, 'scores') else None
        )
        
    except Exception as e:
        logger.error(f"Error in chat endpoint: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/chat/history/{conversation_id}")
async def get_conversation_history(conversation_id: str):
    """Get complete conversation history"""
    history = await mongodb.get_conversation_history(conversation_id)
    
    if not history:
        raise HTTPException(status_code=404, detail="Conversation not found")
        
    return {
        "conversation_id": conversation_id,
        "messages": history
    }

@app.post("/chat/summarize", response_model=SummaryResponse)
async def summarize_conversation(request: SummarizeRequest):
    """Generate a summary of a conversation"""
    try:
        messages = await mongodb.get_messages_for_summary(request.conversation_id)
        
        if not messages:
            raise HTTPException(status_code=404, detail="Conversation not found")
        
        summary = await summarizer.summarize_conversation(
            messages,
            include_metadata=request.include_metadata
        )
        
        return SummaryResponse(**summary)
        
    except Exception as e:
        logger.error(f"Error generating summary: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/chat/feedback/{conversation_id}")
async def submit_feedback(
    conversation_id: str,
    feedback_request: FeedbackRequest
):
    """Submit feedback for a conversation"""
    try:
        # Validate conversation exists
        conversation = await mongodb.get_conversation_metadata(conversation_id)
        if not conversation:
            raise HTTPException(status_code=404, detail="Conversation not found")
        
        # Update feedback
        success = await mongodb.update_feedback(
            conversation_id=conversation_id,
            feedback=feedback_request.feedback,
            rating=feedback_request.rating
        )
        
        if not success:
            raise HTTPException(
                status_code=500,
                detail="Failed to update feedback"
            )
            
        return {
            "status": "success",
            "message": "Feedback submitted successfully",
            "data": {
                "conversation_id": conversation_id,
                "feedback": feedback_request.feedback,
                "rating": feedback_request.format_rating()
            }
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error submitting feedback: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/debug/config")
async def debug_config():
    """Debug endpoint to check configuration"""
    import os
    from config.config import settings
    from pathlib import Path
    
    debug_info = {
        "environment_variables": {
            "OPENAI_API_KEY": "[SET]" if os.getenv('OPENAI_API_KEY') else "[NOT SET]",
            "OPENAI_MODEL": os.getenv('OPENAI_MODEL', '[NOT SET]')
        },
        "settings": {
            "OPENAI_API_KEY": "[SET]" if settings.OPENAI_API_KEY else "[NOT SET]",
            "OPENAI_MODEL": settings.OPENAI_MODEL,
        },
        "files": {
            "env_file_exists": Path('.env').exists(),
            "openai_config_exists": (Path.home() / '.openai' / 'api_key').exists()
        }
    }
    
    if settings.OPENAI_API_KEY:
        key = settings.OPENAI_API_KEY
        debug_info["api_key_info"] = {
            "length": len(key),
            "preview": f"{key[:4]}...{key[-4:]}" if len(key) > 8 else "[INVALID LENGTH]"
        }
        
    return debug_info

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)