contextsage 0.1.0__py3-none-any.whl

contextsage/__init__.py

@@ -0,0 +1,22 @@
+"""ContextSage: intelligent, production-grade context summarization for LangGraph agents.
+
+ContextSage is a narrowly-focused enhancement layer over LangGraph/LangChain's
+built-in ``SummarizationMiddleware``. It analyzes heterogeneous agent context
+(mixed natural language, JSON, logs, errors, tables, code) to decide *what*
+should be preserved, compressed, or left untouched, then delegates the
+actual LLM-based semantic summarization to LangGraph itself.
+
+The public surface of this package is intentionally small — a single class:
+
+    from contextsage import IntelligentSummarizationMiddleware
+
+See ``docs/`` for the full architecture, or the module docstrings under
+``contextsage.middleware``, ``contextsage.planning``, ``contextsage.validation``,
+etc. for internal implementation details.
+"""
+
+from contextsage.middleware.summarization import IntelligentSummarizationMiddleware
+
+__all__ = ["IntelligentSummarizationMiddleware", "__version__"]
+
+__version__ = "0.1.0"

contextsage/budget/__init__.py

@@ -0,0 +1,3 @@
+"""Context budget analysis: token counting and window accounting."""
+
+from __future__ import annotations

contextsage/budget/analyzer.py

