headroom-ai 0.2.13__py3-none-any.whl

headroom/parser.py

@@ -0,0 +1,293 @@
+"""Message parsing utilities for Headroom SDK."""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from typing import TYPE_CHECKING, Any
+
+from .config import Block, WasteSignals
+
+if TYPE_CHECKING:
+    from .tokenizer import Tokenizer
+
+
+# Patterns for detecting waste signals
+HTML_TAG_PATTERN = re.compile(r"<[^>]+>")
+HTML_COMMENT_PATTERN = re.compile(r"<!--[\s\S]*?-->")
+BASE64_PATTERN = re.compile(r"[A-Za-z0-9+/]{50,}={0,2}")
+WHITESPACE_PATTERN = re.compile(r"[ \t]{4,}|\n{3,}")
+JSON_BLOCK_PATTERN = re.compile(r"\{[\s\S]{500,}\}")
+
+# Patterns for RAG detection (best effort)
+RAG_MARKERS = [
+    r"\[Document\s*\d+\]",
+    r"\[Source:\s*",
+    r"<context>",
+    r"<document>",
+    r"Retrieved from:",
+    r"From the knowledge base:",
+]
+RAG_PATTERN = re.compile("|".join(RAG_MARKERS), re.IGNORECASE)
+
+
+def compute_hash(text: str) -> str:
+    """Compute SHA256 hash of text, truncated to 16 chars."""
+    return hashlib.sha256(text.encode()).hexdigest()[:16]
+
+
+def detect_waste_signals(text: str, tokenizer: Tokenizer) -> WasteSignals:
+    """
+    Detect waste signals in text.
+
+    Args:
+        text: The text to analyze.
+        tokenizer: Tokenizer for counting tokens.
+
+    Returns:
+        WasteSignals with detected waste.
+    """
+    signals = WasteSignals()
+
+    if not text:
+        return signals
+
+    # HTML tags and comments
+    html_matches = HTML_TAG_PATTERN.findall(text) + HTML_COMMENT_PATTERN.findall(text)
+    if html_matches:
+        html_text = "".join(html_matches)
+        signals.html_noise_tokens = tokenizer.count_text(html_text)
+
+    # Base64 blobs
+    base64_matches = BASE64_PATTERN.findall(text)
+    if base64_matches:
+        base64_text = "".join(base64_matches)
+        signals.base64_tokens = tokenizer.count_text(base64_text)
+
+    # Excessive whitespace
+    ws_matches = WHITESPACE_PATTERN.findall(text)
+    if ws_matches:
+        # Count tokens that could be saved by normalizing
+        ws_text = "".join(ws_matches)
+        signals.whitespace_tokens = max(0, tokenizer.count_text(ws_text) - len(ws_matches))
+
+    # Large JSON blocks
+    json_matches = JSON_BLOCK_PATTERN.findall(text)
+    if json_matches:
+        for match in json_matches:
+            tokens = tokenizer.count_text(match)
+            if tokens > 500:
+                signals.json_bloat_tokens += tokens
+
+    return signals
+
+
+def is_rag_content(text: str) -> bool:
+    """Check if text appears to be RAG-injected content."""
+    return RAG_PATTERN.search(text) is not None
+
+
+def parse_message_to_blocks(
+    message: dict[str, Any],
+    index: int,
+    tokenizer: Tokenizer,
+) -> list[Block]:
+    """
+    Parse a single message into Block objects.
+
+    Args:
+        message: The message dict to parse.
+        index: Position in the message list.
+        tokenizer: Tokenizer for token counting.
+
+    Returns:
+        List of Block objects (usually 1, but tool_calls may produce multiple).
+    """
+    blocks: list[Block] = []
+    role = message.get("role", "unknown")
+
+    # Handle content
+    content = message.get("content")
+    if content:
+        if isinstance(content, str):
+            text = content
+        elif isinstance(content, list):
+            # Multi-modal - extract text parts
+            text_parts = []
+            for part in content:
+                if isinstance(part, dict) and part.get("type") == "text":
+                    text_parts.append(part.get("text", ""))
+                elif isinstance(part, str):
+                    text_parts.append(part)
+            text = "\n".join(text_parts)
+        else:
+            text = str(content)
+
+        # Determine block kind
+        if role == "system":
+            kind = "system"
+        elif role == "user":
+            # Check if this looks like RAG content
+            kind = "rag" if is_rag_content(text) else "user"
+        elif role == "assistant":
+            kind = "assistant"
+        elif role == "tool":
+            kind = "tool_result"
+        else:
+            kind = "unknown"
+
+        # Build flags
+        flags: dict[str, Any] = {}
+        if role == "tool":
+            flags["tool_call_id"] = message.get("tool_call_id")
+
+        # Detect waste
+        waste = detect_waste_signals(text, tokenizer)
+        if waste.total() > 0:
+            flags["waste_signals"] = waste.to_dict()
+
+        blocks.append(
+            Block(
+                kind=kind,  # type: ignore[arg-type]
+                text=text,
+                tokens_est=tokenizer.count_text(text) + 4,  # Add message overhead
+                content_hash=compute_hash(text),
+                source_index=index,
+                flags=flags,
+            )
+        )
+
+    # Handle tool calls (assistant messages with tool_calls)
+    tool_calls = message.get("tool_calls")
+    if tool_calls:
+        for tc in tool_calls:
+            func = tc.get("function", {})
+            tc_text = f"{func.get('name', 'unknown')}({func.get('arguments', '')})"
+
+            blocks.append(
+                Block(
+                    kind="tool_call",
+                    text=tc_text,
+                    tokens_est=tokenizer.count_text(tc_text) + 10,
+                    content_hash=compute_hash(tc_text),
+                    source_index=index,
+                    flags={
+                        "tool_call_id": tc.get("id"),
+                        "function_name": func.get("name"),
+                    },
+                )
+            )
+
+    # If no content or tool_calls, create a minimal block
+    if not blocks:
+        blocks.append(
+            Block(
+                kind="unknown",
+                text="",
+                tokens_est=4,
+                content_hash=compute_hash(""),
+                source_index=index,
+                flags={},
+            )
+        )
+
+    return blocks
+
+
+def parse_messages(
+    messages: list[dict[str, Any]],
+    tokenizer: Tokenizer,
+) -> tuple[list[Block], dict[str, int], WasteSignals]:
+    """
+    Parse all messages into blocks with analysis.
+
+    Args:
+        messages: List of message dicts.
+        tokenizer: Tokenizer instance for token counting.
+
+    Returns:
+        Tuple of (blocks, block_breakdown, total_waste_signals)
+    """
+    all_blocks: list[Block] = []
+    total_waste = WasteSignals()
+
+    for i, msg in enumerate(messages):
+        blocks = parse_message_to_blocks(msg, i, tokenizer)
+        all_blocks.extend(blocks)
+
+        # Accumulate waste signals
+        for block in blocks:
+            if "waste_signals" in block.flags:
+                ws = block.flags["waste_signals"]
+                total_waste.json_bloat_tokens += ws.get("json_bloat", 0)
+                total_waste.html_noise_tokens += ws.get("html_noise", 0)
+                total_waste.base64_tokens += ws.get("base64", 0)
+                total_waste.whitespace_tokens += ws.get("whitespace", 0)
+                total_waste.dynamic_date_tokens += ws.get("dynamic_date", 0)
+                total_waste.repetition_tokens += ws.get("repetition", 0)
+
+    # Compute block breakdown
+    breakdown: dict[str, int] = {}
+    for block in all_blocks:
+        kind = block.kind
+        breakdown[kind] = breakdown.get(kind, 0) + block.tokens_est
+
+    return all_blocks, breakdown, total_waste
+
+
+def find_tool_units(messages: list[dict[str, Any]]) -> list[tuple[int, list[int]]]:
+    """
+    Find tool call units (assistant with tool_calls + corresponding tool responses).
+
+    A tool unit is atomic - if the assistant message is dropped, all its
+    tool responses must also be dropped.
+
+    Args:
+        messages: List of message dicts.
+
+    Returns:
+        List of (assistant_index, [tool_response_indices]) tuples.
+    """
+    units: list[tuple[int, list[int]]] = []
+
+    # Build map of tool_call_id -> message index for tool responses
+    tool_response_map: dict[str, int] = {}
+    for i, msg in enumerate(messages):
+        if msg.get("role") == "tool":
+            tc_id = msg.get("tool_call_id")
+            if tc_id:
+                tool_response_map[tc_id] = i
+
+    # Find assistant messages with tool_calls
+    for i, msg in enumerate(messages):
+        if msg.get("role") == "assistant" and msg.get("tool_calls"):
+            tool_calls = msg["tool_calls"]
+            response_indices: list[int] = []
+
+            for tc in tool_calls:
+                tc_id = tc.get("id")
+                if tc_id and tc_id in tool_response_map:
+                    response_indices.append(tool_response_map[tc_id])
+
+            if response_indices:
+                units.append((i, sorted(response_indices)))
+
+    return units
+
+
+def get_message_content_text(message: dict[str, Any]) -> str:
+    """Extract text content from a message."""
+    content = message.get("content")
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts = []
+        for part in content:
+            if isinstance(part, dict) and part.get("type") == "text":
+                parts.append(part.get("text", ""))
+            elif isinstance(part, str):
+                parts.append(part)
+        return "\n".join(parts)
+    return str(content)

