Update utils/database.py

utils/database.py

@@ -176,41 +176,105 @@ def create_tables(conn: sqlite3.Connection) -> None:
         st.error(f"Unexpected error in create_tables: {e}")
         raise
 
-def create_chat_tables(conn):
+def create_chat_tables(conn: sqlite3.Connection) -> None:
     """Create necessary tables for chat management."""
     try:
-        cursor = conn.cursor()
-        
-        # Create chats table
-        cursor.execute('''
-            CREATE TABLE IF NOT EXISTS chats (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                title TEXT NOT NULL,
-                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                collection_id INTEGER,
-                FOREIGN KEY (collection_id) REFERENCES collections (id) ON DELETE SET NULL
-            )
-        ''')
-        
-        # Create chat messages table
-        cursor.execute('''
-            CREATE TABLE IF NOT EXISTS chat_messages (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                chat_id INTEGER NOT NULL,
-                role TEXT NOT NULL,
-                content TEXT NOT NULL,
-                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                metadata TEXT,  -- Store metadata as JSON string
-                FOREIGN KEY (chat_id) REFERENCES chats (id) ON DELETE CASCADE
-            )
-        ''')
-        
-        conn.commit()
-        
+        with conn_lock:
+            cursor = conn.cursor()
+            
+            # Create tags table first
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS document_tags (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    document_id INTEGER NOT NULL,
+                    tag TEXT NOT NULL,
+                    confidence FLOAT,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE,
+                    UNIQUE(document_id, tag)
+                )
+            ''')
+            
+            # Create chats table with collection_id
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS chats (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    title TEXT NOT NULL,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    collection_id INTEGER,
+                    FOREIGN KEY (collection_id) REFERENCES collections (id) ON DELETE SET NULL
+                )
+            ''')
+            
+            # Create chat messages table
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS chat_messages (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    chat_id INTEGER NOT NULL,
+                    role TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    metadata TEXT,
+                    FOREIGN KEY (chat_id) REFERENCES chats (id) ON DELETE CASCADE
+                )
+            ''')
+            
+            conn.commit()
+            
     except sqlite3.Error as e:
         st.error(f"Error creating chat tables: {e}")
         raise
+
+async def generate_document_tags(content: str) -> List[str]:
+    """Generate tags for a document using AI."""
+    try:
+        llm = ChatOpenAI(temperature=0.2, model="gpt-3.5-turbo")
+        
+        prompt = """Analyze the following document content and generate relevant tags/keywords. 
+        Focus on key themes, topics, and important terminology.
+        Return only the tags as a comma-separated list.
+        Content: {content}"""
+        
+        response = await llm.ainvoke(prompt.format(content=content[:2000]))  # Use first 2000 chars
+        tags = [tag.strip() for tag in response.split(',')]
+        return tags
+    except Exception as e:
+        st.error(f"Error generating tags: {e}")
+        return []
+
+def add_document_tags(conn: sqlite3.Connection, document_id: int, tags: List[str]) -> bool:
+    """Add tags to a document."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            for tag in tags:
+                cursor.execute('''
+                    INSERT OR IGNORE INTO document_tags (document_id, tag)
+                    VALUES (?, ?)
+                ''', (document_id, tag))
+            conn.commit()
+            return True
+    except sqlite3.Error as e:
+        st.error(f"Error adding tags: {e}")
+        return False
+
+def get_document_tags(conn: sqlite3.Connection, document_id: int) -> List[str]:
+    """Get all tags for a document."""
+    try:
+        with conn_lock:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT tag FROM document_tags
+                WHERE document_id = ?
+                ORDER BY tag
+            ''', (document_id,))
+            return [row[0] for row in cursor.fetchall()]
+    except sqlite3.Error as e:
+        st.error(f"Error retrieving tags: {e}")
+        return []
+
+        
 def search_documents_in_collection(conn: sqlite3.Connection, collection_id: int, query: str) -> List[Dict]:
     """Search for documents within a collection."""
     try: