apsimo 1.3.0__py3-none-any.whl

apsimo/contextgate/chunker.py

@@ -0,0 +1,169 @@
+"""Structure-aware text chunking.
+
+Splits text along natural boundaries — fenced code blocks stay atomic,
+then markdown headings, then blank-line paragraphs — and greedily packs
+blocks into chunks near a target token size, with configurable overlap
+between consecutive chunks so facts straddling a boundary survive
+retrieval.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+from apsimo.contextgate.estimate import estimate_tokens
+
+__all__ = ["Chunk", "chunk_text"]
+
+_FENCE_RE = re.compile(r"^(```|~~~)", re.MULTILINE)
+_HEADING_RE = re.compile(r"^#{1,6}\s", re.MULTILINE)
+
+
+@dataclass
+class Chunk:
+    """One packed chunk of a source text."""
+
+    index: int          # 0-based position in document order
+    text: str           # chunk content (includes any overlap prefix)
+    start: int          # char offset of the core (non-overlap) content
+    end: int            # char offset one past the core content
+    tokens: int = 0     # estimated tokens of ``text``
+
+
+def _split_blocks(text: str) -> list[tuple[int, str]]:
+    """Split *text* into (offset, block) pairs along structural boundaries.
+
+    Fenced code blocks are kept whole. Outside fences, headings start new
+    blocks and blank lines separate paragraphs.
+    """
+    blocks: list[tuple[int, str]] = []
+    pos = 0
+    n = len(text)
+
+    while pos < n:
+        fence = _FENCE_RE.search(text, pos)
+        prose_end = fence.start() if fence else n
+
+        # Prose region: split on blank lines and headings
+        prose = text[pos:prose_end]
+        if prose.strip():
+            offset = pos
+            # Insert split points before headings so each heading opens a block
+            paragraphs = re.split(r"(\n\s*\n)", prose)
+            cursor = 0
+            for part in paragraphs:
+                if part.strip() and not re.fullmatch(r"\n\s*\n", part):
+                    # Further split on headings inside the paragraph run
+                    last = 0
+                    for m in _HEADING_RE.finditer(part):
+                        if m.start() > last and part[last:m.start()].strip():
+                            blocks.append((offset + cursor + last, part[last:m.start()]))
+                        last = m.start()
+                    if part[last:].strip():
+                        blocks.append((offset + cursor + last, part[last:]))
+                cursor += len(part)
+
+        if fence is None:
+            break
+
+        # Fenced block: find the closing fence of the same kind
+        marker = fence.group(1)
+        close = text.find("\n" + marker, fence.end())
+        if close == -1:
+            blocks.append((fence.start(), text[fence.start():]))
+            break
+        # Include through the end of the closing-fence line
+        fend = text.find("\n", close + 1 + len(marker))
+        fend = n if fend == -1 else fend + 1
+        blocks.append((fence.start(), text[fence.start():fend]))
+        pos = fend
+
+    return blocks
+
+
+def _hard_split(offset: int, block: str, target_tokens: int) -> list[tuple[int, str]]:
+    """Split an oversized block on sentence/newline boundaries, then windows."""
+    limit_chars = max(1, target_tokens * 4)
+    pieces: list[tuple[int, str]] = []
+    cursor = 0
+    n = len(block)
+    while cursor < n:
+        end = min(cursor + limit_chars, n)
+        if end < n:
+            # Back up to the last sentence end or newline in the window
+            window = block[cursor:end]
+            cut = max(window.rfind(". "), window.rfind("\n"))
+            if cut > limit_chars // 4:
+                end = cursor + cut + 1
+        pieces.append((offset + cursor, block[cursor:end]))
+        cursor = end
+    return pieces
+
+
+def chunk_text(
+    text: str,
+    target_tokens: int = 1024,
+    overlap_tokens: int = 128,
+) -> list[Chunk]:
+    """Chunk *text* into ~*target_tokens* pieces along structural boundaries.
+
+    Consecutive chunks share an overlap of roughly *overlap_tokens* (the
+    tail of the previous chunk is prepended to the next), so information
+    spanning a boundary remains retrievable. ``start``/``end`` offsets
+    always refer to the core (non-overlap) content.
+    """
+    if not text.strip():
+        return []
+
+    blocks: list[tuple[int, str]] = []
+    for offset, block in _split_blocks(text):
+        if estimate_tokens(block) > target_tokens:
+            blocks.extend(_hard_split(offset, block, target_tokens))
+        else:
+            blocks.append((offset, block))
+
+    chunks: list[Chunk] = []
+    cur_parts: list[tuple[int, str]] = []
+    cur_tokens = 0
+
+    def _flush() -> None:
+        nonlocal cur_parts, cur_tokens
+        if not cur_parts:
+            return
+        start = cur_parts[0][0]
+        last_off, last_text = cur_parts[-1]
+        end = last_off + len(last_text)
+        core = text[start:end]
+        prefix = ""
+        if chunks and overlap_tokens > 0:
+            prev = chunks[-1]
+            tail_chars = overlap_tokens * 4
+            tail = text[max(prev.start, prev.end - tail_chars):prev.end]
+            # Cut at a word boundary so the overlap reads cleanly
+            sp = tail.find(" ")
+            if 0 <= sp < len(tail) - 1:
+                tail = tail[sp + 1:]
+            prefix = tail + "\n"
+        body = prefix + core
+        chunks.append(
+            Chunk(
+                index=len(chunks),
+                text=body,
+                start=start,
+                end=end,
+                tokens=estimate_tokens(body),
+            )
+        )
+        cur_parts = []
+        cur_tokens = 0
+
+    for offset, block in blocks:
+        btok = estimate_tokens(block)
+        if cur_parts and cur_tokens + btok > target_tokens:
+            _flush()
+        cur_parts.append((offset, block))
+        cur_tokens += btok
+    _flush()
+
+    return chunks

apsimo/contextgate/estimate.py

@@ -0,0 +1,54 @@
+"""Cheap token estimation — no tokenizer dependency.
+
+The gate only needs estimates good to ~±20%; the decision headroom
+(default 0.8) absorbs the imprecision. English prose runs ~4 chars per
+token; code and symbol-dense text runs denser (~3 chars per token), so a
+crude density probe adjusts the ratio.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+
+__all__ = ["estimate_tokens"]
+
+_DEFAULT_CHARS_PER_TOKEN = 4.0
+_CODE_CHARS_PER_TOKEN = 3.0
+
+
+def _chars_per_token() -> float:
+    try:
+        v = float(os.environ.get("COLONY_CONTEXT_CHARS_PER_TOKEN", ""))
+        if v > 0:
+            return v
+    except ValueError:
+        pass
+    return _DEFAULT_CHARS_PER_TOKEN
+
+
+def _looks_dense(text: str, sample_limit: int = 20000) -> bool:
+    """True when the text is symbol/whitespace-dense (code, logs, JSON)."""
+    sample = text[:sample_limit]
+    if not sample:
+        return False
+    symbolish = sum(
+        1 for c in sample if not (c.isalpha() or c in " .,;:'\"!?-")
+    )
+    return symbolish / len(sample) > 0.25
+
+
+def estimate_tokens(text: str) -> int:
+    """Estimate the token count of *text*.
+
+    Uses a chars-per-token heuristic (configurable via
+    ``COLONY_CONTEXT_CHARS_PER_TOKEN``), with a denser ratio for
+    code-like input. Deliberately dependency-free; accuracy within
+    ~±20% is sufficient for gating decisions.
+    """
+    if not text:
+        return 0
+    cpt = _chars_per_token()
+    if _looks_dense(text):
+        cpt = min(cpt, _CODE_CHARS_PER_TOKEN)
+    return math.ceil(len(text) / cpt)

