ds-fire-fighter/app.py

import os
from typing import Optional
from fastapi import FastAPI, HTTPException, Security, Depends
from fastapi.security import APIKeyHeader
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from dotenv import load_dotenv
from utils.document_loader import load_document
import json
from pydantic import BaseModel
from src.llm import ai_chat
from langchain_openai import ChatOpenAI
import requests
import tempfile
from scripts.generate_pdf  import create_pdf
from scripts.generate_theme import generate_theme
from scripts.generate_quiz import generate_quiz
from typing import Dict, Any
from fastapi.responses import Response
from datetime import datetime
from fastapi import HTTPException
from pydantic import BaseModel
from typing import Optional, Union, Dict, Any
import os
import requests
import os
from PyPDF2 import PdfReader
from config import QUIZ_TYPES
# Load environment variables
load_dotenv()
API_KEY = os.getenv("API_KEY_ACCESS")

base_path = os.path.join("data", "config_files")
QUESTIONS_PATH = os.path.join(base_path, "questions.json")
THEME_CONTEXT_PATH = os.path.join(base_path, "theme_context.json")

# Load themes at module level
with open(THEME_CONTEXT_PATH, "r") as f:
    themes = json.load(f)

# Initialize FastAPI app
app = FastAPI(
    title="Fire Fighter Interview API",
    description="API For fire fighter",
    version="1.0.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Setup API key authentication
api_key_header = APIKeyHeader(name="Authorization", auto_error=False)



async def get_api_key(api_key_header: str = Security(api_key_header)) -> str:
    """Validate API key from header"""
    if not api_key_header or not api_key_header.startswith('Bearer '):
        raise HTTPException(
            status_code=401,
            detail={"error": "Unauthorized", "message": "API key is missing or invalid."}
        )
    
    token = api_key_header.split(' ')[1]
    if token != API_KEY:
        raise HTTPException(
            status_code=401,
            detail={"error": "Unauthorized", "message": "API key does not match."}
        )
    
    return token

class ChatRequest(BaseModel):
    resume_url: Optional[str] = None
    query: str=None
    conversation_id: str
    theme_id: Optional[int] = 1

class ChatResponse(BaseModel):
    message: str
    end: bool
    error: Optional[str] = None



class GeneratePDFRequest(BaseModel):
    conversation_id: str
    feedback: Optional[str] = None
    previous_results: Optional[Dict[str, Any]] = None
    resume_url: Optional[str] = None
    full_history_url: Optional[str] = None

class QuizRequest(BaseModel):
    pdf_url: str
    quiz_type: int  # 1, 2, or 3 corresponding to QUIZ_TYPES

class QuizResponse(BaseModel):
    success: bool
    message: str
    quiz_data: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    

async def extract_pdf_text(pdf_url: str) -> Union[str, None]:
    """Extract text from PDF and handle potential errors."""
    try:
        response = requests.get(pdf_url)
        response.raise_for_status()
        
        # Create a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_pdf:
            temp_pdf.write(response.content)
            temp_path = temp_pdf.name

        # Extract text from PDF
        reader = PdfReader(temp_path)
        text = "\n\n".join(
            page.extract_text() for page in reader.pages if page.extract_text()
        )

        # Clean up temporary file
        os.unlink(temp_path)

        if not text.strip():
            return None
        return text

    except requests.RequestException as e:
        raise HTTPException(
            status_code=400,
            detail=f"Error downloading PDF: {str(e)}"
        )
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error processing PDF: {str(e)}"
        )

@app.post("/rescue-career/chat", response_model=ChatResponse)
async def chat_endpoint(
    request: ChatRequest,
    api_key: str = Depends(get_api_key)
):
    try:
        # Validate theme
        matching_themes = [t for t in themes if t["id"] == request.theme_id]
        if not matching_themes:
            raise HTTPException(
                status_code=400,
                detail=f"No theme found with ID {request.theme_id}"
            )

        # Only try to load document if resume_url is provided
        resume_docs = ""
        if request.resume_url:
            docs = load_document(request.resume_url)
            if not docs:
                raise HTTPException(
                    status_code=400, 
                    detail="Invalid resume URL: Unable to fetch document"
                )
            resume_docs = "\n".join(f"- {doc.page_content}" for doc in docs)

        # Get AI chat response
        response = ai_chat(
            query=request.query,
            conversation_id=request.conversation_id,
            theme_id=request.theme_id,
            resume=resume_docs
        )

        # Parse response
        try:
            parsed_response = json.loads(response)
            return ChatResponse(
                message=parsed_response.get("message", ""),
                end=parsed_response.get("end", "no") == "yes",
                error=None
            )
        except json.JSONDecodeError:
            return ChatResponse(
                message=response,
                end=False,
                error=None
            )

    except HTTPException as e:
        # Re-raise HTTP exceptions
        raise
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error processing chat request: {str(e)}"
        )
     

    
