model_selector.py

# filename: model_selector.py

from language_models import GPT4Model, ClaudeModel, GeminiModel, LanguageModel
import os
from tiktoken import get_encoding
from log_config import get_logger

logger = get_logger('ModelSelector')

# Load the tokenizer encoding
try:
    encoding = get_encoding("cl100k_base")
except Exception as e:
    logger.error("Failed to load encoding: %s", str(e))
    raise

def count_tokens(text: str) -> int:
    """Count the number of tokens in a text string.
    Args:
        text (str): The text string to count tokens in.
    Returns:
        int: The number of tokens in the text string.
    """
    return len(encoding.encode(text))

def select_optimal_model(raw_text: str) -> LanguageModel:
    """Select the optimal language model for analyzing the given text based on token count.
    Args:
        raw_text (str): The text string to analyze.
    Returns:
        LanguageModel: The selected optimal language model instance.
    """
    if not isinstance(raw_text, str):
        raise ValueError("Input text must be a string.")

    try:
        token_count = count_tokens(raw_text)
        logger.info("Token count for input: %d", token_count)
    except Exception as e:
        logger.error("Error counting tokens: %s", str(e))
        raise

    # Token limits loaded from environment variables, with default values
    model_preferences = {
        GPT4Model: int(os.getenv('GPT4_TOKEN_LIMIT', '128000')),
        ClaudeModel: int(os.getenv('CLAUDE_TOKEN_LIMIT', '199999')),
        GeminiModel: int(os.getenv('GEMINI_TOKEN_LIMIT', '1000000'))
    }

    # Filter models that have sufficient token limit
    available_models = sorted(
        (model for model, limit in model_preferences.items() if token_count <= limit),
        key=lambda model: model_preferences[model]
    )

    # Choose the first available model based on the sorted order
    selected_model_class = available_models[0] if available_models else None

    if selected_model_class:
        selected_model = selected_model_class()  # Create an instance of the selected model class
        logger.info("Selected %s for analysis with token count of %d", selected_model.__class__.__name__, token_count)
    else:
        selected_model = None
        logger.warning("No suitable model found for the given token count.")

    return selected_model

# file: model_selector.py (end)