Spaces:
Paused
Paused
| 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 [] | |