Enhance app.py to improve error handling during email generation. Added try-except blocks in both conversation_predict and update_ui functions to catch exceptions and display error messages in the Streamlit interface. This change ensures a more robust user experience by providing feedback in case of issues during response generation.

app.py

@@ -113,18 +113,20 @@ def conversation_predict(input_text: str, cv_sections: Dict[str, str]):
     prompt = create_email_prompt(input_text, cv_sections)
     
     # Use the streaming API
-    for response in client.text_generation(
-        model="google/gemma-2b-it",
-        prompt=prompt,
-        max_new_tokens=512,
-        temperature=0.7,
-        top_p=0.95,
-        stream=True
-    ):
-        if hasattr(response, 'token'):  # Handle different response formats
-            yield response.token.text
-        else:
-            yield response.generated_text
+    try:
+        for response in client.text_generation(
+            model="google/gemma-2b-it",
+            prompt=prompt,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.95,
+            stream=True
+        ):
+            # The streaming response returns text directly
+            yield response
+    except Exception as e:
+        st.error(f"Error generating response: {str(e)}")
+        yield ""
 
 def respond(
     message: str,
@@ -182,16 +184,19 @@ def update_ui(message, cv_file, cv_sections):
         if message and cv_file and isinstance(cv_sections, dict):
             email_text = ""
             # Stream the response
-            for chunk in conversation_predict(message, cv_sections):
-                if chunk:
-                    email_text += chunk
-                    # Update the text area with each chunk
-                    email_placeholder.text_area(
-                        "Generated Email",
-                        value=email_text,
-                        height=400,
-                        key="email_output"
-                    )
+            try:
+                for chunk in conversation_predict(message, cv_sections):
+                    if chunk:
+                        email_text += chunk
+                        # Update the text area with each chunk
+                        email_placeholder.text_area(
+                            "Generated Email",
+                            value=email_text,
+                            height=400,
+                            key="email_output"
+                        )
+            except Exception as e:
+                st.error(f"Error during email generation: {str(e)}")
         else:
             st.warning("Please upload a CV and enter a job description.")