flashlite 0.1.0__py3-none-any.whl

flashlite/config.py

@@ -0,0 +1,154 @@
+"""Configuration management and environment loading."""
+
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from dotenv import load_dotenv
+
+from .types import ConfigError, RateLimitConfig, RetryConfig
+
+
+@dataclass
+class FlashliteConfig:
+    """Main configuration for Flashlite client."""
+
+    # Default model to use if not specified per-request
+    default_model: str | None = None
+
+    # Default completion parameters
+    default_temperature: float | None = None
+    default_max_tokens: int | None = None
+
+    # Retry configuration
+    retry: RetryConfig = field(default_factory=RetryConfig)
+
+    # Rate limiting configuration
+    rate_limit: RateLimitConfig = field(default_factory=RateLimitConfig)
+
+    # Template directory
+    template_dir: Path | str | None = None
+
+    # Logging
+    log_requests: bool = False
+    log_level: str = "INFO"
+
+    # Default kwargs to pass to all completions
+    default_kwargs: dict[str, Any] = field(default_factory=dict)
+
+    # Timeout in seconds
+    timeout: float = 600.0
+
+    @classmethod
+    def from_env(cls) -> "FlashliteConfig":
+        """Create config from environment variables."""
+        config = cls()
+
+        # Read config from FLASHLITE_ prefixed env vars
+        if model := os.getenv("FLASHLITE_DEFAULT_MODEL"):
+            config.default_model = model
+
+        if temp := os.getenv("FLASHLITE_DEFAULT_TEMPERATURE"):
+            try:
+                config.default_temperature = float(temp)
+            except ValueError:
+                raise ConfigError(f"Invalid FLASHLITE_DEFAULT_TEMPERATURE: {temp}")
+
+        if max_tokens := os.getenv("FLASHLITE_DEFAULT_MAX_TOKENS"):
+            try:
+                config.default_max_tokens = int(max_tokens)
+            except ValueError:
+                raise ConfigError(f"Invalid FLASHLITE_DEFAULT_MAX_TOKENS: {max_tokens}")
+
+        if log_level := os.getenv("FLASHLITE_LOG_LEVEL"):
+            config.log_level = log_level.upper()
+
+        if os.getenv("FLASHLITE_LOG_REQUESTS", "").lower() in ("1", "true", "yes"):
+            config.log_requests = True
+
+        if template_dir := os.getenv("FLASHLITE_TEMPLATE_DIR"):
+            config.template_dir = Path(template_dir)
+
+        if rpm := os.getenv("FLASHLITE_RATE_LIMIT_RPM"):
+            try:
+                config.rate_limit.requests_per_minute = float(rpm)
+            except ValueError:
+                raise ConfigError(f"Invalid FLASHLITE_RATE_LIMIT_RPM: {rpm}")
+
+        if tpm := os.getenv("FLASHLITE_RATE_LIMIT_TPM"):
+            try:
+                config.rate_limit.tokens_per_minute = float(tpm)
+            except ValueError:
+                raise ConfigError(f"Invalid FLASHLITE_RATE_LIMIT_TPM: {tpm}")
+
+        if timeout := os.getenv("FLASHLITE_TIMEOUT"):
+            try:
+                config.timeout = float(timeout)
+            except ValueError:
+                raise ConfigError(f"Invalid FLASHLITE_TIMEOUT: {timeout}")
+
+        return config
+
+
+def load_env_files(
+    env_file: str | Path | None = None,
+    env_files: list[str | Path] | None = None,
+) -> None:
+    """
+    Load environment variables from .env files.
+
+    Args:
+        env_file: Single env file to load
+        env_files: Multiple env files to load (later files override earlier)
+    """
+    files_to_load: list[Path] = []
+
+    if env_files:
+        files_to_load.extend(Path(f) for f in env_files)
+    elif env_file:
+        files_to_load.append(Path(env_file))
+    else:
+        # Default: try to load .env from current directory
+        default_env = Path(".env")
+        if default_env.exists():
+            files_to_load.append(default_env)
+
+    # Load files in order (later overrides earlier)
+    for file_path in files_to_load:
+        if file_path.exists():
+            load_dotenv(file_path, override=True)
+
+
+def validate_api_keys(required_providers: list[str] | None = None) -> dict[str, bool]:
+    """
+    Check which API keys are configured.
+
+    Args:
+        required_providers: If provided, raise error if any are missing
+
+    Returns:
+        Dict mapping provider names to whether their key is set
+    """
+    key_mapping = {
+        "openai": "OPENAI_API_KEY",
+        "anthropic": "ANTHROPIC_API_KEY",
+        "google": "GOOGLE_API_KEY",
+        "cohere": "COHERE_API_KEY",
+        "azure": "AZURE_API_KEY",
+        "bedrock": "AWS_ACCESS_KEY_ID",
+    }
+
+    results = {}
+    for provider, env_var in key_mapping.items():
+        results[provider] = bool(os.getenv(env_var))
+
+    if required_providers:
+        missing = [p for p in required_providers if not results.get(p)]
+        if missing:
+            raise ConfigError(
+                f"Missing API keys for providers: {', '.join(missing)}. "
+                f"Set the corresponding environment variables."
+            )
+
+    return results

