app.py

import streamlit as st
from graph import EssayWriter, RouteQuery, GraphState
from crew import *
import os
import base64

#st.title("Multi-Agent Essay Writing Assistant")
st.set_page_config(page_title="Multi-Agent Essay Writing Assistant", page_icon="🤖🤖🤖✍️", layout = "centered")

# Ensure session state variables are initialized properly
if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]

if "app" not in st.session_state:
    st.session_state["app"] = None

if "chat_active" not in st.session_state:
    st.session_state["chat_active"] = True

# Sidebar with essay settings and user-defined length
with st.sidebar:
    st.subheader("About:")
    st.info(
        "\n\n* This app uses the 'gpt-4o-mini-2024-07-18' model."
        "\n\n* Writing essays may take some time, approximately 1-2 minutes."
    )

    # API Key Retrieval
    openai_key = st.secrets.get("OPENAI_API_KEY", "")

    st.divider()

    # User-defined essay length selection
    st.subheader("📝 Configure Essay Settings:")
    essay_length = st.number_input(
        "Select Essay Length (words):",
        min_value=150,
        max_value=400,
        value=250,
        step=50
    )

    st.divider()

    # Reference section
    st.subheader("📖 References:")
    st.markdown(
        "[1. Multi-Agent System with CrewAI and LangChain](https://discuss.streamlit.io/t/new-project-i-have-build-a-multi-agent-system-with-crewai-and-langchain/84002)",
        unsafe_allow_html=True
    )

# Initialize agents function
def initialize_agents():
    if not openai_key:
        st.error("⚠️ OpenAI API key is missing! Please provide a valid key through Hugging Face Secrets.")
        st.session_state["chat_active"] = True
        return None

    os.environ["OPENAI_API_KEY"] = openai_key
    try:
        # Prevent re-initialization
        if "app" in st.session_state and st.session_state["app"] is not None:
            return st.session_state["app"]
        
        # Initialize the full EssayWriter instance
        essay_writer = EssayWriter()  # ✅ Store the full instance
        st.session_state["app"] = essay_writer  # ✅ Now contains `graph`
        st.session_state["chat_active"] = False  # ✅ Enable chat after successful initialization

        return essay_writer
    except Exception as e:
        st.error(f"❌ Error initializing agents: {e}")
        st.session_state["chat_active"] = True
        return None


# Automatically initialize agents on app load
if st.session_state["app"] is None:
    st.session_state["app"] = initialize_agents()

if st.session_state["app"] is None:
    st.error("⚠️ Failed to initialize agents. Please check your API key and restart the app.")

app = st.session_state["app"]

# Function to invoke the agent and generate a response
def generate_response(topic, length):
    if not app or not hasattr(app, "graph"):
        st.error("⚠️ Agents are not initialized. Please check the system or restart the app.")
        return {"response": "Error: Agents not initialized."}

    return app.graph.invoke(input={"topic": topic, "length": length})  # Ensure `graph` is invoked


# Define Tabs
tab1, tab2 = st.tabs(["📜 Essay Generation", "📊 Workflow Viz"])

# 📜 Tab 1: Essay Generation
with tab1:
    st.subheader("📝 Generate an Essay")

    # Display chat messages from the session
    if "messages" not in st.session_state:
        st.session_state["messages"] = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]

    for message in st.session_state["messages"]:
        with st.chat_message(message["role"]):
            st.markdown(message["content"], unsafe_allow_html=True)

    # Handle user input
    if topic := st.chat_input(placeholder="📝 Ask a question or provide an essay topic...", disabled=st.session_state["chat_active"]):
        st.chat_message("user").markdown(topic)
        st.session_state["messages"].append({"role": "user", "content": topic})

        with st.spinner("⏳ Generating your essay..."):
            response = None
            if app:
                response = app.write_essay({"topic": topic})
            else:
                st.error("⚠️ Agents are not initialized. Please check the system or restart the app.")

        # Handle the assistant's response
        with st.chat_message("assistant"):
            if response and "essay" in response:  # Display essay preview and download link
                essay = response["essay"]
                st.markdown(f"### 📝 Essay Preview ({essay_length} words)")
                st.markdown(f"#### {essay['header']}")
                st.markdown(essay["entry"])
                
                for para in essay["paragraphs"]:
                    st.markdown(f"**{para['sub_header']}**")
                    st.markdown(para["paragraph"])
                
                st.markdown("**🖊️ Conclusion:**")
                st.markdown(essay["conclusion"])

                # Provide download link for the PDF (only if available)
                pdf_name = response.get("pdf_name")
                if pdf_name and os.path.exists(pdf_name):
                    with open(pdf_name, "rb") as pdf_file:
                        b64 = base64.b64encode(pdf_file.read()).decode()
                        href = f"<a href='data:application/octet-stream;base64,{b64}' download='{pdf_name}'>📄 Click here to download the PDF</a>"
                        st.markdown(href, unsafe_allow_html=True)

                # Save response in session state
                st.session_state["messages"].append(
                    {"role": "assistant", "content": f"Here is your {essay_length}-word essay preview and the download link."}
                )
            elif response:  # Other responses (fallback)
                st.markdown(response["response"])
                st.session_state["messages"].append({"role": "assistant", "content": response["response"]})
            else:
                st.error("⚠️ No response received. Please try again.")


# 📊 Tab 2: Workflow Visualization
with tab2:
    st.subheader("📊 Multi-Agent Workflow Visualization")

    try:
        # Load and display the workflow graph
        graph_path = "graph.png"  # Ensure graph.py generates this
        if os.path.exists(graph_path):
            st.image(graph_path, caption="Multi-Agent Workflow", use_column_width=True)
        else:
            st.warning("⚠️ Workflow graph not found. Ensure `graph.py` generates `graph.png`.")

    except Exception as e:
        st.error(f"❌ Error generating workflow visualization: {e}")

# Acknowledgment Section
st.markdown(
    """
    <div style="text-align: center; font-size: 14px; color: #555; padding-top: 250px; margin-top: 250px;">
        <strong>Acknowledgment:</strong> This app is based on Mesut Duman's work: 
        <a href="https://github.com/mesutdmn/Autonomous-Multi-Agent-Systems-with-CrewAI-Essay-Writer/tree/main" 
           target="_blank" style="color: #007BFF; text-decoration: none;">
           CrewAI Essay Writer
        </a>
    </div>
    """,
    unsafe_allow_html=True,
)