powerailabs-contextkit 0.1.0__tar.gz → 0.2.0__tar.gz

{powerailabs_contextkit-0.1.0 → powerailabs_contextkit-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: powerailabs-contextkit
-Version: 0.1.0
+Version: 0.2.0
 Summary: Assemble: declare prioritized, pinnable context blocks; pack them to a token budget with an inspectable receipt.
 Author: Raghav Mishra
 License-Expression: MIT

{powerailabs_contextkit-0.1.0 → powerailabs_contextkit-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "powerailabs-contextkit"
-version = "0.1.0"
+version = "0.2.0"
 description = "Assemble: declare prioritized, pinnable context blocks; pack them to a token budget with an inspectable receipt."
 requires-python = ">=3.11"
 license = "MIT"

{powerailabs_contextkit-0.1.0 → powerailabs_contextkit-0.2.0}/src/powerailabs/contextkit/__init__.py

@@ -65,10 +65,11 @@ class AssemblyReport:
     reserved_output: int
     model: str
     decisions: list[BlockDecision] = field(default_factory=list)
+    order: str = "default"
 
     def __str__(self) -> str:
         lines = [
-            f"AssemblyReport(model={self.model}) "
+            f"AssemblyReport(model={self.model}, order={self.order}) "
             f"budget={self.budget} reserved_output={self.reserved_output} "
             f"used={self.used}/{self.budget - self.reserved_output}",
         ]
@@ -79,12 +80,22 @@ class AssemblyReport:
         return "\n".join(lines)
 
 
-# Render order: system first, conversational/context middle, the user turn last.
+# Default render order: system first, conversational/context middle, the user turn last.
 _ROLE_RANK = {"system": 0, "tool": 1, "assistant": 2, "user": 3}
+_ORDERS = ("default", "attention", "cache")
 
 
 class Context:
-    """A token-budgeted, declarative context assembler. See docs/contextkit.md §3, §5."""
+    """A token-budgeted, declarative context assembler. See docs/contextkit.md §3, §5.
+
+    ``order`` controls how kept blocks are arranged in the final messages (docs/contextkit.md §2):
+
+    - ``"default"`` — role-grouped: system → context/history → the user turn.
+    - ``"attention"`` — "lost-in-the-middle": highest-priority context blocks ride the edges
+      (just after system / just before the user turn), weakest in the dead center.
+    - ``"cache"`` — stable prefix first (pinned, high-priority blocks lead) to maximize provider
+      prompt-cache / KV-cache hits across calls.
+    """
 
     def __init__(
         self,
@@ -92,11 +103,15 @@ class Context:
         model: str,
         reserve_output: int = 0,
         compressor: Any = None,
+        order: str = "default",
     ) -> None:
+        if order not in _ORDERS:
+            raise ValueError(f"order must be one of {_ORDERS}, got {order!r}")
         self.budget_tokens = budget_tokens
         self.model = model
         self.reserve_output = reserve_output
         self._compressor = compressor
+        self.order = order
         self._blocks: list[Block] = []
         self._report: AssemblyReport | None = None
         self._messages: list[dict] = []
@@ -151,7 +166,7 @@ class Context:
 
         used = 0
         decisions: list[BlockDecision] = []
-        kept: list[tuple[int, str, str]] = []  # (insertion_index, role, rendered_content)
+        kept: list[tuple[int, Block, str]] = []  # (insertion_index, block, rendered_content)
 
         for idx, block in order:
             text = block.content if isinstance(block.content, str) else str(block.content)
@@ -160,7 +175,7 @@ class Context:
 
             if before <= remaining:
                 used += before
-                kept.append((idx, block.role, text))
+                kept.append((idx, block, text))
                 decisions.append(BlockDecision(block.role, "kept", before, before))
                 continue
 
@@ -176,17 +191,18 @@ class Context:
                 continue
             after = tokens.count(new_text, self.model)
             used += after
-            kept.append((idx, block.role, new_text))
+            kept.append((idx, block, 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]
+        ordered = _order_blocks(kept, self.order)
+        messages = [{"role": block.role, "content": content} for _, block, content in ordered]
         report = AssemblyReport(
             budget=budget_tokens,
             used=used,
             reserved_output=self.reserve_output,
             model=self.model,
             decisions=decisions,
+            order=self.order,
         )
         if emit:
             bus.emit(report)
@@ -243,6 +259,36 @@ class Context:
             return None
 
 
+def _order_blocks(kept: list[tuple[int, Block, str]], mode: str) -> list[tuple[int, Block, str]]:
+    """Arrange kept blocks for rendering per the chosen strategy. Deterministic."""
+    if mode == "cache":
+        # Stable prefix: pinned, high-priority blocks lead so the prompt prefix is reused.
+        return sorted(kept, key=lambda k: (not k[1].pin, -k[1].priority, k[0]))
+    if mode == "attention":
+        systems = sorted(
+            (k for k in kept if k[1].role == "system"), key=lambda k: (-k[1].priority, k[0])
+        )
+        finals = sorted(
+            (k for k in kept if k[1].role == "user"), key=lambda k: (k[1].priority, k[0])
+        )  # ascending -> highest-priority user turn ends up last (strongest end position)
+        middles = sorted(
+            (k for k in kept if k[1].role not in ("system", "user")),
+            key=lambda k: (-k[1].priority, k[0]),
+        )
+        return [*systems, *_edge_load(middles), *finals]
+    # default: role-grouped, insertion order within a role.
+    return sorted(kept, key=lambda k: (_ROLE_RANK.get(k[1].role, 1), k[0]))
+
+
+def _edge_load(items: list) -> list:
+    """Edge-load a priority-descending list: highest at both edges, lowest in the center."""
+    left: list = []
+    right: list = []
+    for i, item in enumerate(items):
+        (left if i % 2 == 0 else right).append(item)
+    return left + right[::-1]
+
+
 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"):

{powerailabs_contextkit-0.1.0 → powerailabs_contextkit-0.2.0}/tests/test_contextkit.py

@@ -104,6 +104,48 @@ def test_report_emitted_on_bus():
     assert len(seen) == 1 and seen[0].model == "gpt-4o"
 
 
+def test_attention_order_edge_loads_priority():
+    # Lost-in-the-middle: highest-priority context blocks ride the edges, weakest in the center.
+    ctx = Context(budget_tokens=1000, model="gpt-4o", order="attention")
+    ctx.add(Block("SYS", priority=100, pin=True, role="system"))
+    ctx.add(Block("p9", priority=9, role="assistant"))
+    ctx.add(Block("p1", priority=1, role="assistant"))
+    ctx.add(Block("p5", priority=5, role="assistant"))
+    ctx.add(Block("p7", priority=7, role="assistant"))
+    ctx.add(Block("USER", priority=10, pin=True, role="user"))
+    msgs = ctx.assemble()
+    assert msgs[0]["content"] == "SYS"  # system anchored first
+    assert msgs[-1]["content"] == "USER"  # user turn anchored last
+    middle = [m["content"] for m in msgs[1:-1]]
+    # desc by priority = [p9,p7,p5,p1] -> edge-loaded -> [p9,p5,p1,p7]: strongest on the edges
+    assert middle[0] == "p9" and middle[-1] == "p7"
+    assert middle[len(middle) // 2] == "p1"  # weakest in the center
+    assert ctx.report().order == "attention"
+
+
+def test_cache_order_puts_pinned_prefix_first():
+    ctx = Context(budget_tokens=1000, model="gpt-4o", order="cache")
+    ctx.add(Block("volatile", priority=8, pin=False, role="user"))
+    ctx.add(Block("stable-hi", priority=10, pin=True, role="system"))
+    ctx.add(Block("stable-lo", priority=2, pin=True, role="assistant"))
+    msgs = ctx.assemble()
+    # pinned blocks form the stable prefix (priority desc), volatile trails
+    assert [m["content"] for m in msgs] == ["stable-hi", "stable-lo", "volatile"]
+
+
+def test_invalid_order_rejected():
+    with pytest.raises(ValueError):
+        Context(budget_tokens=10, model="gpt-4o", order="bogus")
+
+
+def test_default_order_unchanged():
+    ctx = Context(budget_tokens=1000, model="gpt-4o")  # default
+    ctx.add(Block("u", priority=9, role="user"))
+    ctx.add(Block("s", priority=1, role="system"))
+    assert [m["role"] for m in ctx.assemble()] == ["system", "user"]
+    assert ctx.report().order == "default"
+
+
 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"))