App_Function_Libraries/DB/Prompts_DB.py

# Prompts_DB.py
# Description: Functions to manage the prompts database.
#
# Imports
import sqlite3
import logging
#
# External Imports
import re
from typing import Tuple
#
# Local Imports
from App_Function_Libraries.Utils.Utils import get_database_path
#
#######################################################################################################################
#
# Functions to manage prompts DB

def create_prompts_db():
    logging.debug("create_prompts_db: Creating prompts database.")
    with sqlite3.connect(get_database_path('prompts.db')) as conn:
        cursor = conn.cursor()
        cursor.executescript('''

            CREATE TABLE IF NOT EXISTS Prompts (

                id INTEGER PRIMARY KEY AUTOINCREMENT,

                name TEXT NOT NULL UNIQUE,

                author TEXT,

                details TEXT,

                system TEXT,

                user TEXT

            );

            CREATE TABLE IF NOT EXISTS Keywords (

                id INTEGER PRIMARY KEY AUTOINCREMENT,

                keyword TEXT NOT NULL UNIQUE COLLATE NOCASE

            );

            CREATE TABLE IF NOT EXISTS PromptKeywords (

                prompt_id INTEGER,

                keyword_id INTEGER,

                FOREIGN KEY (prompt_id) REFERENCES Prompts (id),

                FOREIGN KEY (keyword_id) REFERENCES Keywords (id),

                PRIMARY KEY (prompt_id, keyword_id)

            );

            CREATE INDEX IF NOT EXISTS idx_keywords_keyword ON Keywords(keyword);

            CREATE INDEX IF NOT EXISTS idx_promptkeywords_prompt_id ON PromptKeywords(prompt_id);

            CREATE INDEX IF NOT EXISTS idx_promptkeywords_keyword_id ON PromptKeywords(keyword_id);

        ''')

# FIXME - dirty hack that should be removed later...
# Migration function to add the 'author' column to the Prompts table
def add_author_column_to_prompts():
    with sqlite3.connect(get_database_path('prompts.db')) as conn:
        cursor = conn.cursor()
        # Check if 'author' column already exists
        cursor.execute("PRAGMA table_info(Prompts)")
        columns = [col[1] for col in cursor.fetchall()]

        if 'author' not in columns:
            # Add the 'author' column
            cursor.execute('ALTER TABLE Prompts ADD COLUMN author TEXT')
            print("Author column added to Prompts table.")
        else:
            print("Author column already exists in Prompts table.")

add_author_column_to_prompts()

def normalize_keyword(keyword):
    return re.sub(r'\s+', ' ', keyword.strip().lower())


# FIXME - update calls to this function to use the new args
def add_prompt(name, author, details, system=None, user=None, keywords=None):
    logging.debug(f"add_prompt: Adding prompt with name: {name}, author: {author}, system: {system}, user: {user}, keywords: {keywords}")
    if not name:
        logging.error("add_prompt: A name is required.")
        return "A name is required."

    try:
        with sqlite3.connect(get_database_path('prompts.db')) as conn:
            cursor = conn.cursor()
            cursor.execute('''

                INSERT INTO Prompts (name, author, details, system, user)

                VALUES (?, ?, ?, ?, ?)

            ''', (name, author, details, system, user))
            prompt_id = cursor.lastrowid

            if keywords:
                normalized_keywords = [normalize_keyword(k) for k in keywords if k.strip()]
                for keyword in set(normalized_keywords):  # Use set to remove duplicates
                    cursor.execute('''

                        INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)

                    ''', (keyword,))
                    cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,))
                    keyword_id = cursor.fetchone()[0]
                    cursor.execute('''

                        INSERT OR IGNORE INTO PromptKeywords (prompt_id, keyword_id) VALUES (?, ?)

                    ''', (prompt_id, keyword_id))
        return "Prompt added successfully."
    except sqlite3.IntegrityError:
        return "Prompt with this name already exists."
    except sqlite3.Error as e:
        return f"Database error: {e}"


