Update app.py

app.py

@@ -91,7 +91,14 @@ class AIAgent:
         self.current_task: Optional[str] = None
         self.search_engine_url: str = "https://www.google.com/search?q="  # Default search engine
         self.prompts: List[str] = []  # Store prompts for future use
-        self.code_generator = pipeline('text-generation', model='gpt2')  # Initialize code generator
+        self.code_generator = None  # Initialize code generator later
+        self.available_models = [
+            "gpt2",
+            "facebook/bart-large-cnn",
+            "google/flan-t5-xl",
+            "bigscience/T0_3B",
+        ]  # Add more as needed
+        self.selected_model = "gpt2"  # Default model
 
     # --- Search Functionality ---
     def search(self, query: str) -> List[str]:
@@ -110,9 +117,13 @@ class AIAgent:
     def code_generation(self, snippet: str) -> str:
         """Generates code based on the provided snippet or description."""
         try:
-            generated_text = self.code_generator(snippet, max_length=500, num_return_sequences=1)[0][
-                'generated_text'
-            ]
+            if not self.code_generator:
+                self.code_generator = pipeline(
+                    'text-generation', model=self.selected_model
+                )
+            generated_text = self.code_generator(
+                snippet, max_length=500, num_return_sequences=1
+            )[0]['generated_text']
             return generated_text
         except Exception as e:
             raise CodeGenerationError(f"Error during code generation: {e}")
@@ -384,12 +395,70 @@ class AIAgent:
 # --- Streamlit Integration ---
 if __name__ == "__main__":
     agent = AIAgent()
+    st.set_page_config(
+        page_title="AI Agent",
+        page_icon="🤖",
+        layout="wide",
+        initial_sidebar_state="expanded",
+    )
     st.title("AI Agent")
-    st.write("Enter a command for the AI Agent:")
-    input_str = st.text_input("")
-    if input_str:
-        agent.handle_input(input_str)
-    # agent.run()  # Removed to avoid conflict with Streamlit
+    st.sidebar.title("Agent Settings")
+
+    # --- Command Dropdown ---
+    command_options = [
+        "SEARCH",
+        "CODEGEN",
+        "REFINE-CODE",
+        "TEST-CODE",
+        "INTEGRATE-CODE",
+        "TEST-APP",
+        "GENERATE-REPORT",
+        "WORKSPACE-EXPLORER",
+        "ADD_PROMPT",
+        "ACTION_PROMPT",
+        "COMPRESS_HISTORY_PROMPT",
+        "LOG_PROMPT",
+        "LOG_RESPONSE",
+        "MODIFY_PROMPT",
+        "PREFIX",
+        "SEARCH_QUERY",
+        "READ_PROMPT",
+        "TASK_PROMPT",
+        "UNDERSTAND_TEST_RESULTS_PROMPT",
+    ]
+    selected_command = st.sidebar.selectbox("Command", command_options)
+
+    # --- Model Dropdown ---
+    selected_model = st.sidebar.selectbox(
+        "Model", agent.available_models, index=agent.available_models.index(agent.selected_model)
+    )
+    agent.selected_model = selected_model
+
+    # --- Input Field ---
+    input_str = st.text_input(f"Enter input for {selected_command}:")
+
+    # --- Execute Command ---
+    if st.button("Execute"):
+        if input_str:
+            agent.handle_input(f"{selected_command} {input_str}")
+            st.write(f"Output: {agent.task_history[-1]['output']}")
+
+    # --- Task History ---
+    st.subheader("Task History")
+    for task in agent.task_history:
+        st.write(f"**Action:** {task['action']}")
+        st.write(f"**Input:** {task['input']}")
+        st.write(f"**Output:** {task['output']}")
+
+    # --- Workspace Explorer ---
+    st.subheader("Workspace Explorer")
+    with st.expander("Explore Workspace"):
+        try:
+            workspace_output = agent.workspace_explorer()
+            st.write(workspace_output)
+        except WorkspaceExplorerError as e:
+            st.error(f"Error exploring workspace: {e}")
+
 
 # --- Gradio Integration ---
 def gradio_interface(input_text):