Update app.py

app.py

@@ -8,6 +8,7 @@ import openai
 from dotenv import load_dotenv
 from fastapi import FastAPI, Request, UploadFile, File, Form
 from fastapi.responses import HTMLResponse, StreamingResponse
+from contextlib import asynccontextmanager
 from fastrtc import (
     AdditionalOutputs,
     AsyncStreamHandler,
@@ -1493,10 +1494,42 @@ class PersonalAssistantDB:
                 )
             """)
             
+            # Check if table exists before trying to check columns
+            cursor = await db.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='user_memories'")
+            table_exists = await cursor.fetchone()
+            
+            if table_exists:
+                # Check if columns exist and add them if they don't (for migration)
+                cursor = await db.execute("PRAGMA table_info(user_memories)")
+                columns = await cursor.fetchall()
+                column_names = [col[1] for col in columns]
+                
+                if 'is_document' not in column_names:
+                    try:
+                        await db.execute("ALTER TABLE user_memories ADD COLUMN is_document BOOLEAN DEFAULT 0")
+                        print("Added is_document column to user_memories table")
+                    except Exception as e:
+                        print(f"Error adding is_document column: {e}")
+                
+                if 'document_name' not in column_names:
+                    try:
+                        await db.execute("ALTER TABLE user_memories ADD COLUMN document_name TEXT")
+                        print("Added document_name column to user_memories table")
+                    except Exception as e:
+                        print(f"Error adding document_name column: {e}")
+            
+            # Commit the schema changes before creating indexes
+            await db.commit()
+            
             # Create indexes for better performance
             await db.execute("CREATE INDEX IF NOT EXISTS idx_memories_category ON user_memories(category)")
             await db.execute("CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id)")
-            await db.execute("CREATE INDEX IF NOT EXISTS idx_memories_document ON user_memories(is_document)")
+            
+            # Create index on is_document column
+            try:
+                await db.execute("CREATE INDEX IF NOT EXISTS idx_memories_document ON user_memories(is_document)")
+            except Exception as e:
+                print(f"Could not create index on is_document: {e}")
             
             await db.commit()
     
@@ -2318,14 +2351,10 @@ stream = Stream(
     time_limit=300 if get_space() else None,
 )
 
-app = FastAPI()
-
-# Mount stream
-stream.mount(app)
-
-# Initialize database on startup
-@app.on_event("startup")
-async def startup_event():
+# Lifespan event handler
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup
     try:
         await PersonalAssistantDB.init()
         print(f"Database initialized at: {DB_PATH}")
@@ -2346,6 +2375,16 @@ async def startup_event():
         # Try to create directory if it doesn't exist
         os.makedirs(PERSISTENT_DIR, exist_ok=True)
         await PersonalAssistantDB.init()
+    
+    yield  # Application runs here
+    
+    # Shutdown (if needed)
+    print("Application shutting down...")
+
+app = FastAPI(lifespan=lifespan)
+
+# Mount stream
+stream.mount(app)
 
 # Intercept offer to capture settings
 @app.post("/webrtc/offer", include_in_schema=False)