app.py

from jose import JWTError, jwt                                                                                                                                                                                                                
from datetime import datetime, timedelta                                                                                                                                                                                                      
from typing import Optional  
from fastapi import FastAPI, UploadFile, File, HTTPException, Form, BackgroundTasks, Depends                                                                                                                                                  
from fastapi.responses import JSONResponse                                                                                                                                                                                                    
from fastapi.middleware.cors import CORSMiddleware                                                                                                                                                                                            
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm                                                                                                                                                                  
from typing import List, Dict, Optional                                                                                                                                                                                                       
from uuid import uuid4                                                                                                                                                                                                                        
import os                                                                                                                                                                                                                                     
import uvicorn                                                                                                                                                                                                                                
from embedder import process_and_embed_document                                                                                                                                                                                               
from QA import process_excel_and_evaluate, analyze_results, answer_question                                                                                                                                                                   
import dotenv                                                                                                                                                                                                                                 
from sqlalchemy import create_engine                                                                                                                                                                                                          
from sqlalchemy.orm import sessionmaker                                                                                                                                                                                                       
from models import User, Company, File                                                                                                                                                                                                        
from passlib.context import CryptContext    
from dotenv import load_dotenv         
from mistralai.client import MistralClient                                                                                                                                                       


load_dotenv()
mistral_client = MistralClient(api_key=os.getenv("MISTRAL_API_KEY"))
# Database setup                                                                                                                                                                                                                              
DATABASE_URL = os.getenv("DATABASE_URL")                                                                                                                                                      
                                                                                                                                                                                              
# Ensure the DATABASE_URL uses 'postgresql://' for psycopg2 compatibility if it's 'postgres://'                                                                                               
if DATABASE_URL and DATABASE_URL.startswith("postgres://"):                                                                                                                                   
    DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://", 1)                                                                                                                    
                                                                                                                                                                                              
engine = create_engine(                                                                                                                                                                       
    DATABASE_URL,                                                                                                                                                                             
    pool_recycle=3600,  # Recycle connections after 1 hour (3600 seconds) to prevent stale connections                                                                                        
    connect_args={"sslmode": "disable"} # Explicitly disable SSL for local development                                                                                                        
)                                                                                                                                                                                             
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)                                                                                                                                                              

SECRET_KEY = os.getenv("SECRET_KEY")
ALGORITHM = "HS256"                                                                                                                                                                                                                           
ACCESS_TOKEN_EXPIRE_MINUTES = 30  


# Password hashing setup                                                                                                                                                                                                                      
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")                                                                                                                                                                             
                                                                                                                                                                                                                                              
# OAuth2 scheme                                                                                                                                                                                                                               
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")                                                                                                                                                                                        
                                                                                                                                                                                                                                              
class TaskManager:                                                                                                                                                                                                                            
    def __init__(self):                                                                                                                                                                                                                       
        self.tasks: Dict[str, Dict] = {}                                                                                                                                                                                                      
                                                                                                                                                                                                                                              
    def create_task(self, task_func, *args, **kwargs):                                                                                                                                                                                        
        task_id = str(uuid4())                                                                                                                                                                                                                
        self.tasks[task_id] = {                                                                                                                                                                                                               
            "status": "processing",                                                                                                                                                                                                           
            "result": None                                                                                                                                                                                                                    
        }                                                                                                                                                                                                                                     
        background_tasks = BackgroundTasks()                                                                                                                                                                                                  
        background_tasks.add_task(self._run_task, task_id, task_func, *args, **kwargs)                                                                                                                                                        
        return task_id, background_tasks                                                                                                                                                                                                      
                                                                                                                                                                                                                                              
    def _run_task(self, task_id, task_func, *args, **kwargs):                                                                                                                                                                                 
        try:                                                                                                                                                                                                                                  
            result = task_func(*args, **kwargs)                                                                                                                                                                                               
            self.tasks[task_id]["status"] = "completed"                                                                                                                                                                                       
            self.tasks[task_id]["result"] = result                                                                                                                                                                                            
        except Exception as e:                                                                                                                                                                                                                
            self.tasks[task_id]["status"] = "failed"                                                                                                                                                                                          
            self.tasks[task_id]["result"] = str(e)                                                                                                                                                                                            
                                                                                                                                                                                                                                              
    def get_task_status(self, task_id):                                                                                                                                                                                                       
        return self.tasks.get(task_id, {"status": "not_found"})                                                                                                                                                                               
                                                                                                                                                                                                                                              
