Spaces:
Sleeping
Sleeping
File size: 2,078 Bytes
cac5839 |
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 |
from fastapi import APIRouter, UploadFile, File, Query
import os
from PIL import Image
import pytesseract
from pdf2image import convert_from_path
from app.services.vector_store import add_document, query_documents
import uuid
router = APIRouter()
UPLOAD_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'data', 'uploads')
os.makedirs(UPLOAD_DIR, exist_ok=True)
@router.get("/ping")
def ping():
return {"message": "pong"}
@router.post("/upload-document")
def upload_document(file: UploadFile = File(...)):
file_location = os.path.join(UPLOAD_DIR, file.filename)
with open(file_location, "wb") as f:
f.write(file.file.read())
extracted_text = ""
if file.filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
# OCR for images
image = Image.open(file_location)
extracted_text = pytesseract.image_to_string(image)
elif file.filename.lower().endswith('.pdf'):
# OCR for PDFs
try:
pages = convert_from_path(file_location)
text_list = [pytesseract.image_to_string(page) for page in pages]
extracted_text = "\n".join(text_list)
except Exception as e:
extracted_text = f"Error processing PDF: {str(e)}"
else:
extracted_text = "OCR not supported for this file type."
# Save extracted text
text_filename = file.filename + ".txt"
text_path = os.path.join(UPLOAD_DIR, text_filename)
with open(text_path, "w", encoding="utf-8") as tf:
tf.write(extracted_text)
# Store in vector DB
doc_id = str(uuid.uuid4())
add_document(doc_id, extracted_text, metadata={"filename": file.filename, "path": file_location})
return {"filename": file.filename, "saved_to": file_location, "ocr_text_file": text_path, "vector_id": doc_id}
@router.get("/search")
def search_documents(query: str = Query(..., description="Search query"), n_results: int = 3):
results = query_documents(query, n_results)
return results |