search_personalization/ubi_tracker.py

"""
UBI (User Behavior Insights) Tracker Module

This module provides configuration management and utility functions for UBI tracking
in the OpenSearch Product Catalog. It handles loading configuration from environment
variables and provides functions to check if UBI tracking is enabled.
"""

import os
from typing import Dict, Any, Optional
from dotenv import load_dotenv
import streamlit as st
import streamlit.components.v1 as components
import logging

# Import proxy functions
from search_personalization.ubi_proxy import send_ubi_query, send_ubi_event

# Load environment variables from .env file
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def is_ubi_enabled() -> bool:
    """
    Check if UBI tracking is enabled.
    
    UBI tracking is considered enabled if:
    1. UBI_ENABLED environment variable is set to 'true' (case-insensitive)
    2. Both UBI_EVENTS_ENDPOINT and UBI_QUERIES_ENDPOINT are configured
    
    Returns:
        bool: True if UBI tracking is enabled and properly configured, False otherwise
    """
    # Check if UBI is explicitly enabled
    enabled = os.getenv('UBI_ENABLED', 'true').lower() == 'true'
    
    if not enabled:
        return False
    
    # Check if both ingestion endpoints are configured
    events_endpoint = os.getenv('UBI_EVENTS_ENDPOINT')
    queries_endpoint = os.getenv('UBI_QUERIES_ENDPOINT')
    
    if not events_endpoint or not queries_endpoint:
        print("WARNING: UBI tracking is enabled but UBI_EVENTS_ENDPOINT and/or UBI_QUERIES_ENDPOINT are not configured")
        return False
    
    return True


def get_ubi_config() -> Dict[str, Any]:
    """
    Get UBI configuration from environment variables.
    
    Returns a dictionary containing all UBI-related configuration values:
    - events_endpoint: OpenSearch Ingestion Service endpoint for UBI events
    - queries_endpoint: OpenSearch Ingestion Service endpoint for UBI queries
    - aws_region: AWS region for the ingestion service
    - debug_mode: Whether to enable debug logging
    - enabled: Whether UBI tracking is enabled
    - application_name: Name of the application for UBI data
    
    Returns:
        Dict[str, Any]: Dictionary containing UBI configuration
    """
    config = {
        'events_endpoint': os.getenv('UBI_EVENTS_ENDPOINT'),
        'queries_endpoint': os.getenv('UBI_QUERIES_ENDPOINT'),
        'aws_region': os.getenv('AWS_REGION', 'us-east-1'),
        'debug_mode': os.getenv('UBI_DEBUG_MODE', 'false').lower() == 'true',
        'enabled': is_ubi_enabled(),
        'application_name': os.getenv('UBI_APPLICATION_NAME', 'os-product-catalog')
    }
    
    return config


def validate_ubi_config() -> tuple[bool, Optional[str]]:
    """
    Validate UBI configuration.
    
    Checks that all required configuration values are present and valid.
    
    Returns:
        tuple[bool, Optional[str]]: (is_valid, error_message)
            - is_valid: True if configuration is valid, False otherwise
            - error_message: Description of validation error, or None if valid
    """
    config = get_ubi_config()
    
    if not config['enabled']:
        return True, None  # Valid to have UBI disabled
    
    # Check required fields when UBI is enabled
    if not config['events_endpoint']:
        return False, "UBI_EVENTS_ENDPOINT is required when UBI_ENABLED is true"
    
    if not config['queries_endpoint']:
        return False, "UBI_QUERIES_ENDPOINT is required when UBI_ENABLED is true"
    
    # Validate endpoints are HTTPS
    if not config['events_endpoint'].startswith('https://'):
        return False, "UBI_EVENTS_ENDPOINT must be an HTTPS URL"
    
    if not config['queries_endpoint'].startswith('https://'):
        return False, "UBI_QUERIES_ENDPOINT must be an HTTPS URL"
    
    if not config['aws_region']:
        return False, "AWS_REGION is required for UBI tracking"
    
    if not config['application_name']:
        return False, "UBI_APPLICATION_NAME is required for UBI tracking"
    
    return True, None


