Update app.py

app.py

@@ -14,7 +14,6 @@ GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 # Configure Generative AI model
 if GOOGLE_API_KEY:
     genai.configure(api_key=GOOGLE_API_KEY)
-    model = genai.GenerativeModel('gemini-pro')  # Or specify your model
 else:
     st.error(
         "Google AI Studio API key not found. Please add it to your .env file.  "
@@ -25,55 +24,13 @@ else:
 st.title("Prompt Engineering Playground")
 st.subheader("Experiment with Prompting Techniques")
 
-# Sidebar for explanations and resources
-with st.sidebar:
-    st.header("Prompting Concepts")
-    st.markdown(
-        """
-        This app demonstrates various prompt engineering techniques.  Explore how different prompts
-        affect the output of a large language model.
-        """
-    )
-    st.subheader("Key Techniques:")
-    st.markdown(
-        """
-        -   **Clear Instructions**:  Provide explicit and unambiguous directions.
-        -   **Delimiters**:   Use special characters to separate input parts.
-        -   **Structured Output**: Request output in a specific format (JSON).
-        -   **Assumption Checking**:  Verify conditions in the input.
-        -   **Few-Shot Prompting**: Provide input-output examples.
-        -   **Temperature Control**: Adjust output randomness.
-        -   **Chain of Thought (CoT)**: Elicit step-by-step reasoning.
-        -   **Prompt Templates**: Employ pre-defined prompt structures.
-        -   **System Prompt**:  Influence the model's overall behavior.
-        -   **Retrieval Augmentation**: Provide external knowledge.
-        """
-    )
-    st.subheader("Important Considerations:")
-    st.markdown(
-        """
-        -   **Context Window**: Be mindful of the maximum input length the model can handle.
-        -   **Tokenization**:  Understand how text is broken down into tokens.
-        -   **Bias and Safety**:  Be aware of potential biases in the model's output and take steps to mitigate them.
-        -    **Rate Limits**:  The API has usage limits.  The app includes basic handling, but monitor your usage.
-        """
-    )
-    st.subheader("Resources:")
-    st.markdown(
-        """
-        -   [Google Generative AI Course](https://developers.google.com/learn/generative-ai)
-        -   [Prompt Engineering Guide](https://www.promptingguide.ai/)
-        -   [Google AI Platform](https://cloud.google.com/ai-platform)
-        """
-    )
-
 # --- Helper Functions ---
 def code_block(text: str, language: str = "text") -> None:
     """Displays text as a formatted code block in Streamlit."""
     st.markdown(f"```{language}\n{text}\n```", unsafe_allow_html=True)
 
 
-def display_response(response: Any) -> None: # Removed the type hint GenerateContentResponse
+def display_response(response: Any) -> None:
     """Displays the model's response, handling text, and error cases."""
     if response.text:
         st.subheader("Generated Response:")
@@ -85,12 +42,13 @@ def display_response(response: Any) -> None: # Removed the type hint GenerateCon
         st.error(f"Full response object: {response}")  # Print the full response for debugging
 
 
-def generate_with_retry(prompt: str, generation_config: genai.types.GenerationConfig, max_retries: int = 3, delay: int = 5) -> Any: # Removed the type hint
+def generate_with_retry(prompt: str, model_name: str, generation_config: genai.types.GenerationConfig, max_retries: int = 3, delay: int = 5) -> Any:
     """
-    Generates content with retry logic to handle potential API errors (e.g., rate limits).
+    Generates content with retry logic to handle potential API errors (e.g., rate limits, model not found).
 
     Args:
         prompt: The prompt string.
+        model_name: The name of the model to use.
         generation_config: The generation configuration.
         max_retries: Maximum number of retries.
         delay: Delay in seconds between retries.
@@ -99,20 +57,58 @@ def generate_with_retry(prompt: str, generation_config: genai.types.GenerationCo
         The generated response.
 
     Raises:
-        Exception: If the generation fails after maximum retries.
+        Exception: If the generation fails after maximum retries or a critical error occurs.
     """
     for i in range(max_retries):
         try:
+            model = genai.GenerativeModel(model_name)  # Use the selected model name
             response = model.generate_content(prompt, generation_config=generation_config)
-            return response  # Return the response if successful
+            return response
         except Exception as e:
-            st.warning(f"Error during generation (attempt {i + 1}/{max_retries}): {e}")
-            if i < max_retries - 1:
+            error_message = str(e)
+            st.warning(f"Error during generation (attempt {i + 1}/{max_retries}): {error_message}")
+            if "404" in error_message and "not found" in error_message:
+                st.error(
+                    f"Model '{model_name}' is not available or not supported.  Please select a different model."
+                )
+                return None # Return None to indicate failure.  The calling code must handle this.
+            elif i < max_retries - 1:
                 st.info(f"Retrying in {delay} seconds...")
-                time.sleep(delay)  # Use time.sleep for retrying
+                time.sleep(delay)
             else:
                 raise  # Re-raise the exception after the last retry
-    raise Exception("Failed to generate content after maximum retries")  #Should never reach here.
+    raise Exception("Failed to generate content after maximum retries")
+
+
+
+# --- Model Selection and Initialization ---
+available_models = []
+try:
+    available_models = genai.list_models()  # Get the list of available models
+except Exception as e:
+    st.error(f"Error listing models: {e}.  Please check your API key and network connection.")
+    st.stop()
+
+model_names = [
+    model.name for model in available_models if "generateContent" in model.supported_generation_methods
+]  #get models supporting generateContent
+
+if not model_names:
+    st.error(
+        "No models supporting 'generateContent' found.  This application requires a model that supports this method."
+    )
+    st.stop()
+
+default_model = "gemini-pro" if "gemini-pro" in model_names else model_names[0] #select default model
+
+selected_model = st.selectbox("Select a Model:", model_names, index=model_names.index(default_model))  # Let user choose
+
+# Re-initialize the model with the selected name.  This is done *outside* the generate_with_retry loop.
+try:
+    model = genai.GenerativeModel(selected_model)
+except Exception as e:
+    st.error(f"Error initializing model {selected_model}: {e}")
+    st.stop()
 
 # --- Prompting Techniques Section ---
 st.header("Experiment with Prompts")
@@ -129,14 +125,14 @@ prompt_technique = st.selectbox(
         "Chain of Thought (CoT)",
         "Prompt Templates",
         "System Prompt",
-        "Retrieval Augmentation"  # Added Retrieval Augmentation
+        "Retrieval Augmentation"
     ],
     index=0,
 )
 
 prompt_input = st.text_area("Enter your prompt here:", height=150)
 
