AI backed processing

backend/database.py

@@ -1,23 +1,15 @@
 import os
+import urllib.parse
 from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, Boolean, ForeignKey, LargeBinary
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, relationship
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.sql import func
 from dotenv import load_dotenv
 
 # Load environment variables
 load_dotenv()
 
-DB_USER = os.getenv("DB_USER")
-DB_PASSWORD = os.getenv("DB_PASSWORD")
-DB_HOST = os.getenv("DB_HOST")
-DB_PORT = os.getenv("DB_PORT")
-import urllib.parse
-
-# ... (imports)
-
-# Load environment variables
-load_dotenv()
-
 DB_USER = os.getenv("DB_USER")
 DB_PASSWORD = os.getenv("DB_PASSWORD")
 DB_HOST = os.getenv("DB_HOST")
@@ -54,10 +46,33 @@ class Attachment(Base):
     email_id = Column(Integer, ForeignKey("emails.id"))
     filename = Column(String)
     content_type = Column(String)
-    file_content = Column(LargeBinary) # Storing content directly in DB as requested
-    
+    file_path = Column(String, nullable=True) # Path to file on disk
+    file_content = Column(LargeBinary, nullable=True) # Stored in DB (for small files)
+
     email = relationship("Email", back_populates="attachments")
 
+class Vendor(Base):
+    __tablename__ = "vendors"
+    
+    id = Column(Integer, primary_key=True, index=True)
+    name = Column(String, unique=True, index=True)
+    default_model = Column(String, default="text") # 'text' (Gemma) or 'vision' (Qwen)
+    created_at = Column(DateTime(timezone=True), server_default=func.now())
+    
+    documents = relationship("Document", back_populates="vendor")
+
+class Document(Base):
+    __tablename__ = "documents"
+    
+    id = Column(Integer, primary_key=True, index=True)
+    vendor_id = Column(Integer, ForeignKey("vendors.id"), nullable=True)
+    filename = Column(String)
+    upload_date = Column(DateTime(timezone=True), server_default=func.now())
+    status = Column(String, default="pending") # pending, verified
+    processed_data = Column(JSONB) # The final verified JSON
+    
+    vendor = relationship("Vendor", back_populates="documents")
+
 def get_db():
     db = SessionLocal()
     try:

backend/llm_service.py

@@ -0,0 +1,54 @@
+import ollama
+import json
+import base64
+
+def extract_data(text: str = None, image_path: str = None, model_type: str = "text") -> dict:
+    """
+    Extracts structured data using either Text (Gemma) or Vision (Qwen) models.
+    """
+    
+    prompt = """
+    You are an expert data extraction assistant. 
+    Extract the following fields from the provided document and return them as a SINGLE VALID JSON OBJECT:
+    - invoice_number (string)
+    - date (string)
+    - vendor_name (string)
+    - total_amount (string)
+    - currency (string)
+    - line_items (list of objects with: description, quantity, unit_price, total)
+
+    IMPORTANT: 
+    - Return ONLY the JSON. No markdown formatting, no explanations. 
+    - If a field is not found, use null.
+    """
+    
+    messages = [{'role': 'user', 'content': prompt}]
+    model = 'gemma:2b'
+    
+    if model_type == 'vision':
+        if not image_path:
+            return {"error": "Image path required for vision mode"}
+        
+        # Qwen-VL handles images passed in the message
+        model = 'qwen2.5vl:7b' # Using the installed model ID
+        messages[0]['images'] = [image_path]
+        messages[0]['content'] = "Analyze this image. " + prompt
+    else:
+        # Text Mode
+        if not text:
+            return {"error": "Text required for text mode"}
+        messages[0]['content'] += f"\n\n---\n{text}\n---"
+
+    try:
+        response = ollama.chat(model=model, messages=messages)
+        content = response['message']['content']
+        
+        # Clean up markdown
+        content = content.replace("```json", "").replace("```", "").strip()
+        
+        return json.loads(content)
+
+    except Exception as e:
+        print(f"LLM Extraction Error ({model_type}): {e}")
+        return {"error": str(e), "raw_output": content if 'content' in locals() else ""}
+

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}

backend/requirements.txt

@@ -10,3 +10,4 @@ imap-tools
 apscheduler
 python-dotenv
 pdfplumber
+ollama