def create_ubi_tracker(
    ingestion_endpoint: str,
    aws_region: str,
    debug_mode: bool = False
) -> bool:
    """
    Create and initialize the UBI tracking component.
    
    This function embeds a custom HTML/JavaScript component in the Streamlit app
    that handles UBI tracking on the client side. The component manages Client_ID,
    Session_ID, and Query_ID, formats UBI data, and sends it to OpenSearch Ingestion Service.
    
    The component communicates Client_ID and Session_ID back to Python via Streamlit's
    component API, allowing the backend proxy to use these IDs when tracking.
    
    Args:
        ingestion_endpoint: OpenSearch Ingestion Service HTTPS endpoint
        aws_region: AWS region for SigV4 signing
        debug_mode: Enable console logging for debugging (default: False)
    
    Returns:
        bool: True if component initialized successfully, False otherwise
    
    Requirements: 8.4, 8.5, 12.4, 12.5
    """
    # Read the HTML template
    template_path = os.path.join(os.path.dirname(__file__), 'ubi_component_template.html')
    
    try:
        with open(template_path, 'r') as f:
            html_template = f.read()
    except FileNotFoundError:
        logger.error(f"UBI component template not found at {template_path}")
        # Set flag to disable tracking if component fails
        if 'ubi_component_failed' not in st.session_state:
            st.session_state.ubi_component_failed = True
        return False
    except Exception as e:
        logger.error(f"Error reading UBI component template: {e}")
        if 'ubi_component_failed' not in st.session_state:
            st.session_state.ubi_component_failed = True
        return False
    
    try:
        # Inject configuration into the HTML
        config_script = f"""
        <script>
            // Initialize UBI client with configuration
            document.addEventListener('DOMContentLoaded', function() {{
                try {{
                    if (typeof initializeUbiClient === 'function') {{
                        initializeUbiClient({{
                            ingestionEndpoint: '{ingestion_endpoint}',
                            awsRegion: '{aws_region}',
                            debugMode: {str(debug_mode).lower()}
                        }});
                    }} else {{
                        console.error('UBI: initializeUbiClient function not found');
                    }}
                }} catch (error) {{
                    console.error('UBI: Failed to initialize client:', error);
                }}
            }});
        </script>
        """
        
        # Combine template with configuration
        full_html = html_template.replace('</body>', f'{config_script}</body>')
        
        # Embed the component and capture return value
        component_value = components.html(full_html, height=0, scrolling=False)
        
        # Store Client_ID and Session_ID in session state if returned from component
        if component_value and isinstance(component_value, dict):
            if 'client_id' in component_value:
                st.session_state.ubi_client_id = component_value['client_id']
                if debug_mode:
                    logger.info(f"Received Client_ID from component: {component_value['client_id']}")
            
            if 'session_id' in component_value:
                st.session_state.ubi_session_id = component_value['session_id']
                if debug_mode:
                    logger.info(f"Received Session_ID from component: {component_value['session_id']}")
        
        # Mark component as successfully loaded
        if 'ubi_component_failed' in st.session_state:
            del st.session_state.ubi_component_failed
        
        return True
        
    except Exception as e:
        logger.error(f"Error initializing UBI component: {e}")
        # Set flag to disable tracking if component fails
        if 'ubi_component_failed' not in st.session_state:
            st.session_state.ubi_component_failed = True
        return False


def track_query(query_text: str, query_id: str, user_id: Optional[str] = None, query_response_hit_ids: Optional[list] = None) -> None:
    """
    Track a search query by sending it to the backend proxy.
    
    This function formats the query data and sends it to OpenSearch Ingestion Service
    via the backend proxy with AWS SigV4 authentication.
    
    Args:
        query_text: The user's search query
        query_id: Unique identifier for this query
        user_id: User identifier for persona tracking (optional)
        query_response_hit_ids: List of document IDs returned by the search in order (optional)
    
    Requirements: 8.4, 8.5, 9.1, 9.2, 12.4, 12.5
    """
    try:
        # Check if component failed to load
        if st.session_state.get('ubi_component_failed', False):
            logger.debug("UBI component failed to load, skipping query tracking")
            return
        
        # Get configuration
        config = get_ubi_config()
        
        if not config['enabled']:
            logger.debug("UBI tracking is disabled, skipping query tracking")
            return
        
        # Get Client_ID from session state (should be set by component)
        client_id = st.session_state.get('ubi_client_id', 'unknown')
        
        # Format query data
        from datetime import datetime
        query_data = {
            'application': config['application_name'],
            'query_id': query_id,
            'client_id': client_id,
            'user_query': query_text or '',
            'object_id_field': 'id',
            'query_attributes': {},
            'timestamp': datetime.utcnow().isoformat() + 'Z'
        }
        
        # Add user_id if provided
        if user_id:
            query_data['user_id'] = user_id
        
        # Add query_response_hit_ids if provided
        if query_response_hit_ids is not None:
            query_data['query_response_hit_ids'] = query_response_hit_ids
        
        # Send via proxy
        success, error = send_ubi_query(query_data)
        
        if not success:
            logger.error(f"Failed to track query: {error}")
        elif config['debug_mode']:
            logger.info(f"Query tracked successfully: {query_id}")
    
    except Exception as e:
        logger.exception(f"Error tracking query: {e}")
        # Silent failure - do not disrupt user experience