headroom/pricing/__init__.py

@@ -0,0 +1,51 @@
+"""Pricing module for LLM cost estimation.
+
+This module provides pricing information and cost estimation utilities
+for various LLM providers. Uses LiteLLM's community-maintained pricing
+database for up-to-date costs across 100+ models.
+"""
+
+# Legacy imports for backwards compatibility
+from .anthropic_prices import (
+    ANTHROPIC_PRICES,
+    get_anthropic_registry,
+)
+from .anthropic_prices import (
+    LAST_UPDATED as ANTHROPIC_LAST_UPDATED,
+)
+from .litellm_pricing import (
+    LiteLLMModelPricing,
+    estimate_cost,
+    get_litellm_model_cost,
+    get_model_pricing,
+    list_available_models,
+)
+from .openai_prices import (
+    LAST_UPDATED as OPENAI_LAST_UPDATED,
+)
+from .openai_prices import (
+    OPENAI_PRICES,
+    get_openai_registry,
+)
+from .registry import CostEstimate, ModelPricing, PricingRegistry
+
+__all__ = [
+    # LiteLLM-based pricing (preferred)
+    "LiteLLMModelPricing",
+    "estimate_cost",
+    "get_litellm_model_cost",
+    "get_model_pricing",
+    "list_available_models",
+    # Core classes
+    "CostEstimate",
+    "ModelPricing",
+    "PricingRegistry",
+    # Legacy - OpenAI (deprecated, use LiteLLM instead)
+    "OPENAI_LAST_UPDATED",
+    "OPENAI_PRICES",
+    "get_openai_registry",
+    # Legacy - Anthropic (deprecated, use LiteLLM instead)
+    "ANTHROPIC_LAST_UPDATED",
+    "ANTHROPIC_PRICES",
+    "get_anthropic_registry",
+]