flashlite/conversation/__init__.py

@@ -0,0 +1,30 @@
+"""Conversation management module for flashlite."""
+
+from .context import (
+    ContextLimits,
+    ContextManager,
+    check_context_fit,
+    estimate_messages_tokens,
+    estimate_tokens,
+    truncate_messages,
+)
+from .manager import Conversation, ConversationState, Turn
+from .multi_agent import Agent, ChatMessage, MultiAgentChat
+
+__all__ = [
+    # Conversation management
+    "Conversation",
+    "ConversationState",
+    "Turn",
+    # Multi-agent conversations
+    "MultiAgentChat",
+    "Agent",
+    "ChatMessage",
+    # Context management
+    "ContextManager",
+    "ContextLimits",
+    "estimate_tokens",
+    "estimate_messages_tokens",
+    "check_context_fit",
+    "truncate_messages",
+]

flashlite/conversation/context.py

@@ -0,0 +1,319 @@
+"""Context window management for conversations."""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+from ..types import Messages
+
+logger = logging.getLogger(__name__)
+
+
+# Approximate token counts per character for different models
+# These are rough estimates - actual tokenization varies
+CHARS_PER_TOKEN_ESTIMATE = 4
+
+
+@dataclass
+class ContextLimits:
+    """Context window limits for a model."""
+
+    max_tokens: int
+    recommended_max: int | None = None  # Leave room for response
+
+    @classmethod
+    def for_model(cls, model: str) -> "ContextLimits":
+        """Get context limits for a model (approximate)."""
+        model_lower = model.lower()
+
+        # GPT-4 variants
+        if "gpt-4o" in model_lower:
+            return cls(max_tokens=128_000, recommended_max=120_000)
+        if "gpt-4-turbo" in model_lower or "gpt-4-1106" in model_lower:
+            return cls(max_tokens=128_000, recommended_max=120_000)
+        if "gpt-4-32k" in model_lower:
+            return cls(max_tokens=32_768, recommended_max=30_000)
+        if "gpt-4" in model_lower:
+            return cls(max_tokens=8_192, recommended_max=7_000)
+
+        # GPT-3.5 variants
+        if "gpt-3.5-turbo-16k" in model_lower:
+            return cls(max_tokens=16_384, recommended_max=15_000)
+        if "gpt-3.5" in model_lower:
+            return cls(max_tokens=16_384, recommended_max=15_000)
+
+        # Claude variants
+        if "claude-3" in model_lower or "claude-sonnet-4" in model_lower:
+            return cls(max_tokens=200_000, recommended_max=190_000)
+        if "claude-2" in model_lower:
+            return cls(max_tokens=100_000, recommended_max=95_000)
+        if "claude" in model_lower:
+            return cls(max_tokens=100_000, recommended_max=95_000)
+
+        # Gemini
+        if "gemini-1.5" in model_lower:
+            return cls(max_tokens=1_000_000, recommended_max=900_000)
+        if "gemini" in model_lower:
+            return cls(max_tokens=32_768, recommended_max=30_000)
+
+        # Mistral
+        if "mistral-large" in model_lower:
+            return cls(max_tokens=128_000, recommended_max=120_000)
+        if "mistral" in model_lower:
+            return cls(max_tokens=32_768, recommended_max=30_000)
+
+        # Default conservative estimate
+        return cls(max_tokens=8_192, recommended_max=7_000)
+
+
+def estimate_tokens(text: str) -> int:
+    """
+    Estimate the number of tokens in text.
+
+    This is a rough approximation. For accurate counts, use tiktoken
+    or the provider's tokenizer.
+
+    Args:
+        text: The text to estimate tokens for
+
+    Returns:
+        Estimated token count
+    """
+    return len(text) // CHARS_PER_TOKEN_ESTIMATE + 1
+
+
+def estimate_messages_tokens(messages: Messages) -> int:
+    """
+    Estimate total tokens in a messages list.
+
+    Args:
+        messages: List of messages
+
+    Returns:
+        Estimated token count
+    """
+    total = 0
+    for msg in messages:
+        # Add overhead for message structure
+        total += 4  # Approximate overhead per message
+
+        content = msg.get("content", "")
+        if isinstance(content, str):
+            total += estimate_tokens(content)
+        elif isinstance(content, list):
+            # Handle multi-part content (e.g., images)
+            for part in content:
+                if isinstance(part, dict) and "text" in part:
+                    total += estimate_tokens(part["text"])
+                else:
+                    total += 100  # Rough estimate for non-text content
+
+    return total
+
+
+def check_context_fit(
+    messages: Messages,
+    model: str,
+    max_response_tokens: int = 4096,
+) -> tuple[bool, dict[str, Any]]:
+    """
+    Check if messages fit within the model's context window.
+
+    Args:
+        messages: The messages to check
+        model: The model name
+        max_response_tokens: Expected max tokens in response
+
+    Returns:
+        Tuple of (fits, info_dict) where info_dict contains:
+        - estimated_tokens: Estimated input tokens
+        - max_tokens: Model's max context
+        - remaining: Tokens remaining for response
+        - warning: Optional warning message
+    """
+    limits = ContextLimits.for_model(model)
+    estimated = estimate_messages_tokens(messages)
+    remaining = limits.max_tokens - estimated - max_response_tokens
+
+    info: dict[str, Any] = {
+        "estimated_tokens": estimated,
+        "max_tokens": limits.max_tokens,
+        "remaining": remaining,
+    }
+
+    if remaining < 0:
+        info["warning"] = (
+            f"Messages ({estimated} tokens) + response ({max_response_tokens}) "
+            f"exceed context limit ({limits.max_tokens})"
+        )
+        return False, info
+
+    if limits.recommended_max and estimated > limits.recommended_max:
+        info["warning"] = (
+            f"Messages ({estimated} tokens) exceed recommended max "
+            f"({limits.recommended_max}). Consider truncating."
+        )
+
+    return True, info
+
+
+def truncate_messages(
+    messages: Messages,
+    max_tokens: int,
+    strategy: str = "oldest",
+    keep_system: bool = True,
+) -> Messages:
+    """
+    Truncate messages to fit within a token budget.
+
+    Args:
+        messages: The messages to truncate
+        max_tokens: Maximum tokens to keep
+        strategy: Truncation strategy - "oldest" removes oldest messages first
+        keep_system: Whether to always keep system messages
+
+    Returns:
+        Truncated messages list
+    """
+    if strategy != "oldest":
+        raise ValueError(f"Unknown truncation strategy: {strategy}")
+
+    messages_list = list(messages)
+    current_tokens = estimate_messages_tokens(messages_list)
+
+    if current_tokens <= max_tokens:
+        return messages_list
+
+    # Separate system messages if we're keeping them
+    system_messages = []
+    other_messages = []
+
+    for msg in messages_list:
+        if keep_system and msg.get("role") == "system":
+            system_messages.append(msg)
+        else:
+            other_messages.append(msg)
+
+    system_tokens = estimate_messages_tokens(system_messages)
+    available_tokens = max_tokens - system_tokens
+
+    # Remove oldest messages until we fit
+    while other_messages and estimate_messages_tokens(other_messages) > available_tokens:
+        other_messages.pop(0)
+
+    result = system_messages + other_messages
+
+    removed_count = len(messages_list) - len(result)
+    if removed_count > 0:
+        logger.info(
+            f"Truncated {removed_count} messages to fit context window "
+            f"({current_tokens} -> {estimate_messages_tokens(result)} tokens)"
+        )
+
+    return result
+
+
+class ContextManager:
+    """
+    Manages context window for a conversation.
+
+    Provides automatic truncation when approaching limits and
+    warnings when context is getting full.
+
+    Example:
+        ctx = ContextManager(model="gpt-4o", max_response_tokens=4096)
+
+        # Check if messages fit
+        fits, info = ctx.check(messages)
+        if not fits:
+            messages = ctx.truncate(messages)
+
+        # Or use auto mode
+        messages = ctx.prepare(messages)  # Automatically truncates if needed
+    """
+
+    def __init__(
+        self,
+        model: str,
+        max_response_tokens: int = 4096,
+        auto_truncate: bool = True,
+        truncation_strategy: str = "oldest",
+        keep_system: bool = True,
+        warn_threshold: float = 0.8,
+    ):
+        """
+        Initialize context manager.
+
+        Args:
+            model: Model name to get context limits for
+            max_response_tokens: Expected max tokens in response
+            auto_truncate: Whether to automatically truncate when needed
+            truncation_strategy: Strategy for truncation ("oldest")
+            keep_system: Whether to preserve system messages during truncation
+            warn_threshold: Warn when context usage exceeds this ratio (0-1)
+        """
+        self._model = model
+        self._limits = ContextLimits.for_model(model)
+        self._max_response_tokens = max_response_tokens
+        self._auto_truncate = auto_truncate
+        self._truncation_strategy = truncation_strategy
+        self._keep_system = keep_system
+        self._warn_threshold = warn_threshold
+
+    def check(self, messages: Messages) -> tuple[bool, dict[str, Any]]:
+        """Check if messages fit within context limits."""
+        return check_context_fit(messages, self._model, self._max_response_tokens)
+
+    def truncate(self, messages: Messages) -> Messages:
+        """Truncate messages to fit within limits."""
+        max_input = self._limits.max_tokens - self._max_response_tokens
+        return truncate_messages(
+            messages,
+            max_input,
+            strategy=self._truncation_strategy,
+            keep_system=self._keep_system,
+        )
+
+    def prepare(self, messages: Messages) -> Messages:
+        """
+        Prepare messages for completion.
+
+        Checks fit and truncates if needed (when auto_truncate is enabled).
+        Logs warnings when approaching limits.
+
+        Args:
+            messages: The messages to prepare
+
+        Returns:
+            Messages ready for completion (possibly truncated)
+        """
+        fits, info = self.check(messages)
+
+        # Check warning threshold
+        usage_ratio = info["estimated_tokens"] / self._limits.max_tokens
+        if usage_ratio > self._warn_threshold:
+            logger.warning(
+                f"Context usage at {usage_ratio:.0%} "
+                f"({info['estimated_tokens']}/{self._limits.max_tokens} tokens)"
+            )
+
+        if not fits:
+            if self._auto_truncate:
+                logger.warning(f"Context exceeded, truncating: {info.get('warning')}")
+                return self.truncate(messages)
+            else:
+                raise ValueError(
+                    f"Messages exceed context limit: {info.get('warning')}"
+                )
+
+        return list(messages)
+
+    @property
+    def model(self) -> str:
+        """The model this manager is configured for."""
+        return self._model
+
+    @property
+    def max_tokens(self) -> int:
+        """Maximum context tokens for the model."""
+        return self._limits.max_tokens