58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
import os
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
|
|
DATABASE_URL = os.environ.get("DATABASE_URL")
|
|
|
|
def get_conn():
|
|
return psycopg2.connect(DATABASE_URL, cursor_factory=RealDictCursor)
|
|
|
|
def init_db():
|
|
with get_conn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS files (
|
|
file_hash TEXT PRIMARY KEY,
|
|
filename TEXT,
|
|
uploaded_at TIMESTAMP DEFAULT NOW()
|
|
);
|
|
""")
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS rag_cache (
|
|
file_hash TEXT PRIMARY KEY,
|
|
chunks_json TEXT,
|
|
faiss_index BYTEA,
|
|
vectors BYTEA,
|
|
meta_json TEXT,
|
|
created_at TIMESTAMP DEFAULT NOW()
|
|
);
|
|
""")
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS sessions (
|
|
session_id TEXT PRIMARY KEY,
|
|
file_hashes TEXT[],
|
|
doc_type TEXT,
|
|
created_at TIMESTAMP DEFAULT NOW()
|
|
);
|
|
""")
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS outlines (
|
|
session_id TEXT PRIMARY KEY,
|
|
outline_text TEXT,
|
|
created_at TIMESTAMP DEFAULT NOW()
|
|
);
|
|
""")
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS briefings (
|
|
session_id TEXT PRIMARY KEY,
|
|
briefing_json TEXT,
|
|
html TEXT,
|
|
created_at TIMESTAMP DEFAULT NOW()
|
|
);
|
|
""")
|
|
conn.commit()
|
|
print("[DB]
|
|
珥湲고
|
|
猷", flush=True)
|