apsimo/contextgate/gate.py

@@ -0,0 +1,313 @@
+"""Gate decision + context preparation service.
+
+``decide()`` picks a strategy from (estimated size, budget, task kind);
+``prepare_context()`` executes it and returns the final context text plus
+metadata. See the package docstring for the full model.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Awaitable, Callable, Optional
+
+from apsimo.contextgate.chunker import Chunk, chunk_text
+from apsimo.contextgate.estimate import estimate_tokens
+from apsimo.contextgate.retrieve import EmbedFn, rank_chunks
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "GateConfig",
+    "GateDecision",
+    "PreparedContext",
+    "classify_task",
+    "decide",
+    "prepare_context",
+]
+
+SummarizeFn = Callable[[str, str], Awaitable[str]]  # (chunk_text, query) -> summary
+
+
+class GateDecision(str, Enum):
+    PASS_THROUGH = "pass_through"
+    RETRIEVE = "retrieve"
+    MAP_REDUCE = "map_reduce"
+
+
+@dataclass
+class GateConfig:
+    """Tunables for the context gate. All values env-overridable."""
+
+    mode: str = "auto"                 # auto | on | off
+    headroom: float = 0.8              # gate when est > budget * headroom
+    default_budget_tokens: int = 0     # used when caller/tier give none (0 = don't gate)
+    chunk_tokens: int = 1024
+    overlap_tokens: int = 128
+    min_score: float = 0.05            # drop retrieval chunks scoring below this
+
+    @classmethod
+    def from_env(cls) -> "GateConfig":
+        cfg = cls()
+        mode = os.environ.get("COLONY_CONTEXT_GATE", "").strip().lower()
+        if mode in ("auto", "on", "off"):
+            cfg.mode = mode
+        for attr, env, cast in (
+            ("headroom", "COLONY_CONTEXT_GATE_HEADROOM", float),
+            ("default_budget_tokens", "COLONY_CONTEXT_GATE_BUDGET", int),
+            ("chunk_tokens", "COLONY_CONTEXT_CHUNK_TOKENS", int),
+            ("overlap_tokens", "COLONY_CONTEXT_OVERLAP_TOKENS", int),
+        ):
+            raw = os.environ.get(env, "")
+            if raw:
+                try:
+                    setattr(cfg, attr, cast(raw))
+                except ValueError:
+                    logger.warning("Invalid %s=%r — using default", env, raw)
+        return cfg
+
+
+@dataclass
+class PreparedContext:
+    """Result of ``prepare_context``."""
+
+    text: str
+    decision: GateDecision
+    est_tokens_in: int
+    est_tokens_out: int
+    budget_tokens: int
+    chunks_total: int = 0
+    chunks_used: int = 0
+    coverage: float = 1.0   # fraction of source chars represented in output
+    task_kind: str = ""
+
+
+# ---------------------------------------------------------------------------
+# Task classification — the smarter-than-size signal
+# ---------------------------------------------------------------------------
+
+_HOLISTIC_RE = re.compile(
+    r"\b(summar[iy][sz]e|overview|review|rewrite|rephrase|translate|proofread|"
+    r"critique|tl;?dr|digest|condense|outline|abstract)\b",
+    re.IGNORECASE,
+)
+_RETRIEVAL_RE = re.compile(
+    r"(\?|\b(what|when|where|who|whom|whose|which|why|how|find|look ?up|search|"
+    r"locate|extract|quote|list all|did|does|do|is|are|was|were)\b)",
+    re.IGNORECASE,
+)
+
+
+def classify_task(query: str) -> str:
+    """Classify the caller's intent as ``retrieval`` or ``holistic``.
+
+    Retrieval tasks (needle questions, lookups) benefit from ranked-chunk
+    RAG; holistic tasks (summarize/review the whole document) need
+    coverage of everything and get map-reduce instead. Callers that know
+    their intent should pass ``task_kind`` explicitly — this heuristic is
+    only the fallback.
+    """
+    q = (query or "").strip()
+    if not q:
+        return "holistic"
+    if _HOLISTIC_RE.search(q):
+        return "holistic"
+    if _RETRIEVAL_RE.search(q):
+        return "retrieval"
+    # A short, specific query usually names the thing to find.
+    return "retrieval" if len(q) < 200 else "holistic"
+
+
+# ---------------------------------------------------------------------------
+# Decision
+# ---------------------------------------------------------------------------
+
+def decide(
+    est_tokens: int,
+    budget_tokens: int,
+    query: str = "",
+    task_kind: Optional[str] = None,
+    config: Optional[GateConfig] = None,
+) -> GateDecision:
+    """Pick a strategy for content of *est_tokens* against *budget_tokens*.
+
+    ``mode=off`` or an unknown budget (0) always passes through; otherwise
+    content within ``budget * headroom`` passes through, and oversized
+    content is routed to RETRIEVE or MAP_REDUCE by task kind.
+    """
+    cfg = config or GateConfig.from_env()
+    if cfg.mode == "off":
+        return GateDecision.PASS_THROUGH
+    budget = budget_tokens or cfg.default_budget_tokens
+    if budget <= 0:
+        return GateDecision.PASS_THROUGH
+    if est_tokens <= budget * cfg.headroom:
+        return GateDecision.PASS_THROUGH
+    kind = task_kind if task_kind in ("retrieval", "holistic") else classify_task(query)
+    return GateDecision.RETRIEVE if kind == "retrieval" else GateDecision.MAP_REDUCE
+
+
+# ---------------------------------------------------------------------------
+# Preparation
+# ---------------------------------------------------------------------------
+
+def _assemble(
+    selected: list[Chunk],
+    total: int,
+    source_chars: int,
+    label: str,
+) -> tuple[str, float]:
+    """Join *selected* chunks (document order) with provenance markers."""
+    selected = sorted(selected, key=lambda c: c.index)
+    covered = sum(c.end - c.start for c in selected)
+    coverage = min(1.0, covered / source_chars) if source_chars else 1.0
+    parts = [
+        f"[context gate: {label} — {len(selected)} of {total} chunks, "
+        f"~{coverage:.0%} of the source shown]"
+    ]
+    for c in selected:
+        parts.append(f"--- chunk {c.index + 1}/{total} (chars {c.start}-{c.end}) ---")
+        parts.append(c.text)
+    return "\n".join(parts), coverage
+
+
+def _pack_to_budget(
+    ranked: list[tuple[Chunk, float]],
+    budget_tokens: int,
+    min_score: float,
+) -> list[Chunk]:
+    selected: list[Chunk] = []
+    used = 0
+    for chunk, score in ranked:
+        if score < min_score and selected:
+            break
+        if used + chunk.tokens > budget_tokens:
+            if not selected:
+                selected.append(chunk)  # always include at least the best chunk
+            break
+        selected.append(chunk)
+        used += chunk.tokens
+    return selected
+
+
+def _sample_evenly(chunks: list[Chunk], budget_tokens: int) -> list[Chunk]:
+    """Pick evenly-spaced chunks so the selection spans the whole source."""
+    if not chunks:
+        return []
+    avg = max(1, sum(c.tokens for c in chunks) // len(chunks))
+    k = max(1, min(len(chunks), budget_tokens // avg))
+    if k >= len(chunks):
+        return list(chunks)
+    step = len(chunks) / k
+    picked = []
+    used = 0
+    for i in range(k):
+        c = chunks[int(i * step)]
+        if used + c.tokens > budget_tokens and picked:
+            break
+        picked.append(c)
+        used += c.tokens
+    return picked
+
+
+async def prepare_context(
+    content: str,
+    query: str = "",
+    budget_tokens: int = 0,
+    task_kind: Optional[str] = None,
+    config: Optional[GateConfig] = None,
+    embed_fn: Optional[EmbedFn] = None,
+    summarize_fn: Optional[SummarizeFn] = None,
+) -> PreparedContext:
+    """Prepare *content* for a model call within *budget_tokens*.
+
+    Returns the content unchanged when it fits (or gating is off).
+    Otherwise chunks it and either retrieves the chunks most relevant to
+    *query* (retrieval tasks) or map-reduces via *summarize_fn* /
+    coverage-samples (holistic tasks). Never raises on ranking or
+    summarization failure — degrades toward coverage sampling.
+    """
+    cfg = config or GateConfig.from_env()
+    est = estimate_tokens(content)
+    budget = budget_tokens or cfg.default_budget_tokens
+    decision = decide(est, budget, query, task_kind, cfg)
+    kind = task_kind if task_kind in ("retrieval", "holistic") else classify_task(query)
+
+    if decision == GateDecision.PASS_THROUGH:
+        return PreparedContext(
+            text=content,
+            decision=decision,
+            est_tokens_in=est,
+            est_tokens_out=est,
+            budget_tokens=budget,
+            task_kind=kind,
+        )
+
+    chunks = chunk_text(content, cfg.chunk_tokens, cfg.overlap_tokens)
+    total = len(chunks)
+    effective_budget = max(1, int(budget * cfg.headroom))
+
+    if decision == GateDecision.RETRIEVE:
+        ranked = await rank_chunks(chunks, query, embed_fn)
+        selected = _pack_to_budget(ranked, effective_budget, cfg.min_score)
+        text, coverage = _assemble(
+            selected, total, len(content), "chunks selected for relevance to the query"
+        )
+    else:  # MAP_REDUCE
+        if summarize_fn is not None:
+            summaries: list[Chunk] = []
+            for c in chunks:
+                try:
+                    s = await summarize_fn(c.text, query)
+                except Exception:
+                    logger.warning(
+                        "summarize_fn failed on chunk %d — using head of chunk",
+                        c.index, exc_info=True,
+                    )
+                    s = c.text[: cfg.chunk_tokens]
+                summaries.append(
+                    Chunk(
+                        index=c.index,
+                        text=s,
+                        start=c.start,
+                        end=c.end,
+                        tokens=estimate_tokens(s),
+                    )
+                )
+            # Keep summaries within budget (they normally fit; sample if not)
+            if sum(s.tokens for s in summaries) > effective_budget:
+                summaries = _sample_evenly(summaries, effective_budget)
+            text, coverage = _assemble(
+                summaries, total, len(content), "per-chunk summaries (map-reduce)"
+            )
+        else:
+            selected = _sample_evenly(chunks, effective_budget)
+            text, coverage = _assemble(
+                selected, total, len(content), "evenly-spaced coverage sample"
+            )
+
+    prepared = PreparedContext(
+        text=text,
+        decision=decision,
+        est_tokens_in=est,
+        est_tokens_out=estimate_tokens(text),
+        budget_tokens=budget,
+        chunks_total=total,
+        chunks_used=text.count("--- chunk "),
+        coverage=coverage,
+        task_kind=kind,
+    )
+    logger.info(
+        "context gate: %s (%s) %d -> %d est tokens, %d/%d chunks, %.0f%% coverage",
+        prepared.decision.value,
+        prepared.task_kind,
+        prepared.est_tokens_in,
+        prepared.est_tokens_out,
+        prepared.chunks_used,
+        prepared.chunks_total,
+        prepared.coverage * 100,
+    )
+    return prepared

apsimo/contextgate/retrieve.py

@@ -0,0 +1,115 @@
+"""Chunk ranking for the context gate.
+
+Embeddings when the caller provides an ``embed_fn`` (any async
+``list[str] -> list[list[float]]``); otherwise a dependency-free lexical
+TF-IDF cosine ranker, so the gate works even where no embedding stack is
+configured.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+import re
+from collections import Counter
+from typing import Awaitable, Callable, Optional, Sequence
+
+from apsimo.contextgate.chunker import Chunk
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["rank_chunks", "lexical_scores"]
+
+EmbedFn = Callable[[list[str]], Awaitable[list[list[float]]]]
+
+_WORD_RE = re.compile(r"[a-z0-9]+")
+
+# Minimal English stopword set — enough to stop function words from
+# dominating TF-IDF; deliberately tiny to stay language-tolerant.
+_STOPWORDS = frozenset(
+    "a an and are as at be but by for from has have if in into is it its of on "
+    "or that the their there these they this to was were what when where which "
+    "who will with you your".split()
+)
+
+
+def _terms(text: str) -> list[str]:
+    return [w for w in _WORD_RE.findall(text.lower()) if w not in _STOPWORDS]
+
+
+def lexical_scores(chunks: Sequence[Chunk], query: str) -> list[float]:
+    """TF-IDF cosine similarity of each chunk against *query* (0..1)."""
+    q_terms = _terms(query)
+    if not q_terms or not chunks:
+        return [0.0] * len(chunks)
+
+    chunk_terms = [_terms(c.text) for c in chunks]
+    n = len(chunks)
+    df: Counter[str] = Counter()
+    for terms in chunk_terms:
+        df.update(set(terms))
+
+    def _idf(term: str) -> float:
+        return math.log(1.0 + n / (1.0 + df.get(term, 0)))
+
+    q_tf = Counter(q_terms)
+    q_vec = {t: (1.0 + math.log(c)) * _idf(t) for t, c in q_tf.items()}
+    q_norm = math.sqrt(sum(v * v for v in q_vec.values())) or 1.0
+
+    scores: list[float] = []
+    for terms in chunk_terms:
+        tf = Counter(terms)
+        dot = 0.0
+        norm_sq = 0.0
+        for t, c in tf.items():
+            w = (1.0 + math.log(c)) * _idf(t)
+            norm_sq += w * w
+            qv = q_vec.get(t)
+            if qv:
+                dot += w * qv
+        norm = math.sqrt(norm_sq) or 1.0
+        scores.append(dot / (norm * q_norm))
+    return scores
+
+
+def _cosine(a: Sequence[float], b: Sequence[float]) -> float:
+    dot = sum(x * y for x, y in zip(a, b))
+    na = math.sqrt(sum(x * x for x in a)) or 1.0
+    nb = math.sqrt(sum(y * y for y in b)) or 1.0
+    return dot / (na * nb)
+
+
+async def rank_chunks(
+    chunks: Sequence[Chunk],
+    query: str,
+    embed_fn: Optional[EmbedFn] = None,
+) -> list[tuple[Chunk, float]]:
+    """Rank *chunks* by relevance to *query*, best first.
+
+    Uses *embed_fn* (batched: query + chunks in one call) when provided,
+    falling back to lexical TF-IDF on any embedding failure so ranking
+    never hard-fails.
+    """
+    if not chunks:
+        return []
+
+    if embed_fn is not None and query:
+        try:
+            vectors = await embed_fn([query] + [c.text for c in chunks])
+            if len(vectors) == len(chunks) + 1:
+                q_vec = vectors[0]
+                scored = [
+                    (c, _cosine(q_vec, v)) for c, v in zip(chunks, vectors[1:])
+                ]
+                scored.sort(key=lambda cs: cs[1], reverse=True)
+                return scored
+            logger.warning(
+                "embed_fn returned %d vectors for %d texts — falling back to lexical",
+                len(vectors), len(chunks) + 1,
+            )
+        except Exception:
+            logger.warning("embed_fn failed — falling back to lexical ranking", exc_info=True)
+
+    scored = list(zip(chunks, lexical_scores(chunks, query)))
+    scored.sort(key=lambda cs: cs[1], reverse=True)
+    return scored

apsimo/delivery/__init__.py

@@ -0,0 +1,16 @@
+"""Colony proactive delivery system.
+
+Bridges the autonomy loop's initiative/insight generation with the gateway's
+messaging adapters so Colony can proactively reach users when it has something
+worth saying.
+
+Components:
+- ProactiveDeliveryBridge: Queues and manages pending deliveries
+- RateLimiter: Per-person rate limiting (max 3/day, quiet hours, 2h cooldown)
+"""
+
+from apsimo.delivery.bridge import ProactiveDeliveryBridge
+from apsimo.delivery.rate_limiter import DeliveryRateLimiter
+from apsimo.delivery.channels import ChannelRegistry
+
+__all__ = ["ProactiveDeliveryBridge", "DeliveryRateLimiter", "ChannelRegistry"]