def get_db():                                                                                                                                                                                                                                 
    db = SessionLocal()                                                                                                                                                                                                                       
    try:                                                                                                                                                                                                                                      
        yield db                                                                                                                                                                                                                              
    finally:                                                                                                                                                                                                                                  
        db.close()                                                                                                                                                                                                                            
                                                                                                                                                                                                                                              
def verify_password(plain_password, hashed_password):                                                                                                                                                                                         
    return pwd_context.verify(plain_password, hashed_password)                                                                                                                                                                                
                                                                                                                                                                                                                                              
def get_user(db, username: str):                                                                                                                                                                                                              
    return db.query(User).filter(User.username == username).first()                                                                                                                                                                           
                                                                                                                                                                                                                                              
def authenticate_user(db, username: str, password: str):                                                                                                                                                                                      
    user = get_user(db, username)                                                                                                                                                                                                             
    if not user:                                                                                                                                                                                                                              
        return False                                                                                                                                                                                                                          
    if not verify_password(password, user.password):                                                                                                                                                                                          
        return False                                                                                                                                                                                                                          
    return user                                                                                                                                                                                                                               
                                                                                                                                                                                                                                              
def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):                                                                                                                                                               
    to_encode = data.copy()                                                                                                                                                                                                                   
    if expires_delta:                                                                                                                                                                                                                         
        expire = datetime.utcnow() + expires_delta                                                                                                                                                                                            
    else:                                                                                                                                                                                                                                     
        expire = datetime.utcnow() + timedelta(minutes=15)                                                                                                                                                                                    
    to_encode.update({"exp": expire})                                                                                                                                                                                                         
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)                                                                                                                                                                      
    return encoded_jwt                                                                                                                                                                                                                        
                                                                                                                                                                                                                                              
async def get_current_user(token: str = Depends(oauth2_scheme), db: SessionLocal = Depends(get_db)):                                                                                                                                          
    credentials_exception = HTTPException(                                                                                                                                                                                                    
        status_code=401,                                                                                                                                                                                                                      
        detail="Could not validate credentials",                                                                                                                                                                                              
        headers={"WWW-Authenticate": "Bearer"},                                                                                                                                                                                               
    )                                                                                                                                                                                                                                         
    try:                                                                                                                                                                                                                                      
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])                                                                                                                                                                       
        username: str = payload.get("sub")                                                                                                                                                                                                    
        if username is None:                                                                                                                                                                                                                  
            raise credentials_exception                                                                                                                                                                                                       
    except JWTError:                                                                                                                                                                                                                          
        raise credentials_exception                                                                                                                                                                                                           
                                                                                                                                                                                                                                              
    user = get_user(db, username=username)                                                                                                                                                                                                    
    if user is None:                                                                                                                                                                                                                          
        raise credentials_exception                                                                                                                                                                                                           
    return user                                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                              
