app.py

from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn
import os
from dotenv import load_dotenv
from typing import List, Optional
import json

# Load environment variables
load_dotenv()

# Initialize FastAPI app
app = FastAPI(title="EY Research Assistant API")

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

@app.get("/")
async def root():
    return {"message": "Welcome to EY Research Assistant API"}

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

import shutil
from utils.document_processor import process_document
from utils.vector_store import VectorStore
from utils.ai_service import AIService
from utils.database import Database
import uuid
import datetime

# Initialize services
try:
    vector_store = VectorStore()
    ai_service = AIService()
    db = Database()
except Exception as e:
    print(f"Error initializing services: {e}")

@app.post("/upload-document")
async def upload_document(
    file: UploadFile = File(...),
    user_id: str = Form(None),
    document_type: str = Form(None)
):
    try:
        # Create a unique ID for the document
        document_id = str(uuid.uuid4())
        
        # Save the file to disk
        file_location = f"uploads/{document_id}_{file.filename}"
        with open(file_location, "wb") as file_object:
            shutil.copyfileobj(file.file, file_object)
        
        # Process the document to extract text
        document_text = process_document(file_location)
        
        if not document_text:
            return JSONResponse(
                status_code=400,
                content={"detail": "Could not extract text from document"}
            )
        
        # Split the text into chunks
        chunks = ai_service.text_splitter.split_text(document_text)
        
        # Get embeddings for chunks
        embeddings = ai_service.get_embeddings(chunks)
        
        # Prepare vectors for Pinecone
        vectors = []
        for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
            vectors.append({
                "id": f"{document_id}_{i}",
                "values": embedding,
                "metadata": {
                    "document_id": document_id,
                    "chunk_index": i,
                    "text": chunk,
                    "source": file.filename
                }
            })
        
        # Upsert vectors to Pinecone
        vector_store.upsert_vectors(vectors)
        
        # Store document metadata in Supabase
        document_data = {
            "id": document_id,
            "filename": file.filename,
            "user_id": user_id,
            "document_type": document_type,
            "upload_date": datetime.datetime.now().isoformat(),
            "chunk_count": len(chunks)
        }
        
        db.insert_document(document_data)
        
        return {
            "document_id": document_id,
            "filename": file.filename,
            "chunk_count": len(chunks),
            "message": "Document processed successfully"
        }
        
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"detail": f"Error processing document: {str(e)}"}
        )

@app.post("/query")
async def query(
    question: str = Form(...),
    user_id: str = Form(None),
    model: str = Form("gpt-4")  # Options: gpt-4, claude
):
    try:
        # Store query in database
        query_data = {
            "id": str(uuid.uuid4()),
            "user_id": user_id,
            "question": question,
            "model": model,
            "created_at": datetime.datetime.now().isoformat()
        }
        
        db.insert_query(query_data)
        
        # Generate embedding for the question
        question_embedding = ai_service.get_embeddings([question])[0]
        
        # Query Pinecone for relevant chunks
        query_response = vector_store.query(question_embedding)
        
        if not query_response.matches:
            return {"answer": "I couldn't find relevant information to answer your question."}
        
        # Extract the text from the matched chunks
        context = "\n\n".join([match.metadata["text"] for match in query_response.matches])
        
        # Generate the answer based on the model choice
        if model == "claude" and ai_service.anthropic_client:
            prompt = f"""
            Question: {question}
            
            Context information:
            {context}
            
            Please answer the question based only on the provided context. If the answer cannot be determined from the context, say so.
            """
            
            answer = ai_service.query_anthropic(prompt)
        else:
            # Use OpenAI via langchain
            prompt = f"""
            Answer the following question based on the context provided. If the answer cannot be determined from the context, say so.
            
            Question: {question}
            
            Context:
            {context}
            """
            
            answer = ai_service.openai_llm.predict(prompt)
        
        # Update query with the answer
        db.supabase.table("queries").update({"answer": answer}).eq("id", query_data["id"]).execute()
        
        return {
            "query_id": query_data["id"],
            "answer": answer,
            "sources": [{"document": match.metadata["source"], "relevance": match.score} for match in query_response.matches[:3]]
        }
        
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"detail": f"Error processing query: {str(e)}"}
        )

@app.get("/documents")
async def get_documents(user_id: Optional[str] = None):
    """Get all documents, optionally filtered by user_id."""
    try:
        response = db.get_documents(user_id)
        return {"documents": response.data}
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"detail": f"Error getting documents: {str(e)}"}
        )

@app.get("/queries")
async def get_queries(user_id: Optional[str] = None, limit: int = 10):
    """Get recent queries, optionally filtered by user_id."""
    try:
        response = db.get_queries(user_id, limit)
        return {"queries": response.data}
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"detail": f"Error getting queries: {str(e)}"}
        )

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