run_agentic_pipeline.py

import gradio as gr
import pandas as pd
import json
import logging
from typing import Dict, Any, List, Optional

# Assuming these functions are in ui.insights_ui_generator
# Make sure to have this file updated with the new `format_report_for_display` function
try:
    from ui.insights_ui_generator import (
        format_report_for_display,
        extract_key_results_for_selection,
        format_single_okr_for_display
    )
    AGENTIC_MODULES_LOADED = True
except ImportError as e:
    logging.error(f"Could not import agentic pipeline display modules: {e}. Tabs 3 and 4 will be disabled.")
    AGENTIC_MODULES_LOADED = False
    def format_report_for_display(report_data): return "Agentic modules not loaded. Report unavailable."
    def extract_key_results_for_selection(okrs_dict): return []
    def format_single_okr_for_display(okr_data, **kwargs): return "Agentic modules not loaded. OKR display unavailable."


logger = logging.getLogger(__name__)

def load_and_display_agentic_results(
    token_state: dict, 
    orchestration_raw_results_st: Optional[dict], 
    selected_key_result_ids_st: List[str], 
    key_results_for_selection_st: List[dict]
):
    """
    Loads pre-computed agentic results from the state, populates the report library dropdown,
    and displays the latest report and its associated OKRs by default.

    This function is designed to work with the UI defined in app.py and expects a specific
    order of outputs.

    Args:
        token_state: The main state dictionary containing the bubble_agentic_analysis_data DataFrame.
        orchestration_raw_results_st: The state holding the raw JSON/dict of the currently displayed report.
        selected_key_result_ids_st: The state for the IDs of selected Key Results.
        key_results_for_selection_st: The state holding the list of all available Key Results for selection.

    Returns:
        A tuple of Gradio updates matching the `agentic_display_outputs` list in `app.py`.
    """
    # Default empty/initial return values that match the output components list
    # The order is critical:
    # 1. agentic_report_display_md
    # 2. report_selector_dd
    # 3. key_results_cbg
    # 4. okr_detail_display_md
    # 5. orchestration_raw_results_st
    # 6. selected_key_result_ids_st
    # 7. key_results_for_selection_st
    # 8. agentic_pipeline_status_md
    initial_updates = (
        "No agentic analysis data found in Bubble.",
        gr.update(choices=[], value=None, interactive=False),
        gr.update(choices=[], value=[], interactive=False),
        "No OKRs to display.",
        None,
        [],
        [],
        "Status: No agentic analysis data found."
    )

    if not AGENTIC_MODULES_LOADED:
        error_updates = list(initial_updates)
        error_updates[7] = "Status: Critical module import error."
        return tuple(error_updates)

    agentic_df = token_state.get("bubble_agentic_analysis_data")

    if agentic_df is None or agentic_df.empty:
        logger.warning("Agentic analysis DataFrame is missing or empty in the state.")
        return initial_updates

    try:
        # --- 1. Prepare Report Library ---
        if 'Created Date' not in agentic_df.columns or '_id' not in agentic_df.columns:
            raise KeyError("Required columns ('Created Date', '_id') not found in agentic data.")
        
        # Ensure 'Created Date' is datetime, then sort to get the latest report first
        agentic_df['Created Date'] = pd.to_datetime(agentic_df['Created Date'])
        agentic_df = agentic_df.sort_values(by='Created Date', ascending=False).reset_index(drop=True)
        
        # Create choices for the dropdown: (Display Name, Unique ID)
        report_choices = [
            (f"{row.get('report_type', 'Report')} - {row['Created Date'].strftime('%Y-%m-%d %H:%M')}", row['_id'])
            for _, row in agentic_df.iterrows()
        ]
        
        if not report_choices:
            return initial_updates

        # --- 2. Load and Display the Latest Report by Default ---
        latest_report_series = agentic_df.iloc[0]
        latest_report_id = latest_report_series['_id']
        
        # Format the latest report's content for the Markdown display
        report_display_md = format_report_for_display(latest_report_series)
        
        # Create the update for the report library dropdown
        report_selector_update = gr.update(choices=report_choices, value=latest_report_id, interactive=True)

        # --- 3. Load and Prepare OKRs from the Latest Report ---
        raw_results_state = None
        okr_details_md = "No OKRs found in the latest report."
        key_results_cbg_update = gr.update(choices=[], value=[], interactive=False)
        all_krs_state = []
        
        # Assumption: The full JSON from the agent is stored in 'orchestration_results'.
        if 'orchestration_results' in latest_report_series and pd.notna(latest_report_series['orchestration_results']):
            try:
                raw_results_state = json.loads(latest_report_series['orchestration_results'])
            except json.JSONDecodeError:
                logger.error(f"Failed to parse 'orchestration_results' JSON for report ID {latest_report_id}")
                raw_results_state = {} # Avoid crashing, proceed with empty data
        else:
            raw_results_state = {}

        actionable_okrs_dict = raw_results_state.get("actionable_okrs", {})
        
        if actionable_okrs_dict:
            all_krs_state = extract_key_results_for_selection(actionable_okrs_dict)
            if all_krs_state:
                kr_choices = [(kr['kr_description'], kr['unique_kr_id']) for kr in all_krs_state]
                key_results_cbg_update = gr.update(choices=kr_choices, value=[], interactive=True)
                
                # Format all OKRs for initial display
                okrs_list = actionable_okrs_dict.get("okrs", [])
                output_md_parts = [
                    format_single_okr_for_display(okr_data, okr_main_index=okr_idx)
                    for okr_idx, okr_data in enumerate(okrs_list)
                ]
                okr_details_md = "\n\n---\n\n".join(output_md_parts) if output_md_parts else okr_details_md

        status_update = f"Status: Loaded {len(agentic_df)} reports. Displaying the latest from {latest_report_series['Created Date'].strftime('%Y-%m-%d')}."
        
        return (
            report_display_md, 
            report_selector_update,
            key_results_cbg_update, 
            okr_details_md, 
            raw_results_state, 
            [], # Reset selected KRs
            all_krs_state,
            status_update
        )

    except Exception as e:
        logger.error(f"Failed to process and display agentic results: {e}", exc_info=True)
        error_updates = list(initial_updates)
        error_updates[0] = f"An error occurred while loading reports: {e}"
        error_updates[7] = f"Status: Error - {e}"
        return tuple(error_updates)