def process_files_task(file_paths, company_name):                                                                                                                                                                                             
    try:                                                                                                                                                                                                                                      
        # Process and embed the documents                                                                                                                                                                                                     
        for file_path in file_paths:                                                                                                                                                                                                          
            print(f"Processing file: {file_path}")                                                                                                                                                                                            
            process_and_embed_document(file_path, company_name)                                                                                                                                                                               
                                                                                                                                                                                                                                              
        # Remove the temporary files                                                                                                                                                                                                          
        for file_path in file_paths:                                                                                                                                                                                                          
            os.remove(file_path)                                                                                                                                                                                                              
                                                                                                                                                                                                                                              
        print("All files processed and embedded successfully.")                                                                                                                                                                               
        print("Starting compliance evaluation...")                                                                                                                                                                                            
                                                                                                                                                                                                                                              
        # Process the compliance evaluation                                                                                                                                                                                                   
        # input_file = "/home/abeltagy/Optomatica/FRA_Project/data/expanded_queries_Baselines_optomatica_20250812_101757_copy.csv"                                                                                                                                                            
        input_file = '/app/data/expanded_queries_Baselines_optomatica_20250812_101757.csv'  # Adjusted for Docker                                                                                                                      
        results = process_excel_and_evaluate(input_file, company_name)                                                                                                                                                                        
                                                                                                                                                                                                                                              
        print("Compliance evaluation completed.")                                                                                                                                                                                             
        print(f"Number of results: {len(results)}")                                                                                                                                                                                           
                                                                                                                                                                                                                                              
        # Analyze the results                                                                                                                                                                                                                 
        analysis = analyze_results(results)                                                                                                                                                                                                   
                                                                                                                                                                                                                                              
        # Filter results to only include non-compliant or insufficient information items                                                                                                                                                      
        filtered_results = [                                                                                                                                                                                                                  
            {                                                                                                                                                                                                                                 
                "original_description": r["original_description"],                                                                                                                                                                            
                "compliance_status": r["compliance_status"],                                                                                                                                                                                  
                "recommendations": r["recommendations"]                                                                                                                                                                                       
            }                                                                                                                                                                                                                                 
            for r in results                                                                                                                                                                                                                  
            if r["compliance_status"] in ["Non-Compliant", "Insufficient Information"]                                                                                                                                                        
        ]                                                                                                                                                                                                                                     
                                                                                                                                                                                                                                              
        return {                                                                                                                                                                                                                              
            "results": filtered_results,                                                                                                                                                                                                      
            "analysis": analysis                                                                                                                                                                                                              
        }                                                                                                                                                                                                                                     
    except Exception as e:                                                                                                                                                                                                                    
        raise HTTPException(status_code=500, detail=str(e))                                                                                                                                                                                   
                                                                                                                                                                                                                                              
task_manager = TaskManager()                                                                                                                                                                                                                  
                                                                                                                                                                                                                                              
app = FastAPI(title="FRA API", description="API for document processing")                                                                                                                                                                     
                                                                                                                                                                                                                                              
app.add_middleware(                                                                                                                                                                                                                           
    CORSMiddleware,                                                                                                                                                                                                                           
    allow_origins=["*"],                                                                                                                                                                                               
    allow_credentials=True,                                                                                                                                                                                                                   
    allow_methods=["*"],                                                                                                                                                                                               
    allow_headers=["*"],                                                                                                                                                                                            
)                                                                                                                                                                                                                                             
                                                                                                                                                                                                                                              
# Load environment variables                                                                                                                                                                                                                  
dotenv.load_dotenv()                                                                                                                                                                                                                          
                                                                                                                                                                                                                                              
@app.get("/")                                                                                                                                                                                                                                 
def read_root():                                                                                                                                                                                                                              
    return {"message": "Welcome to the Financial Regulatory Authority API"}                                                                                                                                                                   
                                                                                                                                                                                                                                              
