pal-code 0.1.0__py3-none-any.whl

pal/context_manager/components/assembler.py

@@ -0,0 +1,75 @@
+"""One token budget for everything that goes into a request besides the user's own message.
+
+History, the exact details of old turns, remembered notes, document excerpts and code-graph excerpts each have their own cap, and together
+they could add up to a lot. The assembler keeps the total under `budget` by dropping the least useful first: remembered notes, then
+the facts, then the oldest turns (down to the last two), then document excerpts, then code excerpts, then the remaining turns
+(never the last one). What the user asks about right now (excerpts that matched it) outranks older history.
+"""
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+from pal.context_manager.components.turns import Turn
+from pal.tokens import estimate_tokens
+
+# the order blocks are dropped in, least useful first
+DROP_ORDER = ("memory", "facts", "documents", "graph")
+
+
+@dataclass
+class Block:
+    source: str            # "facts" | "memory" | "documents" | "graph"
+    text: str
+    tokens: int = 0
+
+
+@dataclass
+class ContextReport:
+    budget: int
+    used: int
+    by_source: dict = field(default_factory=dict)      # tokens per source, after fitting
+    dropped: list = field(default_factory=list)        # block sources dropped to fit
+    turns_dropped: int = 0
+    tokens_dropped_from_turns: int = 0
+
+
+class ContextAssembler:
+    def __init__(self, budget: int, count: Callable[[str], int] = estimate_tokens):
+        self.budget = budget
+        self.count = count
+
+    def fit(self, turns: list[Turn], blocks: list[Block], turn_cost: Callable[[Turn], int]) -> tuple[list[Turn], list[Block], ContextReport]:
+        blocks = [Block(b.source, b.text, b.tokens or self.count(b.text)) for b in blocks]
+        turns = list(turns)
+        used = sum(turn_cost(t) for t in turns) + sum(b.tokens for b in blocks)
+        dropped: list[str] = []
+        lost = 0
+
+        def drop_block(source: str) -> None:
+            nonlocal used
+            for b in [b for b in blocks if b.source == source]:
+                blocks.remove(b)
+                used -= b.tokens
+                dropped.append(source)
+
+        def drop_turns(keep: int) -> None:
+            nonlocal used, lost
+            while used > self.budget and len(turns) > keep:
+                cost = turn_cost(turns.pop(0))
+                used -= cost
+                lost += cost
+
+        before = len(turns)
+        for source in ("memory", "facts"):
+            if used > self.budget:
+                drop_block(source)
+        drop_turns(keep=2)
+        for source in ("documents", "graph"):
+            if used > self.budget:
+                drop_block(source)
+        drop_turns(keep=1)
+        by_source = {"history": sum(turn_cost(t) for t in turns)}
+        for b in blocks:
+            by_source[b.source] = by_source.get(b.source, 0) + b.tokens
+        return turns, blocks, ContextReport(self.budget, used, by_source, dropped, before - len(turns), lost)

pal/context_manager/components/context_cache.py

@@ -0,0 +1,45 @@
+import time
+from dataclasses import dataclass
+
+from pal.context_manager.utils.hashing import hash_segment
+from pal.context_manager.utils.storage import InMemoryStorage, StorageBackend
+
+
+@dataclass
+class CacheEntry:
+    content: str
+    created_at: float
+    hits: int = 0
+
+
+class ContextCache:
+    def __init__(self, storage: StorageBackend = None, ttl: float = 3600):
+        self.storage = storage or InMemoryStorage()
+        self.ttl = ttl
+        self.hits = 0
+        self.misses = 0
+
+    def get_or_store(self, content: str) -> tuple[bool, str]:
+        key = hash_segment(content)
+        entry = self.storage.get(key)
+
+        if entry and (time.time() - entry.created_at) < self.ttl:
+            entry.hits += 1
+            self.storage.set(key, entry)
+            self.hits += 1
+            return True, key
+
+        self.storage.set(key, CacheEntry(content=content, created_at=time.time()))
+        self.misses += 1
+        return False, key
+
+    def resolve(self, key: str) -> str | None:
+        entry = self.storage.get(key)
+        return entry.content if entry else None
+
+    def hit_rate(self) -> float:
+        total = self.hits + self.misses
+        return self.hits / total if total else 0.0
+
+    def size_bytes(self) -> int:
+        return sum(len(self.storage.get(k).content.encode("utf-8")) for k in self.storage.keys())

