Update app.py

app.py

@@ -1,44 +1,69 @@
-import os
 import streamlit as st
-import torch
-from transformers import GPTNeoForCausalLM, GPT2Tokenizer
-
-# Function to reverse engineer code
-def reverse_prompt_engineer(input_code):
-    # Load the tokenizer and model
-    tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
-    model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B").to("cuda")
+from streamlit_chat import message
+from src.langchain_agent import init, init_agent
+
+def main():
     
-    # Generate a prompt using the analyzed code
-    prompt = "Reverse engineer the following code:\n\n" + input_code
+    #initialise agent and streamlit page
+    init()
+    agent_executor = init_agent()
     
-    # Tokenize the prompt and generate similar code using the model
-    input_ids = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
-    generated_ids = model.generate(input_ids, max_length=100, do_sample=True)
-    generated_code = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    # initialize message history
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    # store agent in memory
+    if "clarina" not in st.session_state:
+        st.session_state.clarina = agent_executor
+
+    # store generated responses in memory
+    if 'generated' not in st.session_state:
+        st.session_state.generated = []
     
-    return generated_code
-
-# Set Streamlit page configuration
-st.set_page_config(
-    page_title="Code Generator",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
-
-# Code Generation Page
-st.title("Code Generator")
-
-# User input code area
-user_input = st.text_area("Input Code")
-
-# Generate Code button
-if st.button("Generate Code"):
-    if not user_input:
-        st.error("Input field is empty!")
-    else:
-        try:
-            generated_code = reverse_prompt_engineer(user_input)
-            st.code(generated_code)
-        except Exception as e:
-            st.error(f"An error occurred: {str(e)}")
+    # define function to generate response
+    def generate_response(user_input):
+        # handle user input
+        if user_input:
+            # save user input
+            st.session_state.messages.append(user_input) 
+            # get response from agent
+            with st.spinner("Thinking..."):
+                response = st.session_state.clarina.reverse_prompt_engineer(user_input)
+            # save response
+            st.session_state.messages.append(response) 
+            st.session_state.generated.append(response)    
+
+    # container for chat history
+    response_container = st.container()
+    # container for text box
+    container = st.container()
+
+    with container:
+
+        # initialize session state to clear input text box after user enters input
+        if "temp" not in st.session_state:
+            st.session_state.temp = ""
+
+        def clear_text():
+            """callback function to clear input text box"""
+            st.session_state.temp = st.session_state.user_input
+            st.session_state.user_input = ""
+
+        st.text_input("user input",key="user_input",placeholder = "Enter your code here", label_visibility="hidden",on_change=clear_text) # get user input
+        generate_response(st.session_state.temp) # generate response
+
+    # display message history
+    if st.session_state.generated:
+        
+        with response_container:
+            messages = st.session_state.get('messages', []) 
+            for i, msg in enumerate(messages):
+                if i % 2 == 0:
+                    # display user input
+                    message(msg, is_user=True, key=str(i) + '_user')
+                else:
+                    # display response
+                    message(msg, is_user=False, key=str(i) + '_ai')    
+
+if __name__ == '__main__':
+    main()