Create database.py

database.py

@@ -0,0 +1,37 @@
+import sqlite3
+
+def initialize_database():
+    """Initialize the SQLite database and create the 'documents' table if it doesn't exist."""
+    # Connect to the SQLite database (or create it if it doesn't exist)
+    conn = sqlite3.connect('dataset.db')
+    cursor = conn.cursor()
+
+    # Create the 'documents' table if it doesn't exist
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS documents (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            text TEXT NOT NULL,
+            topics TEXT
+        )
+    ''')
+
+    # Commit changes and close the connection
+    conn.commit()
+    conn.close()
+
+def save_to_db(chunks, topics=None):
+    """Save chunks to the SQLite database."""
+    # Ensure the database and table are initialized
+    initialize_database()
+
+    # Connect to the database
+    conn = sqlite3.connect('dataset.db')
+    cursor = conn.cursor()
+
+    # Insert chunks into the database
+    for chunk in chunks:
+        cursor.execute('INSERT INTO documents (text, topics) VALUES (?, ?)', (chunk, topics))
+
+    # Commit changes and close the connection
+    conn.commit()
+    conn.close()