def fetch_prompt_details(name):
    logging.debug(f"fetch_prompt_details: Fetching details for prompt: {name}")
    with sqlite3.connect(get_database_path('prompts.db')) as conn:
        cursor = conn.cursor()
        cursor.execute('''

            SELECT p.name, p.author, p.details, p.system, p.user, GROUP_CONCAT(k.keyword, ', ') as keywords

            FROM Prompts p

            LEFT JOIN PromptKeywords pk ON p.id = pk.prompt_id

            LEFT JOIN Keywords k ON pk.keyword_id = k.id

            WHERE p.name = ?

            GROUP BY p.id

        ''', (name,))
        return cursor.fetchone()


def list_prompts(page=1, per_page=10):
    logging.debug(f"list_prompts: Listing prompts for page {page} with {per_page} prompts per page.")
    offset = (page - 1) * per_page
    with sqlite3.connect(get_database_path('prompts.db')) as conn:
        cursor = conn.cursor()
        cursor.execute('SELECT name FROM Prompts LIMIT ? OFFSET ?', (per_page, offset))
        prompts = [row[0] for row in cursor.fetchall()]

        # Get total count of prompts
        cursor.execute('SELECT COUNT(*) FROM Prompts')
        total_count = cursor.fetchone()[0]

    total_pages = (total_count + per_page - 1) // per_page
    return prompts, total_pages, page


def insert_prompt_to_db(title, author, description, system_prompt, user_prompt, keywords=None):
    return add_prompt(title, author, description, system_prompt, user_prompt, keywords)


def get_prompt_db_connection():
    prompt_db_path = get_database_path('prompts.db')
    return sqlite3.connect(prompt_db_path)


