app.py

import streamlit as st
import matplotlib.pyplot as plt
from qiskit import QuantumCircuit
from qiskit.visualization import circuit_drawer
from openai import OpenAI
import io
import base64

# Set page configuration
st.set_page_config(page_title="AI-Powered Quantum Algorithm Designer", layout="wide")

# Set API Key (replace with your own in Streamlit secrets)
API_KEY = st.secrets["AIML_API_KEY"]
BASE_URL = "https://api.aimlapi.com/"

# Initialize OpenAI client
client = OpenAI(
    api_key=API_KEY,
    base_url=BASE_URL,
)

def design_algorithm(problem_description, goals, language, hardware_info, error_rates, model):
    # Construct the prompt for the API
    prompt = f"""
You are an AI assistant specialized in designing quantum algorithms. Design an algorithm for the following problem:

Problem Description:
{problem_description}

Optimization Goals:
{goals}

Target Quantum Programming Language:
{language}

Target Quantum Hardware:
{hardware_info}

Gate Error Rates:
{error_rates}

Provide the quantum algorithm in the specified language with the following format:
Algorithm Code:
<Your code in {language}>
Explanation:
<Brief explanation of the algorithm's logic and how it addresses the specified problem>
Analysis:
<Analysis of the algorithm's performance, complexity, and how it fits the optimization goals>
"""

    # Call the OpenAI API
    chat_completion = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "user", "content": prompt},
        ],
        max_tokens=2000,
    )

    response = chat_completion.choices[0].message.content
    return response, chat_completion.usage.total_tokens

def extract_response_sections(response):
    # Use simple parsing to extract sections from the response
    try:
        code = response.split("Algorithm Code:")[1].split("Explanation:")[0].strip()
        explanation = response.split("Explanation:")[1].split("Analysis:")[0].strip()
        analysis = response.split("Analysis:")[1].strip()
    except IndexError:
        raise ValueError("The model's response was not in the expected format. Please try again.")
    
    return code, explanation, analysis

def visualize_circuit(circuit_code):
    try:
        # Convert the code to a QuantumCircuit object if the language is Qiskit
        qc = QuantumCircuit.from_qasm_str(circuit_code)
        circuit_drawing = circuit_drawer(qc, output='mpl', style={'backgroundcolor': '#EEEEEE'})
        
        # Save the drawing to a bytes buffer
        buf = io.BytesIO()
        circuit_drawing.savefig(buf, format='png')
        buf.seek(0)
        
        # Encode the bytes buffer to base64
        encoded_image = base64.b64encode(buf.getvalue()).decode()
        return f"data:image/png;base64,{encoded_image}"
    except Exception:
        return None

# Streamlit UI
st.title("🧠 AI-Powered Quantum Algorithm Designer")
st.markdown("""
This tool uses OpenAI's o1 models to design quantum algorithms based on user-defined problems, optimization goals, 
and the desired programming language. It leverages deep reasoning capabilities to produce efficient and tailored 
quantum solutions.
""")

col1, col2 = st.columns(2)

with col1:
    st.subheader("Input Parameters")
    problem_description = st.text_area("Describe your computational problem:", 
                                       height=150, 
                                       placeholder="e.g., Solve the Traveling Salesman Problem using quantum techniques.")
    goals = st.text_area("Enter optimization goals:", 
                         height=100, 
                         placeholder="e.g., minimize qubit count, reduce circuit depth")
    language = st.selectbox("Choose target quantum programming language:", 
                            ["Qiskit", "PyQuil", "Cirq", "Qmod"], 
                            help="The language in which the algorithm should be generated.")
    hardware_info = st.text_input("Target quantum hardware:", 
                                  placeholder="e.g., IBM Q System One")
    error_rates = st.text_input("Gate error rates:", 
                                placeholder="e.g., 'X: 0.001, CNOT: 0.01'")
    model = st.selectbox("Choose o1 model:", 
                         ["o1-mini", "o1-preview"], 
                         help="o1-preview offers deeper reasoning but is more expensive")

submit_button = st.button("Generate Algorithm")

if submit_button and problem_description:
    with st.spinner("Designing algorithm... This may take a moment."):
        try:
            response, tokens_used = design_algorithm(problem_description, goals, language, hardware_info, error_rates, model)
            st.success("Algorithm Generated Successfully!")
            
            # Extract code, explanation, and analysis
            try:
                code, explanation, analysis = extract_response_sections(response)
                
                st.subheader("Generated Algorithm")
                st.code(code, language="python" if language == "Qiskit" else language.lower())
                
                st.subheader("Explanation")
                st.write(explanation)
                
                st.subheader("Performance Analysis")
                st.write(analysis)
                
                # Attempt visualization if the language is Qiskit and the code is in QASM format
                if language == "Qiskit" and "OPENQASM" in code:
                    st.subheader("Visualized Quantum Circuit")
                    circuit_image = visualize_circuit(code)
                    if circuit_image:
                        st.image(circuit_image, use_column_width=True)
                    else:
                        st.warning("Could not visualize the circuit. Please review the code output.")
                
            except ValueError as ve:
                st.error(str(ve))
                st.write("Here is the raw response for debugging:")
                st.code(response, language="text")
            
            st.info(f"Tokens Used: {tokens_used}")
            
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
else:
    st.info("Enter your problem details and click 'Generate Algorithm' to begin.")

st.sidebar.title("About")
st.sidebar.info("""
This Quantum Algorithm Designer uses OpenAI's o1 models to provide advanced algorithm design capabilities.

- **o1-preview**: Ideal for complex design tasks requiring deep reasoning.
- **o1-mini**: Faster and more cost-effective for simpler tasks.

The app translates user requirements into customized quantum algorithms in various programming languages.
""")

st.sidebar.title("Limitations")
st.sidebar.warning("""
- No real-time web data access.
- External data upload not supported.
- Visualization is limited to Qiskit-based QASM outputs.
- The generated algorithms may require user review for edge cases.
""")

st.sidebar.title("Usage Tips")
st.sidebar.info("""
1. Provide detailed problem descriptions and goals for better results.
2. Select the language and hardware that match your needs.
3. Start with o1-mini for faster prototyping; use o1-preview for intricate tasks.
4. Review the code and explanation for insights into the design.
""")