headroom/pricing/anthropic_prices.py

@@ -0,0 +1,81 @@
+"""Anthropic model pricing information."""
+
+from datetime import date
+
+from .registry import ModelPricing, PricingRegistry
+
+# Last verified date for pricing information
+LAST_UPDATED = date(2025, 1, 6)
+
+# Official pricing page
+SOURCE_URL = "https://www.anthropic.com/pricing"
+
+# All prices are in USD per 1 million tokens
+ANTHROPIC_PRICES: dict[str, ModelPricing] = {
+    "claude-3-5-sonnet-20241022": ModelPricing(
+        model="claude-3-5-sonnet-20241022",
+        provider="anthropic",
+        input_per_1m=3.00,
+        output_per_1m=15.00,
+        cached_input_per_1m=0.30,
+        batch_input_per_1m=1.50,
+        batch_output_per_1m=7.50,
+        context_window=200_000,
+        notes="Most intelligent Claude model, best for complex tasks",
+    ),
+    "claude-3-5-sonnet-latest": ModelPricing(
+        model="claude-3-5-sonnet-latest",
+        provider="anthropic",
+        input_per_1m=3.00,
+        output_per_1m=15.00,
+        cached_input_per_1m=0.30,
+        batch_input_per_1m=1.50,
+        batch_output_per_1m=7.50,
+        context_window=200_000,
+        notes="Alias for claude-3-5-sonnet-20241022",
+    ),
+    "claude-3-5-haiku-20241022": ModelPricing(
+        model="claude-3-5-haiku-20241022",
+        provider="anthropic",
+        input_per_1m=0.80,
+        output_per_1m=4.00,
+        cached_input_per_1m=0.08,
+        batch_input_per_1m=0.40,
+        batch_output_per_1m=2.00,
+        context_window=200_000,
+        notes="Fast and cost-effective for simple tasks",
+    ),
+    "claude-3-opus-20240229": ModelPricing(
+        model="claude-3-opus-20240229",
+        provider="anthropic",
+        input_per_1m=15.00,
+        output_per_1m=75.00,
+        cached_input_per_1m=1.50,
+        batch_input_per_1m=7.50,
+        batch_output_per_1m=37.50,
+        context_window=200_000,
+        notes="Previous generation powerful model for complex tasks",
+    ),
+    "claude-3-haiku-20240307": ModelPricing(
+        model="claude-3-haiku-20240307",
+        provider="anthropic",
+        input_per_1m=0.25,
+        output_per_1m=1.25,
+        cached_input_per_1m=0.03,
+        context_window=200_000,
+        notes="Previous generation fastest and most compact model",
+    ),
+}
+
+
+def get_anthropic_registry() -> PricingRegistry:
+    """Create and return an Anthropic pricing registry.
+
+    Returns:
+        PricingRegistry configured with Anthropic model prices.
+    """
+    return PricingRegistry(
+        last_updated=LAST_UPDATED,
+        source_url=SOURCE_URL,
+        prices=ANTHROPIC_PRICES.copy(),
+    )