def search_prompts(query):
    logging.debug(f"search_prompts: Searching prompts with query: {query}")
    try:
        with get_prompt_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""

                SELECT p.name, p.details, p.system, p.user, GROUP_CONCAT(k.keyword, ', ') as keywords

                FROM Prompts p

                LEFT JOIN PromptKeywords pk ON p.id = pk.prompt_id

                LEFT JOIN Keywords k ON pk.keyword_id = k.id

                WHERE p.name LIKE ? OR p.details LIKE ? OR p.system LIKE ? OR p.user LIKE ? OR k.keyword LIKE ?

                GROUP BY p.id

                ORDER BY p.name

            """, (f'%{query}%', f'%{query}%', f'%{query}%', f'%{query}%', f'%{query}%'))
            return cursor.fetchall()
    except sqlite3.Error as e:
        logging.error(f"Error searching prompts: {e}")
        return []


def search_prompts_by_keyword(keyword, page=1, per_page=10):
    logging.debug(f"search_prompts_by_keyword: Searching prompts by keyword: {keyword}")
    normalized_keyword = normalize_keyword(keyword)
    offset = (page - 1) * per_page
    with sqlite3.connect(get_database_path('prompts.db')) as conn:
        cursor = conn.cursor()
        cursor.execute('''

            SELECT DISTINCT p.name

            FROM Prompts p

            JOIN PromptKeywords pk ON p.id = pk.prompt_id

            JOIN Keywords k ON pk.keyword_id = k.id

            WHERE k.keyword LIKE ?

            LIMIT ? OFFSET ?

        ''', ('%' + normalized_keyword + '%', per_page, offset))
        prompts = [row[0] for row in cursor.fetchall()]

        # Get total count of matching prompts
        cursor.execute('''

            SELECT COUNT(DISTINCT p.id)

            FROM Prompts p

            JOIN PromptKeywords pk ON p.id = pk.prompt_id

            JOIN Keywords k ON pk.keyword_id = k.id

            WHERE k.keyword LIKE ?

        ''', ('%' + normalized_keyword + '%',))
        total_count = cursor.fetchone()[0]

    total_pages = (total_count + per_page - 1) // per_page
    return prompts, total_pages, page


def update_prompt_keywords(prompt_name, new_keywords):
    logging.debug(f"update_prompt_keywords: Updating keywords for prompt: {prompt_name}")
    try:
        with sqlite3.connect(get_database_path('prompts.db')) as conn:
            cursor = conn.cursor()

            cursor.execute('SELECT id FROM Prompts WHERE name = ?', (prompt_name,))
            prompt_id = cursor.fetchone()
            if not prompt_id:
                return "Prompt not found."
            prompt_id = prompt_id[0]

            cursor.execute('DELETE FROM PromptKeywords WHERE prompt_id = ?', (prompt_id,))

            normalized_keywords = [normalize_keyword(k) for k in new_keywords if k.strip()]
            for keyword in set(normalized_keywords):  # Use set to remove duplicates
                cursor.execute('INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)', (keyword,))
                cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,))
                keyword_id = cursor.fetchone()[0]
                cursor.execute('INSERT INTO PromptKeywords (prompt_id, keyword_id) VALUES (?, ?)',
                               (prompt_id, keyword_id))

            # Remove unused keywords
            cursor.execute('''

                DELETE FROM Keywords

                WHERE id NOT IN (SELECT DISTINCT keyword_id FROM PromptKeywords)

            ''')
        return "Keywords updated successfully."
    except sqlite3.Error as e:
        return f"Database error: {e}"


def add_or_update_prompt(title, author, description, system_prompt, user_prompt, keywords=None):
    logging.debug(f"add_or_update_prompt: Adding or updating prompt: {title}")
    if not title:
        return "Error: Title is required."

    existing_prompt = fetch_prompt_details(title)
    if existing_prompt:
        # Update existing prompt
        result = update_prompt_in_db(title, author, description, system_prompt, user_prompt)
        if "successfully" in result:
            # Update keywords if the prompt update was successful
            keyword_result = update_prompt_keywords(title, keywords or [])
            result += f" {keyword_result}"
    else:
        # Insert new prompt
        result = insert_prompt_to_db(title, author, description, system_prompt, user_prompt, keywords)

    return result


def load_prompt_details(selected_prompt):
    logging.debug(f"load_prompt_details: Loading prompt details for {selected_prompt}")
    if selected_prompt:
        details = fetch_prompt_details(selected_prompt)
        if details:
            return details[0], details[1], details[2], details[3], details[4], details[5]
    return "", "", "", "", "", ""


def update_prompt_in_db(title, author, description, system_prompt, user_prompt):
    logging.debug(f"update_prompt_in_db: Updating prompt: {title}")
    try:
        with sqlite3.connect(get_database_path('prompts.db')) as conn:
            cursor = conn.cursor()
            cursor.execute(
                "UPDATE Prompts SET author = ?, details = ?, system = ?, user = ? WHERE name = ?",
                (author, description, system_prompt, user_prompt, title)
            )
            if cursor.rowcount == 0:
                return "No prompt found with the given title."
        return "Prompt updated successfully!"
    except sqlite3.Error as e:
        return f"Error updating prompt: {e}"


def delete_prompt(prompt_id):
    logging.debug(f"delete_prompt: Deleting prompt with ID: {prompt_id}")
    try:
        with sqlite3.connect(get_database_path('prompts.db')) as conn:
            cursor = conn.cursor()

            # Delete associated keywords
            cursor.execute("DELETE FROM PromptKeywords WHERE prompt_id = ?", (prompt_id,))

            # Delete the prompt
            cursor.execute("DELETE FROM Prompts WHERE id = ?", (prompt_id,))

            if cursor.rowcount == 0:
                return f"No prompt found with ID {prompt_id}"
            else:
                conn.commit()
                return f"Prompt with ID {prompt_id} has been successfully deleted"
    except sqlite3.Error as e:
        return f"An error occurred: {e}"


def delete_prompt_keyword(keyword: str) -> str:
    """

    Delete a keyword and its associations from the prompts database.



    Args:

        keyword (str): The keyword to delete



    Returns:

        str: Success/failure message

    """
    logging.debug(f"delete_prompt_keyword: Deleting keyword: {keyword}")
    try:
        with sqlite3.connect(get_database_path('prompts.db')) as conn:
            cursor = conn.cursor()

            # First normalize the keyword
            normalized_keyword = normalize_keyword(keyword)

            # Get the keyword ID
            cursor.execute("SELECT id FROM Keywords WHERE keyword = ?", (normalized_keyword,))
            result = cursor.fetchone()

            if not result:
                return f"Keyword '{keyword}' not found."

            keyword_id = result[0]

            # Delete keyword associations from PromptKeywords
            cursor.execute("DELETE FROM PromptKeywords WHERE keyword_id = ?", (keyword_id,))

            # Delete the keyword itself
            cursor.execute("DELETE FROM Keywords WHERE id = ?", (keyword_id,))

            # Get the number of affected prompts
            affected_prompts = cursor.rowcount

            conn.commit()

            logging.info(f"Keyword '{keyword}' deleted successfully")
            return f"Successfully deleted keyword '{keyword}' and removed it from {affected_prompts} prompts."

    except sqlite3.Error as e:
        error_msg = f"Database error deleting keyword: {str(e)}"
        logging.error(error_msg)
        return error_msg
    except Exception as e:
        error_msg = f"Error deleting keyword: {str(e)}"
        logging.error(error_msg)
        return error_msg


def export_prompt_keywords_to_csv() -> Tuple[str, str]:
    """

    Export all prompt keywords to a CSV file with associated metadata.



    Returns:

        Tuple[str, str]: (status_message, file_path)

    """
    import csv
    import tempfile
    import os
    from datetime import datetime

    logging.debug("export_prompt_keywords_to_csv: Starting export")
    try:
        # Create a temporary file with a specific name in the system's temp directory
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        temp_dir = tempfile.gettempdir()
        file_path = os.path.join(temp_dir, f'prompt_keywords_export_{timestamp}.csv')

        with sqlite3.connect(get_database_path('prompts.db')) as conn:
            cursor = conn.cursor()

            # Get keywords with related prompt information
            query = '''

            SELECT 

                k.keyword,

                GROUP_CONCAT(p.name, ' | ') as prompt_names,

                COUNT(DISTINCT p.id) as num_prompts,

                GROUP_CONCAT(DISTINCT p.author, ' | ') as authors

            FROM Keywords k

            LEFT JOIN PromptKeywords pk ON k.id = pk.keyword_id

            LEFT JOIN Prompts p ON pk.prompt_id = p.id

            GROUP BY k.id, k.keyword

            ORDER BY k.keyword

            '''

            cursor.execute(query)
            results = cursor.fetchall()

            # Write to CSV
            with open(file_path, 'w', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)
                writer.writerow([
                    'Keyword',
                    'Associated Prompts',
                    'Number of Prompts',
                    'Authors'
                ])

                for row in results:
                    writer.writerow([
                        row[0],  # keyword
                        row[1] if row[1] else '',  # prompt_names (may be None)
                        row[2],  # num_prompts
                        row[3] if row[3] else ''  # authors (may be None)
                    ])

        status_msg = f"Successfully exported {len(results)} prompt keywords to CSV."
        logging.info(status_msg)

        return status_msg, file_path

    except sqlite3.Error as e:
        error_msg = f"Database error exporting keywords: {str(e)}"
        logging.error(error_msg)
        return error_msg, "None"
    except Exception as e:
        error_msg = f"Error exporting keywords: {str(e)}"
        logging.error(error_msg)
        return error_msg, "None"


def view_prompt_keywords() -> str:
    """

    View all keywords currently in the prompts database.



    Returns:

        str: Markdown formatted string of all keywords

    """
    logging.debug("view_prompt_keywords: Retrieving all keywords")
    try:
        with sqlite3.connect(get_database_path('prompts.db')) as conn:
            cursor = conn.cursor()
            cursor.execute("""

                SELECT k.keyword, COUNT(DISTINCT pk.prompt_id) as prompt_count 

                FROM Keywords k

                LEFT JOIN PromptKeywords pk ON k.id = pk.keyword_id

                GROUP BY k.id, k.keyword

                ORDER BY k.keyword

            """)

            keywords = cursor.fetchall()
            if keywords:
                keyword_list = [f"- {k[0]} ({k[1]} prompts)" for k in keywords]
                return "### Current Prompt Keywords:\n" + "\n".join(keyword_list)
            return "No keywords found."

    except Exception as e:
        error_msg = f"Error retrieving keywords: {str(e)}"
        logging.error(error_msg)
        return error_msg


def export_prompts(

        export_format='csv',

        filter_keywords=None,

        include_system=True,

        include_user=True,

        include_details=True,

        include_author=True,

        include_keywords=True,

        markdown_template=None

) -> Tuple[str, str]:
    """

    Export prompts to CSV or Markdown with configurable options.



    Args:

        export_format (str): 'csv' or 'markdown'

        filter_keywords (List[str], optional): Keywords to filter prompts by

        include_system (bool): Include system prompts in export

        include_user (bool): Include user prompts in export

        include_details (bool): Include prompt details/descriptions

        include_author (bool): Include author information

        include_keywords (bool): Include associated keywords

        markdown_template (str, optional): Template for markdown export



    Returns:

        Tuple[str, str]: (status_message, file_path)

    """
    import csv
    import tempfile
    import os
    import zipfile
    from datetime import datetime

    try:
        # Get prompts data
        with get_prompt_db_connection() as conn:
            cursor = conn.cursor()

            # Build query based on included fields
            select_fields = ['p.name']
            if include_author:
                select_fields.append('p.author')
            if include_details:
                select_fields.append('p.details')
            if include_system:
                select_fields.append('p.system')
            if include_user:
                select_fields.append('p.user')

            query = f"""

                SELECT DISTINCT {', '.join(select_fields)}

                FROM Prompts p

            """

            # Add keyword filtering if specified
            if filter_keywords:
                placeholders = ','.join(['?' for _ in filter_keywords])
                query += f"""

                    JOIN PromptKeywords pk ON p.id = pk.prompt_id

                    JOIN Keywords k ON pk.keyword_id = k.id

                    WHERE k.keyword IN ({placeholders})

                """

            cursor.execute(query, filter_keywords if filter_keywords else ())
            prompts = cursor.fetchall()

            # Get keywords for each prompt if needed
            if include_keywords:
                prompt_keywords = {}
                for prompt in prompts:
                    cursor.execute("""

                        SELECT k.keyword

                        FROM Keywords k

                        JOIN PromptKeywords pk ON k.id = pk.keyword_id

                        JOIN Prompts p ON pk.prompt_id = p.id

                        WHERE p.name = ?

                    """, (prompt[0],))
                    prompt_keywords[prompt[0]] = [row[0] for row in cursor.fetchall()]

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        if export_format == 'csv':
            # Export as CSV
            temp_file = os.path.join(tempfile.gettempdir(), f'prompts_export_{timestamp}.csv')
            with open(temp_file, 'w', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)

                # Write header
                header = ['Name']
                if include_author:
                    header.append('Author')
                if include_details:
                    header.append('Details')
                if include_system:
                    header.append('System Prompt')
                if include_user:
                    header.append('User Prompt')
                if include_keywords:
                    header.append('Keywords')
                writer.writerow(header)

                # Write data
                for prompt in prompts:
                    row = list(prompt)
                    if include_keywords:
                        row.append(', '.join(prompt_keywords.get(prompt[0], [])))
                    writer.writerow(row)

            return f"Successfully exported {len(prompts)} prompts to CSV.", temp_file

        else:
            # Export as Markdown files in ZIP
            temp_dir = tempfile.mkdtemp()
            zip_path = os.path.join(tempfile.gettempdir(), f'prompts_export_{timestamp}.zip')

            # Define markdown templates
            templates = {
                "Basic Template": """# {title}

