Spaces:
Paused
Paused
File size: 4,258 Bytes
2e91995 1804a7a 2e91995 1804a7a 2e91995 1804a7a 2e91995 1804a7a 2e91995 1804a7a 2e91995 f73c3de 2e91995 f73c3de 2e91995 f73c3de 1804a7a 2e91995 f73c3de 2e91995 1804a7a 2e91995 f73c3de 2e91995 f73c3de 2e91995 f73c3de 1804a7a 2e91995 f73c3de 2e91995 f73c3de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import json
from sqlalchemy import create_engine, text
from src.core.config import Config
from datetime import datetime
class MemoryManager:
def __init__(self):
self.config = Config()
try:
self.engine = create_engine(self.config.DB_URL)
except:
self.engine = None
def get_conn(self): return self.engine.connect()
def save_attachment(self, user_id, workspace_id, filename, filetype, analysis):
try:
with self.get_conn() as conn:
# Get or Create Session
row = conn.execute(text("SELECT id FROM chat_sessions WHERE user_id = :uid ORDER BY last_active DESC LIMIT 1"), {"uid": str(user_id)}).fetchone()
if row:
sid = row[0]
else:
res = conn.execute(text("INSERT INTO chat_sessions (user_id, workspace_id, title) VALUES (:uid, :wid, 'New Chat') RETURNING id"), {"uid": str(user_id), "wid": str(workspace_id)}).fetchone()
sid = res[0]
conn.execute(text("INSERT INTO chat_attachments (session_id, file_name, file_type, analysis_summary) VALUES (:sid, :f, :t, :a)"), {"sid": sid, "f": filename, "t": filetype, "a": analysis})
conn.execute(text("UPDATE chat_sessions SET last_active = CURRENT_TIMESTAMP WHERE id = :sid"), {"sid": sid})
conn.commit()
except Exception as e: print(f"DB Error: {e}")
def add_message(self, user_id, workspace_id, role, content):
try:
with self.get_conn() as conn:
# Simple lookup for demo
sid_res = conn.execute(text("SELECT id FROM chat_sessions WHERE user_id = :uid ORDER BY last_active DESC LIMIT 1"), {"uid": str(user_id)}).fetchone()
sid = sid_res[0] if sid_res else 1
conn.execute(text("INSERT INTO chat_messages (session_id, role, content) VALUES (:sid, :role, :content)"), {"sid": sid, "role": role, "content": str(content)})
conn.commit()
except: pass
def get_context_string(self, user_id, limit=6):
try:
with self.get_conn() as conn:
# 1. Get Messages
rows = conn.execute(text("SELECT m.role, m.content FROM chat_messages m JOIN chat_sessions s ON m.session_id = s.id WHERE s.user_id = :uid ORDER BY m.created_at DESC LIMIT :lim"), {"uid": str(user_id), "lim": limit}).fetchall()
history = "\n".join([f"{r[0]}: {r[1]}" for r in reversed(rows)])
# 2. Get Attachments (SMARTER LOGIC)
# We fetch the *Latest* one separately to mark it as CURRENT
att_rows = conn.execute(text("SELECT a.file_name, a.analysis_summary, a.created_at FROM chat_attachments a JOIN chat_sessions s ON a.session_id = s.id WHERE s.user_id = :uid ORDER BY a.created_at DESC LIMIT 3"), {"uid": str(user_id)}).fetchall()
vision_context = ""
if att_rows:
# The first one is the latest
latest = att_rows[0]
vision_context += f"\n[CURRENTLY LOOKING AT IMAGE]:\nFile: {latest[0]}\nDescription: {latest[1]}\n"
# The rest are history
if len(att_rows) > 1:
vision_context += "\n[PREVIOUS IMAGES (Context only)]:\n"
for r in att_rows[1:]:
vision_context += f"- {r[0]}: {r[1]}\n"
return vision_context + "\n" + history
except Exception as e:
return f"Error: {e}"
def save_workflow(self, workspace_id, name, json_data):
# (Same as before)
with self.get_conn() as conn:
conn.execute(text("INSERT INTO scenarios (workspace_id, name, description, steps, status, created_at) VALUES (:wid, :name, 'AI Generated', :steps, 'active', :time)"),
{"wid": workspace_id, "name": name, "steps": json.dumps(json_data), "time": datetime.now().isoformat()})
conn.commit()
return 1
# Required for server.py imports
def get_user_workspaces(self, uid): return []
|