@app.post("/rescue-career/generate-theme")
async def generate_pdf_endpoint(
    request: GeneratePDFRequest,
    api_key: str = Depends(get_api_key)
):
    
    try:
        # Here you would fetch the conversation data using the conversation_id
        # This is a placeholder - replace with your actual conversation data fetching logic
        conversation_data = await get_conversation_data(request.conversation_id)
        
        if not conversation_data:
            raise HTTPException(
                status_code=404,
                detail=f"No conversation found with ID {request.conversation_id}"
            )
        
        resume_docs = ""
        if request.resume_url:
            docs = load_document(request.resume_url)
            if not docs:
                raise HTTPException(
                    status_code=400, 
                    detail="Invalid resume URL: Unable to fetch document"
                )
            resume_docs = "\n".join(f"- {doc.page_content}" for doc in docs)


        full_history_docs = ""
        if request.full_history_url:
            docs = load_document(request.full_history_url)
            if not docs:
                raise HTTPException(
                    status_code=400, 
                    detail="Invalid full_history URL: Unable to fetch document"
                )
            full_history_docs = "\n".join(f"- {doc.page_content}" for doc in docs)
        # Generate theme data using the generate_theme function
        theme_data = generate_theme(
            conversation_data=conversation_data,
            feedback=request.feedback,
            previous_result=request.previous_results,
            resume = resume_docs,
            full_history = full_history_docs
        )  
        
        if not theme_data:
            raise HTTPException(
                status_code=500,
                detail="Failed to generate theme data"
            )
        
        # Generate the PDF using the create_pdf function
        pdf_content = create_pdf(theme_data)
        
        # Create filename with timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"theme_{timestamp}.pdf"
        
        # Return the PDF as a response
        return Response(
            content=pdf_content,
            media_type="application/pdf",
            headers={
                "Content-Disposition": f'attachment; filename="{filename}"'
            }
        )
    
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error generating PDF: {str(e)}"
        )
        
@app.post("/rescue-career/generate-quiz", response_model=QuizResponse)
async def generate_quiz_endpoint(
    request: QuizRequest,
    api_key: str = Depends(get_api_key)
):
    """Generate quiz based on PDF content and quiz type."""
    # Validate quiz type
    if request.quiz_type not in QUIZ_TYPES:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid quiz type. Must be one of: {list(QUIZ_TYPES)}"
        )

    try:
        # Extract text from PDF
        pdf_text = await extract_pdf_text(request.pdf_url)
        if not pdf_text:
            return QuizResponse(
                success=False,
                message="PDF extraction completed but no text content found",
                error="Empty PDF content"
            )
            
        # Generate quiz using the existing function
        quiz_data = generate_quiz(
            startpop_pdf=pdf_text,
            quiz_type=request.quiz_type
        )

        if not quiz_data:
            return QuizResponse(
                success=False,
                message="Quiz generation failed",
                error="Unable to generate quiz from the provided content"
            )

        return QuizResponse(
            success=True,
            message="Quiz generated successfully",
            quiz_data=quiz_data
        )

    except HTTPException as he:
        raise he
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Unexpected error during quiz generation: {str(e)}"
        )
        
        
        
async def get_conversation_data(conversation_id: str) -> dict:
    """
    Fetch conversation data using the conversation ID.
    Replace this with your actual implementation to fetch conversation data.
    """
    try:
        storage_path = "conversations.json"
        with open(storage_path, 'r') as f:
                convs =  json.load(f)
        convs_id = convs[conversation_id]
        return convs_id
    except Exception as e:
        print(f"Error fetching conversation data: {e}")
        return None
    
    
@app.on_event("startup")
async def startup_event():
    """Initialize required components on startup"""
    pass

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