app.py

"""
Streamlit app for running a multi-agent health science research workflow.

This app provides a user interface for configuring, starting, monitoring, 
and downloading reports generated by a multi-agent AI system. Users can 
select LLM type, set workflow parameters, and view real-time progress and logs.
"""

import asyncio
import datetime
import logging
import os
import queue
import threading
import time
from dataclasses import asdict
from pathlib import Path
from typing import List, Optional

from dotenv import load_dotenv
import streamlit as st

# Import your existing classes
from multi_agent_workflow import LLMType, MultiAgentWorkflow, WorkflowConfig

load_dotenv()

# Configure Streamlit page
st.set_page_config(
    page_title="Health Science Research Agent",
    page_icon="🔬",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        text-align: center;
        padding: 1rem 0;
        background: linear-gradient(90deg, #008080 0%, #764ba2 100%);
        color: white;
        border-radius: 10px;
        margin-bottom: 2rem;
    }
    .agent-status {
        padding: 0.5rem;
        border-radius: 5px;
        margin: 0.5rem 0;
        border-left: 4px solid #667eea;
        background-color: #f0f2f6;
        color: #222;
        font-weight: 500;
    }
    .tool-result {
        background-color: #e8f4fd;
        padding: 0.5rem;
        border-radius: 5px;
        border-left: 4px solid #1f77b4;
        margin: 0.5rem 0;
        color: #222;
        font-weight: 500;
    }
    .success-box {
        background-color: #d4edda;
        color: #155724;
        padding: 1rem;
        border-radius: 5px;
        border: 1px solid #c3e6cb;
        margin: 1rem 0;
    }
    .error-box {
        background-color: #f8d7da;
        color: #721c24;
        padding: 1rem;
        border-radius: 5px;
        border: 1px solid #f5c6cb;
        margin: 1rem 0;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'workflow_running' not in st.session_state:
    st.session_state.workflow_running = False
if 'workflow_results' not in st.session_state:
    st.session_state.workflow_results = []
if 'config' not in st.session_state:
    st.session_state.config = None


def create_sidebar() -> Optional[WorkflowConfig]:
    """Create the configuration sidebar"""
    st.sidebar.title("🔧 Configuration")

    # LLM Configuration
    llm_options = [llm_type.name for llm_type in LLMType]
    llm_name = st.sidebar.selectbox(
        "LLM Name",
        options=llm_options,
        index=0,
        help="Select the language model to use"
    )
    # Convert back to enum for config
    llm_enum = LLMType[llm_name]

    # Report Configuration
    st.sidebar.subheader("Report Settings")
    target_word_count = st.sidebar.slider(
        "Target Word Count",
        min_value=1000,
        max_value=10000,
        value=5000,
        step=500,
        help="Target length for the generated report"
    )

    min_developments = st.sidebar.slider(
        "Minimum Developments",
        min_value=1,
        max_value=10,
        value=5,
        help="Minimum number of health developments to include"
    )

    max_developments = st.sidebar.slider(
        "Maximum Developments",
        min_value=1,
        max_value=15,
        value=7,
        help="Maximum number of health developments to include"
    )

    # File Configuration
    st.sidebar.subheader("File Settings")
    docs_dir = st.sidebar.text_input(
        "Output Directory",
        value="./docs",
        help="Directory where reports will be saved"
    )

    report_filename = st.sidebar.text_input(
        "Report Filename",
        value="health_report.md",
        help="Base filename for the report (timestamp will be added)"
    )

    # Workflow Configuration
    st.sidebar.subheader("Workflow Settings")
    max_iterations = st.sidebar.slider(
        "Max Iterations",
        min_value=1,
        max_value=20,
        value=10,
        help="Maximum number of iterations before stopping"
    )

    timeout_seconds = st.sidebar.slider(
        "Timeout (minutes)",
        min_value=5,
        max_value=60,
        value=30,
        help="Maximum time to run the workflow"
    ) * 60  # Convert to seconds

    # Logging Configuration
    log_level = st.sidebar.selectbox(
        "Log Level",
        options=["DEBUG", "INFO", "WARNING", "ERROR"],
        index=1,
        help="Logging verbosity level"
    )

    # Trusted Sources Configuration
    st.sidebar.subheader("Trusted Sources")
    with st.sidebar.expander("Edit Trusted Sources"):
        default_sources = [
            "PubMed (pubmed.ncbi.nlm.nih.gov)",
            "Google Scholar",
            "The Lancet (thelancet.com)",
            "Nature Medicine (nature.com/nm)",
            "ScienceDirect (sciencedirect.com)",
            "ScienceDaily (sciencedaily.com)",
            "Medical News Today (medicalnewstoday.com)",
            "NIH (nih.gov)",
            "WHO (who.int)",
            "CDC (cdc.gov)"
        ]

        sources_text = st.text_area(
            "Sources (one per line)",
            value="\n".join(default_sources),
            height=200,
            help="List of trusted sources, one per line"
        )
        trusted_sources = [s.strip()
                           for s in sources_text.split('\n') if s.strip()]

    # Create configuration object
    try:
        config = WorkflowConfig(
            tavily_api_key=os.getenv("TAVILY_API_KEY", ""),
            llm_type=llm_enum,  # Pass enum value as before
            docs_dir=docs_dir,
            default_report_filename=report_filename,
            target_word_count=target_word_count,
            min_developments=min_developments,
            max_developments=max_developments,
            max_iterations=max_iterations,
            timeout_seconds=timeout_seconds,
            trusted_sources=trusted_sources,
            log_level=log_level
        )
        st.session_state.config = config
        return config
    except ValueError as e:
        st.sidebar.error(f"Configuration Error: {e}")
        return None


def display_workflow_status(event_queue):
    """Display real-time workflow status"""
    status_container = st.container()

    with status_container:
        if not event_queue.empty():
            try:
                while not event_queue.empty():
                    event_data = event_queue.get_nowait()

                    if event_data['type'] == 'agent_change':
                        st.markdown(f"""
                        <div class="agent-status">
                            🤖 <strong>Agent:</strong> {event_data['agent']} 
                            (Iteration {event_data['iteration']})
                        </div>
                        """, unsafe_allow_html=True)

                    elif event_data['type'] == 'agent_output':
                        if event_data['content']:
                            st.write("📤 **Output:**", event_data['content'])
                        if event_data['tool_calls']:
                            st.write("🛠️ **Planning to use tools:**",
                                     event_data['tool_calls'])

                    elif event_data['type'] == 'tool_result':
                        st.markdown(f"""
                        <div class="tool-result">
                            🔧 <strong>Tool Result ({event_data['tool_name']}):</strong><br>
                            <strong>Arguments:</strong> {event_data['arguments']}<br>
                            <strong>Output:</strong> {event_data['output'][:200]}...
                        </div>
                        """, unsafe_allow_html=True)

                    elif event_data['type'] == 'tool_call':
                        st.write(
                            f"🔨 **Calling Tool:** {event_data['tool_name']}")
                        st.write(f"**Arguments:** {event_data['arguments']}")

                    # Store in session state for persistence
                    if not st.session_state.workflow_results or st.session_state.workflow_results[-1] != event_data:
                        st.session_state.workflow_results.append(event_data)

            except queue.Empty:
                pass


async def run_workflow_async(config, event_queue):
    """Run the workflow asynchronously and capture events"""
    try:
        # Custom workflow class to capture events
        class StreamlitMultiAgentWorkflow(MultiAgentWorkflow):
            def __init__(self, config, event_queue):
                super().__init__(config)
                self.event_queue = event_queue

            async def run(self):
                try:
                    self.logger.info("Starting multi-agent workflow")
                    prompt = self.config.get_prompt_template()

                    handler = self.agent_workflow.run(user_msg=prompt)

                    current_agent = None
                    iteration_count = 0

                    async for event in handler.stream_events():
                        if iteration_count >= self.config.max_iterations:
                            self.logger.warning("Reached max iterations")
                            break

                        # Capture agent changes
                        if (hasattr(event, "current_agent_name") and
                                event.current_agent_name != current_agent):
                            current_agent = event.current_agent_name
                            iteration_count += 1

                            self.event_queue.put({
                                'type': 'agent_change',
                                'agent': current_agent,
                                'iteration': iteration_count,
                                'timestamp': datetime.datetime.now()
                            })

                        # Capture agent output
                        elif hasattr(event, 'response') and hasattr(event.response, 'content'):
                            tool_calls = []
                            if hasattr(event, 'tool_calls') and event.tool_calls:
                                tool_calls = [
                                    call.tool_name for call in event.tool_calls]

                            self.event_queue.put({
                                'type': 'agent_output',
                                'content': event.response.content,
                                'tool_calls': tool_calls,
                                'timestamp': datetime.datetime.now()
                            })

                        # Capture tool results
                        elif hasattr(event, 'tool_name') and hasattr(event, 'tool_output'):
                            self.event_queue.put({
                                'type': 'tool_result',
                                'tool_name': event.tool_name,
                                'arguments': getattr(event, 'tool_kwargs', {}),
                                'output': str(event.tool_output),
                                'timestamp': datetime.datetime.now()
                            })

                        # Capture tool calls
                        elif hasattr(event, 'tool_name') and hasattr(event, 'tool_kwargs'):
                            self.event_queue.put({
                                'type': 'tool_call',
                                'tool_name': event.tool_name,
                                'arguments': event.tool_kwargs,
                                'timestamp': datetime.datetime.now()
                            })

                    self.event_queue.put({
                        'type': 'completion',
                        'message': 'Workflow completed successfully',
                        'timestamp': datetime.datetime.now()
                    })

                except Exception as e:
                    self.event_queue.put({
                        'type': 'error',
                        'message': str(e),
                        'timestamp': datetime.datetime.now()
                    })

        workflow = StreamlitMultiAgentWorkflow(config, event_queue)
        await workflow.run()

    except Exception as e:
        event_queue.put({
            'type': 'error',
            'message': f"Workflow initialization failed: {str(e)}",
            'timestamp': datetime.datetime.now()
        })


def run_workflow_in_thread(config, event_queue):
    """Run the async workflow in a separate thread"""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        loop.run_until_complete(run_workflow_async(config, event_queue))
    finally:
        # Cancel all pending tasks before closing the loop
        pending = asyncio.all_tasks(loop)
        for task in pending:
            task.cancel()
        loop.run_until_complete(asyncio.gather(
            *pending, return_exceptions=True))
        loop.close()


def main():
    """Main Streamlit application"""
    # Header
    st.markdown("""
    <div class="main-header">
        <h1>🔬 Health Science Research Agent</h1>
        <p>Multi-Agent AI System for Health Science Research & Report Generation</p>
    </div>
    """, unsafe_allow_html=True)

    # Create sidebar configuration
    config = create_sidebar()

    # Main content area
    col1, col2 = st.columns([2, 1])

    with col1:
        st.header("🚀 Workflow Control")

        # Display current configuration
        if config:
            with st.expander("📋 Current Configuration", expanded=False):
                config_dict = asdict(config)
                for key, value in config_dict.items():
                    if key != 'tavily_api_key':  # Don't display API key
                        st.write(f"**{key}:** {value}")

        # Control buttons
        col_start, col_stop, col_clear = st.columns(3)

        with col_start:
            start_button = st.button(
                "▶️ Start Workflow",
                disabled=st.session_state.workflow_running or config is None,
                help="Start the multi-agent research workflow"
            )

        with col_stop:
            stop_button = st.button(
                "⏹️ Stop Workflow",
                disabled=not st.session_state.workflow_running,
                help="Stop the currently running workflow"
            )

        with col_clear:
            clear_button = st.button(
                "🗑️ Clear Results",
                help="Clear all workflow results and logs"
            )

        # Handle button clicks
        if start_button and config:
            st.session_state.workflow_running = True
            st.session_state.workflow_results = []

            # Create event queue for communication
            event_queue = queue.Queue()

            # Start workflow in background thread
            workflow_thread = threading.Thread(
                target=run_workflow_in_thread,
                args=(config, event_queue)
            )
            workflow_thread.daemon = True
            workflow_thread.start()

            # Store thread and queue in session state
            st.session_state.workflow_thread = workflow_thread
            st.session_state.event_queue = event_queue

            st.success("🚀 Workflow started! Monitor progress below.")
            st.rerun()

        if stop_button:
            st.session_state.workflow_running = False
            st.warning("⏹️ Workflow stopped by user.")
            st.rerun()

        if clear_button:
            st.session_state.workflow_results = []
            st.info("🗑️ Results cleared.")
            st.rerun()

    with col2:
        st.header("📊 Status")

        if st.session_state.workflow_running:
            with st.spinner("Workflow is running. Please wait..."):
                st.markdown("🟢 **Status:** Running")
                st.markdown("⏱️ **Started:** " +
                            datetime.datetime.now().strftime("%H:%M:%S"))
        else:
            st.markdown("🔴 **Status:** Stopped")

        # Display report files
        if config and Path(config.docs_dir).exists():
            st.subheader("📁 Generated Reports")
            report_files = list(Path(config.docs_dir).glob("*.md"))

            if report_files:
                for file_path in sorted(report_files, reverse=True):
                    file_name = file_path.name
                    file_size = file_path.stat().st_size
                    mod_time = datetime.datetime.fromtimestamp(
                        file_path.stat().st_mtime)

                    with st.expander(f"📄 {file_name}"):
                        st.write(f"**Size:** {file_size} bytes")
                        st.write(
                            f"**Modified:** {mod_time.strftime('%Y-%m-%d %H:%M:%S')}")

                        # Download button
                        with open(file_path, 'r', encoding='utf-8') as f:
                            content = f.read()

                        st.download_button(
                            label="⬇️ Download",
                            data=content,
                            file_name=file_name,
                            mime="text/markdown"
                        )
            else:
                st.info("No reports generated yet")

    # Workflow Results Section
    st.header("📈 Workflow Progress")

    # Auto-refresh when workflow is running
    if st.session_state.workflow_running and hasattr(st.session_state, 'event_queue'):
        display_workflow_status(st.session_state.event_queue)

        # Check if workflow thread is still alive
        if hasattr(st.session_state, 'workflow_thread'):
            if not st.session_state.workflow_thread.is_alive():
                st.session_state.workflow_running = False
                st.success("✅ Workflow completed!")

        # Auto-refresh every 2 seconds when running
        time.sleep(2)
        st.rerun()

    # Display historical results
    if st.session_state.workflow_results:
        st.subheader("📋 Workflow Log")

        # Display results in reverse chronological order
        # Show last 20
        for i, result in enumerate(reversed(st.session_state.workflow_results[-20:])):
            timestamp = result.get(
                'timestamp', datetime.datetime.now()).strftime("%H:%M:%S")

            if result['type'] == 'agent_change':
                st.markdown(f"""
                <div class="agent-status">
                    <small>{timestamp}</small><br>
                    🤖 <strong>Agent:</strong> {result['agent']} (Iteration {result['iteration']})
                </div>
                """, unsafe_allow_html=True)

            elif result['type'] == 'completion':
                st.markdown(f"""
                <div class="success-box">
                    <small>{timestamp}</small><br>
                    ✅ <strong>Completed:</strong> {result['message']}
                </div>
                """, unsafe_allow_html=True)

            elif result['type'] == 'error':
                st.markdown(f"""
                <div class="error-box">
                    <small>{timestamp}</small><br>
                    ❌ <strong>Error:</strong> {result['message']}
                </div>
                """, unsafe_allow_html=True)


if __name__ == "__main__":
    main()