Code Editor for app.py

import streamlit as st
import pandas as pd
import os
import json
import uuid
from datetime import datetime
import plotly.express as px
from openai import OpenAI

# Page configuration
st.set_page_config(
    page_title="Prompt Optimization Tool",
    page_icon="📝",
    layout="wide"
)

# Initialize session state variables
if "prompts" not in st.session_state:
    st.session_state.prompts = []
if "current_prompt" not in st.session_state:
    st.session_state.current_prompt = {
        "id": str(uuid.uuid4()),
        "title": "",
        "content": "",
        "tags": [],
        "created_at": datetime.now().isoformat(),
        "updated_at": datetime.now().isoformat(),
        "folder": "Default",
        "versions": []
    }
if "folders" not in st.session_state:
    st.session_state.folders = ["Default"]
if "filter_tag" not in st.session_state:
    st.session_state.filter_tag = "All"
if "filter_folder" not in st.session_state:
    st.session_state.filter_folder = "All"
if "all_tags" not in st.session_state:
    st.session_state.all_tags = []
if "search_query" not in st.session_state:
    st.session_state.search_query = ""
if "last_api_response" not in st.session_state:
    st.session_state.last_api_response = None

# Function to save prompts to file
def save_prompts_to_file():
    with open("prompts.json", "w") as f:
        json.dump(st.session_state.prompts, f)

# Function to load prompts from file
def load_prompts_from_file():
    try:
        with open("prompts.json", "r") as f:
            st.session_state.prompts = json.load(f)
    except FileNotFoundError:
        st.session_state.prompts = []
    
    # Extract all unique tags and folders
    all_tags = []
    all_folders = ["Default"]
    for prompt in st.session_state.prompts:
        all_tags.extend(prompt.get("tags", []))
        folder = prompt.get("folder", "Default")
        if folder not in all_folders:
            all_folders.append(folder)
    
    st.session_state.all_tags = list(set(all_tags))
    st.session_state.folders = all_folders

# Load prompts at the start
load_prompts_from_file()

# App title and description
st.title("📝 Prompt Optimization Tool")
st.markdown("Create, optimize, and manage your AI prompts in one place.")

# Main navigation
tabs = st.tabs(["Create & Optimize", "Browse Prompts", "Analytics"])

# Tab 1: Create & Optimize Prompts
with tabs[0]:
    col1, col2 = st.columns([3, 2])
    
    with col1:
        st.subheader("Prompt Editor")
        
        # Prompt title input
        st.session_state.current_prompt["title"] = st.text_input("Prompt Title", st.session_state.current_prompt["title"])
        
        # Prompt content textarea
        st.session_state.current_prompt["content"] = st.text_area("Prompt Content", st.session_state.current_prompt["content"], height=250)
        
        # Tags and folder
        tag_input = st.text_input("Tags (comma separated)")
        if tag_input:
            st.session_state.current_prompt["tags"] = [tag.strip() for tag in tag_input.split(",")]
        
        # Folder selection or creation
        folder_col1, folder_col2 = st.columns([3, 1])
        
        with folder_col1:
            folder_option = st.selectbox("Select Folder", st.session_state.folders)
            if folder_option:
                st.session_state.current_prompt["folder"] = folder_option
        
        with folder_col2:
            new_folder = st.text_input("Or Create New Folder")
            if new_folder and new_folder not in st.session_state.folders:
                st.session_state.folders.append(new_folder)
                st.session_state.current_prompt["folder"] = new_folder
        
        # Save button
        if st.button("Save Prompt"):
            if not st.session_state.current_prompt["title"]:
                st.error("Please provide a title for your prompt.")
            elif not st.session_state.current_prompt["content"]:
                st.error("Please provide content for your prompt.")
            else:
                # Update timestamp
                st.session_state.current_prompt["updated_at"] = datetime.now().isoformat()
                
                # Check if this is an edit of an existing prompt
                existing_prompt_index = None
                for i, prompt in enumerate(st.session_state.prompts):
                    if prompt["id"] == st.session_state.current_prompt["id"]:
                        existing_prompt_index = i
                        break
                
                if existing_prompt_index is not None:
                    # Save the current content as a new version
                    current_version = {
                        "content": st.session_state.prompts[existing_prompt_index]["content"],
                        "timestamp": st.session_state.prompts[existing_prompt_index]["updated_at"]
                    }
                    if "versions" not in st.session_state.current_prompt:
                        st.session_state.current_prompt["versions"] = []
                    st.session_state.current_prompt["versions"].append(current_version)
                    # Update the prompt
                    st.session_state.prompts[existing_prompt_index] = st.session_state.current_prompt.copy()
                else:
                    # Add new prompt
                    st.session_state.prompts.append(st.session_state.current_prompt.copy())
                
                # Save to file
                save_prompts_to_file()
                
                # Reset current prompt
                st.session_state.current_prompt = {
                    "id": str(uuid.uuid4()),
                    "title": "",
                    "content": "",
                    "tags": [],
                    "created_at": datetime.now().isoformat(),
                    "updated_at": datetime.now().isoformat(),
                    "folder": "Default",
                    "versions": []
                }
                
                st.success("Prompt saved successfully!")
                st.rerun()
    
    with col2:
        st.subheader("Prompt Optimization")
        
        optimization_option = st.selectbox(
            "Optimization Goal",
            ["Make it clearer", "Make it more concise", "Make it more detailed", "Improve readability", "Custom"]
        )
        
        custom_optimization = ""
        if optimization_option == "Custom":
            custom_optimization = st.text_input("Describe your optimization goal")
        
        if st.button("Optimize Prompt"):
            if not st.session_state.current_prompt["content"]:
                st.error("Please enter a prompt to optimize.")
            else:
                with st.spinner("Optimizing your prompt..."):
                    try:
                        # Setup OpenAI client
                        base_url = "https://editor.ploomber.io/_api/ai-proxy/v1"
                        api_key = os.environ["PLOOMBER_API_KEY"]
                        client = OpenAI(base_url=base_url, api_key=api_key)
                        
                        # Prepare the optimization instruction
                        if optimization_option == "Custom":
                            instruction = custom_optimization
                        else:
                            instruction = optimization_option
                        
                        # Call the API
                        response = client.chat.completions.create(
                            model="gpt-4o-mini",
                            messages=[
                                {"role": "system", "content": "You are a prompt engineering expert. Your task is to optimize the provided prompt according to the specified goal."},
                                {"role": "user", "content": f"Please optimize the following prompt. Goal: {instruction}\n\nPrompt:\n{st.session_state.current_prompt['content']}"}
                            ]
                        )
                        
                        # Store the optimized prompt
                        optimized_prompt = response.choices[0].message.content
                        st.session_state.last_api_response = optimized_prompt
                        
                        # Add as a new version in the current prompt
                        current_version = {
                            "content": st.session_state.current_prompt["content"],
                            "timestamp": datetime.now().isoformat()
                        }
                        if "versions" not in st.session_state.current_prompt:
                            st.session_state.current_prompt["versions"] = []
                        st.session_state.current_prompt["versions"].append(current_version)
                        
                        # Update the current prompt with the optimized version
                        st.session_state.current_prompt["content"] = optimized_prompt
                        st.session_state.current_prompt["updated_at"] = datetime.now().isoformat()
                        
                    except Exception as e:
                        st.error(f"Error during optimization: {str(e)}")
                st.rerun()
        
        # Display optimization result
        if st.session_state.last_api_response:
            st.subheader("Optimization Result")
            st.write(st.session_state.last_api_response)
            
            if st.button("Apply Optimization"):
                st.session_state.current_prompt["content"] = st.session_state.last_api_response
                st.session_state.last_api_response = None
                st.rerun()
        
        # Prompt versions
        if st.session_state.current_prompt.get("versions"):
            st.subheader("Previous Versions")
            for i, version in enumerate(reversed(st.session_state.current_prompt["versions"])):
                with st.expander(f"Version {len(st.session_state.current_prompt['versions']) - i} - {version['timestamp'][:16].replace('T', ' ')}"):
                    st.write(version["content"])
                    if st.button(f"Restore Version {len(st.session_state.current_prompt['versions']) - i}"):
                        st.session_state.current_prompt["content"] = version["content"]
                        st.rerun()

# Tab 2: Browse Prompts
with tabs[1]:
    st.subheader("Browse & Search Prompts")
    
    # Search and filter options
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.session_state.search_query = st.text_input("Search prompts", st.session_state.search_query)
    
    with col2:
        tag_options = ["All"] + st.session_state.all_tags
        st.session_state.filter_tag = st.selectbox("Filter by tag", tag_options)
    
    with col3:
        folder_options = ["All"] + st.session_state.folders
        st.session_state.filter_folder = st.selectbox("Filter by folder", folder_options)
    
    # Apply filters
    filtered_prompts = st.session_state.prompts
    
    # Apply search
    if st.session_state.search_query:
        filtered_prompts = [p for p in filtered_prompts if 
                            st.session_state.search_query.lower() in p["title"].lower() or 
                            st.session_state.search_query.lower() in p["content"].lower()]
    
    # Apply tag filter
    if st.session_state.filter_tag != "All":
        filtered_prompts = [p for p in filtered_prompts if st.session_state.filter_tag in p.get("tags", [])]
    
    # Apply folder filter
    if st.session_state.filter_folder != "All":
        filtered_prompts = [p for p in filtered_prompts if p.get("folder", "Default") == st.session_state.filter_folder]
    
    # Display prompts
    if not filtered_prompts:
        st.info("No prompts found. Create a new prompt in the 'Create & Optimize' tab.")
    else:
        for prompt in filtered_prompts:
            with st.expander(f"{prompt['title']} - {prompt.get('folder', 'Default')}"):
                st.write(prompt["content"])
                
                # Display tags
                if prompt.get("tags"):
                    st.write("Tags: " + ", ".join(prompt["tags"]))
                
                # Display timestamps
                created_at = datetime.fromisoformat(prompt["created_at"]).strftime("%Y-%m-%d %H:%M")
                updated_at = datetime.fromisoformat(prompt["updated_at"]).strftime("%Y-%m-%d %H:%M")
                st.write(f"Created: {created_at} | Last updated: {updated_at}")
                
                # Edit and delete buttons
                col1, col2 = st.columns(2)
                
                with col1:
                    if st.button(f"Edit '{prompt['title']}'"):
                        st.session_state.current_prompt = prompt.copy()
                        st.switch_page("app.py")  # This will refresh to the first tab
                
                with col2:
                    if st.button(f"Delete '{prompt['title']}'"):
                        if st.session_state.current_prompt["id"] == prompt["id"]:
                            st.session_state.current_prompt = {
                                "id": str(uuid.uuid4()),
                                "title": "",
                                "content": "",
                                "tags": [],
                                "created_at": datetime.now().isoformat(),
                                "updated_at": datetime.now().isoformat(),
                                "folder": "Default",
                                "versions": []
                            }
                        
                        st.session_state.prompts.remove(prompt)
                        save_prompts_to_file()
                        st.success(f"Prompt '{prompt['title']}' deleted.")
                        st.rerun()

# Tab 3: Analytics
with tabs[2]:
    st.subheader("Prompt Analytics")
    
    if not st.session_state.prompts:
        st.info("No prompts available for analytics. Create some prompts first.")
    else:
        # Prepare data for analytics
        folder_counts = {}
        tag_counts = {}
        creation_dates = []
        prompt_lengths = []
        
        for prompt in st.session_state.prompts:
            # Folder counts
            folder = prompt.get("folder", "Default")
            folder_counts[folder] = folder_counts.get(folder, 0) + 1
            
            # Tag counts
            for tag in prompt.get("tags", []):
                tag_counts[tag] = tag_counts.get(tag, 0) + 1
            
            # Creation dates
            creation_dates.append({
                "title": prompt["title"],
                "date": datetime.fromisoformat(prompt["created_at"]).strftime("%Y-%m-%d")
            })
            
            # Prompt lengths
            prompt_lengths.append({
                "title": prompt["title"],
                "length": len(prompt["content"])
            })
        
        col1, col2 = st.columns(2)
        
        with col1:
            # Folder distribution
            st.subheader("Prompts by Folder")
            folder_df = pd.DataFrame({
                "Folder": folder_counts.keys(),
                "Count": folder_counts.values()
            })
            if not folder_df.empty:
                fig = px.pie(folder_df, values="Count", names="Folder", title="Prompt Distribution by Folder")
                st.plotly_chart(fig, use_container_width=True)
            
            # Creation timeline
            st.subheader("Prompt Creation Timeline")
            date_df = pd.DataFrame(creation_dates)
            if not date_df.empty:
                date_df["date"] = pd.to_datetime(date_df["date"])
                date_count = date_df.groupby("date").size().reset_index(name="count")
                fig = px.line(date_count, x="date", y="count", title="Prompts Created Over Time")
                st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            # Tag distribution
            st.subheader("Prompts by Tag")
            if tag_counts:
                tag_df = pd.DataFrame({
                    "Tag": tag_counts.keys(),
                    "Count": tag_counts.values()
                })
                fig = px.bar(tag_df, x="Tag", y="Count", title="Prompt Count by Tag")
                st.plotly_chart(fig, use_container_width=True)
            else:
                st.info("No tags have been added to prompts yet.")
            
            # Prompt length distribution
            st.subheader("Prompt Length Distribution")
            length_df = pd.DataFrame(prompt_lengths)
            if not length_df.empty:
                fig = px.histogram(length_df, x="length", title="Distribution of Prompt Lengths (characters)")
                st.plotly_chart(fig, use_container_width=True)
        
        # Overall statistics
        st.subheader("Summary Statistics")
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            st.metric("Total Prompts", len(st.session_state.prompts))
        
        with col2:
            st.metric("Total Folders", len(folder_counts))
        
        with col3:
            st.metric("Total Tags", len(tag_counts))
        
        with col4:
            if prompt_lengths:
                avg_length = sum(item["length"] for item in prompt_lengths) / len(prompt_lengths)
                st.metric("Average Prompt Length", f"{int(avg_length)} chars")
            else:
                st.metric("Average Prompt Length", "N/A")

# Footer
st.markdown("---")
st.markdown("© 2023 Prompt Optimization Tool | Made with Streamlit")