Update app.py

app.py

@@ -185,27 +185,30 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 def format_json_output(raw_output):
-    """Format agent output into proper JSON"""
+    """Format CrewOutput or raw string into proper JSON structure"""
     try:
-        # Try direct JSON parsing
-        return json.loads(raw_output)
-    except json.JSONDecodeError:
-        # If direct parsing fails, try to extract JSON block
+        # Handle CrewOutput object
+        if hasattr(raw_output, 'raw_output'):
+            raw_text = str(raw_output.raw_output)
+        else:
+            raw_text = str(raw_output)
+
+        # Try to find and parse JSON structure
         json_pattern = r"\{[\s\S]*\}"
-        match = re.search(json_pattern, raw_output)
+        match = re.search(json_pattern, raw_text)
         if match:
             try:
                 return json.loads(match.group())
             except:
                 pass
-        
-        # If JSON extraction fails, create structured format
+
+        # If no JSON found, create structured format
         return {
-            "exec_summary": "Error formatting report",
-            "detailed_report": raw_output,
-            "sources": []
+            "exec_summary": extract_section(raw_text, "Executive Summary"),
+            "detailed_report": raw_text,
+            "sources": extract_sources(raw_text)
         }
-
+        
 def display_thinking_animation():
     """Display thinking animation dots"""
     return """
@@ -359,12 +362,33 @@ def run_market_research(topic: str, progress_container, chat_container):
         return format_json_output(result)
 
     except Exception as e:
-        st.error(f"Error: {str(e)}")
-        return None
+        st.error(f"Error formatting output: {str(e)}")
+        return {
+            "exec_summary": "Error formatting report",
+            "detailed_report": raw_text if 'raw_text' in locals() else str(raw_output),
+            "sources": []
+        }
+
+def extract_section(text, section_name):
+    """Extract a section from the text"""
+    pattern = f"{section_name}.*?\n(.*?)(?=\n\n|$)"
+    match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
+    return match.group(1).strip() if match else ""
+
+def extract_sources(text):
+    """Extract sources from the text"""
+    sources = []
+    source_pattern = r"Source:.*?(?:\n|$)|\[.*?\]|\(https?://.*?\)"
+    matches = re.finditer(source_pattern, text, re.MULTILINE)
+    return [match.group().strip() for match in matches]
 
 def main():
     st.title("🤖 AI Market Research Generator")
     
+    # Initialize session state for button control
+    if 'generating' not in st.session_state:
+        st.session_state.generating = False
+    
     tab1, tab2 = st.tabs(["Generate Report", "View Reports"])
     
     with tab1:
@@ -377,48 +401,75 @@ def main():
                 placeholder="e.g., Electric Vehicles Market"
             )
             
-            if st.button("Generate Report", type="primary"):
-                if not topic:
-                    st.error("Please enter a research topic")
-                    return
+            # Dynamic button state
+            if st.session_state.generating:
+                st.button("Generating Report...", disabled=True)
+            else:
+                if st.button("Generate Report", type="primary"):
+                    if not topic:
+                        st.error("Please enter a research topic")
+                    else:
+                        st.session_state.generating = True
+                        st.rerun()
+        
+        # Only show chat window if generating
+        if st.session_state.generating:
+            with col2:
+                st.markdown('<div class="chat-window">', unsafe_allow_html=True)
+                chat_container = st.empty()
+                st.markdown('</div>', unsafe_allow_html=True)
                 
-                # Create containers for progress and chat
                 progress_container = st.container()
                 
-                with col2:
-                    st.markdown('<div class="chat-window">', unsafe_allow_html=True)
-                    chat_container = st.container()
-                    st.markdown('</div>', unsafe_allow_html=True)
-                
-                result = run_market_research(topic, progress_container, chat_container)
-                
-                if result:
-                    st.session_state.current_report = result
-                    st.success("Report generated successfully! View it in the Reports tab.")
+                try:
+                    result = run_market_research(topic, progress_container, chat_container)
+                    
+                    if result:
+                        st.session_state.current_report = result
+                        st.success("Report generated successfully! View it in the Reports tab.")
+                except Exception as e:
+                    st.error(f"Error generating report: {str(e)}")
+                finally:
+                    st.session_state.generating = False
+                    st.rerun()
     
     with tab2:
         if hasattr(st.session_state, 'current_report'):
-            st.subheader("📊 Generated Report")
+            report = st.session_state.current_report
             
-            # Display Executive Summary
-            with st.expander("Executive Summary", expanded=True):
-                st.markdown(st.session_state.current_report['exec_summary'])
+            # Display Executive Summary in a card
+            st.markdown("""
+                <div style='background-color: white; padding: 20px; border-radius: 10px; margin: 10px 0;'>
+                    <h3>Executive Summary</h3>
+                    <p>{}</p>
+                </div>
+            """.format(report['exec_summary']), unsafe_allow_html=True)
             
-            # Display Detailed Report
-            with st.expander("Detailed Report", expanded=True):
-                st.markdown(st.session_state.current_report['detailed_report'])
+            # Display Detailed Report in a card
+            st.markdown("""
+                <div style='background-color: white; padding: 20px; border-radius: 10px; margin: 10px 0;'>
+                    <h3>Detailed Report</h3>
+                    <p>{}</p>
+                </div>
+            """.format(report['detailed_report']), unsafe_allow_html=True)
             
-            # Display Sources
-            with st.expander("Sources", expanded=True):
-                for source in st.session_state.current_report['sources']:
-                    st.markdown(f"- {source}")
+            # Display Sources in a card
+            st.markdown("""
+                <div style='background-color: white; padding: 20px; border-radius: 10px; margin: 10px 0;'>
+                    <h3>Sources</h3>
+                    <ul>
+                        {}
+                    </ul>
+                </div>
+            """.format('\n'.join([f"<li>{source}</li>" for source in report['sources']])), 
+            unsafe_allow_html=True)
             
             # Download buttons
             col1, col2 = st.columns(2)
             with col1:
                 st.download_button(
                     "Download Full Report",
-                    data=st.session_state.current_report['detailed_report'],
+                    data=report['detailed_report'],
                     file_name=f"{topic}_report.md",
                     mime="text/markdown"
                 )