beforeLocal

api/database.py

@@ -0,0 +1,44 @@
+import os
+from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, text
+from sqlalchemy.orm import declarative_base, sessionmaker
+from pgvector.sqlalchemy import Vector
+
+POSTGRES_USER = os.environ.get("POSTGRES_USER", "allmail")
+POSTGRES_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "postgres")
+POSTGRES_DB = os.environ.get("POSTGRES_DB", "emails_db")
+DB_HOST = os.environ.get("DB_HOST", "localhost")
+
+DATABASE_URL = f"postgresql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{DB_HOST}/{POSTGRES_DB}"
+
+engine = create_engine(DATABASE_URL)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+Base = declarative_base()
+
+# Gemini gemini-embedding-001 with output_dimensionality=768
+EMBEDDING_DIMENSIONS = 768
+
+class Email(Base):
+    __tablename__ = "emails"
+
+    id = Column(Integer, primary_key=True, index=True)
+    message_id = Column(String, unique=True, index=True)
+    subject = Column(Text)
+    sender = Column(String)
+    date = Column(DateTime)
+    content = Column(Text)
+    embedding = Column(Vector(EMBEDDING_DIMENSIONS))
+
+def init_db():
+    # Install pgvector extension if not exists
+    with engine.connect() as conn:
+        conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
+        conn.commit()
+    Base.metadata.create_all(bind=engine)
+
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()