app.py

# Modify the init_app function in app.py to ensure proper initialization and handling of errors:

def init_app():
    """Initialize everything just once using st.session_state."""

    # Initialize empty collections to prevent KeyError
    if "chat_history" not in st.session_state:
        st.session_state["chat_history"] = []
    
    # Create a flag for initialization attempts
    if "initialization_attempted" not in st.session_state:
        st.session_state["initialization_attempted"] = False

    if "initialized" not in st.session_state:
        st.session_state["initialization_attempted"] = True
        try:
            # Add key presence checks before starting initialization
            st.session_state["initialized"] = False  # Mark as not initialized until complete
            logger.info("Starting app initialization")
            
            # Load settings
            settings = Settings()
            st.session_state["settings"] = settings
            
            # (1) instrumentation (if needed)
            with st.spinner("Setting up instrumentation..."):
                tracer_provider = setup_instrumentation()
                st.session_state["tracer"] = tracer_provider.get_tracer("llamaindex_app")
                logger.info("Instrumentation setup complete")

            # (2) OpenAI client
            with st.spinner("Initializing OpenAI client..."):
                openai_client = init_openai_client()
                st.session_state["openai_client"] = openai_client
                logger.info("OpenAI client initialized")

            # (3) index manager & query engine
            with st.spinner("Loading index and query engine..."):
                # Create these objects explicitly
                index_manager = IndexManager(openai_client=openai_client)
                query_engine = index_manager.get_query_engine()
                
                # Store them in session_state
                st.session_state["index_manager"] = index_manager
                st.session_state["query_engine"] = query_engine
                logger.info("Index and query engine loaded")

            # (4) classifier
            with st.spinner("Initializing query classifier..."):
                # Create classifier only when we are sure query_engine exists
                if "query_engine" in st.session_state:
                    classifier = QueryClassifier(
                        query_engine=st.session_state["query_engine"],
                        openai_client=openai_client
                    )
                    st.session_state["classifier"] = classifier
                    logger.info("Query classifier initialized")
                else:
                    raise ValueError("Query engine not initialized properly")

            st.session_state["initialized"] = True  # Now mark as fully initialized
            logger.info("App initialization complete")
            st.success("App initialized successfully!")
            
        except Exception as e:
            st.error(f"Failed to initialize app: {str(e)}")
            logger.error(f"Initialization error: {str(e)}", exc_info=True)
            st.session_state["initialization_error"] = str(e)
            return False
        
        return True
    
    return "initialized" in st.session_state and st.session_state["initialized"]


# Modify the main function to handle errors more gracefully:

def main():
    """Main Streamlit app function."""
    st.title("Assurant 10-K Analysis & Risk Assessment App")
    
    # Sidebar with app info
    with st.sidebar:
        st.subheader("About this app")
        st.write("""
        This application analyzes Assurant's 10-K reports and provides risk assessment.
        
        You can ask questions about:
        - Assurant's financial performance
        - Risk factors
        - Business operations
        - Market trends
        """)
        
        # Debug section in sidebar
        if st.checkbox("Show debug info"):
            st.subheader("Debug Information")
            st.write("Session State Keys:")
            for key in st.session_state:
                st.write(f"- {key}: {'Present' if st.session_state[key] is not None else 'None'}")
            
            st.write("Python Path:")
            for path in sys.path:
                st.write(f"- {path}")
            
            # Display environment variables (without sensitive values)
            st.write("Environment Variables:")
            safe_vars = {
                k: (v if not any(secret in k.lower() for secret in ["key", "secret", "password", "token"]) else "****") 
                for k, v in os.environ.items()
                if k.startswith(("OPENAI_", "LLAMAINDEX_"))
            }
            for k, v in safe_vars.items():
                st.write(f"- {k}: {v}")
            
            # Also display settings from the Settings class
            if "settings" in st.session_state:
                settings = st.session_state["settings"]
                st.write("App Settings:")
                st.write(f"- OpenAI Model: {settings.OPENAI_MODEL}")
                st.write(f"- OpenAI Base URL: {settings.OPENAI_BASE_URL if settings.OPENAI_BASE_URL else 'Default API URL'}")
                if settings.OPENAI_ORG_ID:
                    st.write(f"- OpenAI Organization: {settings.OPENAI_ORG_ID}")

    # Initialize the app
    initialization_status = init_app()
    
    if initialization_status is False:
        st.error(f"Initialization failed: {st.session_state.get('initialization_error', 'Unknown error')}")
        st.warning("Please check your environment variables and connection settings.")
        return
    
    # Safety check - if we've attempted initialization but still don't have key components
    if st.session_state.get("initialization_attempted", False) and not st.session_state.get("initialized", False):
        st.warning("The application is still initializing or had problems during initialization. Please wait or refresh the page.")
        return
    
    # Query section - Only show if properly initialized
    if st.session_state.get("initialized", False):
        st.subheader("Ask a question about Assurant's 10-K")
        user_question = st.text_input("Your question:", placeholder="e.g., How did Assurant perform last quarter?")
        submit = st.button("Submit Question")

        if submit and user_question.strip():
            # Safely retrieve references from st.session_state with error handling
            try:
                # Check that all required components exist
                required_keys = ["query_engine", "classifier", "tracer"]
                missing_keys = [key for key in required_keys if key not in st.session_state]
                
                if missing_keys:
                    st.error(f"Missing required components: {', '.join(missing_keys)}")
                    return
                
                query_engine = st.session_state["query_engine"]
                classifier = st.session_state["classifier"]
                tracer = st.session_state["tracer"]
                session_id = str(uuid.uuid4())

                with st.spinner("Analyzing your question..."):
                    # Log the question being processed
                    logger.info(f"Processing query: {user_question}")
                    
                    response, error = process_interaction(
                        query_engine,
                        classifier,
                        tracer,
                        user_question,
                        session_id
                    )

                    # Store in st.session_state so we can display entire conversation
                    if error:
                        st.session_state["chat_history"].append(("user", user_question))
                        st.session_state["chat_history"].append(("assistant", f"Error: {error}"))
                        st.error(f"Error: {error}")
                        logger.error(f"Query processing error: {error}")
                    else:
                        st.session_state["chat_history"].append(("user", user_question))
                        st.session_state["chat_history"].append(("assistant", response.response))
                        logger.info("Query processed successfully")

                        # If there are any sources
                        if getattr(response, "source_nodes", None):
                            source_text = "\n".join(
                                f"- {node.metadata.get('file_name', 'Unknown source')}"
                                for node in response.source_nodes
                            )
                            st.session_state["chat_history"].append(("assistant_sources", source_text))
                            logger.info(f"Found {len(response.source_nodes)} source nodes")
            except Exception as e:
                st.error(f"Error processing your question: {str(e)}")
                logger.error(f"Process interaction error: {str(e)}", exc_info=True)
                st.session_state["chat_history"].append(("user", user_question))
                st.session_state["chat_history"].append(("assistant", f"Error processing your question: {str(e)}"))

    # Display the chat history
    st.subheader("Conversation")
    history_container = st.container()
    
    with history_container:
        for idx, (role, content) in enumerate(st.session_state.get("chat_history", [])):
            if role == "user":
                st.markdown(f"**You**: {content}")
            elif role == "assistant":
                st.markdown(f"**Assistant**: {content}")
            elif role == "assistant_sources":
                with st.expander("View Sources"):
                    st.markdown(content)
            
            # Add a separator between conversations
            if idx < len(st.session_state.get("chat_history", [])) - 1 and role == "assistant":
                st.markdown("---")