@@ -0,0 +1,100 @@
+"""Context budget analysis.
+
+The :class:`BudgetAnalyzer` computes an explicit, auditable accounting of
+the model's context window: how many tokens are already spoken for by the
+system prompt, tool schemas, reserved output, and safety margin, and
+therefore how many tokens are actually available for conversation input.
+
+ContextSage never assumes "characters / 4" is an authoritative token count; it
+delegates to a :class:`~contextsage.budget.tokens.TokenCounter` and clearly
+tracks whether the resulting numbers are estimates.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from contextsage.budget.tokens import TokenCounter
+from contextsage.core.models import BudgetReport
+from contextsage.exceptions import BudgetError
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Sequence
+
+    from langchain_core.messages import BaseMessage
+
+_DEFAULT_MAX_CONTEXT_TOKENS = 128_000
+_DEFAULT_SAFETY_MARGIN_FRACTION = 0.10
+_DEFAULT_SUMMARIZATION_OVERHEAD_TOKENS = 2_000
+
+
+@dataclass(slots=True)
+class BudgetConfig:
+    """User/derived configuration for budget analysis."""
+
+    maximum_context_tokens: int = _DEFAULT_MAX_CONTEXT_TOKENS
+    reserved_output_tokens: int = 4_000
+    safety_margin_fraction: float = _DEFAULT_SAFETY_MARGIN_FRACTION
+    summarization_overhead_tokens: int = _DEFAULT_SUMMARIZATION_OVERHEAD_TOKENS
+
+    def __post_init__(self) -> None:
+        if self.maximum_context_tokens <= 0:
+            raise BudgetError("maximum_context_tokens must be positive")
+        if not 0.0 <= self.safety_margin_fraction < 1.0:
+            raise BudgetError("safety_margin_fraction must be in [0, 1)")
+        if self.reserved_output_tokens < 0:
+            raise BudgetError("reserved_output_tokens must not be negative")
+
+
+class BudgetAnalyzer:
+    """Computes a :class:`BudgetReport` for a given message list."""
+
+    def __init__(self, token_counter: TokenCounter, config: BudgetConfig | None = None) -> None:
+        self._token_counter = token_counter
+        self._config = config or BudgetConfig()
+
+    def analyze(
+        self,
+        messages: Sequence[BaseMessage],
+        *,
+        system_prompt: str | None = None,
+        tool_schemas: Iterable[str] = (),
+    ) -> BudgetReport:
+        try:
+            current_input_tokens = self._token_counter.count_messages(messages)
+            system_tokens = self._token_counter.count_text(system_prompt) if system_prompt else 0
+            tool_schema_tokens = sum(
+                self._token_counter.count_text(schema) for schema in tool_schemas
+            )
+        except Exception as exc:
+            raise BudgetError(f"failed to compute budget: {exc}") from exc
+
+        safety_margin_tokens = round(
+            self._config.maximum_context_tokens * self._config.safety_margin_fraction
+        )
+
+        return BudgetReport(
+            maximum_context_tokens=self._config.maximum_context_tokens,
+            reserved_output_tokens=self._config.reserved_output_tokens,
+            system_tokens=system_tokens,
+            tool_schema_tokens=tool_schema_tokens,
+            current_input_tokens=current_input_tokens,
+            safety_margin_tokens=safety_margin_tokens,
+            summarization_overhead_tokens=self._config.summarization_overhead_tokens,
+        )
+
+
+def infer_max_context_tokens(model: object, default: int = _DEFAULT_MAX_CONTEXT_TOKENS) -> int:
+    """Best-effort extraction of a model's max input tokens from its profile.
+
+    Falls back to ``default`` when the model exposes no profile information;
+    this mirrors LangGraph's own graceful-degradation approach for models
+    without published context-window metadata.
+    """
+    profile = getattr(model, "profile", None)
+    if isinstance(profile, dict):
+        value = profile.get("max_input_tokens")
+        if isinstance(value, int) and value > 0:
+            return value
+    return default

contextsage/budget/tokens.py

@@ -0,0 +1,216 @@
+"""Provider/model-aware token counting.
+
+ContextSage needs token counts in two places with different cost profiles:
+
+* **Budget analysis** (whole-conversation accounting) — should match, as
+  closely as possible, whatever the wrapped LangGraph
+  ``SummarizationMiddleware`` itself uses, so ContextSage's "is summarization
+  required" decision is consistent with the mechanics that ultimately run.
+* **Per-unit accounting during decomposition** — invoked many times over a
+  single large tool output, so it must be cheap and cache-friendly (this is
+  why every counter is wrapped in :class:`CachingTokenCounter`).
+
+We never hardcode "characters / 4" as if it were authoritative.
+:class:`TiktokenTokenCounter` self-reports that it is an authoritative,
+provider-verified count via :attr:`TokenCounter.is_authoritative`.
+"""
+
+from __future__ import annotations
+
+import socket
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+from contextsage.exceptions import TokenCountingError
+
+if TYPE_CHECKING:
+    import tiktoken
+    from langchain_core.language_models.chat_models import BaseChatModel
+    from langchain_core.messages import BaseMessage
+
+_TEXT_CACHE_SIZE = 4096
+
+#: Used when no model is supplied (and no ``encoding_name`` override) —
+#: ``cl100k_base`` is the most broadly compatible encoding across
+#: OpenAI-compatible and many open-weight tokenizer families, and is used
+#: directly with no model-name lookup, so the zero-configuration default
+#: never depends on network access to resolve an uncached encoding.
+_DEFAULT_ENCODING_NAME = "cl100k_base"
+
+#: Tiktoken downloads uncached encoding files over plain HTTP with no
+#: explicit timeout (see ``tiktoken.load.read_file``). To keep ContextSage
+#: from hanging indefinitely in network-restricted/offline environments —
+#: the same production-reliability concern that ruled out on-demand grammar
+#: downloads elsewhere in this package — resolving an *explicit* model name
+#: to its encoding is bounded by a short socket timeout, with a safe local
+#: fallback to ``cl100k_base`` on failure.
+_ENCODING_DOWNLOAD_TIMEOUT_SECONDS = 5.0
+
+
+class TokenCounter(ABC):
+    """Abstract base class for every token counting strategy in ContextSage.
+
+    ``tiktoken`` is the default, always-installed implementation
+    (:class:`TiktokenTokenCounter`), but token counting remains pluggable:
+    anyone integrating a provider-specific tokenizer
+    (Anthropic, Gemini, a local model's own tokenizer, ...) subclasses
+    ``TokenCounter`` and passes an instance to
+    ``IntelligentSummarizationMiddleware(token_counter=...)`` — no other
+    ContextSage code needs to change.
+    """
+
+    #: Whether this counter reports an authoritative, provider-verified count
+    #: (``True``) or an estimate (``False``).
+    is_authoritative: bool
+
+    @abstractmethod
+    def count_text(self, text: str) -> int:
+        """Return the (approximate or authoritative) token count for raw text."""
+
+    @abstractmethod
+    def count_messages(self, messages: Iterable[BaseMessage]) -> int:
+        """Return the total token count for a sequence of messages."""
+
+
+def _resolve_encoding(
+    model: str,
+    *,
+    fallback_encoding_name: str = _DEFAULT_ENCODING_NAME,
+    download_timeout_seconds: float = _ENCODING_DOWNLOAD_TIMEOUT_SECONDS,
+) -> tiktoken.Encoding:
+    """Resolve an explicit ``model`` name to a tiktoken encoding, bounded and fail-safe.
+
+    Tries the model-specific encoding first (bounded by a short socket
+    timeout so an uncached encoding never hangs indefinitely on a
+    network-restricted host); falls back to ``fallback_encoding_name`` if
+    that fails for any reason.
+    """
+    import tiktoken  # noqa: PLC0415
+
+    previous_timeout = socket.getdefaulttimeout()
+    socket.setdefaulttimeout(download_timeout_seconds)
+    try:
+        try:
+            return tiktoken.encoding_for_model(model)
+        except Exception:
+            return tiktoken.get_encoding(fallback_encoding_name)
+    finally:
+        socket.setdefaulttimeout(previous_timeout)
+
+
+class TiktokenTokenCounter(TokenCounter):
+    """Authoritative-ish counting using ``tiktoken``.
+
+    ``tiktoken`` is a direct, always-installed dependency, and this is the
+    only built-in :class:`TokenCounter` implementation: it is the closest
+    thing ContextSage has to a provider-verified count without requiring a
+    live API round-trip, for both OpenAI-compatible and most open-weight
+    tokenizer families (via ``cl100k_base``/``o200k_base``-family BPE).
+
+    With no arguments, uses ``cl100k_base`` directly — no model-name lookup,
+    so the zero-configuration default has no network dependency. Pass
+    ``model`` to resolve the encoding for a specific model actually in use
+    (e.g. ``"gpt-4o"``, ``"gpt-4.1"``); pass ``encoding_name`` to select an
+    encoding directly. ``fallback_encoding_name``/``download_timeout_seconds``
+    tune the ``model``-based resolution path itself (see :func:`_resolve_encoding`).
+    Remains fully pluggable for non-GPT-family providers: pass any other
+    :class:`TokenCounter` subclass via
+    ``IntelligentSummarizationMiddleware(token_counter=...)``.
+    """
+
+    is_authoritative = True
+
+    def __init__(
+        self,
+        model: str | None = None,
+        encoding_name: str | None = None,
+        *,
+        fallback_encoding_name: str = _DEFAULT_ENCODING_NAME,
+        download_timeout_seconds: float = _ENCODING_DOWNLOAD_TIMEOUT_SECONDS,
+    ) -> None:
+        import tiktoken  # noqa: PLC0415
+
+        self.model = model
+        try:
+            if encoding_name is not None:
+                self._encoding = tiktoken.get_encoding(encoding_name)
+            elif model is not None:
+                self._encoding = _resolve_encoding(
+                    model,
+                    fallback_encoding_name=fallback_encoding_name,
+                    download_timeout_seconds=download_timeout_seconds,
+                )
+            else:
+                self._encoding = tiktoken.get_encoding(fallback_encoding_name)
+        except Exception as exc:
+            raise TokenCountingError(f"tiktoken failed to load an encoding: {exc}") from exc
+
+    def count_text(self, text: str) -> int:
+        if not text:
+            return 0
+        return len(self._encoding.encode(text, disallowed_special=()))
+
+    def count_messages(self, messages: Iterable[BaseMessage]) -> int:
+        return sum(self.count_text(_message_text(m)) for m in messages)
+
+
+class CachingTokenCounter(TokenCounter):
+    """Wraps another counter with a bounded LRU cache over raw text."""
+
+    def __init__(self, inner: TokenCounter, maxsize: int = _TEXT_CACHE_SIZE) -> None:
+        self._inner = inner
+        self.is_authoritative = inner.is_authoritative
+        self._cached_count_text = lru_cache(maxsize=maxsize)(self._inner.count_text)
+
+    def count_text(self, text: str) -> int:
+        return self._cached_count_text(text)
+
+    def count_messages(self, messages: Iterable[BaseMessage]) -> int:
+        return self._inner.count_messages(messages)
+
+
+def _message_text(message: BaseMessage) -> str:
+    content = getattr(message, "content", "")
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts = []
+        for block in content:
+            if isinstance(block, str):
+                parts.append(block)
+            elif isinstance(block, dict):
+                parts.append(str(block.get("text", block)))
+        return "\n".join(parts)
+    return str(content)
+
+
+def _extract_model_name(model: BaseChatModel | None) -> str | None:
+    """Best-effort extraction of a model name string from a chat model instance."""
+    if model is None:
+        return None
+    for attr in ("model_name", "model"):
+        value = getattr(model, attr, None)
+        if isinstance(value, str) and value:
+            return value
+    profile = getattr(model, "profile", None)
+    if isinstance(profile, dict):
+        value = profile.get("name")
+        if isinstance(value, str) and value:
+            return value
+    return None
+
+
+def create_default_token_counter(model: BaseChatModel | None = None) -> TokenCounter:
+    """Build the default token counter used when the caller does not supply one.
+
+    Resolves the tokenizer to whatever model is actually configured on the
+    middleware (e.g. ``"gpt-4.1"``, ``"gpt-4o-mini"``); falls back to the
+    universal ``cl100k_base`` encoding directly (no network-dependent
+    model-name lookup) when no model name can be determined. The counter
+    remains fully pluggable: pass any other :class:`TokenCounter` subclass
+    via ``IntelligentSummarizationMiddleware(token_counter=...)``.
+    """
+    model_name = _extract_model_name(model)
+    return CachingTokenCounter(TiktokenTokenCounter(model=model_name))

contextsage/classification/__init__.py

@@ -0,0 +1,3 @@
+"""Structural signal detection for heterogeneous context regions."""
+
+from __future__ import annotations

contextsage/classification/signals.py

@@ -0,0 +1,131 @@
+"""Structural signal detection.
+
+Detection here is intentionally lightweight and conservative: this module
+answers "what does this chunk of text look like?" with a confidence score,
+not "parse this into a full AST". Low-confidence detections are treated as
+plain text by callers — false structural classification is worse than no
+classification.
+
+Kind detection itself is delegated to a pluggable
+:class:`~contextsage.parsers.base.ParserRegistry` (see :mod:`contextsage.parsers`):
+:class:`StructuralSignalDetector` owns the *always-on* enrichment (identifier
+extraction, error-keyword/severity tagging) that applies regardless of which
+kind was detected, and never needs a mature parsing library itself — only
+the Python standard library is used here, to keep the dependency footprint
+minimal.
+
+Classification is logically distinct from parsing (a parser decides *what
+kind* a region is; the detector's own enrichment applies uniformly no
+matter which kind was decided) and is therefore overridable independently:
+all three patterns below are overridable at construction time (no
+subclassing required), the same way parsers expose their own default
+patterns — pass ``severity_pattern``/``error_keywords_pattern``/
+``identifier_patterns`` to recognize an application's own severity
+vocabulary or identifier conventions via the public API.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Sequence
+from dataclasses import replace
+
+from contextsage.core.models import ContentSignals, RegionKind
+from contextsage.parsers.base import ParserRegistry, default_registry
+
+#: The built-in defaults, exposed so callers can override just one of them
+#: (or compose with them) without having to reconstruct the others.
+DEFAULT_SEVERITY_PATTERN = re.compile(r"\b(DEBUG|INFO|WARN(?:ING)?|ERROR|CRITICAL|FATAL|TRACE)\b")
+DEFAULT_ERROR_KEYWORDS_PATTERN = re.compile(
+    r"\b(traceback|exception|error|failed|failure|stack trace|panic)\b", re.IGNORECASE
+)
+DEFAULT_IDENTIFIER_PATTERNS: tuple[re.Pattern[str], ...] = (
+    re.compile(r"\b[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\b"),
+    re.compile(r"\b[A-Z]{2,6}-\d{3,}\b"),
+    re.compile(r"\b(?:req|request|trace|correlation|txn|transaction)[_-]?id\s*[:=]\s*[\w-]+", re.I),
+    re.compile(r"\bcustomer\s+\d+\b", re.I),
+)
+
+
+class StructuralSignalDetector:
+    """Detects the structural nature of a chunk of text.
+
+    ``registry`` is an optional, advanced extension point: pass a custom
+    :class:`~contextsage.parsers.base.ParserRegistry` (built from your own
+    :class:`~contextsage.parsers.base.ContentParser` subclasses, alongside or
+    instead of the built-in ones) to plug in domain-specific detection
+    without forking ContextSage. The default registry (built-in JSON/table/
+    code/log/error parsers) is used when ``registry`` is omitted, so ordinary
+    usage of :class:`contextsage.IntelligentSummarizationMiddleware` never
+    needs to know this exists.
+
+    ``severity_pattern``, ``error_keywords_pattern``, and
+    ``identifier_patterns`` override the always-on enrichment applied on top
+    of whatever the registry detected; they default to
+    :data:`DEFAULT_SEVERITY_PATTERN`, :data:`DEFAULT_ERROR_KEYWORDS_PATTERN`,
+    and :data:`DEFAULT_IDENTIFIER_PATTERNS` respectively.
+    """
+
+    def __init__(
+        self,
+        registry: ParserRegistry | None = None,
+        *,
+        severity_pattern: re.Pattern[str] = DEFAULT_SEVERITY_PATTERN,
+        error_keywords_pattern: re.Pattern[str] = DEFAULT_ERROR_KEYWORDS_PATTERN,
+        identifier_patterns: Sequence[re.Pattern[str]] = DEFAULT_IDENTIFIER_PATTERNS,
+    ) -> None:
+        self._registry = registry or default_registry()
+        self._severity_pattern = severity_pattern
+        self._error_keywords_pattern = error_keywords_pattern
+        self._identifier_patterns = tuple(identifier_patterns)
+
+    @property
+    def registry(self) -> ParserRegistry:
+        """The configured parser registry.
+
+        Exposed so :class:`~contextsage.context.decomposition.ContextDecomposer`
+        can derive its own segmentation-time patterns from whichever
+        parsers are actually registered (default or overridden via
+        ``IntelligentSummarizationMiddleware(parsers=[...])``), rather than
+        maintaining independent, possibly-stale copies.
+        """
+        return self._registry
+
+    def detect(self, text: str) -> ContentSignals:
+        stripped = text.strip()
+        if not stripped:
+            return ContentSignals(kind=RegionKind.TEXT, confidence=1.0)
+
+        signals = self._registry.detect(stripped) or ContentSignals(
+            kind=RegionKind.TEXT, confidence=0.5
+        )
+        identifiers = self._extract_identifiers(text)
+        has_error = bool(self._error_keywords_pattern.search(text))
+        severity_match = self._severity_pattern.search(text)
+        return replace(
+            signals,
+            identifiers=identifiers,
+            has_error=signals.has_error or has_error,
+            severity=signals.severity
+            or (severity_match.group(0).upper() if severity_match else None),
+        )
+
+    def _extract_identifiers(self, text: str) -> tuple[str, ...]:
+        found: list[str] = []
+        for pattern in self._identifier_patterns:
+            found.extend(m.group(0) for m in pattern.finditer(text))
+        # De-duplicate while preserving order.
+        seen: set[str] = set()
+        unique = []
+        for item in found:
+            if item not in seen:
+                seen.add(item)
+                unique.append(item)
+        return tuple(unique)
+
+
+def normalize_for_repetition(line: str) -> str:
+    """Collapse volatile substrings (numbers, hex ids, timestamps) for dedup grouping."""
+    normalized = re.sub(r"\b[0-9a-fA-F]{8,}\b", "<ID>", line)
+    normalized = re.sub(r"\d+", "<N>", normalized)
+    return normalized.strip()

contextsage/context/__init__.py

@@ -0,0 +1,3 @@
+"""Context observation and heterogeneous decomposition into ContextUnits."""
+
+from __future__ import annotations