src/gepa_optimizer/__init__.py

"""
GEPA Universal Prompt Optimizer

A modern, modular Python library for universal prompt optimization powered by GEPA.

Quick Start (No custom evaluator needed!):
    
    from gepa_optimizer import quick_optimize
    
    result = await quick_optimize(
        seed_prompt="Your initial prompt",
        dataset=[
            {"input": "task1", "output": "expected1"},
            {"input": "task2", "output": "expected2"},
        ],
        model="openai/gpt-4o"  # or any: "google/gemini-1.5-pro", "anthropic/claude-3-5-sonnet-20241022"
    )
    print(result.optimized_prompt)
"""

# Core functionality
from .core import GepaOptimizer
from .core.base_adapter import BaseGepaAdapter
from .core.universal_adapter import UniversalGepaAdapter

# Configuration and models
from .models import OptimizationConfig, OptimizationResult, OptimizedResult, ModelConfig

# Data processing
from .data import UniversalConverter, DataLoader, DataValidator
from .data.scroll_dataset_loader import ScrollDatasetLoader, load_scroll_dataset
from .data.validation_dataset_loader import ValidationDatasetLoader, load_validation_dataset, load_validation_split
from .data.index_caching_loader import IndexCachingDatasetLoader, load_index_caching_dataset, load_index_caching_split

# LLM clients
from .llms import VisionLLMClient
from .llms.base_llm import BaseLLMClient
from .llms.batch_llm import BatchLLMClient

# Evaluators - including Universal Semantic Evaluator (works for ANY task!)
from .evaluation import (
    BaseEvaluator,
    UniversalSemanticEvaluator,
    create_universal_evaluator,
    UITreeEvaluator,
    ScrollElementEvaluator,
    ValidationEvaluator,
    IndexCachingEvaluator
)

# LLEGO Genetic Operators
from .operators import (
    # Base interfaces
    BaseGeneticOperator,
    BaseCrossoverOperator,
    BaseMutationOperator,
    # Concrete operators
    FitnessGuidedCrossover,
    DiversityGuidedMutation,
    LLEGOIntegrationLayer,
    # Data models
    PromptCandidate,
    PromptMetadata
)

# Utilities
from .utils import setup_logging, calculate_metrics, sanitize_prompt, APIKeyManager
from .utils.exceptions import GepaOptimizerError, GepaDependencyError, InvalidInputError, DatasetError

# Logging infrastructure
from .infrastructure.logging import get_logger, configure_logging, LogContext

# Type definitions (for type hints in user code)
from .types import (
    DatasetItem,
    EvaluationResult,
    LLMResponse,
    CandidateDict,
    LLMClientProtocol,
    EvaluatorProtocol,
)

__version__ = "0.1.0"


# ═══════════════════════════════════════════════════════════════════════════════
# CONVENIENCE FUNCTION: quick_optimize
# No evaluator needed - uses Universal Semantic Evaluator automatically
# ═══════════════════════════════════════════════════════════════════════════════

async def quick_optimize(
    seed_prompt: str,
    dataset: list,
    model: str,
    max_iterations: int = 5,
    max_metric_calls: int = 50,
    batch_size: int = 4,
    use_llego: bool = True,
    verbose: bool = True
) -> OptimizedResult:
    """
    🚀 Quick prompt optimization - no custom evaluator needed!
    
    Uses Universal Semantic Evaluator that works for ANY task.
    
    Args:
        seed_prompt: Your initial prompt to optimize
        dataset: List of dicts with 'input' and 'output' (expected) keys
                 Can also include 'image' key for multi-modal tasks
        model: LLM model to use in format "provider/model-name" (REQUIRED)
               Examples:
                 - "google/gemini-1.5-pro"
                 - "google/gemini-2.5-flash-preview-05-20"
                 - "openai/gpt-4o"
                 - "openai/gpt-4-turbo"
                 - "anthropic/claude-3-5-sonnet-20241022"
        max_iterations: Maximum optimization iterations (default: 5)
        max_metric_calls: Maximum evaluation calls (default: 50)
        batch_size: Samples per evaluation batch (default: 4)
        use_llego: Enable LLEGO genetic operators (default: True)
        verbose: Show progress logs (default: True)
    
    Returns:
        OptimizedResult with optimized prompt and improvement metrics
    
    Example:
        >>> result = await quick_optimize(
        ...     seed_prompt="Count the objects in the image",
        ...     dataset=[
        ...         {"input": "image1.jpg", "output": "5 objects", "image": "base64..."},
        ...         {"input": "image2.jpg", "output": "3 objects", "image": "base64..."},
        ...     ],
        ...     model="openai/gpt-4o",  # or "google/gemini-1.5-pro", etc.
        ...     max_iterations=3
        ... )
        >>> print(result.optimized_prompt)
    """
    import logging
    
    if verbose:
        logging.basicConfig(level=logging.INFO)
    
    # Create LLM client
    llm_client = VisionLLMClient.from_model_string(model)
    
    # Create Universal Semantic Evaluator (uses same LLM for analysis)
    evaluator = UniversalSemanticEvaluator(
        llm_client=llm_client,
        use_llm_analysis=True
    )
    
    # Create configuration
    config = OptimizationConfig(
        model=model,
        reflection_model=model,
        max_iterations=max_iterations,
        max_metric_calls=max_metric_calls,
        batch_size=batch_size,
        use_llego_operators=use_llego,
        enable_gepa_reflection_with_llego=use_llego,
        num_gepa_reflection_candidates=3,
        n_crossover=2,
        n_mutation=2,
        verbose=verbose
    )
    
    # Create optimizer
    optimizer = GepaOptimizer(
        config=config,
        llm_client=llm_client,
        evaluator=evaluator
    )
    
    # Run optimization
    result = await optimizer.train(
        seed_prompt=seed_prompt,
        dataset=dataset
    )
    
    return result


