Spaces:
Running
Running
File size: 5,966 Bytes
ec96972 e15840d ec96972 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
import os
import warnings
import logging
# Set up cache directory for HuggingFace models
cache_dir = os.path.join(os.getcwd(), ".cache")
os.makedirs(cache_dir, exist_ok=True)
os.environ['HF_HOME'] = cache_dir
os.environ['TRANSFORMERS_CACHE'] = cache_dir
# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['TF_LOGGING_LEVEL'] = 'ERROR'
os.environ['TF_ENABLE_DEPRECATION_WARNINGS'] = '0'
# Suppress specific TensorFlow deprecation warnings
warnings.filterwarnings('ignore', category=DeprecationWarning, module='tensorflow')
logging.getLogger('tensorflow').setLevel(logging.ERROR)
from fastapi import FastAPI, Request, HTTPException, Depends, Header
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from parser import parse_pdf_from_url, parse_pdf_from_file
from embedder import build_faiss_index
from retriever import retrieve_chunks
from llm import query_gemini
import uvicorn
app = FastAPI(title="HackRx Insurance Policy Assistant", version="1.0.0")
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/")
async def root():
return {"message": "HackRx Insurance Policy Assistant API is running!"}
@app.get("/health")
async def health_check():
return {"status": "healthy", "message": "API is ready to process requests"}
class QueryRequest(BaseModel):
documents: str
questions: list[str]
class LocalQueryRequest(BaseModel):
document_path: str
questions: list[str]
def verify_token(authorization: str = Header(None)):
if not authorization or not authorization.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Invalid authorization header")
token = authorization.replace("Bearer ", "")
# For demo purposes, accept any token. In production, validate against a database
if not token:
raise HTTPException(status_code=401, detail="Invalid token")
return token
@app.post("/api/v1/hackrx/run")
async def run_query(request: QueryRequest, token: str = Depends(verify_token)):
try:
print(f"Processing {len(request.questions)} questions...")
text_chunks = parse_pdf_from_url(request.documents)
print(f"Extracted {len(text_chunks)} text chunks from PDF")
index, texts = build_faiss_index(text_chunks)
# Get relevant chunks for all questions at once
all_chunks = set()
for question in request.questions:
top_chunks = retrieve_chunks(index, texts, question)
all_chunks.update(top_chunks)
# Process all questions in a single LLM call
print(f"Processing all {len(request.questions)} questions in batch...")
response = query_gemini(request.questions, list(all_chunks))
# Extract answers from the JSON response
if isinstance(response, dict) and "answers" in response:
answers = response["answers"]
# Ensure we have the right number of answers
while len(answers) < len(request.questions):
answers.append("Not Found")
answers = answers[:len(request.questions)]
else:
# Fallback if response is not in expected format
answers = [response] if isinstance(response, str) else []
# Ensure we have the right number of answers
while len(answers) < len(request.questions):
answers.append("Not Found")
answers = answers[:len(request.questions)]
print(f"Generated {len(answers)} answers")
return { "answers": answers }
except Exception as e:
print(f"Error: {str(e)}")
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
@app.post("/api/v1/hackrx/local")
async def run_local_query(request: LocalQueryRequest):
try:
print(f"Processing local document: {request.document_path}")
print(f"Processing {len(request.questions)} questions...")
# Parse local PDF file
text_chunks = parse_pdf_from_file(request.document_path)
print(f"Extracted {len(text_chunks)} text chunks from local PDF")
index, texts = build_faiss_index(text_chunks)
# Get relevant chunks for all questions at once
all_chunks = set()
for question in request.questions:
top_chunks = retrieve_chunks(index, texts, question)
all_chunks.update(top_chunks)
# Process all questions in a single LLM call
print(f"Processing all {len(request.questions)} questions in batch...")
response = query_gemini(request.questions, list(all_chunks))
# Extract answers from the JSON response
if isinstance(response, dict) and "answers" in response:
answers = response["answers"]
# Ensure we have the right number of answers
while len(answers) < len(request.questions):
answers.append("Not Found")
answers = answers[:len(request.questions)]
else:
# Fallback if response is not in expected format
answers = [response] if isinstance(response, str) else []
# Ensure we have the right number of answers
while len(answers) < len(request.questions):
answers.append("Not Found")
answers = answers[:len(request.questions)]
print(f"Generated {len(answers)} answers")
return { "answers": answers }
except Exception as e:
print(f"Error: {str(e)}")
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
if __name__ == "__main__":
port = int(os.environ.get("PORT", 7860))
uvicorn.run("app:app", host="0.0.0.0", port=port) |