-# Temperature slider (common to several techniques)
+# Temperature slider
 temperature = st.slider(
     "Temperature:",
     min_value=0.0,
@@ -151,14 +147,15 @@ if st.button("Generate Response"):
         st.warning("Please enter a prompt.")
     else:
         with st.spinner("Generating..."):
-            generation_config = genai.types.GenerationConfig(temperature=temperature)  # Create it once
+            generation_config = genai.types.GenerationConfig(temperature=temperature)
 
             try:
                 if prompt_technique == "Using Delimiters":
                     delimiter = st.text_input("Enter your delimiter (e.g., ###, ---):", "###")
                     processed_prompt = f"Here is the input, with parts separated by '{delimiter}':\n{prompt_input}\n Please process each part separately."
-                    response = generate_with_retry(processed_prompt, generation_config)
-                    display_response(response)
+                    response = generate_with_retry(processed_prompt, selected_model, generation_config)
+                    if response:
+                        display_response(response)
 
                 elif prompt_technique == "Structured Output (JSON)":
                     json_format = st.text_input(
@@ -166,22 +163,24 @@ if st.button("Generate Response"):
                         "{'key1': type, 'key2': type}",
                     )
                     processed_prompt = f"Please provide the output in JSON format, following this structure: {json_format}.  Here is the information: {prompt_input}"
-                    response = generate_with_retry(processed_prompt, generation_config)
-                    try:
-                        json_output = json.loads(response.text)
-                        st.subheader("Generated JSON Output:")
-                        st.json(json_output)
-                    except json.JSONDecodeError:
-                        st.error("Failed to decode JSON. Raw response:")
-                        code_block(response.text, "json")
+                    response = generate_with_retry(processed_prompt, selected_model, generation_config)
+                    if response:
+                        try:
+                            json_output = json.loads(response.text)
+                            st.subheader("Generated JSON Output:")
+                            st.json(json_output)
+                        except json.JSONDecodeError:
+                            st.error("Failed to decode JSON. Raw response:")
+                            code_block(response.text, "json")
 
                 elif prompt_technique == "Checking Assumptions":
                     assumption = st.text_input(
                         "State the assumption you want the model to check:", "The text is about a historical event."
                     )
                     processed_prompt = f"First, check if the following assumption is true: '{assumption}'.  Then, answer the prompt: {prompt_input}"
-                    response = generate_with_retry(processed_prompt, generation_config)
-                    display_response(response)
+                    response = generate_with_retry(processed_prompt, selected_model, generation_config)
+                    if response:
+                        display_response(response)
 
                 elif prompt_technique == "Few-Shot Prompting":
                     example1_input = st.text_area("Example 1 Input:", height=50)
@@ -195,18 +194,20 @@ if st.button("Generate Response"):
                         processed_prompt += f"Input: {example2_input}\nOutput: {example2_output}\n"
                     processed_prompt += f"\nNow, answer the following:\nInput: {prompt_input}"
 
-                    response = generate_with_retry(processed_prompt, generation_config)
-                    display_response(response)
+                    response = generate_with_retry(processed_prompt, selected_model, generation_config)
+                    if response:
+                        display_response(response)
 
                 elif prompt_technique == "Temperature Control":
-                    # The temperature slider is already handled
-                    response = generate_with_retry(prompt_input, generation_config)
-                    display_response(response)
+                    response = generate_with_retry(prompt_input, selected_model, generation_config)
+                    if response:
+                        display_response(response)
 
                 elif prompt_technique == "Chain of Thought (CoT)":
                     cot_prompt = f"Let's think step by step. {prompt_input}"
-                    response = generate_with_retry(cot_prompt, generation_config)
-                    display_response(response)
+                    response = generate_with_retry(cot_prompt, selected_model, generation_config)
+                    if response:
+                        display_response(response)
 
                 elif prompt_technique == "Prompt Templates":
                     template_name = st.selectbox(
@@ -224,35 +225,43 @@ if st.button("Generate Response"):
                     else:
                         processed_prompt = prompt_input
 
-                    response = generate_with_retry(processed_prompt, generation_config)
-                    display_response(response)
+                    response = generate_with_retry(processed_prompt, selected_model, generation_config)
+                    if response:
+                        display_response(response)
 
                 elif prompt_technique == "System Prompt":
                     system_prompt_text = st.text_area(
                         "Enter system prompt:", "You are a helpful and informative assistant.", height=100
                     )
-                    user_prompt = f"{prompt_input}"  # User's input
+                    user_prompt = f"{prompt_input}"
 
                     response = generate_with_retry(
                         contents=[
-                            genai.Content(role="system", parts=[genai.Part(text=system_prompt_text)]), #send system prompt
-                            genai.Content(role="user", parts=[genai.Part(text=user_prompt)]), # send user prompt
+                            genai.Content(role="system", parts=[genai.Part(text=system_prompt_text)]),
+                            genai.Content(role="user", parts=[genai.Part(text=user_prompt)]),
                         ],
+                        model=selected_model, # Pass the model name here as well
                         generation_config=generation_config,
                     )
-                    display_response(response)
-                
+                    if response:
+                        display_response(response)
+
                 elif prompt_technique == "Retrieval Augmentation":
-                    context_text = st.text_area("Enter context text (knowledge base):",
-                                                "This is the context the model can use to answer the question.",
-                                                height=150)
+                    context_text = st.text_area(
+                        "Enter context text (knowledge base):",
+                        "This is the context the model can use to answer the question.",
+                        height=150,
+                    )
                     processed_prompt = f"Given the following context: \n\n {context_text} \n\n Answer the following question: {prompt_input}"
-                    response = generate_with_retry(processed_prompt, generation_config)
-                    display_response(response)
+                    response = generate_with_retry(processed_prompt, selected_model, generation_config)
+                    if response:
+                        display_response(response)
 
                 else:  # Simple Instruction
-                    response = generate_with_retry(prompt_input, generation_config)
-                    display_response(response)
+                    response = generate_with_retry(prompt_input, selected_model, generation_config)
+                    if response:
+                        display_response(response)
 
             except Exception as e:
                 st.error(f"An error occurred: {e}")
+