"""Shared utilities for Headroom SDK.""" from __future__ import annotations import hashlib import json import re import uuid from datetime import datetime from typing import Any # Marker format for Headroom modifications MARKER_PREFIX = " str: """Generate a unique request ID.""" return str(uuid.uuid4()) def compute_hash(data: str | bytes) -> str: """Compute SHA256 hash, returning hex string.""" if isinstance(data, str): data = data.encode("utf-8") return hashlib.sha256(data).hexdigest() def compute_short_hash(data: str | bytes, length: int = 16) -> str: """Compute truncated SHA256 hash.""" return compute_hash(data)[:length] def compute_messages_hash(messages: list[dict[str, Any]]) -> str: """Compute hash of messages list for deduplication.""" # Serialize deterministically serialized = json.dumps(messages, sort_keys=True, separators=(",", ":")) return compute_short_hash(serialized) def compute_prefix_hash(messages: list[dict[str, Any]], prefix_count: int | None = None) -> str: """ Compute hash of message prefix for cache alignment. Args: messages: List of messages. prefix_count: Number of messages to include (default: all system messages + 1). Returns: Hash of the prefix content. """ if not messages: return compute_short_hash("") if prefix_count is None: # Default: system messages + first non-system prefix_count = 1 for i, msg in enumerate(messages): if msg.get("role") == "system": prefix_count = i + 2 else: break prefix_messages = messages[:prefix_count] serialized = json.dumps(prefix_messages, sort_keys=True, separators=(",", ":")) return compute_short_hash(serialized) def format_timestamp(dt: datetime | None = None) -> str: """Format datetime as ISO8601 string.""" if dt is None: dt = datetime.utcnow() return dt.isoformat() + "Z" def parse_timestamp(ts: str) -> datetime: """Parse ISO8601 timestamp string.""" # Handle both with and without Z suffix ts = ts.rstrip("Z") return datetime.fromisoformat(ts) def create_marker(marker_type: str, **kwargs: Any) -> str: """ Create a Headroom marker string. Args: marker_type: Type of marker (e.g., "tool_digest", "dropped_context"). **kwargs: Attributes to include in the marker. Returns: Formatted marker string. """ attrs = " ".join(f'{k}="{v}"' for k, v in kwargs.items()) if attrs: return f"{MARKER_PREFIX}{marker_type} {attrs}{MARKER_SUFFIX}" return f"{MARKER_PREFIX}{marker_type}{MARKER_SUFFIX}" def create_tool_digest_marker(original_hash: str) -> str: """Create marker for crushed tool output.""" return create_marker("tool_digest", sha256=original_hash) def create_dropped_context_marker(reason: str, count: int | None = None) -> str: """Create marker for dropped context.""" if count is not None: return create_marker("dropped_context", reason=reason, count=str(count)) return create_marker("dropped_context", reason=reason) def create_truncated_marker(original_length: int, truncated_to: int) -> str: """Create marker for truncated content.""" return create_marker( "truncated", original=str(original_length), truncated_to=str(truncated_to), ) def extract_markers(text: str) -> list[dict[str, Any]]: """ Extract Headroom markers from text. Returns: List of dicts with marker_type and attributes. """ pattern = re.compile(r"]*)>") markers = [] for match in pattern.finditer(text): marker_type = match.group(1) attrs_str = match.group(2).strip() # Parse attributes attrs: dict[str, str] = {} if attrs_str: attr_pattern = re.compile(r'(\w+)="([^"]*)"') for attr_match in attr_pattern.finditer(attrs_str): attrs[attr_match.group(1)] = attr_match.group(2) markers.append({"type": marker_type, "attributes": attrs}) return markers def safe_json_loads(text: str) -> tuple[Any | None, bool]: """ Safely parse JSON, returning (result, success). Args: text: JSON string to parse. Returns: Tuple of (parsed_result or None, success_bool). """ try: return json.loads(text), True except (json.JSONDecodeError, ValueError): return None, False def safe_json_dumps(obj: Any, **kwargs: Any) -> str: """ Safely serialize to JSON with defaults. Args: obj: Object to serialize. **kwargs: Additional json.dumps arguments. Returns: JSON string. """ kwargs.setdefault("ensure_ascii", False) kwargs.setdefault("separators", (",", ":")) # Compact by default return json.dumps(obj, **kwargs) def estimate_cost( input_tokens: int, output_tokens: int, model: str, cached_tokens: int = 0, provider: Any = None, ) -> float | None: """ Estimate API cost in USD using provider. Args: input_tokens: Number of input tokens. output_tokens: Number of output tokens. model: Model name. cached_tokens: Number of cached input tokens. provider: Provider instance for cost estimation. Returns: Estimated cost in USD, or None if not available. """ if provider is None: return None result = provider.estimate_cost(input_tokens, output_tokens, model, cached_tokens) return float(result) if result is not None else None def format_cost(cost: float) -> str: """Format cost as human-readable string.""" if cost < 0.01: return f"${cost:.4f}" return f"${cost:.2f}" def deep_copy_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: """Create a deep copy of messages list.""" result: list[dict[str, Any]] = json.loads(json.dumps(messages)) return result