def track_event(
    action_name: str,
    query_id: str,
    object_id: Optional[str] = None,
    event_attributes: Optional[Dict[str, Any]] = None,
    message: Optional[str] = None,
    message_type: Optional[str] = None,
    user_id: Optional[str] = None
) -> None:
    """
    Track a user event by sending it to the backend proxy.
    
    This function formats the event data and sends it to OpenSearch Ingestion Service
    via the backend proxy with AWS SigV4 authentication.
    
    Args:
        action_name: Type of event (search, add_to_cart, hover, etc.)
        query_id: Query ID to associate with this event
        object_id: Product ID (for product-related events)
        event_attributes: Additional event metadata
        message: Human-readable event description
        message_type: Message type (INFO, QUERY, etc.) - defaults to QUERY for search, INFO otherwise
        user_id: User identifier for persona tracking (optional)
    
    Requirements: 8.4, 8.5, 9.1, 9.2, 12.4, 12.5
    """
    try:
        # Check if component failed to load
        if st.session_state.get('ubi_component_failed', False):
            logger.debug("UBI component failed to load, skipping event tracking")
            return
        
        # Get configuration
        config = get_ubi_config()
        
        if not config['enabled']:
            logger.debug("UBI tracking is disabled, skipping event tracking")
            return
        
        # Get Client_ID and Session_ID from session state
        client_id = st.session_state.get('ubi_client_id', 'unknown')
        session_id = st.session_state.get('ubi_session_id', 'unknown')
        
        # Determine message_type if not provided
        if message_type is None:
            message_type = 'QUERY' if action_name == 'search' else 'INFO'
        
        # Format event data
        from datetime import datetime
        event_data = {
            'application': config['application_name'],
            'action_name': action_name,
            'query_id': query_id,
            'client_id': client_id,
            'session_id': session_id,
            'user_id': user_id or '',
            'timestamp': datetime.utcnow().isoformat() + 'Z',
            'message_type': message_type,
            'message': message or ''
        }
        
        # Add event_attributes if provided
        if event_attributes:
            event_data['event_attributes'] = event_attributes
            
            # Ensure object_id_field is set if object exists
            if event_attributes.get('object') and event_attributes['object'].get('object_id'):
                if not event_attributes['object'].get('object_id_field'):
                    event_attributes['object']['object_id_field'] = 'id'
                
                # Ensure position.ordinal exists for events with object_id
                if not event_attributes.get('position'):
                    event_attributes['position'] = {'ordinal': 0}
                elif not isinstance(event_attributes['position'].get('ordinal'), int):
                    event_attributes['position']['ordinal'] = 0
        elif object_id:
            # If objectId is provided but no eventAttributes, create minimal structure
            event_data['event_attributes'] = {
                'object': {
                    'object_id': object_id,
                    'object_id_field': 'id'
                },
                'position': {
                    'ordinal': 0
                }
            }
        
        # Send via proxy
        success, error = send_ubi_event(event_data)
        
        if not success:
            logger.error(f"Failed to track event: {error}")
        elif config['debug_mode']:
            logger.info(f"Event tracked successfully: {action_name}")
    
    except Exception as e:
        logger.exception(f"Error tracking event: {e}")
        # Silent failure - do not disrupt user experience