pal/context_manager/components/kv_bridge.py

@@ -0,0 +1,16 @@
+from pal.llm_wrapper.components.base import LLMWrapper
+
+
+class KVBridge:
+    """Delegates dedup to the active LLM core's native KV cache when
+    available (e.g. llama.cpp's prompt-prefix cache, Ollama's automatic
+    prefix reuse), instead of duplicating that logic here."""
+
+    def __init__(self, core: LLMWrapper):
+        self.core = core
+
+    def active(self) -> bool:
+        return self.core.supports_kv_cache()
+
+    def reset(self) -> None:
+        self.core.reset_kv_cache()

pal/context_manager/components/pager.py

@@ -0,0 +1,48 @@
+import time
+from dataclasses import dataclass, field
+
+from pal.context_manager.utils.relevance import recency_score, text_similarity
+
+
+@dataclass
+class Segment:
+    id: str
+    content: str
+    role: str = "user"
+    timestamp: float = field(default_factory=time.time)
+
+
+class ContextPager:
+    def __init__(self, hot_window: int = 6, cold_threshold: float = 0.35):
+        self.hot_window = hot_window
+        self.cold_threshold = cold_threshold
+        self.cold_store: dict[str, Segment] = {}
+
+    def page_out(self, history: list[Segment]) -> tuple[list[Segment], list[str]]:
+        if len(history) <= self.hot_window:
+            return history, []
+
+        hot = history[-self.hot_window:]
+        cold = history[: -self.hot_window]
+        cold_refs = []
+        for seg in cold:
+            self.cold_store[seg.id] = seg
+            cold_refs.append(seg.id)
+        return hot, cold_refs
+
+    def page_in(self, query: str, cold_refs: list[str], top_k: int = 2) -> list[Segment]:
+        candidates = [self.cold_store[r] for r in cold_refs if r in self.cold_store]
+        if not candidates:
+            return []
+
+        texts = [c.content for c in candidates]
+        sims = text_similarity(query, texts)  # embeddings if available, keyword overlap otherwise
+        now = time.time()
+
+        scored = []
+        for seg, sim in zip(candidates, sims):
+            score = 0.7 * sim + 0.3 * recency_score(seg.timestamp, now)
+            scored.append((score, seg))
+
+        scored.sort(key=lambda x: x[0], reverse=True)
+        return [seg for score, seg in scored[:top_k] if score >= self.cold_threshold]

pal/context_manager/components/simple_context_manager.py

@@ -0,0 +1,56 @@
+"""SimpleContextManager (AIOS Appendix A.4): pause a long generation when its time slice ends and resume it later without
+losing the work done.
+
+Two ways, chosen by what the core can give us. Text-based (every cloud API and local server): the text decoded so far is the
+snapshot, and resuming continues from it. Logits-based (an in-process Hugging Face model): the beam-search state is the
+snapshot. Both are done by the interrupt controller already used by the Round Robin scheduler; this class gives it the
+per-caller save, load and clear interface the adapter needs.
+"""
+from __future__ import annotations
+
+import threading
+from typing import Any
+
+from pal.llm_wrapper.components.base import LLMSyscall, LLMWrapper
+
+
+class SimpleContextManager:
+    def __init__(self) -> None:
+        self._contexts: dict[Any, Any] = {}
+        self._lock = threading.Lock()
+        from pal.orchestrator.components.context_interrupt import (
+            ContextInterruptController,  # lazy: avoids a cycle
+        )
+        self._controller = ContextInterruptController()
+
+    def load_context(self, pid: Any) -> Any:
+        with self._lock:
+            return self._contexts.get(pid)
+
+    def clear_context(self, pid: Any) -> None:
+        with self._lock:
+            self._contexts.pop(pid, None)
+
+    def generate_response_with_interruption(
+        self, model_name: str, model: LLMWrapper, messages: list[dict], tools: list[dict] | None, message_return_type: str,
+        temperature: float | None, max_tokens: int | None, pid: Any, time_limit: float, response_format: dict | None = None,
+    ) -> tuple[str, bool]:
+        """Generate for at most `time_limit` seconds. Returns (text so far, finished). If not finished, the state is kept under
+        `pid`; calling again with the same pid continues from it. Plain text only: tool calls and JSON output are not interrupted."""
+        from pal.orchestrator.components.context_interrupt import ContextSnapshot, InterruptResult
+
+        base = LLMSyscall(messages=messages, model=model_name or None, temperature=temperature, max_tokens=max_tokens)
+        snap = self.load_context(pid)
+        prior = ""
+        if snap is not None and snap.mode == "text":
+            prior = snap.text_so_far
+            partial = InterruptResult(complete=False, text_so_far=prior, elapsed=0.0, snapshot=snap)
+            base = self._controller.build_continuation(base, partial)
+        res = self._controller.run_slice(model, base, time_limit, resume_snapshot=snap if snap and snap.mode == "logits" else None)
+        text = res.text_so_far if res.cumulative else prior + res.text_so_far
+        if res.complete:
+            self.clear_context(pid)
+            return text, True
+        with self._lock:
+            self._contexts[pid] = ContextSnapshot(mode=res.snapshot.mode, text_so_far=text, beam_state=res.snapshot.beam_state)
+        return text, False

