AI backed processing

backend/main.py

@@ -44,7 +44,7 @@ def extract_text_from_pdf(file_bytes: bytes) -> str:
         return ""
 
 # Internal modules
-from database import get_db, Email
+from database import get_db, Email, Vendor, Document
 from scheduler import start_scheduler, stop_scheduler
 from mail_service import fetch_and_store_emails
 
@@ -87,7 +87,7 @@ class LoginResponse(BaseModel):
 
 class NERResponse(BaseModel):
     text: str
-
+    file_path: str
 
 
 def extract_text_from_image(file_bytes: bytes) -> str:
@@ -104,13 +104,41 @@ async def extract_text(file: UploadFile = File(...)):
     content = await file.read()
     filename = file.filename.lower()
     
+    # Save file for Vision mode
+    file_path = f"uploads/{file.filename}"
+    with open(file_path, "wb") as f:
+        f.write(content)
+    
     extracted_text = ""
     
     if filename.endswith(".pdf"):
         # Try text extraction first
-        extracted_text = extract_text_from_pdf(content)
-        
-        # If text is empty, it might be a scanned PDF. 
+        with pdfplumber.open(io.BytesIO(content)) as pdf:
+            try:
+                text = ""
+                for page in pdf.pages:
+                    page_text = page.extract_text(layout=True)
+                    if page_text:
+                        text += page_text + "\n"
+                if text.strip():
+                    extracted_text = text.strip()
+            except Exception:
+                pass
+
+        if not extracted_text:
+             try:
+                 # Fallback to pypdf
+                 reader = PdfReader(io.BytesIO(content))
+                 text = ""
+                 for page in reader.pages:
+                    page_text = page.extract_text()
+                    if page_text:
+                        text += page_text + "\n"
+                 extracted_text = text.strip()
+             except:
+                 pass
+
+        # If text is still empty, it might be a scanned PDF. 
         if not extracted_text.strip():
              try:
                  images = convert_from_bytes(content)
@@ -125,7 +153,35 @@ async def extract_text(file: UploadFile = File(...)):
     else:
         raise HTTPException(status_code=400, detail="Unsupported file type")
     
-    return NERResponse(text=extracted_text)
+    return NERResponse(text=extracted_text, file_path=file_path)
+
+# 3. AI Extraction Module
+from llm_service import extract_data
+from pdf2image import convert_from_path
+
+class AITextRequest(BaseModel):
+    text: Optional[str] = None
+    file_path: Optional[str] = None
+    model_type: str = "text"
+
+@app.post("/api/extract/ai")
+def extract_ai_data(request: AITextRequest):
+    final_image_path = request.file_path
+    
+    if request.model_type == "vision" and request.file_path and request.file_path.endswith(".pdf"):
+        # Convert PDF first page to image
+        try:
+            images = convert_from_path(request.file_path)
+            if images:
+                # Save temp image
+                temp_img_path = request.file_path + ".jpg"
+                images[0].save(temp_img_path, "JPEG")
+                final_image_path = temp_img_path
+        except Exception as e:
+            print(f"Error converting PDF for vision: {e}")
+            
+    data = extract_data(text=request.text, image_path=final_image_path, model_type=request.model_type)
+    return data
 
 import zipfile
 import mimetypes
@@ -255,3 +311,36 @@ def sync_emails():
 @app.get("/")
 def read_root():
     return {"message": "OCR Backend API is running"}
+
+class DocumentSaveRequest(BaseModel):
+    vendor_name: str
+    file_path: str
+    model_type: str
+    data: dict
+
+@app.post("/api/documents/save")
+def save_document(request: DocumentSaveRequest, db: Session = Depends(get_db)):
+    # 1. Find or Create Vendor
+    vendor = db.query(Vendor).filter(Vendor.name == request.vendor_name).first()
+    if not vendor:
+        vendor = Vendor(name=request.vendor_name, default_model=request.model_type)
+        db.add(vendor)
+        db.commit()
+        db.refresh(vendor)
+    else:
+        # Update preference
+        vendor.default_model = request.model_type
+        db.commit()
+    
+    # 2. Save Document
+    filename = request.file_path.split('/')[-1]
+    doc = Document(
+        vendor_id=vendor.id,
+        filename=filename,
+        status="verified",
+        processed_data=request.data
+    )
+    db.add(doc)
+    db.commit()
+    
+    return {"message": "Document saved and Vendor preference updated", "vendor_id": vendor.id}