headroom/pricing/litellm_pricing.py

@@ -0,0 +1,113 @@
+"""LiteLLM-based pricing for model cost estimation.
+
+Uses LiteLLM's community-maintained model cost database instead of
+hardcoded values. This provides up-to-date pricing for 100+ models.
+
+See: https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import litellm
+
+
+@dataclass
+class LiteLLMModelPricing:
+    """Pricing information from LiteLLM's database.
+
+    All costs are in USD per 1 million tokens.
+    """
+
+    model: str
+    input_cost_per_1m: float
+    output_cost_per_1m: float
+    max_tokens: int | None = None
+    max_input_tokens: int | None = None
+    max_output_tokens: int | None = None
+    supports_vision: bool = False
+    supports_function_calling: bool = False
+
+
+def get_litellm_model_cost() -> dict[str, Any]:
+    """Get LiteLLM's full model cost dictionary.
+
+    Returns:
+        Dictionary mapping model names to their pricing/capability info.
+    """
+    return litellm.model_cost
+
+
+def get_model_pricing(model: str) -> LiteLLMModelPricing | None:
+    """Get pricing for a model from LiteLLM's database.
+
+    Args:
+        model: Model name (e.g., 'gpt-4o', 'claude-3-5-sonnet-20241022').
+
+    Returns:
+        LiteLLMModelPricing if found, None otherwise.
+    """
+    cost_data = litellm.model_cost
+
+    # Try exact match first
+    info = cost_data.get(model)
+
+    # Try common provider prefixes if not found
+    if info is None:
+        for prefix in ["openai/", "anthropic/", "google/", "mistral/", "deepseek/"]:
+            if f"{prefix}{model}" in cost_data:
+                info = cost_data[f"{prefix}{model}"]
+                break
+
+    if info is None:
+        return None
+
+    # LiteLLM stores cost per token, convert to per 1M
+    input_per_token = info.get("input_cost_per_token", 0) or 0
+    output_per_token = info.get("output_cost_per_token", 0) or 0
+
+    return LiteLLMModelPricing(
+        model=model,
+        input_cost_per_1m=input_per_token * 1_000_000,
+        output_cost_per_1m=output_per_token * 1_000_000,
+        max_tokens=info.get("max_tokens"),
+        max_input_tokens=info.get("max_input_tokens"),
+        max_output_tokens=info.get("max_output_tokens"),
+        supports_vision=info.get("supports_vision", False),
+        supports_function_calling=info.get("supports_function_calling", False),
+    )
+
+
+def estimate_cost(
+    model: str,
+    input_tokens: int = 0,
+    output_tokens: int = 0,
+) -> float | None:
+    """Estimate cost for a model using LiteLLM's pricing.
+
+    Args:
+        model: Model name.
+        input_tokens: Number of input tokens.
+        output_tokens: Number of output tokens.
+
+    Returns:
+        Estimated cost in USD, or None if model not found.
+    """
+    pricing = get_model_pricing(model)
+    if pricing is None:
+        return None
+
+    input_cost = (input_tokens / 1_000_000) * pricing.input_cost_per_1m
+    output_cost = (output_tokens / 1_000_000) * pricing.output_cost_per_1m
+    return input_cost + output_cost
+
+
+def list_available_models() -> list[str]:
+    """List all models with pricing info in LiteLLM's database.
+
+    Returns:
+        List of model names.
+    """
+    return list(litellm.model_cost.keys())

headroom/pricing/openai_prices.py

@@ -0,0 +1,91 @@
+"""OpenAI model pricing information."""
+
+from datetime import date
+
+from .registry import ModelPricing, PricingRegistry
+
+# Last verified date for pricing information
+LAST_UPDATED = date(2025, 1, 6)
+
+# Official pricing page
+SOURCE_URL = "https://openai.com/api/pricing/"
+
+# All prices are in USD per 1 million tokens
+OPENAI_PRICES: dict[str, ModelPricing] = {
+    "gpt-4o": ModelPricing(
+        model="gpt-4o",
+        provider="openai",
+        input_per_1m=2.50,
+        output_per_1m=10.00,
+        cached_input_per_1m=1.25,
+        context_window=128_000,
+        notes="Most capable GPT-4o model",
+    ),
+    "gpt-4o-mini": ModelPricing(
+        model="gpt-4o-mini",
+        provider="openai",
+        input_per_1m=0.15,
+        output_per_1m=0.60,
+        cached_input_per_1m=0.075,
+        context_window=128_000,
+        notes="Affordable small model for fast, lightweight tasks",
+    ),
+    "o1": ModelPricing(
+        model="o1",
+        provider="openai",
+        input_per_1m=15.00,
+        output_per_1m=60.00,
+        cached_input_per_1m=7.50,
+        context_window=200_000,
+        notes="Reasoning model for complex, multi-step tasks",
+    ),
+    "o1-mini": ModelPricing(
+        model="o1-mini",
+        provider="openai",
+        input_per_1m=1.10,
+        output_per_1m=4.40,
+        cached_input_per_1m=0.55,
+        context_window=128_000,
+        notes="Smaller reasoning model, cost-effective for coding tasks",
+    ),
+    "o3-mini": ModelPricing(
+        model="o3-mini",
+        provider="openai",
+        input_per_1m=1.10,
+        output_per_1m=4.40,
+        cached_input_per_1m=0.55,
+        context_window=200_000,
+        notes="Latest small reasoning model",
+    ),
+    "gpt-4-turbo": ModelPricing(
+        model="gpt-4-turbo",
+        provider="openai",
+        input_per_1m=10.00,
+        output_per_1m=30.00,
+        cached_input_per_1m=5.00,
+        context_window=128_000,
+        notes="Previous generation GPT-4 Turbo model",
+    ),
+    "gpt-3.5-turbo": ModelPricing(
+        model="gpt-3.5-turbo",
+        provider="openai",
+        input_per_1m=0.50,
+        output_per_1m=1.50,
+        cached_input_per_1m=0.25,
+        context_window=16_385,
+        notes="Fast, inexpensive model for simple tasks",
+    ),
+}
+
+
+def get_openai_registry() -> PricingRegistry:
+    """Create and return an OpenAI pricing registry.
+
+    Returns:
+        PricingRegistry configured with OpenAI model prices.
+    """
+    return PricingRegistry(
+        last_updated=LAST_UPDATED,
+        source_url=SOURCE_URL,
+        prices=OPENAI_PRICES.copy(),
+    )