Source code for aidputils.agents.tools.http.response_optimizer

"""
Response Optimizer for HTTP Tool.

Provides LLM token optimization using pluggable parser/optimizer architecture.
"""

import logging
from typing import Any, Dict

from aidputils.agents.tools.http.response_registry import (
    ResponseParserRegistry,
    ResponseOptimizerRegistry,
)

logger = logging.getLogger(__name__)




class OptimizationError(ValueError):
    """Exception raised for response optimization errors."""
    pass





def optimize_json_response(
    data: Any,
    field_containing_data: str = None,
    include_fields: list = None,
    exclude_fields: list = None
) -> Any:
    """
    Optimize JSON response for LLM consumption.
    
    Args:
        data: The JSON data to optimize
        field_containing_data: Dot-notation path to extract
        include_fields: Fields to include in result
        exclude_fields: Fields to exclude from result
        
    Returns:
        Optimized JSON data
    """
    from aidputils.agents.tools.http.optimizers.json_optimizer import JsonResponseOptimizer
    
    optimizer = JsonResponseOptimizer()
    config = {
        'field_containing_data': field_containing_data,
        'include_fields': include_fields,
        'exclude_fields': exclude_fields,
    }
    config = {k: v for k, v in config.items() if v is not None}
    return optimizer.optimize(data, config)





def extract_html_content(
    html: str,
    css_selector: str = None,
    strip_tags: bool = False
) -> str:
    """
    Extract content from HTML using CSS selector.
    
    Args:
        html: The HTML content
        css_selector: CSS selector to extract
        strip_tags: Whether to strip HTML tags from result
        
    Returns:
        Extracted content
    """
    from aidputils.agents.tools.http.optimizers.html_optimizer import HtmlResponseOptimizer
    
    optimizer = HtmlResponseOptimizer()
    config = {
        'css_selector': css_selector,
        'strip_tags': strip_tags,
    }
    return optimizer.optimize(html, config)





def truncate_text(text: str, max_characters: int) -> str:
    """
    Truncate text to a maximum number of characters.
    
    Args:
        text: The text to truncate
        max_characters: Maximum characters allowed
        
    Returns:
        Truncated text (with '...' if truncated)
    """
    from aidputils.agents.tools.http.optimizers.text_optimizer import TextResponseOptimizer
    
    optimizer = TextResponseOptimizer()
    config = {'max_characters': max_characters}
    return optimizer.optimize(text, config)





def optimize_response(
    response_data: Any,
    content_type: str,
    optimization_config: Dict[str, Any]
) -> Any:
    """
    Optimize a response based on configuration using pluggable architecture.
    
    Args:
        response_data: The response data to optimize
        content_type: Content-Type header value
        optimization_config: Optimization settings
        
    Returns:
        Optimized response data
    """
    if not optimization_config:
        return response_data
    
    response_type = optimization_config.get("response_type", "").lower()
    
    # Auto-detect response type from content-type if not specified
    if not response_type:
        parser = ResponseParserRegistry.get_parser(content_type)
        response_type = parser.get_response_type()
    
    # Get optimizer for response type
    try:
        optimizer = ResponseOptimizerRegistry.get_optimizer(response_type)
    except KeyError:
        logger.warning(f"No optimizer for response type '{response_type}', returning data as-is")
        return response_data
    
    # Apply optimization
    result = optimizer.optimize(response_data, optimization_config)
    
    # Apply text truncation if specified (works for any text output)
    max_characters = optimization_config.get("max_characters")
    if max_characters and isinstance(result, str):
        result = truncate_text(result, max_characters)
    
    return result