@app.post("/register")                                                                                                                                                                                                                        
async def register_user(username: str = Form(...), password: str = Form(...), db: SessionLocal = Depends(get_db)):                                                                                                                            
    # Check if user already exists                                                                                                                                                                                                            
    db_user = get_user(db, username)                                                                                                                                                                                                          
    if db_user:                                                                                                                                                                                                                               
        raise HTTPException(status_code=400, detail="Username already registered")                                                                                                                                                            
                                                                                                                                                                                                                                              
    # Hash the password                                                                                                                                                                                                                       
    hashed_password = pwd_context.hash(password)                                                                                                                                                                                              
                                                                                                                                                                                                                                              
    # Create new user                                                                                                                                                                                                                         
    new_user = User(username=username, password=hashed_password)                                                                                                                                                                              
    db.add(new_user)                                 

    db.commit()                                                                                                                                                                                                                               
    db.refresh(new_user)                                                                                                                                                                                                                      
                                                                                                                                                                                                                                              
    # Generate access token                                                                                                                                                                                                                   
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)                                                                                                                                                                     
    access_token = create_access_token(                                                                                                                                                                                                       
        data={"sub": new_user.username}, expires_delta=access_token_expires                                                                                                                                                                   
    )                                                                                                                                                                                                                                         
                                                                                                                                                                                                                                              
    return {                                                                                                                                                                                                                                  
        "message": "User created successfully",                                                                                                                                                                                               
        "access_token": access_token,                                                                                                                                                                                                         
        "token_type": "bearer"                                                                                                                                                                                                                
    }                                                                                                                                                                                            
                                                                                                                                                                                                                                              
@app.post("/login")                                                                                                                                                                                                                           
async def login(form_data: OAuth2PasswordRequestForm = Depends(), db: SessionLocal = Depends(get_db)):                                                                                                                                        
    user = authenticate_user(db, form_data.username, form_data.password)                                                                                                                                                                      
    if not user:                                                                                                                                                                                                                              
        raise HTTPException(status_code=401, detail="Incorrect username or password")                                                                                                                                                         
                                                                                                                                                                                                                                              
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)                                                                                                                                                                     
    access_token = create_access_token(                                                                                                                                                                                                       
        data={"sub": user.username}, expires_delta=access_token_expires                                                                                                                                                                       
    )                                                                                                                                                                                                                                         
    return {"access_token": access_token, "token_type": "bearer"}                                                                                                                                                                                                  

@app.post("/createCompany")                                                                                                                                                                                                                      
async def create_company(                                                                                                                                                                                                                     
    company_name: str = Form(...),                                                                                                                                                                                                            
    current_user: User = Depends(get_current_user),                                                                                                                                                                                           
    db: SessionLocal = Depends(get_db)):                                                                                                                                                                                                                                            
    """                                                                                                                                                                                                                                       
    Create a new company and associate it with the current user.                                                                                                                                                                              
    """                                                                                                                                                                                                                                       
    try:                                                                                                                                                                                                                                      
        # Check if company already exists                                                                                                                                                                                                     
        db_company = db.query(Company).filter(Company.name == company_name).first()                                                                                                                                                           
        if db_company:                                                                                                                                                                                                                        
            # If company exists, check if user is already associated                                                                                                                                                                          
            if current_user in db_company.users:                                                                                                                                                                                              
                raise HTTPException(status_code=400, detail="Company already exists and you are already associated with it")                                                                                                                  
            # Associate user with existing company                                                                                                                                                                                            
            db_company.users.append(current_user)                                                                                                                                                                                             
        else:                                                                                                                                                                                                                                 
            # Create new company                                                                                                                                                                                                              
            db_company = Company(name=company_name)                                                                                                                                                                                           
            db_company.users.append(current_user)                                                                                                                                                                                             
            db.add(db_company)                                                                                                                                                                                                                
                                                                                                                                                                                                                                              
        db.commit()                                                                                                                                                                                                                           
        db.refresh(db_company)                                                                                                                                                                                                                
                                                                                                                                                                                                                                              
        return {"message": "Company created/associated successfully", "company_id": db_company.id}                                                                                                                                            
    except Exception as e:                                                                                                                                                                                                                    
        db.rollback()                                                                                                                                                                                                                         
        raise HTTPException(status_code=500, detail=str(e))                                                                                                                                                                                   
                                                                                                                                                                                                                                              