{author_section}

{details_section}

{system_section}

{user_section}

{keywords_section}

""",
                "Detailed Template": """# {title}



## Author

{author_section}



## Description

{details_section}



## System Prompt

{system_section}



## User Prompt

{user_section}



## Keywords

{keywords_section}

"""
            }

            template = templates.get(markdown_template, markdown_template or templates["Basic Template"])

            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for prompt in prompts:
                    # Create markdown content
                    md_content = template.format(
                        title=prompt[0],
                        author_section=f"Author: {prompt[1]}" if include_author else "",
                        details_section=prompt[2] if include_details else "",
                        system_section=prompt[3] if include_system else "",
                        user_section=prompt[4] if include_user else "",
                        keywords_section=', '.join(prompt_keywords.get(prompt[0], [])) if include_keywords else ""
                    )

                    # Create safe filename
                    safe_filename = re.sub(r'[^\w\-_\. ]', '_', prompt[0])
                    md_path = os.path.join(temp_dir, f"{safe_filename}.md")

                    # Write markdown file
                    with open(md_path, 'w', encoding='utf-8') as f:
                        f.write(md_content)

                    # Add to ZIP
                    zipf.write(md_path, os.path.basename(md_path))

            return f"Successfully exported {len(prompts)} prompts to Markdown files.", zip_path

    except Exception as e:
        error_msg = f"Error exporting prompts: {str(e)}"
        logging.error(error_msg)
        return error_msg, "None"


create_prompts_db()

#
# End of Propmts_DB.py
#######################################################################################################################