sci/research_page.py

import streamlit as st
from datetime import datetime
from colorprinter.print_color import *

from _base_class import StreamlitBaseClass
from projects_page import Project
from agent_research import ResearchReport, MasterAgent, StructureAgent, ToolAgent, ArchiveAgent, process_step
import os
import json


class ResearchPage(StreamlitBaseClass):
    """
    ResearchPage - A Streamlit interface for deep research using AI agents.
    
    This class provides a user interface for conducting in-depth research using
    multiple specialized AI agents working together. It allows users to input
    research questions, track progress, and view detailed research reports.
    
    Attributes:
        username (str): The username of the current user.
        project_name (str): Name of the selected project.
        project (Project): Project instance the research is associated with.
        page_name (str): Name of the current page ("Research").
        research_state (dict): Dictionary tracking the current state of research.
        report (ResearchReport): Instance for tracking research progress and results.
    
    Methods:
        run(): Main method to render the research interface and handle interactions.
        sidebar_actions(): Renders sidebar elements for selecting projects and research options.
        start_new_research(): Initiates a new research session.
        view_saved_reports(): Displays a list of saved research reports.
        display_report(): Renders a research report in the Streamlit interface.
        show_research_progress(): Displays the current research progress.
    """
    def __init__(self, username):
        super().__init__(username=username)
        self.project_name = None
        self.project = None
        self.page_name = "Research"
        
        # Research state tracking
        self.research_state = {
            "in_progress": False,
            "completed": False,
            "question": None,
            "started_at": None,
            "report": None,
            "current_step": None,
            "steps_completed": 0,
            "total_steps": 0
        }
        
        self.report = None
        
        # Initialize attributes from session state if available
        if self.page_name in st.session_state:
            for k, v in st.session_state[self.page_name].items():
                setattr(self, k, v)
                
        # Create reports directory if it doesn't exist
        os.makedirs(f"/home/lasse/sci/reports", exist_ok=True)

    def run(self):
        self.update_current_page("Research")
        self.sidebar_actions()
        
        st.title("Deep Research")
        
        if not self.project:
            st.warning("Please select a project to start researching.")
            return
            
        # Main interface
        if self.research_state["in_progress"]:
            self.show_research_progress()
        elif self.research_state["completed"]:
            self.display_report(self.research_state["report"])
        else:
            # Input for new research
            st.subheader(f"New Research for Project: {self.project_name}")
            with st.form("research_form"):
                question = st.text_area("Enter your research question:", 
                                       help="Be specific about what you want to research. Complex questions will be broken down into sub-questions.")
                start_button = st.form_submit_button("Start Research")
                
                if start_button and question:
                    self.start_new_research(question)
                    st.rerun()
            
            # Option to view saved reports
            with st.expander("View Saved Reports"):
                self.view_saved_reports()
    
    def sidebar_actions(self):
        with st.sidebar:
            with st.form("select_project"):
                self.project = self.choose_project("Project for research:")
                submitted = st.form_submit_button("Select Project")
                
                if submitted and self.project:
                    self.project_name = self.project.name
                    st.success(f"Selected project: {self.project_name}")

            if self.research_state["in_progress"]:
                st.info(f"Research in progress: {self.research_state['question']}")
                if st.button("Cancel Research"):
                    self.research_state["in_progress"] = False
                    st.rerun()
                    
            elif self.research_state["completed"]:
                if st.button("Start New Research"):
                    self.research_state["completed"] = False
                    self.research_state["report"] = None
                    st.rerun()

    def start_new_research(self, question):
        """Initiates a new research session with the given question"""
        self.research_state["question"] = question
        self.research_state["in_progress"] = True
        self.research_state["completed"] = False
        self.research_state["started_at"] = datetime.now().isoformat()
        
        # Initialize the research report
        self.report = ResearchReport(
            question=question, 
            username=self.username, 
            project_name=self.project_name
        )
        
        # Save current state
        st.session_state[self.page_name] = {
            "project_name": self.project_name,
            "project": self.project,
            "research_state": self.research_state,
            "report": self.report
        }
        
        # Start a new thread to run the research process
        # In a production environment, you might want to use a background job
        # For now, we'll run it in the main thread with streamlit spinner
        with st.spinner("Research in progress... This may take several minutes."):
            try:
                # Initialize agents
                master_agent = MasterAgent(
                    username=self.username, 
                    project=self.project, 
                    report=self.report, 
                    chat=True
                )
                structure_agent = StructureAgent(
                    username=self.username, 
                    model="small", 
                    report=self.report
                )
                tool_agent = ToolAgent(
                    username=self.username,
                    model="tools",
                    system_message="You are an assistant with tools. Always choose a tool to help with the task.",
                    report=self.report,
                    project=self.project,
                    chat=True
                )
                archive_agent = ArchiveAgent(
                    username=self.username,
                    report=self.report,
                    project=self.project,
                    system_message="You are an assistant specialized in reading and summarizing research information.",
                    chat=True
                )
                
                # Track the research state in the master agent
                master_agent.research_state["original_question"] = question
                
                # Execute the research workflow
                # 1. Create research plan
                st.text("Creating research plan...")
                research_plan = master_agent.make_plan(question)
                self.report.log_plan(research_plan)
                
                # 2. Structure the plan
                st.text("Structuring research plan...")
                structured_plan = structure_agent.make_structured(research_plan, question)
                self.report.log_plan(research_plan, structured_plan.model_dump())
                
                # Update total steps count
                self.research_state["total_steps"] = len(structured_plan.steps)
                
                # 3. Execute the plan step by step
                execution_results = {}
                
                for step_name, tasks in structured_plan.steps.items():
                    st.text(f"Processing step: {step_name}")
                    self.research_state["current_step"] = step_name
                    self.research_state["steps_completed"] += 1
                    
                    # Collect all task descriptions in this step
                    step_tasks = [
                        {"task_name": task_name, "task_description": task_description}
                        for task_name, task_description in tasks
                    ]
                    
                    # Process the entire step
                    step_result = process_step(
                        step_name, step_tasks, master_agent, tool_agent, archive_agent
                    )
                    execution_results[step_name] = step_result
                
                # 4. Evaluate if more steps are needed
                st.text("Evaluating research plan...")
                plan_evaluation = master_agent.evaluate_plan(execution_results)
                self.report.log_plan_evaluation(plan_evaluation)
                
                # 5. Write the final report
                st.text("Writing final report...")
                final_report = master_agent.write_report(execution_results)
                self.report.log_final_report(final_report)
                
                # 6. Save the reports
                timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
                report_path = f"/home/lasse/sci/reports/research_report_{self.username}_{timestamp}"
                
                # Save JSON report
                json_path = f"{report_path}.json"
                with open(json_path, "w") as f:
                    json.dump(self.report.get_full_report(), f, indent=2)
                
                # Save markdown report
                markdown_report = self.report.get_markdown_report()
                markdown_path = f"{report_path}.md"
                with open(markdown_path, "w") as f:
                    f.write(markdown_report)
                
                # Update research state
                self.research_state["in_progress"] = False
                self.research_state["completed"] = True
                self.research_state["report"] = {
                    "json_path": json_path,
                    "markdown_path": markdown_path,
                    "report_data": self.report.get_full_report(),
                    "markdown_content": markdown_report
                }
                
            except Exception as e:
                st.error(f"An error occurred during research: {str(e)}")
                import traceback
                st.code(traceback.format_exc())
                self.research_state["in_progress"] = False
        
        # Update session state
        st.session_state[self.page_name] = {
            "project_name": self.project_name,
            "project": self.project,
            "research_state": self.research_state,
            "report": self.report
        }

    def view_saved_reports(self):
        """Displays a list of saved research reports"""
        reports_dir = "/home/lasse/sci/reports"
        if not os.path.exists(reports_dir):
            st.info("No saved reports found.")
            return
            
        # Get all report files
        json_files = [f for f in os.listdir(reports_dir) if f.endswith('.json') and f.startswith('research_report')]
        
        if not json_files:
            st.info("No saved reports found.")
            return
            
        for file in sorted(json_files, reverse=True):
            file_path = os.path.join(reports_dir, file)
            try:
                with open(file_path, 'r') as f:
                    report_data = json.load(f)
                    
                # Extract basic info
                question = report_data.get("metadata", {}).get("question", "Unknown question")
                project = report_data.get("metadata", {}).get("project_name", "No project")
                started_at = report_data.get("metadata", {}).get("started_at", "Unknown time")
                
                # Format the date
                try:
                    date_obj = datetime.fromisoformat(started_at)
                    date_str = date_obj.strftime("%Y-%m-%d %H:%M")
                except:
                    date_str = started_at
                
                # Create an expandable section for each report
                st.markdown(f"_{question} ({project} - {date_str})_")
                st.markdown(f"**Project:** {project}")
                st.markdown(f"**Date:** {date_str}")
                
                # Button to view full report
                if st.button("View Full Report", key=f"view_{file}"):
                    # Load corresponding markdown file if it exists
                    md_file = file.replace('.json', '.md')
                    md_path = os.path.join(reports_dir, md_file)
                    
                    if os.path.exists(md_path):
                        with open(md_path, 'r') as f:
                            markdown_content = f.read()
                    else:
                        markdown_content = None
                        
                    self.research_state["completed"] = True
                    self.research_state["report"] = {
                        "json_path": file_path,
                        "markdown_path": md_path if os.path.exists(md_path) else None,
                        "report_data": report_data,
                        "markdown_content": markdown_content
                    }
                    st.rerun()
                        
            except Exception as e:
                st.error(f"Error loading report {file}: {str(e)}")

    def display_report(self, report_data):
        """Renders a research report in the Streamlit interface"""
        if not report_data:
            st.warning("No report data available.")
            return
            
        st.title("Research Report")
        
        # Get report data
        markdown_content = report_data.get("markdown_content")
        json_data = report_data.get("report_data")
        
        if markdown_content:
            # Display the markdown report
            st.markdown(markdown_content)
        elif json_data:
            # Fallback to displaying JSON data in a more readable format
            question = json_data.get("metadata", {}).get("question", "Unknown question")
            st.header(f"Research on: {question}")
            
            # Display metadata
            st.subheader("Metadata")
            metadata = json_data.get("metadata", {})
            st.markdown(f"**Project:** {metadata.get('project_name', 'None')}")
            st.markdown(f"**Started:** {metadata.get('started_at', 'Unknown')}")
            st.markdown(f"**Finished:** {metadata.get('finished_at', 'Unknown')}")
            
            # Display final report
            st.subheader("Research Findings")
            st.markdown(json_data.get("final_report", "No final report available."))
            
            # Display steps
            st.subheader("Research Steps")
            steps = json_data.get("steps", {})
            for step_name, step_data in steps.items():
                with st.expander(step_name):
                    st.markdown(f"**Summary:** {step_data.get('summary', 'No summary available.')}")
                    
                    # Display tools used
                    st.markdown("**Tools used:**")
                    for tool in step_data.get("tools_used", []):
                        st.markdown(f"- {tool.get('tool', 'Unknown tool')} with query: _{tool.get('args', {}).get('query', 'No query')}_")
            
        else:
            st.error("No report content available to display.")
        
        # Download buttons
        col1, col2 = st.columns(2)
        with col1:
            if report_data.get("markdown_path") and os.path.exists(report_data["markdown_path"]):
                with open(report_data["markdown_path"], "r") as f:
                    markdown_content = f.read()
                    st.download_button(
                        label="Download as Markdown",
                        data=markdown_content,
                        file_name=os.path.basename(report_data["markdown_path"]),
                        mime="text/markdown"
                    )
                    
        with col2:
            if report_data.get("json_path") and os.path.exists(report_data["json_path"]):
                with open(report_data["json_path"], "r") as f:
                    json_content = f.read()
                    st.download_button(
                        label="Download as JSON",
                        data=json_content,
                        file_name=os.path.basename(report_data["json_path"]),
                        mime="application/json"
                    )

    def show_research_progress(self):
        """Displays the current research progress"""
        st.subheader("Research in Progress")
        st.markdown(f"**Question:** {self.research_state['question']}")
        
        # Show progress bar
        progress = 0
        if self.research_state["total_steps"] > 0:
            progress = self.research_state["steps_completed"] / self.research_state["total_steps"]
        
        st.progress(progress)
        
        # Show current step
        current_step = self.research_state.get("current_step", "Planning")
        st.markdown(f"**Current step:** {current_step}")
        
        st.info("Research is ongoing. This may take several minutes depending on the complexity of the question.")
        st.warning("Please do not navigate away from this page while research is in progress.")