@app.get("/getCompanies")                                                                                                                                                                                                                       
async def get_user_companies(                                                                                                                                                                                                                 
    current_user: User = Depends(get_current_user),                                                                                                                                                                                           
    db: SessionLocal = Depends(get_db)):                                                                                                                                                                                                                                            
    """                                                                                                                                                                                                                                       
    Get all companies associated with the current user.                                                                                                                                                                                       
    """                                                                                                                                                                                                                                       
    try:                                                                                                                                                                                                                                      
        # Get all companies associated with the current user                                                                                                                                                                                  
        companies = db.query(Company).filter(Company.users.any(id=current_user.id)).all()                                                                                                                                                     
                                                                                                                                                                                                                                              
        # Convert companies to a list of dictionaries                                                                                                                                                                                         
        company_list = [{"id": company.id, "name": company.name} for company in companies]                                                                                                                                                    
                                                                                                                                                                                                                                              
        return {"companies": company_list}                                                                                                                                                                                                    
    except Exception as e:                                                                                                                                                                                                                    
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/process_files/")                                                                                                                                                                                                                  
async def process_files(                                                                                                                                                                                                                      
    files: List[UploadFile] = File(),                                                                                                                                                                                                   
    company_name: str = Form(...),                                                                                                                                                                                                            
    background_tasks: BackgroundTasks = BackgroundTasks(),                                                                                                                                                                                    
    current_user: User = Depends(get_current_user)):                                                                                                                                                                                          
    try:                                                                                                                                                                                                                                      
        # Create a temporary directory to store uploaded files                                                                                                                                                                                
        os.makedirs("temp", exist_ok=True)                                                                                                                                                                                                    
                                                                                                                                                                                                                                              
        # Save the uploaded files temporarily                                                                                                                                                                                                 
        file_paths = []                                                                                                                                                                                                                       
        for file in files:                                                                                                                                                                                                                    
            file_path = os.path.join("temp", file.filename)                                                                                                                                                                                   
            with open(file_path, "wb") as buffer:                                                                                                                                                                                             
                buffer.write(await file.read())                                                                                                                                                                                               
            file_paths.append(file_path)                                                                                                                                                                                                      
                                                                                                                                                                                                                                              
        # Create a task to process the files                                                                                                                                                                                                  
        task_id, bg_tasks = task_manager.create_task(                                                                                                                                                                                         
            process_files_task,                                                                                                                                                                                                               
            file_paths,                                                                                                                                                                                                                       
            company_name                                                                                                                                                                                                                      
        )                                                                                                                                                                                                                                     
        background_tasks.add_task(bg_tasks)                                                                                                                                                                                                   
                                                                                                                                                                                                                                              
        return JSONResponse(content={"task_id": task_id})                                                                                                                                                                                     
    except Exception as e:                                                                                                                                                                                                                    
        raise HTTPException(status_code=500, detail=str(e))                                                                                                                                                                                    
                                                                                                                                                                                                                                              
@app.get("/task_status/{task_id}")                                                                                                                                                                                                            
async def get_task_status(task_id: str):                                                                                                                                                                                                      
    task_info = task_manager.get_task_status(task_id)                                                                                                                                                                                         
    return JSONResponse(content=task_info)                                                                                                                                                                                                    
                                                                                                                                                                                                                                              
@app.post("/ask")                                                                                                                                                                                                                             
async def ask_question(query: str = Form(), company_name: str = Form(), results: str = Form(), current_user: User = Depends(get_current_user)):                                                                                               
    try:                                                                                                                                                                                                                                      
        print(f"Received question: {query} for company: {company_name}")                                                                                                                                                                      
        response = answer_question(query, company_name, results)                                                                                                                                                                              
        return {"answer": response}                                                                                                                                                                                                           
    except Exception as e:                                                                                                                                                                                                                    
        raise HTTPException(status_code=500, detail=f"Error processing question: {str(e)}")                                                                                                                                                   
                                                                                                                                                                                                                                              
if __name__ == "__main__":                                                                                                                                                                                                                    
    uvicorn.run(app, host="0.0.0.0", port=8000)