database.py

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()