pal/context_manager/components/turns.py

@@ -0,0 +1,88 @@
+"""Paging for chat history, one turn (user message + reply) at a time.
+
+The last `hot_turns` turns are always sent. Older turns are "cold": they are dropped from the
+request unless they look relevant to the current question, in which case up to `recall_turns` of
+them are paged back in (best match first, sent in chronological order). A token cap on the hot
+window trims it further if a few turns are very long.
+"""
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+from pal.context_manager.components.pager import ContextPager, Segment
+from pal.tokens import estimate_tokens
+
+
+@dataclass
+class Turn:
+    user: str
+    assistant: str
+    local: bool = False                       # answered by the local LLM
+    ts: float = field(default_factory=time.time)
+    id: str = ""
+    files: tuple = ()                         # files an agent saved in this turn (the content lives in the workspace)
+    # cached by the compressor, never part of the turn's identity
+    _keys: tuple[frozenset, frozenset, list] | None = field(default=None, repr=False, compare=False)
+    _facts: tuple[list[str], list[str]] | None = field(default=None, repr=False, compare=False)
+
+    def text(self) -> str:
+        return f"{self.user}\n{self.assistant}"
+
+
+@dataclass
+class PagedHistory:
+    turns: list[Turn]             # what to send, chronological
+    recalled: int                 # cold turns paged back in
+    full_tokens: int              # cost of sending every turn
+    sent_tokens: int              # cost of what is sent
+    dropped_turns: int
+
+    @property
+    def saved_tokens(self) -> int:
+        return max(0, self.full_tokens - self.sent_tokens)
+
+
+class TurnPager:
+    def __init__(self, hot_turns: int = 6, recall_turns: int = 2, recall_threshold: float = 0.35,
+                 max_history_tokens: int = 6000, count: Callable[[str], int] = estimate_tokens):
+        self.hot_turns = hot_turns
+        self.recall_turns = recall_turns
+        self.recall_threshold = recall_threshold
+        self.max_history_tokens = max_history_tokens
+        self.count = count
+
+    def _cost(self, turn: Turn) -> int:
+        return self.count(turn.user) + self.count(turn.assistant) + 8
+
+    turn_cost = _cost     # public name for the assembler
+
+    def select(self, turns: list[Turn], query: str) -> PagedHistory:
+        full = sum(self._cost(t) for t in turns)
+        if len(turns) <= 1 and full <= self.max_history_tokens:
+            return PagedHistory(list(turns), 0, full, full, 0)
+
+        index = {str(i): t for i, t in enumerate(turns)}
+        segments = [Segment(id=str(i), content=t.text(), role="turn", timestamp=t.ts) for i, t in enumerate(turns)]
+        pager = ContextPager(hot_window=self.hot_turns, cold_threshold=self.recall_threshold)
+        hot_segs, cold_refs = pager.page_out(segments)
+        hot = [index[s.id] for s in hot_segs]
+
+        # token cap: push the oldest hot turns to cold while over budget (always keep the newest)
+        while len(hot) > 1 and sum(self._cost(t) for t in hot) > self.max_history_tokens:
+            oldest = hot.pop(0)
+            cold_refs.append(next(k for k, v in index.items() if v is oldest))
+
+        recalled: list[Turn] = []
+        if self.recall_turns and cold_refs and query.strip():
+            room = self.max_history_tokens - sum(self._cost(t) for t in hot)
+            for seg in pager.page_in(query, cold_refs, top_k=self.recall_turns):
+                t = index[seg.id]
+                if self._cost(t) <= room:
+                    recalled.append(t)
+                    room -= self._cost(t)
+        pos = {id(t): i for i, t in enumerate(turns)}
+        chosen = sorted(recalled + hot, key=lambda t: pos[id(t)])
+        sent = sum(self._cost(t) for t in chosen)
+        return PagedHistory(chosen, len(recalled), full, sent, len(turns) - len(chosen))

