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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: powerailabs-contextkit
-Version: 0.2.0
+Version: 0.4.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.2.0 → powerailabs_contextkit-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "powerailabs-contextkit"
-version = "0.2.0"
+version = "0.4.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.2.0 → powerailabs_contextkit-0.4.0}/src/powerailabs/contextkit/__init__.py

@@ -9,6 +9,7 @@ import each other; ``squeeze`` plugs in by shape via the ``contextkit[squeeze]``
 
 from __future__ import annotations
 
+import inspect
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import Any
@@ -104,6 +105,7 @@ class Context:
         reserve_output: int = 0,
         compressor: Any = None,
         order: str = "default",
+        image_tokens: int = 0,
     ) -> None:
         if order not in _ORDERS:
             raise ValueError(f"order must be one of {_ORDERS}, got {order!r}")
@@ -112,6 +114,7 @@ class Context:
         self.reserve_output = reserve_output
         self._compressor = compressor
         self.order = order
+        self.image_tokens = image_tokens  # token cost charged per image part in multimodal blocks
         self._blocks: list[Block] = []
         self._report: AssemblyReport | None = None
         self._messages: list[dict] = []
@@ -143,6 +146,17 @@ class Context:
         _, report = self._pack(budget_tokens, emit=False)
         return report
 
+    async def aassemble(self) -> list[dict]:
+        """Async assemble — like :meth:`assemble` but awaits ``async`` summarize callbacks.
+
+        Use this when a block's ``summarizer`` is a coroutine (e.g. an LLM summarizer). The sync
+        :meth:`assemble` falls back to truncation for async summarizers.
+        """
+        messages, report = await self._apack(self.budget_tokens, emit=True)
+        self._messages = messages
+        self._report = report
+        return messages
+
     def for_anthropic(self) -> tuple[str, list[dict]]:
         """Anthropic adapter: split system blocks out (the Messages API takes ``system`` apart).
 
@@ -154,38 +168,150 @@ class Context:
         rest = [m for m in self._messages if m["role"] != "system"]
         return system, rest
 
+    def for_gemini(self) -> tuple[str, list[dict]]:
+        """Gemini adapter: returns ``(system_instruction, contents)``.
+
+        ``contents`` are ``{"role": "user"|"model", "parts": [{"text": ...}]}`` (Gemini uses
+        ``model``, not ``assistant``); system blocks become the separate ``system_instruction``.
+        """
+        if not self._messages:
+            self.assemble()
+        system = "\n\n".join(m["content"] for m in self._messages if m["role"] == "system")
+        contents = [
+            {
+                "role": "model" if m["role"] == "assistant" else "user",
+                "parts": [{"text": m["content"]}],
+            }
+            for m in self._messages
+            if m["role"] != "system"
+        ]
+        return system, contents
+
+    def for_bedrock(self) -> tuple[list[dict], list[dict]]:
+        """Bedrock Converse adapter: returns ``(system, messages)``.
+
+        ``system`` is ``[{"text": ...}]`` (or empty); ``messages`` are
+        ``{"role": "user"|"assistant", "content": [{"text": ...}]}`` — Bedrock allows only those
+        two roles, so non-user blocks map to ``assistant``.
+        """
+        if not self._messages:
+            self.assemble()
+        system_text = "\n\n".join(m["content"] for m in self._messages if m["role"] == "system")
+        system = [{"text": system_text}] if system_text else []
+        messages = [
+            {
+                "role": "user" if m["role"] == "user" else "assistant",
+                "content": [{"text": m["content"]}],
+            }
+            for m in self._messages
+            if m["role"] != "system"
+        ]
+        return system, messages
+
     # ------------------------------------------------------------------ internals
 
-    def _pack(self, budget_tokens: int, *, emit: bool) -> tuple[list[dict], AssemblyReport]:
-        effective = max(0, budget_tokens - self.reserve_output)
+    def _ordered_blocks(self) -> list[tuple[int, Block]]:
         # (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]),
+        return sorted(
+            enumerate(self._blocks), key=lambda iv: (not iv[1].pin, -iv[1].priority, iv[0])
         )
 
+    def _block_tokens(self, block: Block) -> int:
+        """Token cost of a block, charging ``image_tokens`` per image part in multimodal content."""
+        content = block.content
+        if isinstance(content, list):
+            text = "".join(
+                p.get("text", "") for p in content if isinstance(p, dict) and "text" in p
+            )
+            images = sum(
+                1
+                for p in content
+                if isinstance(p, dict) and p.get("type") in ("image", "image_url")
+            )
+            return tokens.count(text, self.model) + images * self.image_tokens
+        return tokens.count(str(content), self.model)
+
+    def _finish(
+        self, budget_tokens: int, used: int, decisions: list, kept: list, *, emit: bool
+    ) -> tuple[list[dict], AssemblyReport]:
+        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)
+        return messages, report
+
+    def _pack(self, budget_tokens: int, *, emit: bool) -> tuple[list[dict], AssemblyReport]:
+        effective = max(0, budget_tokens - self.reserve_output)
         used = 0
         decisions: list[BlockDecision] = []
-        kept: list[tuple[int, Block, str]] = []  # (insertion_index, block, rendered_content)
+        kept: list[tuple[int, Block, Any]] = []  # (insertion_index, block, 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)
+        for idx, block in self._ordered_blocks():
+            before = self._block_tokens(block)
             remaining = effective - used
 
             if before <= remaining:
                 used += before
-                kept.append((idx, block, text))
+                kept.append((idx, block, block.content))
                 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})"
+                )
+            if not isinstance(block.content, str):  # can't shrink a multimodal/list block
+                decisions.append(
+                    BlockDecision(block.role, "dropped", before, 0, "multimodal: too large")
+                )
+                continue
+
+            new_text, action, note = self._evict(block, block.content, 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, new_text))
+            decisions.append(BlockDecision(block.role, action, before, after, note))
+
+        return self._finish(budget_tokens, used, decisions, kept, emit=emit)
+
+    async def _apack(self, budget_tokens: int, *, emit: bool) -> tuple[list[dict], AssemblyReport]:
+        effective = max(0, budget_tokens - self.reserve_output)
+        used = 0
+        decisions: list[BlockDecision] = []
+        kept: list[tuple[int, Block, Any]] = []
+
+        for idx, block in self._ordered_blocks():
+            before = self._block_tokens(block)
+            remaining = effective - used
 
+            if before <= remaining:
+                used += before
+                kept.append((idx, block, block.content))
+                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})"
                 )
+            if not isinstance(block.content, str):
+                decisions.append(
+                    BlockDecision(block.role, "dropped", before, 0, "multimodal: too large")
+                )
+                continue
 
-            new_text, action, note = self._evict(block, text, remaining)
+            new_text, action, note = await self._aevict(block, block.content, remaining)
             if new_text is None:
                 decisions.append(BlockDecision(block.role, "dropped", before, 0, note))
                 continue
@@ -194,19 +320,20 @@ class Context:
             kept.append((idx, block, new_text))
             decisions.append(BlockDecision(block.role, action, before, after, note))
 
-        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)
-        return messages, report
+        return self._finish(budget_tokens, used, decisions, kept, emit=emit)
+
+    async def _aevict(self, block: Block, text: str, remaining: int) -> tuple[str | None, str, str]:
+        """Async eviction: await an async summarizer; delegate everything else to ``_evict``."""
+        if (
+            block.evict == "summarize"
+            and block.summarizer is not None
+            and inspect.iscoroutinefunction(block.summarizer)
+        ):
+            summary = await block.summarizer(text, remaining)
+            if tokens.count(summary, self.model) > remaining:
+                summary = _truncate_to_tokens(summary, remaining, self.model)
+            return summary, "summarized", ""
+        return self._evict(block, text, remaining)
 
     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)``."""