def quick_optimize_sync(
    seed_prompt: str,
    dataset: list,
    model: str,
    max_iterations: int = 5,
    max_metric_calls: int = 50,
    batch_size: int = 4,
    use_llego: bool = True,
    verbose: bool = True
) -> OptimizedResult:
    """
    🚀 Synchronous version of quick_optimize.
    
    Same as quick_optimize but runs synchronously (blocks until complete).
    
    Args:
        model: LLM model to use in format "provider/model-name" (REQUIRED)
               Examples: "openai/gpt-4o", "google/gemini-1.5-pro", "anthropic/claude-3-5-sonnet-20241022"
    
    See quick_optimize for full documentation.
    """
    import asyncio
    return asyncio.run(quick_optimize(
        seed_prompt=seed_prompt,
        dataset=dataset,
        model=model,
        max_iterations=max_iterations,
        max_metric_calls=max_metric_calls,
        batch_size=batch_size,
        use_llego=use_llego,
        verbose=verbose
    ))


__all__ = [
    # 🚀 Quick Start (recommended for new users)
    "quick_optimize",
    "quick_optimize_sync",
    
    # Core functionality
    "GepaOptimizer",
    "BaseGepaAdapter",
    "UniversalGepaAdapter",
    
    # Configuration
    "OptimizationConfig", 
    "OptimizationResult",
    "OptimizedResult",
    "ModelConfig",
    
    # Data processing
    "UniversalConverter",
    "DataLoader",
    "DataValidator",
    
    # Dataset loaders
    "ScrollDatasetLoader",
    "load_scroll_dataset",
    "ValidationDatasetLoader",
    "load_validation_dataset",
    "load_validation_split",
    "IndexCachingDatasetLoader",
    "load_index_caching_dataset",
    "load_index_caching_split",
    
    # LLM clients
    "VisionLLMClient",
    "BaseLLMClient",
    "BatchLLMClient",
    
    # Evaluators (Universal recommended for general use)
    "UniversalSemanticEvaluator",
    "create_universal_evaluator",
    "BaseEvaluator",
    "UITreeEvaluator",
    "ScrollElementEvaluator",
    "ValidationEvaluator",
    "IndexCachingEvaluator",
    
    # LLEGO Genetic Operators - Base interfaces
    "BaseGeneticOperator",
    "BaseCrossoverOperator",
    "BaseMutationOperator",
    # LLEGO Genetic Operators - Concrete implementations
    "FitnessGuidedCrossover",
    "DiversityGuidedMutation",
    "LLEGOIntegrationLayer",
    "PromptCandidate",
    "PromptMetadata",
    
    # Utilities
    "APIKeyManager",
    "GepaOptimizerError",
    "GepaDependencyError",
    "InvalidInputError",
    "DatasetError",
    "setup_logging",
    "calculate_metrics",
    "sanitize_prompt",
    
    # Logging infrastructure
    "get_logger",
    "configure_logging",
    "LogContext",
    
    # Type definitions
    "DatasetItem",
    "EvaluationResult",
    "LLMResponse",
    "CandidateDict",
    "LLMClientProtocol",
    "EvaluatorProtocol",
]