pal/context_manager/pipeline.py

@@ -0,0 +1,27 @@
+from pal.context_manager.components.context_cache import ContextCache
+from pal.context_manager.components.kv_bridge import KVBridge
+from pal.context_manager.components.pager import ContextPager, Segment
+from pal.llm_wrapper.components.base import LLMWrapper
+
+
+class ContextManagerPipeline:
+    """Ties system-prompt caching, hot/cold history paging, and the
+    backend's native KV cache together behind one small interface."""
+
+    def __init__(self, core: LLMWrapper, cache_ttl: float = 3600, hot_window: int = 6, cold_threshold: float = 0.35):
+        self.cache = ContextCache(ttl=cache_ttl)
+        self.pager = ContextPager(hot_window=hot_window, cold_threshold=cold_threshold)
+        self.kv_bridge = KVBridge(core)
+
+    def cache_system_prompt(self, system_prompt: str) -> tuple[bool, str]:
+        return self.cache.get_or_store(system_prompt)
+
+    def page_history(self, history: list[Segment], query: str) -> list[Segment]:
+        """Returns hot + relevant paged-in cold segments, in order, ready
+        to append to a request's message list."""
+        hot, cold_refs = self.pager.page_out(history)
+        recalled = self.pager.page_in(query, cold_refs)
+        return recalled + hot
+
+    def size_bytes(self) -> int:
+        return self.cache.size_bytes()

pal/context_manager/utils/__init__.py

@@ -0,0 +1,18 @@
+from pal.context_manager.utils.hashing import hash_segment
+from pal.context_manager.utils.relevance import (
+    embedding_similarity,
+    keyword_overlap,
+    recency_score,
+    text_similarity,
+)
+from pal.context_manager.utils.storage import InMemoryStorage, StorageBackend
+
+__all__ = [
+    "hash_segment",
+    "embedding_similarity",
+    "keyword_overlap",
+    "text_similarity",
+    "recency_score",
+    "InMemoryStorage",
+    "StorageBackend",
+]

pal/context_manager/utils/hashing.py

@@ -0,0 +1,5 @@
+import hashlib
+
+
+def hash_segment(content: str) -> str:
+    return hashlib.sha256(content.encode("utf-8")).hexdigest()

pal/context_manager/utils/relevance.py

@@ -0,0 +1,59 @@
+import re
+import time
+
+try:
+    import numpy as np
+    from sentence_transformers import SentenceTransformer
+    _EMBEDDER = None
+    _EMBED_AVAILABLE = True
+except ImportError:
+    _EMBED_AVAILABLE = False
+
+
+def _get_embedder():
+    global _EMBEDDER
+    if _EMBEDDER is None:
+        _EMBEDDER = SentenceTransformer("all-MiniLM-L6-v2")
+    return _EMBEDDER
+
+
+def embedding_similarity(query: str, segments: list[str]) -> list[float]:
+    if not _EMBED_AVAILABLE or not segments:
+        return [0.0] * len(segments)
+    model = _get_embedder()
+    q_vec = model.encode([query])[0]
+    s_vecs = model.encode(segments)
+    sims = s_vecs @ q_vec / (np.linalg.norm(s_vecs, axis=1) * np.linalg.norm(q_vec) + 1e-8)
+    return sims.tolist()
+
+
+def recency_score(timestamp: float, now: float = None) -> float:
+    now = now or time.time()
+    age = max(0.0, now - timestamp)
+    return 1.0 / (1.0 + age / 60.0)
+
+
+_WORD = re.compile(r"[^\W_]+", re.UNICODE)
+_STOP = frozenset("a an the and or but if of to in on at for with by from as is are was were be been it this that "
+                  "these those i you he she we they me my your our their what which who how why when where do does did "
+                  "can could should would will please".split())
+
+
+def _keywords(text: str) -> set[str]:
+    return {w for w in (m.lower() for m in _WORD.findall(text)) if w not in _STOP and (len(w) > 1 or w.isdigit())}
+
+
+def keyword_overlap(query: str, text: str) -> float:
+    """Share of the query's meaningful words that also appear in `text` (0 to 1). No model needed."""
+    q, t = _keywords(query), _keywords(text)
+    if not q or not t:
+        return 0.0
+    return len(q & t) / len(q)
+
+
+def text_similarity(query: str, texts: list[str]) -> list[float]:
+    """Embedding similarity when sentence-transformers is installed, keyword overlap otherwise."""
+    sims = embedding_similarity(query, texts)
+    if texts and all(s == 0.0 for s in sims):
+        sims = [keyword_overlap(query, t) for t in texts]
+    return sims