@@ -220,15 +347,20 @@ class Context:
             return _truncate_to_tokens(text, remaining, self.model), "truncated", ""
 
         if strategy == "summarize":
-            if block.summarizer is not None:
+            if block.summarizer is not None and not inspect.iscoroutinefunction(block.summarizer):
                 summary = block.summarizer(text, remaining)
                 if tokens.count(summary, self.model) > remaining:
                     summary = _truncate_to_tokens(summary, remaining, self.model)
                 return summary, "summarized", ""
+            note = (
+                "async summarizer needs aassemble()"
+                if block.summarizer is not None
+                else "no summarizer"
+            )
             return (
                 _truncate_to_tokens(text, remaining, self.model),
                 "truncated",
-                ("no summarizer; fell back to truncate"),
+                f"{note}; truncated",
             )
 
         if strategy == "compress":

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

@@ -153,3 +153,78 @@ def test_for_anthropic_splits_system():
     system, messages = ctx.for_anthropic()
     assert system == "you are helpful"
     assert all(m["role"] != "system" for m in messages)
+
+
+def test_multimodal_image_token_cost():
+    ctx = Context(budget_tokens=1000, model="gpt-4o", image_tokens=85)
+    block = Block(
+        [{"type": "text", "text": "look"}, {"type": "image", "image_url": "..."}],
+        priority=9,
+        pin=True,
+        role="user",
+    )
+    ctx.add(block)
+    ctx.assemble()
+    d = ctx.report().decisions[0]
+    # text("look") ~1 tok + 1 image * 85 = ~86
+    assert d.tokens_before >= 85
+    # multimodal content is preserved as a list in the rendered message
+    assert isinstance(ctx.assemble()[0]["content"], list)
+
+
+def test_multimodal_block_dropped_when_too_large():
+    ctx = Context(budget_tokens=20, model="gpt-4o", image_tokens=1000)
+    ctx.add(Block("keep", priority=10, role="system"))
+    ctx.add(Block([{"type": "image"}], priority=1, role="user", evict="drop_oldest"))
+    ctx.assemble()
+    dropped = [d for d in ctx.report().decisions if d.action == "dropped"]
+    assert len(dropped) == 1
+
+
+async def test_async_summarizer_via_aassemble():
+    calls = {"n": 0}
+
+    async def summarizer(text, target):
+        calls["n"] += 1
+        return "async summary"
+
+    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="summarize", summarizer=summarizer))
+    msgs = await ctx.aassemble()
+    assert calls["n"] == 1  # the async summarizer ran
+    assert "async summary" in [m["content"] for m in msgs]
+    assert any(d.action == "summarized" for d in ctx.report().decisions)
+
+
+def test_sync_assemble_falls_back_for_async_summarizer():
+    async def summarizer(text, target):
+        return "nope"
+
+    ctx = Context(budget_tokens=12, model="gpt-4o")
+    ctx.add(Block("z" * 400, priority=1, role="user", evict="summarize", summarizer=summarizer))
+    ctx.assemble()  # sync path can't await -> truncates with a note
+    d = ctx.report().decisions[0]
+    assert d.action == "truncated" and "aassemble" in d.note
+
+
+def test_for_gemini_adapter():
+    ctx = Context(budget_tokens=1000, model="gpt-4o")
+    ctx.add(Block("be helpful", priority=10, pin=True, role="system"))
+    ctx.add(Block("prior reply", priority=5, role="assistant"))
+    ctx.add(Block("question", priority=9, pin=True, role="user"))
+    system, contents = ctx.for_gemini()
+    assert system == "be helpful"
+    roles = [c["role"] for c in contents]
+    assert "model" in roles and "user" in roles and "system" not in roles  # assistant -> model
+    assert contents[0]["parts"] == [{"text": "be helpful"}] or contents[0]["parts"][0]["text"]
+
+
+def test_for_bedrock_adapter():
+    ctx = Context(budget_tokens=1000, model="gpt-4o")
+    ctx.add(Block("be helpful", priority=10, pin=True, role="system"))
+    ctx.add(Block("question", priority=9, pin=True, role="user"))
+    system, messages = ctx.for_bedrock()
+    assert system == [{"text": "be helpful"}]
+    assert messages == [{"role": "user", "content": [{"text": "question"}]}]
+    assert all(m["role"] in ("user", "assistant") for m in messages)