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)
+        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 text is empty, it might be a scanned PDF. 
+        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

frontend/src/app/ocr.service.ts

@@ -15,4 +15,12 @@ export class OcrService {
     formData.append('file', file);
     return this.http.post(`${this.apiUrl}/extract`, formData);
   }
+
+  extractWithAI(text: string, filePath: string | null, modelType: string): Observable<any> {
+    return this.http.post(`http://localhost:8000/api/extract/ai`, { text, file_path: filePath, model_type: modelType });
+  }
+
+  saveDocument(data: any): Observable<any> {
+    return this.http.post(`http://localhost:8000/api/documents/save`, data);
+  }
 }

frontend/src/app/ocr/ocr.component.ts

@@ -9,6 +9,11 @@ import { FileUploadModule } from 'primeng/fileupload';
 import { ProgressBarModule } from 'primeng/progressbar';
 import { InputTextareaModule } from 'primeng/inputtextarea';
 import { ToastModule } from 'primeng/toast';
+import { ButtonModule } from 'primeng/button';
+import { TableModule } from 'primeng/table';
+import { CardModule } from 'primeng/card';
+import { RadioButtonModule } from 'primeng/radiobutton';
+import { InputTextModule } from 'primeng/inputtext';
 
 @Component({
   selector: 'app-ocr',
@@ -19,16 +24,18 @@ import { ToastModule } from 'primeng/toast';
     FileUploadModule, 
     ProgressBarModule, 
     InputTextareaModule, 
-    ToastModule
+    ToastModule,
+    ButtonModule, 
+    TableModule,
+    CardModule,
+    RadioButtonModule,
+    InputTextModule
   ],
   providers: [MessageService],
   template: `
     <div class="card">
         <h2>OCR Extraction</h2>
-        <!-- 
-           Note: "customUpload" mode in PrimeNG FileUpload requires "uploadHandler".
-           "mode='advanced'" gives the sleek UI.
-        -->
+        
         <p-fileUpload mode="advanced" 
                       chooseLabel="Select PDF or Image"
                       uploadLabel="Extract Text"
@@ -44,15 +51,120 @@ import { ToastModule } from 'primeng/toast';
             <p-progressBar mode="indeterminate" [style]="{'height': '6px'}"></p-progressBar>
         </div>
 
-        <div class="mt-4" *ngIf="extractedText !== null">
-            <h3>Extracted Text Result:</h3>
-            <textarea pInputTextarea 
-                      [autoResize]="true" 
-                      [(ngModel)]="extractedText" 
-                      readonly 
-                      class="w-full" 
-                      style="min-height: 300px; width: 100%; border-color: #d1d5db; font-family: monospace;">
-            </textarea>
+        <div class="mt-4 grid" *ngIf="extractedText !== null">
+            <div class="col-12 md:col-6">
+                <h3>Extracted Text Result:</h3>
+                <textarea pInputTextarea 
+                          [autoResize]="true" 
+                          [(ngModel)]="extractedText" 
+                          readonly 
+                          class="w-full" 
+                          style="min-height: 300px; width: 100%; border-color: #d1d5db; font-family: monospace;">
+                </textarea>
+                
+                <div class="mt-3">
+                    <div class="flex flex-column gap-2 mb-3">
+                        <label>AI Analysis Mode:</label>
+                        <div class="flex align-items-center">
+                            <p-radioButton name="model" value="text" [(ngModel)]="modelType" inputId="mod1"></p-radioButton>
+                            <label for="mod1" class="ml-2">Text Analysis (Fast - Gemma)</label>
+                        </div>
+                        <div class="flex align-items-center">
+                            <p-radioButton name="model" value="vision" [(ngModel)]="modelType" inputId="mod2"></p-radioButton>
+                            <label for="mod2" class="ml-2">Vision Analysis (Accurate - Qwen)</label>
+                        </div>
+                    </div>
+                
+                    <p-button label="Process with AI" 
+                              icon="pi pi-bolt" 
+                              [loading]="aiLoading" 
+                              (onClick)="processWithAI()">
+                    </p-button>
+                </div>
+            </div>
+            
+            <div class="col-12 md:col-6" *ngIf="aiResult">
+                <div class="flex justify-content-between align-items-center">
+                    <h3>AI Analysis Result:</h3>
+                    <p-button label="Save & Verify" icon="pi pi-check" styleClass="p-button-success" [loading]="saveLoading" (onClick)="saveDocument()"></p-button>
+                </div>
+                
+                <p-card class="mb-3">
+                    <div class="grid">
+                        <div class="col-6">
+                            <label class="block text-sm font-bold mb-1">Vendor</label>
+                            <input pInputText [(ngModel)]="aiResult.vendor_name" class="w-full" />
+                        </div>
+                        <div class="col-6">
+                            <label class="block text-sm font-bold mb-1">Date</label>
+                            <input pInputText [(ngModel)]="aiResult.date" class="w-full" />
+                        </div>
+                        <div class="col-6 mt-2">
+                            <label class="block text-sm font-bold mb-1">Invoice #</label>
+                            <input pInputText [(ngModel)]="aiResult.invoice_number" class="w-full" />
+                        </div>
+                        <div class="col-6 mt-2">
+                            <label class="block text-sm font-bold mb-1">Total</label>
+                            <input pInputText [(ngModel)]="aiResult.total_amount" class="w-full" />
+                        </div>
+                    </div>
+                </p-card>
+                
+                <p-table [value]="aiResult.line_items" styleClass="p-datatable-sm" [scrollable]="true" scrollHeight="200px">
+                    <ng-template pTemplate="header">
+                        <tr>
+                            <th>Description</th>
+                            <th>Qty</th>
+                            <th>Price</th>
+                            <th>Total</th>
+                        </tr>
+                    </ng-template>
+                    <ng-template pTemplate="body" let-item>
+                        <tr>
+                            <td pEditableColumn>
+                                <p-cellEditor>
+                                    <ng-template pTemplate="input">
+                                        <input pInputText type="text" [(ngModel)]="item.description">
+                                    </ng-template>
+                                    <ng-template pTemplate="output">
+                                        {{item.description}}
+                                    </ng-template>
+                                </p-cellEditor>
+                            </td>
+                            <td pEditableColumn>
+                                <p-cellEditor>
+                                    <ng-template pTemplate="input">
+                                        <input pInputText type="text" [(ngModel)]="item.quantity">
+                                    </ng-template>
+                                    <ng-template pTemplate="output">
+                                        {{item.quantity}}
+                                    </ng-template>
+                                </p-cellEditor>
+                            </td>
+                            <td pEditableColumn>
+                                <p-cellEditor>
+                                    <ng-template pTemplate="input">
+                                        <input pInputText type="text" [(ngModel)]="item.unit_price">
+                                    </ng-template>
+                                    <ng-template pTemplate="output">
+                                        {{item.unit_price}}
+                                    </ng-template>
+                                </p-cellEditor>
+                            </td>
+                            <td pEditableColumn>
+                                <p-cellEditor>
+                                    <ng-template pTemplate="input">
+                                        <input pInputText type="text" [(ngModel)]="item.total">
+                                    </ng-template>
+                                    <ng-template pTemplate="output">
+                                        {{item.total}}
+                                    </ng-template>
+                                </p-cellEditor>
+                            </td>
+                        </tr>
+                    </ng-template>
+                </p-table>
+            </div>
         </div>
         <p-toast></p-toast>
     </div>
@@ -66,15 +178,26 @@ export class OcrComponent {
   extractedText: string | null = null;
   loading: boolean = false;
   
+  aiLoading: boolean = false;
+  saveLoading: boolean = false;
+  aiResult: any = null;
+  
+  // Hybrid AI Props
+  modelType: string = 'text';
+  filePath: string | null = null;
+
   constructor(private ocrService: OcrService, private messageService: MessageService) {}
 
   onUpload(event: any) {
     this.loading = true;
+    this.aiResult = null; // Reset AI result on new upload
+    this.filePath = null;
     const file = event.files[0];
     
     this.ocrService.extractText(file).subscribe({
         next: (res) => {
             this.extractedText = res.text;
+            this.filePath = res.file_path;
             this.loading = false;
             this.messageService.add({severity:'success', summary:'Success', detail:'Text Extracted Successfully'});
         },
@@ -88,5 +211,50 @@ export class OcrComponent {
 
   onClear() {
       this.extractedText = null;
+      this.aiResult = null;
+      this.filePath = null;
+  }
+  
+  processWithAI() {
+      if (!this.extractedText) return;
+      
+      this.aiLoading = true;
+      // Pass text, filePath, and modelType
+      this.ocrService.extractWithAI(this.extractedText, this.filePath, this.modelType).subscribe({
+          next: (res) => {
+              this.aiResult = res;
+              this.aiLoading = false;
+              this.messageService.add({severity:'success', summary:'AI Processing Complete', detail:'Data Extracted'});
+          },
+          error: (err) => {
+             console.error(err);
+             this.aiLoading = false;
+             this.messageService.add({severity:'error', summary:'AI Error', detail:'Could not process with AI'});
+          }
+      });
+  }
+  
+  saveDocument() {
+      if (!this.aiResult || !this.filePath) return;
+      
+      this.saveLoading = true;
+      const payload = {
+          vendor_name: this.aiResult.vendor_name || 'Unknown Vendor',
+          file_path: this.filePath,
+          model_type: this.modelType,
+          data: this.aiResult
+      };
+      
+      this.ocrService.saveDocument(payload).subscribe({
+          next: (res) => {
+              this.saveLoading = false;
+              this.messageService.add({severity:'success', summary:'Saved & Verified', detail:'Document and rules saved'});
+          },
+          error: (err) => {
+              console.error(err);
+              this.saveLoading = false;
+              this.messageService.add({severity:'error', summary:'Save Error', detail:'Failed to save document'});
+          }
+      });
   }
 }