powerailabs-contextkit 0.1.0__tar.gz

powerailabs_contextkit-0.1.0/.gitignore

@@ -0,0 +1,16 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+build/
+dist/
+.venv/
+.uv/
+.ruff_cache/
+.pytest_cache/
+.mypy_cache/
+.coverage
+htmlcov/
+.idea/
+.vscode/
+.DS_Store
+*.log

powerailabs_contextkit-0.1.0/PKG-INFO

@@ -0,0 +1,48 @@
+Metadata-Version: 2.4
+Name: powerailabs-contextkit
+Version: 0.1.0
+Summary: Assemble: declare prioritized, pinnable context blocks; pack them to a token budget with an inspectable receipt.
+Author: Raghav Mishra
+License-Expression: MIT
+Requires-Python: >=3.11
+Requires-Dist: powerailabs-core<0.2,>=0.1
+Provides-Extra: squeeze
+Requires-Dist: powerailabs-squeeze<0.2,>=0.1; extra == 'squeeze'
+Description-Content-Type: text/markdown
+
+# powerailabs-contextkit
+
+Treat the context window like a packed suitcase, not a string you concatenate. Declare blocks
+with priorities and eviction rules; contextkit fits them to a token budget and tells you exactly
+what it kept, shrank, and dropped.
+
+**Every assembled prompt comes with a receipt.**
+
+![status](https://img.shields.io/badge/status-building-yellow) ![license](https://img.shields.io/badge/license-MIT-blue)
+
+🚧 building (v0) · `pip install powerailabs-contextkit` · `from powerailabs.contextkit import Context, Block`
+
+```python
+from powerailabs.contextkit import Context, Block
+
+ctx = Context(budget_tokens=8000, model="claude-opus-4-8", reserve_output=1000)
+ctx.add(Block(system_prompt, priority=10, pin=True, role="system"))
+ctx.add(Block(retrieved_docs, priority=5, evict="compress"))   # uses squeeze if installed
+ctx.add(Block(chat_history, priority=3, evict="drop_oldest"))
+ctx.add(Block(user_msg, priority=9, pin=True, role="user"))
+
+messages = ctx.assemble()      # provider-ready messages, guaranteed within budget
+print(ctx.report())            # the receipt: kept / truncated / dropped + token math
+#   [kept      ] system    42->42tok
+#   [compressed] user      3120->980tok
+#   [dropped   ] user      610->0tok   # chat_history didn't fit
+preview = ctx.whatif(budget_tokens=4000)   # same inputs, tighter budget, no commit
+```
+
+**Inbound:** call contextkit *before* the model call to build the messages you send — it applies
+whenever you assemble the prompt yourself. Assembly is deterministic (stable sort by pinned →
+priority → insertion). `evict="compress"` wires in `powerailabs-contextkit[squeeze]` by shape (no
+import); without it, `compress` falls back to `truncate`. The `report()` decisions flow onto
+core's event stream, so `acttrace` records what context the model actually saw.
+
+See [`docs/contextkit.md`](../../docs/contextkit.md). *Part of the PowerAI Labs stack — github.com/PowerAI-Labs/powerailabs.*

powerailabs_contextkit-0.1.0/README.md

@@ -0,0 +1,36 @@
+# powerailabs-contextkit
+
+Treat the context window like a packed suitcase, not a string you concatenate. Declare blocks
+with priorities and eviction rules; contextkit fits them to a token budget and tells you exactly
+what it kept, shrank, and dropped.
+
+**Every assembled prompt comes with a receipt.**
+
+![status](https://img.shields.io/badge/status-building-yellow) ![license](https://img.shields.io/badge/license-MIT-blue)
+
+🚧 building (v0) · `pip install powerailabs-contextkit` · `from powerailabs.contextkit import Context, Block`
+
+```python
+from powerailabs.contextkit import Context, Block
+
+ctx = Context(budget_tokens=8000, model="claude-opus-4-8", reserve_output=1000)
+ctx.add(Block(system_prompt, priority=10, pin=True, role="system"))
+ctx.add(Block(retrieved_docs, priority=5, evict="compress"))   # uses squeeze if installed
+ctx.add(Block(chat_history, priority=3, evict="drop_oldest"))
+ctx.add(Block(user_msg, priority=9, pin=True, role="user"))
+
+messages = ctx.assemble()      # provider-ready messages, guaranteed within budget
+print(ctx.report())            # the receipt: kept / truncated / dropped + token math
+#   [kept      ] system    42->42tok
+#   [compressed] user      3120->980tok
+#   [dropped   ] user      610->0tok   # chat_history didn't fit
+preview = ctx.whatif(budget_tokens=4000)   # same inputs, tighter budget, no commit
+```
+
+**Inbound:** call contextkit *before* the model call to build the messages you send — it applies
+whenever you assemble the prompt yourself. Assembly is deterministic (stable sort by pinned →
+priority → insertion). `evict="compress"` wires in `powerailabs-contextkit[squeeze]` by shape (no
+import); without it, `compress` falls back to `truncate`. The `report()` decisions flow onto
+core's event stream, so `acttrace` records what context the model actually saw.
+
+See [`docs/contextkit.md`](../../docs/contextkit.md). *Part of the PowerAI Labs stack — github.com/PowerAI-Labs/powerailabs.*

powerailabs_contextkit-0.1.0/pyproject.toml

@@ -0,0 +1,20 @@
+[project]
+name = "powerailabs-contextkit"
+version = "0.1.0"
+description = "Assemble: declare prioritized, pinnable context blocks; pack them to a token budget with an inspectable receipt."
+requires-python = ">=3.11"
+license = "MIT"
+authors = [{ name = "Raghav Mishra" }]
+readme = "README.md"
+dependencies = ["powerailabs-core>=0.1,<0.2"]
+
+[project.optional-dependencies]
+# Wires squeeze in as the evict="compress" strategy (via core's Compressor protocol).
+squeeze = ["powerailabs-squeeze>=0.1,<0.2"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/powerailabs"]     # contributes powerailabs/contextkit only — NEVER add src/powerailabs/__init__.py

powerailabs_contextkit-0.1.0/src/powerailabs/contextkit/__init__.py

@@ -0,0 +1,266 @@
+"""powerailabs.contextkit — assemble context within a token budget, with a receipt.
+
+Treat the context window like a packed suitcase: declare ``Block``s with priority, pin, and a
+per-block eviction rule; :meth:`Context.assemble` packs them to a token budget (deterministically)
+and :meth:`Context.report` returns the receipt — what was kept, shrunk, or dropped, with the token
+math. Depends only on ``powerailabs-core`` (``tokens`` + the ``Compressor`` protocol). Tools never
+import each other; ``squeeze`` plugs in by shape via the ``contextkit[squeeze]`` extra.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from powerailabs.core import bus, tokens
+
+__all__ = ["Block", "Context", "AssemblyReport", "BlockDecision", "BudgetError"]
+
+EvictStrategy = str  # "drop_oldest" | "truncate" | "summarize" | "compress"
+
+
+class BudgetError(Exception):
+    """Raised when pinned blocks alone exceed the budget (they are never evicted)."""
+
+
+@dataclass
+class Block:
+    """A unit of context with packing intent. See docs/contextkit.md §6.
+
+    Attributes:
+        content: The block text (multimodal lists are a roadmap item; start text-only).
+        priority: Higher is admitted first; ties break by insertion order (deterministic).
+        pin: Pinned blocks are never evicted (assembly raises if pinned blocks alone overflow).
+        evict: Strategy when this block overflows the remaining budget.
+        role: Provider message role: ``system`` | ``user`` | ``assistant`` | ``tool``.
+        summarizer: Callback ``(content, target_tokens) -> str`` used when ``evict="summarize"``.
+    """
+
+    content: str | list
+    priority: int = 0
+    pin: bool = False
+    evict: EvictStrategy = "drop_oldest"
+    role: str = "user"
+    summarizer: Callable[[str, int], str] | None = None
+
+
+@dataclass
+class BlockDecision:
+    """What happened to one block during assembly (a line on the receipt)."""
+
+    role: str
+    action: str  # "kept" | "truncated" | "summarized" | "compressed" | "dropped"
+    tokens_before: int
+    tokens_after: int
+    note: str = ""
+
+
+@dataclass
+class AssemblyReport:
+    """The receipt: budget math + per-block decisions. See docs/contextkit.md §6."""
+
+    budget: int
+    used: int
+    reserved_output: int
+    model: str
+    decisions: list[BlockDecision] = field(default_factory=list)
+
+    def __str__(self) -> str:
+        lines = [
+            f"AssemblyReport(model={self.model}) "
+            f"budget={self.budget} reserved_output={self.reserved_output} "
+            f"used={self.used}/{self.budget - self.reserved_output}",
+        ]
+        for d in self.decisions:
+            arrow = f"{d.tokens_before}->{d.tokens_after}tok"
+            note = f"  # {d.note}" if d.note else ""
+            lines.append(f"  [{d.action:<10}] {d.role:<9} {arrow}{note}")
+        return "\n".join(lines)
+
+
+# Render order: system first, conversational/context middle, the user turn last.
+_ROLE_RANK = {"system": 0, "tool": 1, "assistant": 2, "user": 3}
+
+
+class Context:
+    """A token-budgeted, declarative context assembler. See docs/contextkit.md §3, §5."""
+
+    def __init__(
+        self,
+        budget_tokens: int,
+        model: str,
+        reserve_output: int = 0,
+        compressor: Any = None,
+    ) -> None:
+        self.budget_tokens = budget_tokens
+        self.model = model
+        self.reserve_output = reserve_output
+        self._compressor = compressor
+        self._blocks: list[Block] = []
+        self._report: AssemblyReport | None = None
+        self._messages: list[dict] = []
+
+    def add(self, block: Block) -> Context:
+        """Add a block. Returns ``self`` for chaining."""
+        self._blocks.append(block)
+        return self
+
+    def assemble(self) -> list[dict]:
+        """Pack blocks within the budget; return provider-ready messages (OpenAI/Foundry shape).
+
+        Deterministic: stable sort by ``(pinned, priority, insertion order)``. Emits the
+        :class:`AssemblyReport` onto core's bus so ``acttrace`` records what the model saw.
+        """
+        messages, report = self._pack(self.budget_tokens, emit=True)
+        self._messages = messages
+        self._report = report
+        return messages
+
+    def report(self) -> AssemblyReport:
+        """Return the receipt for the most recent :meth:`assemble`. TODO until first assemble."""
+        if self._report is None:
+            raise RuntimeError("call assemble() before report()")
+        return self._report
+
+    def whatif(self, budget_tokens: int) -> AssemblyReport:
+        """Preview the assembly at a different budget without committing (no bus emit)."""
+        _, report = self._pack(budget_tokens, emit=False)
+        return report
+
+    def for_anthropic(self) -> tuple[str, list[dict]]:
+        """Anthropic adapter: split system blocks out (the Messages API takes ``system`` apart).
+
+        Returns ``(system_text, messages)`` from the most recent :meth:`assemble`.
+        """
+        if not self._messages:
+            self.assemble()
+        system = "\n\n".join(m["content"] for m in self._messages if m["role"] == "system")
+        rest = [m for m in self._messages if m["role"] != "system"]
+        return system, rest
+
+    # ------------------------------------------------------------------ internals
+
+    def _pack(self, budget_tokens: int, *, emit: bool) -> tuple[list[dict], AssemblyReport]:
+        effective = max(0, budget_tokens - self.reserve_output)
+        # (not pin) -> pinned (False) sorts first; then priority desc; then insertion order.
+        order = sorted(
+            enumerate(self._blocks),
+            key=lambda iv: (not iv[1].pin, -iv[1].priority, iv[0]),
+        )
+
+        used = 0
+        decisions: list[BlockDecision] = []
+        kept: list[tuple[int, str, str]] = []  # (insertion_index, role, rendered_content)
+
+        for idx, block in order:
+            text = block.content if isinstance(block.content, str) else str(block.content)
+            before = tokens.count(text, self.model)
+            remaining = effective - used
+
+            if before <= remaining:
+                used += before
+                kept.append((idx, block.role, text))
+                decisions.append(BlockDecision(block.role, "kept", before, before))
+                continue
+
+            if block.pin:
+                raise BudgetError(
+                    f"pinned block(s) exceed budget: need >{before} tokens, "
+                    f"{remaining} of {effective} remaining (reserve_output={self.reserve_output})"
+                )
+
+            new_text, action, note = self._evict(block, text, remaining)
+            if new_text is None:
+                decisions.append(BlockDecision(block.role, "dropped", before, 0, note))
+                continue
+            after = tokens.count(new_text, self.model)
+            used += after
+            kept.append((idx, block.role, new_text))
+            decisions.append(BlockDecision(block.role, action, before, after, note))
+
+        kept.sort(key=lambda k: (_ROLE_RANK.get(k[1], 1), k[0]))
+        messages = [{"role": role, "content": content} for _, role, content in kept]
+        report = AssemblyReport(
+            budget=budget_tokens,
+            used=used,
+            reserved_output=self.reserve_output,
+            model=self.model,
+            decisions=decisions,
+        )
+        if emit:
+            bus.emit(report)
+        return messages, report
+
+    def _evict(self, block: Block, text: str, remaining: int) -> tuple[str | None, str, str]:
+        """Apply a block's eviction strategy. Returns ``(content_or_None, action, note)``."""
+        strategy = block.evict
+
+        if strategy == "drop_oldest" or remaining <= 0:
+            note = "no room" if remaining <= 0 and strategy != "drop_oldest" else ""
+            return None, "dropped", note
+
+        if strategy == "truncate":
+            return _truncate_to_tokens(text, remaining, self.model), "truncated", ""
+
+        if strategy == "summarize":
+            if block.summarizer is not None:
+                summary = block.summarizer(text, remaining)
+                if tokens.count(summary, self.model) > remaining:
+                    summary = _truncate_to_tokens(summary, remaining, self.model)
+                return summary, "summarized", ""
+            return (
+                _truncate_to_tokens(text, remaining, self.model),
+                "truncated",
+                ("no summarizer; fell back to truncate"),
+            )
+
+        if strategy == "compress":
+            compressor = self._get_compressor()
+            if compressor is not None:
+                small = _call_compressor(compressor, text, remaining, self.model)
+                if tokens.count(small, self.model) > remaining:
+                    small = _truncate_to_tokens(small, remaining, self.model)
+                return small, "compressed", ""
+            return (
+                _truncate_to_tokens(text, remaining, self.model),
+                "truncated",
+                ("squeeze not installed; fell back to truncate"),
+            )
+
+        # Unknown strategy: be safe, drop.
+        return None, "dropped", f"unknown evict strategy {strategy!r}"
+
+    def _get_compressor(self) -> Any:
+        if self._compressor is not None:
+            return self._compressor
+        # Optional plugin, discovered at runtime via the contextkit[squeeze] extra.
+        import importlib
+
+        try:
+            return importlib.import_module("powerailabs.squeeze").compress
+        except ModuleNotFoundError:
+            return None
+
+
+def _call_compressor(compressor: Any, text: str, target: int, model: str) -> str:
+    """Call either a Compressor-protocol object or a ``squeeze.compress``-style callable."""
+    if hasattr(compressor, "compress"):
+        small, _handle = compressor.compress(text, target_tokens=target, model=model)
+    else:
+        small, _handle = compressor(text, target_tokens=target)
+    return small
+
+
+def _truncate_to_tokens(text: str, target: int, model: str) -> str:
+    """Trim ``text`` to at most ``target`` tokens (deterministic char-ratio shrink)."""
+    if target <= 0:
+        return ""
+    current = tokens.count(text, model)
+    if current <= target:
+        return text
+    ratio = max(1, len(text)) / max(1, current)
+    cut = text[: int(target * ratio)]
+    while cut and tokens.count(cut, model) > target:
+        cut = cut[: int(len(cut) * 0.9)]
+    return cut

powerailabs_contextkit-0.1.0/tests/test_contextkit.py

@@ -0,0 +1,113 @@
+"""Budgeted assembly + the receipt. Deterministic, offline (heuristic token counts)."""
+
+import pytest
+from powerailabs.contextkit import Block, BudgetError, Context
+from powerailabs.core import bus, tokens
+
+
+@pytest.fixture(autouse=True)
+def _heuristic_tokens(monkeypatch):
+    # Force the offline heuristic so token math is deterministic regardless of tiktoken.
+    monkeypatch.setattr(tokens, "_tiktoken_encoding", lambda model: None)
+    yield
+
+
+def test_public_api_present():
+    import powerailabs.contextkit as ck
+
+    for name in ("Block", "Context", "AssemblyReport", "BlockDecision", "BudgetError"):
+        assert hasattr(ck, name)
+
+
+def test_block_defaults():
+    b = Block("hi", priority=5, pin=True, role="system")
+    assert b.evict == "drop_oldest" and b.role == "system" and b.pin
+
+
+def test_assemble_keeps_everything_under_budget():
+    ctx = Context(budget_tokens=1000, model="gpt-4o")
+    ctx.add(Block("system prompt", priority=10, pin=True, role="system"))
+    ctx.add(Block("the question", priority=9, pin=True, role="user"))
+    messages = ctx.assemble()
+    assert [m["role"] for m in messages] == ["system", "user"]  # system first, user last
+    assert all(d.action == "kept" for d in ctx.report().decisions)
+
+
+def test_drop_oldest_evicts_low_priority_when_tight():
+    ctx = Context(budget_tokens=8, model="gpt-4o")  # ~8 tokens of room
+    ctx.add(Block("x" * 4, priority=10, role="system"))  # ~1 tok, kept
+    ctx.add(Block("y" * 200, priority=1, role="user", evict="drop_oldest"))  # too big -> dropped
+    messages = ctx.assemble()
+    roles = [m["role"] for m in messages]
+    assert "user" not in roles  # low-priority block was dropped
+    dropped = [d for d in ctx.report().decisions if d.action == "dropped"]
+    assert len(dropped) == 1 and dropped[0].role == "user"
+
+
+def test_truncate_shrinks_to_fit():
+    ctx = Context(budget_tokens=12, model="gpt-4o")
+    ctx.add(Block("s", priority=10, role="system"))
+    ctx.add(Block("z" * 400, priority=1, role="user", evict="truncate"))
+    ctx.assemble()
+    decision = next(d for d in ctx.report().decisions if d.role == "user")
+    assert decision.action == "truncated"
+    assert decision.tokens_after < decision.tokens_before
+    assert ctx.report().used <= ctx.report().budget - ctx.report().reserved_output
+
+
+def test_pinned_overflow_raises():
+    ctx = Context(budget_tokens=5, model="gpt-4o")
+    ctx.add(Block("w" * 400, priority=10, pin=True, role="system"))
+    with pytest.raises(BudgetError):
+        ctx.assemble()
+
+
+def test_reserve_output_reduces_usable_budget():
+    ctx = Context(budget_tokens=100, model="gpt-4o", reserve_output=96)
+    ctx.add(Block("a" * 200, priority=1, role="user", evict="truncate"))
+    ctx.assemble()
+    # only ~4 tokens usable -> truncated small
+    assert ctx.report().used <= 4
+
+
+def test_whatif_does_not_commit():
+    ctx = Context(budget_tokens=1000, model="gpt-4o")
+    ctx.add(Block("hello there", priority=5, role="user"))
+    ctx.assemble()
+    committed_used = ctx.report().used
+    preview = ctx.whatif(budget_tokens=3)
+    assert preview.budget == 3
+    assert ctx.report().used == committed_used  # committed report unchanged
+
+
+def test_assembly_is_deterministic():
+    def build():
+        c = Context(budget_tokens=50, model="gpt-4o")
+        c.add(Block("alpha", priority=5, role="user"))
+        c.add(Block("beta", priority=5, role="assistant"))
+        c.add(Block("gamma", priority=9, role="system"))
+        return c.assemble()
+
+    assert build() == build()
+
+
+def test_report_emitted_on_bus():
+    bus._reset()
+    seen = []
+    bus.subscribe(seen.append)
+    try:
+        ctx = Context(budget_tokens=100, model="gpt-4o")
+        ctx.add(Block("hi", role="user"))
+        ctx.assemble()
+    finally:
+        bus._reset()
+    assert len(seen) == 1 and seen[0].model == "gpt-4o"
+
+
+def test_for_anthropic_splits_system():
+    ctx = Context(budget_tokens=1000, model="claude-opus-4-8")
+    ctx.add(Block("you are helpful", priority=10, pin=True, role="system"))
+    ctx.add(Block("hello", priority=9, pin=True, role="user"))
+    system, messages = ctx.for_anthropic()
+    assert system == "you are helpful"
+    assert all(m["role"] != "system" for m in messages)