im-amrith/Rag-engine-backend
0
1import os2import psycopg23from dotenv import load_dotenv4from sentence_transformers import SentenceTransformer5from typing import List, Dict6import json7import urllib.parse8 9load_dotenv()10 11class RAGEngine:12 def __init__(self, db_url=None):13 self.db_url = db_url or os.getenv("DATABASE_URL")14 if not self.db_url:15 raise ValueError("DATABASE_URL environment variable is not set")16 17 try:18 self.conn = psycopg2.connect(self.db_url)19 except Exception as e:20 print(f"Direct connection failed: {e}, attempting manual parsing...")21 url = urllib.parse.urlparse(self.db_url)22 self.conn = psycopg2.connect(23 dbname=url.path[1:],24 user=url.username,25 password=url.password,26 host=url.hostname,27 port=url.port,28 sslmode='require'29 )30 self.conn.autocommit = True31 self.model = SentenceTransformer('all-MiniLM-L6-v2')32 33 self._init_db()34 35 def _init_db(self):36 with self.conn.cursor() as cur:37 cur.execute("CREATE EXTENSION IF NOT EXISTS vector")38 cur.execute("""39 CREATE TABLE IF NOT EXISTS users (40 id SERIAL PRIMARY KEY,41 email VARCHAR(255) UNIQUE NOT NULL,42 hashed_password VARCHAR(255) NOT NULL,43 created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP44 );45 CREATE TABLE IF NOT EXISTS documents (46 id SERIAL PRIMARY KEY,47 user_id INTEGER REFERENCES users(id),48 content TEXT,49 metadata JSONB,50 embedding vector(384)51 );52 CREATE TABLE IF NOT EXISTS chat_history (53 id SERIAL PRIMARY KEY,54 user_id INTEGER REFERENCES users(id),55 user_message TEXT,56 ai_message TEXT,57 timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP58 );59 """)60 # Auto-migration for existing tables61 try:62 cur.execute("ALTER TABLE documents ADD COLUMN IF NOT EXISTS user_id INTEGER REFERENCES users(id)")63 cur.execute("ALTER TABLE chat_history ADD COLUMN IF NOT EXISTS user_id INTEGER REFERENCES users(id)")64 except Exception as e:65 print(f"Migration warning: {e}")66 67 def create_user(self, email, hashed_password):68 with self.conn.cursor() as cur:69 try:70 cur.execute("INSERT INTO users (email, hashed_password) VALUES (%s, %s) RETURNING id", (email, hashed_password))71 return cur.fetchone()[0]72 except psycopg2.IntegrityError:73 self.conn.rollback()74 return None75 76 def get_user(self, email):77 with self.conn.cursor() as cur:78 cur.execute("SELECT id, email, hashed_password FROM users WHERE email = %s", (email,))79 return cur.fetchone()80 81 def add_document(self, text: str, metadata: Dict, user_id: int):82 embedding = self.model.encode(text).tolist()83 84 with self.conn.cursor() as cur:85 cur.execute("""86 INSERT INTO documents (content, metadata, embedding, user_id)87 VALUES (%s, %s, %s, %s)88 """, (text, json.dumps(metadata), embedding, user_id))89 90 def query(self, query_text: str, user_id: int, n_results: int = 5):91 query_embedding = self.model.encode(query_text).tolist()92 93 with self.conn.cursor() as cur:94 cur.execute("""95 SELECT content, metadata, 1 - (embedding <=> %s::vector) as similarity96 FROM documents97 WHERE user_id = %s98 ORDER BY embedding <=> %s::vector99 LIMIT %s100 """, (query_embedding, user_id, query_embedding, n_results))101 102 rows = cur.fetchall()103 104 results = {105 "documents": [[row[0] for row in rows]],106 "metadatas": [[row[1] for row in rows]],107 "distances": [[1 - row[2] for row in rows]]108 }109 return results110 111 def list_documents(self, user_id: int, limit: int = 100):112 with self.conn.cursor() as cur:113 # Get total unique documents (files)114 cur.execute("SELECT count(DISTINCT metadata->>'source') FROM documents WHERE user_id = %s", (user_id,))115 count = cur.fetchone()[0]116 117 # Get unique documents by source118 cur.execute("""119 SELECT DISTINCT ON (metadata->>'source') 120 id, 121 metadata, 122 left(content, 200) 123 FROM documents 124 WHERE user_id = %s125 LIMIT %s126 """, (user_id, limit,))127 128 rows = cur.fetchall()129 130 docs = []131 for row in rows:132 docs.append({133 "id": row[0],134 "metadata": row[1],135 "preview": row[2] + "..."136 })137 return {"count": count, "documents": docs}138 139 def save_chat(self, user_message: str, ai_message: str, user_id: int):140 with self.conn.cursor() as cur:141 cur.execute("""142 INSERT INTO chat_history (user_message, ai_message, user_id)143 VALUES (%s, %s, %s)144 """, (user_message, ai_message, user_id))145 146 def get_chat_history(self, user_id: int, limit: int = 50):147 with self.conn.cursor() as cur:148 cur.execute("""149 SELECT id, user_message, ai_message, timestamp 150 FROM chat_history 151 WHERE user_id = %s152 ORDER BY timestamp DESC 153 LIMIT %s154 """, (user_id, limit))155 rows = cur.fetchall()156 return [157 {"id": row[0], "user": row[1], "ai": row[2], "timestamp": row[3].isoformat()} 158 for row in rows159 ]160 161 def get_chat_item(self, chat_id: int, user_id: int):162 with self.conn.cursor() as cur:163 cur.execute("""164 SELECT id, user_message, ai_message, timestamp165 FROM chat_history166 WHERE id = %s AND user_id = %s167 """, (chat_id, user_id))168 row = cur.fetchone()169 if row:170 return {"id": row[0], "user": row[1], "ai": row[2], "timestamp": row[3].isoformat()}171 return None172 173 def keep_alive(self):174 try:175 with self.conn.cursor() as cur:176 cur.execute("SELECT 1")177 print("Pinged DB to keep alive")178 except Exception as e:179 print(f"Keep-alive ping failed: {e}")180 181rag_engine = RAGEngine()182