NightPrince/CodeGenBot-RAG-Assistant

1import streamlit as st2from pipeline.pipeline import CodeGenPipeline3 4st.set_page_config(page_title="CodeGenBot", page_icon="🤖")5st.title("💻 CodeGenBot")6 7# Initialize pipeline only once (cache in session state)8if "pipeline" not in st.session_state:9    st.session_state.pipeline = CodeGenPipeline(10        "hf://datasets/openai/openai_humaneval/openai_humaneval/test-00000-of-00001.parquet"11    )12 13# Memory for chat messages14if "messages" not in st.session_state:15    st.session_state.messages = []16 17# Display previous messages18for msg in st.session_state.messages:19    if msg["role"] == "assistant":20        st.chat_message("assistant").markdown(msg["content"])21    else:22        st.chat_message(msg["role"]).write(msg["content"])23 24# User input25user_input = st.chat_input("Ask CodeGenBot to generate Python code...")26 27if user_input:28    # Save user message29    st.session_state.messages.append({"role": "user", "content": user_input})30    st.chat_message("user").write(user_input)31 32    # Generate code using your pipeline33    with st.spinner("Generating code..."):34        try:35            code_output = st.session_state.pipeline.generate_code_from_prompt(user_input)36        except Exception as e:37            code_output = f"Error: {e}"38 39    # Save assistant's reply and rerun to display it immediately40    st.session_state.messages.append({"role": "assistant", "content": code_output})41    st.rerun()42