File size: 5,712 Bytes
8d9119f |
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 157 158 |
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from starlette.requests import Request
import os
import time
from pathlib import Path
from typing import Optional, List
# Import your backend modules
from backend.file_handler import save_upload
from backend.extractors import extract_text
from backend.qa_engine import QAEngine
from backend.image_processor import ImageProcessor
from backend.response_formatter import ResponseFormatter
app = FastAPI(
title="Intelligent QA Service",
description="Question answering for documents and images"
)
# Try to set Hugging Face token if available in environment
huggingface_token = os.environ.get("HF_TOKEN")
if huggingface_token:
from huggingface_hub import login
login(token=huggingface_token)
# Initialize models with fallback options
try:
qa_engine = QAEngine(model_name="distilbert-base-cased-distilled-squad") # Use a public model
except Exception as e:
print(f"Error initializing QA engine: {str(e)}")
# Fallback to a simpler implementation if needed
from backend.qa_engine import SimpleQAEngine
qa_engine = SimpleQAEngine()
try:
image_processor = ImageProcessor()
except Exception as e:
print(f"Error initializing Image Processor: {str(e)}")
# Create a fallback image processor if needed
from backend.image_processor import SimpleImageProcessor
image_processor = SimpleImageProcessor()
formatter = ResponseFormatter()
# Mount static files and templates
templates = Jinja2Templates(directory="frontend/templates")
app.mount("/static", StaticFiles(directory="frontend/static"), name="static")
@app.get("/")
async def read_root(request: Request):
"""Render the main page"""
return templates.TemplateResponse("index.html", {"request": request})
@app.post("/api/document-qa")
async def document_qa(
file: UploadFile = File(...),
question: str = Form(...)
):
"""Process document and answer question"""
try:
# Save the uploaded file
file_id, file_name = save_upload(file)
file_path = Path(f"/tmp/uploads/{file_name}")
# Extract text from document
document_text = extract_text(str(file_path))
# Get answer from QA engine
if isinstance(document_text, dict):
# Handle structured document text
# This is a simplistic approach - you'd need to convert the
# structured content to plain text for the QA engine
if "content" in document_text:
if isinstance(document_text["content"], list):
if isinstance(document_text["content"][0], dict):
# Handle docx structure
text = " ".join([p["text"] for p in document_text["content"]])
else:
# Handle txt structure
text = " ".join(document_text["content"])
else:
text = str(document_text["content"])
else:
text = str(document_text)
else:
# Plain text from PDF or PPTX
text = document_text
qa_result = qa_engine.answer_question(text, question)
qa_result["timestamp"] = time.time()
# Format response
response = formatter.format_document_qa_response(qa_result, file.filename)
return JSONResponse(content=response)
except Exception as e:
error_response = formatter.format_error_response(str(e))
return JSONResponse(content=error_response, status_code=error_response["status_code"])
@app.post("/api/image-qa")
async def image_qa(
file: UploadFile = File(...),
question: str = Form(...)
):
"""Process image and answer question"""
try:
print(f"Received image: {file.filename}, size: {file.size}, question: {question}")
# Validate file is an image
if not file.content_type.startswith('image/'):
print(f"Invalid content type: {file.content_type}")
return JSONResponse(
content={"error": "File must be an image", "status_code": 400},
status_code=400
)
# Save the uploaded file
file_id, file_name = save_upload(file)
file_path = Path(f"/tmp/uploads/{file_name}")
print(f"Saved image to: {file_path}")
if not file_path.exists():
print(f"File not saved properly at {file_path}")
return JSONResponse(
content={"error": "File could not be saved", "status_code": 500},
status_code=500
)
# Process the image
vqa_result = image_processor.answer_image_question(str(file_path), question)
vqa_result["timestamp"] = time.time()
# Format response
response = formatter.format_image_qa_response(vqa_result, file.filename)
return JSONResponse(content=response)
except Exception as e:
import traceback
print(f"Error in image_qa: {str(e)}")
print(traceback.format_exc())
error_response = formatter.format_error_response(str(e))
return JSONResponse(content=error_response, status_code=error_response.get("status_code", 500))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860) |