pal/context_manager/utils/storage.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+
+class StorageBackend(ABC):
+    @abstractmethod
+    def get(self, key: str) -> Optional[Any]:
+        ...
+
+    @abstractmethod
+    def set(self, key: str, value: Any) -> None:
+        ...
+
+    @abstractmethod
+    def delete(self, key: str) -> None:
+        ...
+
+    @abstractmethod
+    def keys(self) -> list[str]:
+        ...
+
+
+class InMemoryStorage(StorageBackend):
+    def __init__(self):
+        self._store: dict[str, Any] = {}
+
+    def get(self, key):
+        return self._store.get(key)
+
+    def set(self, key, value):
+        self._store[key] = value
+
+    def delete(self, key):
+        self._store.pop(key, None)
+
+    def keys(self):
+        return list(self._store.keys())

pal/documents/__init__.py

@@ -0,0 +1,4 @@
+from pal.documents.components.chunker import Chunk, chunk_markdown
+from pal.documents.components.index import DocInfo, DocumentIndex, Hit
+
+__all__ = ["Chunk", "DocInfo", "DocumentIndex", "Hit", "chunk_markdown"]

pal/documents/components/chunker.py

@@ -0,0 +1,79 @@
+"""Split markdown into chunks of about `target_tokens`, at heading and paragraph boundaries.
+Pure and deterministic: the same text always gives the same chunks, so part numbers are stable."""
+from __future__ import annotations
+
+import re
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from pal.tokens import estimate_tokens
+
+_HEADING = re.compile(r"^#{1,6}\s+(.*\S)\s*$")
+
+
+@dataclass(frozen=True)
+class Chunk:
+    index: int        # 0-based
+    heading: str      # nearest heading above the chunk ("" if none)
+    text: str
+    tokens: int
+
+
+def _split_long(paragraph: str, target: int, count: Callable[[str], int]) -> list[str]:
+    """Break one oversized paragraph on line, then sentence, then character boundaries."""
+    if count(paragraph) <= target:
+        return [paragraph]
+    pieces: list[str] = []
+    for line in paragraph.split("\n"):
+        pieces.extend(re.split(r"(?<=[.!?])\s+", line) if count(line) > target else [line])
+    out: list[str] = []
+    cur = ""
+    for piece in pieces:
+        while count(piece) > target:  # a single huge "sentence": hard cut
+            cut = target * 4
+            out.append(piece[:cut])
+            piece = piece[cut:]
+        if cur and count(cur + "\n" + piece) > target:
+            out.append(cur)
+            cur = piece
+        else:
+            cur = f"{cur}\n{piece}" if cur else piece
+    if cur:
+        out.append(cur)
+    return out
+
+
+def chunk_markdown(text: str, target_tokens: int = 350, count: Callable[[str], int] = estimate_tokens) -> list[Chunk]:
+    text = text.replace("\r\n", "\n").strip()
+    if not text:
+        return []
+    chunks: list[Chunk] = []
+    heading = ""
+    buf: list[str] = []
+    buf_heading = ""
+
+    def flush() -> None:
+        nonlocal buf
+        if buf:
+            body = "\n\n".join(buf).strip()
+            if body:
+                chunks.append(Chunk(len(chunks), buf_heading, body, count(body)))
+        buf = []
+
+    for para in re.split(r"\n\s*\n", text):
+        para = para.strip("\n")
+        if not para.strip():
+            continue
+        m = _HEADING.match(para.split("\n", 1)[0])
+        if m:
+            flush()  # a heading starts a new chunk
+            heading = m.group(1)
+        for piece in _split_long(para, target_tokens, count):
+            heading_only = len(buf) == 1 and _HEADING.match(buf[0]) is not None
+            if buf and not heading_only and count("\n\n".join(buf + [piece])) > target_tokens:
+                flush()
+            if not buf:
+                buf_heading = heading
+            buf.append(piece)
+    flush()
+    return chunks