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: