Code Editor for app.py

import streamlit as st
import os
from openai import OpenAI
from pygments import highlight
from pygments.lexers import get_lexer_by_name
from pygments.formatters import HtmlFormatter
import base64

st.set_page_config(page_title="Chat with Code", layout="wide")

# Initialize OpenAI client with proxy
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)

# Set up the Streamlit page
st.title("💬 Chat with Code")

# Initialize session state for messages if it doesn't exist
if "messages" not in st.session_state:
    st.session_state.messages = []

# Initialize session state for code if it doesn't exist
if "code" not in st.session_state:
    st.session_state.code = ""

# Function to highlight code
def highlight_code(code, language):
    lexer = get_lexer_by_name(language, stripall=True)
    formatter = HtmlFormatter(style="github-dark")
    highlighted = highlight(code, lexer, formatter)
    css = formatter.get_style_defs('.highlight')
    html = f"""
    <style>
    {css}
    .highlight {{ background-color: #1e1e1e; padding: 10px; border-radius: 5px; }}
    </style>
    {highlighted}
    """
    return html

# Create a two-column layout
col1, col2 = st.columns([1, 1])

with col1:
    st.subheader("Chat")
    
    # Display chat messages from history on app rerun
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # Accept user input
    if prompt := st.chat_input("Ask about the code..."):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message in chat message container
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Prepare messages for API call
        messages = [
            {"role": "system", "content": "You are a helpful assistant that explains code."}
        ]
        
        # Add the code as context
        if st.session_state.code:
            messages.append({"role": "user", "content": f"Here is the code I want to ask about:\n```\n{st.session_state.code}\n```"})
        
        # Add all previous messages
        for message in st.session_state.messages:
            messages.append({"role": message["role"], "content": message["content"]})
        
        # Get response from OpenAI
        with st.chat_message("assistant"):
            response_placeholder = st.empty()
            full_response = ""
            
            # Make API call
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                stream=True
            )
            
            # Stream the response
            for chunk in response:
                if chunk.choices[0].delta.content:
                    full_response += chunk.choices[0].delta.content
                    response_placeholder.markdown(full_response + "▌")
            
            response_placeholder.markdown(full_response)
        
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": full_response})

with col2:
    st.subheader("Code Editor")
    
    # Code editor
    code_input = st.text_area("Enter your code here", height=400, value=st.session_state.code)
    
    # Language selection
    language = st.selectbox("Language", ["python", "javascript", "java", "cpp", "csharp", "go", "rust"])
    
    # Update code in session state
    if code_input != st.session_state.code:
        st.session_state.code = code_input
    
    # Display highlighted code
    if st.session_state.code:
        st.markdown(highlight_code(st.session_state.code, language), unsafe_allow_html=True)
    
    # Clear chat button
    if st.button("Clear Chat"):
        st.session_state.messages = []
        st.rerun()