contextos-memory-runtime 1.0.0rc2__py3-none-any.whl

contextos/services/token_counter.py

@@ -0,0 +1,222 @@
+"""Token counting service for ContextOS.
+
+Supports exact provider usage, exact local tokenizers (tiktoken), model-family
+profiles (Claude-like, Qwen-like), and deterministic approximation fallbacks.
+Every token count is explicitly labeled with its measurement source:
+- PROVIDER_REPORTED
+- TOKENIZER_COUNTED
+- APPROXIMATED
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Protocol, runtime_checkable
+
+import tiktoken
+
+from contextos.core.enums import TokenMeasurementSource
+
+
+@runtime_checkable
+class TokenCounter(Protocol):
+    """Protocol for token counting."""
+
+    def count(self, text: str) -> int: ...
+
+    def count_batch(self, texts: list[str]) -> list[int]: ...
+
+    @property
+    def encoding_name(self) -> str: ...
+
+    @property
+    def measurement_source(self) -> TokenMeasurementSource: ...
+
+
+class TiktokenCounter:
+    """Token counter using tiktoken (OpenAI models).
+
+    Implements TokenCounter with TOKENIZER_COUNTED source.
+    """
+
+    def __init__(self, encoding_name: str = "cl100k_base") -> None:
+        self._encoding_name = encoding_name
+        self._encoder = tiktoken.get_encoding(encoding_name)
+
+    def count(self, text: str) -> int:
+        """Count tokens in text."""
+        if not text:
+            return 0
+        return len(self._encoder.encode(text))
+
+    def count_batch(self, texts: list[str]) -> list[int]:
+        """Count tokens for a batch of texts."""
+        return [self.count(t) for t in texts]
+
+    @property
+    def encoding_name(self) -> str:
+        return self._encoding_name
+
+    @property
+    def measurement_source(self) -> TokenMeasurementSource:
+        return TokenMeasurementSource.TOKENIZER_COUNTED
+
+
+class ClaudeProfileTokenCounter:
+    """Deterministic token counter profile modeled after Claude/Anthropic tokenization.
+
+    Uses Anthropic-like subword heuristics and whitespace/punctuation splitting.
+    Implements TokenCounter with APPROXIMATED source (calibrated synthetic profile).
+    """
+
+    _STD_PATTERN = re.compile(
+        r"[A-Za-z0-9]+(?:'[A-Za-z]+)?|[^\w\s]|\s+",
+    )
+
+    def __init__(self, model_name: str = "claude-profile") -> None:
+        self._model_name = model_name
+
+    def count(self, text: str) -> int:
+        if not text:
+            return 0
+        # Deterministic Claude-like subword accounting:
+        # Standard words, punctuation, and subword chunks for long technical terms
+        matches = self._STD_PATTERN.findall(text)
+        total = 0
+        for token in matches:
+            if not token.strip():
+                # Whitespace tokens
+                total += max(1, len(token) // 4)
+            elif len(token) > 6:
+                # Subword split for long tokens
+                total += (len(token) + 3) // 4
+            else:
+                total += 1
+        return total
+
+    def count_batch(self, texts: list[str]) -> list[int]:
+        return [self.count(t) for t in texts]
+
+    @property
+    def encoding_name(self) -> str:
+        return f"profile-{self._model_name}"
+
+    @property
+    def measurement_source(self) -> TokenMeasurementSource:
+        return TokenMeasurementSource.APPROXIMATED
+
+
+class QwenProfileTokenCounter:
+    """Deterministic token counter profile modeled after Qwen 152k vocabulary.
+
+    Qwen's vocabulary compresses technical terms and code more densely while
+    treating punctuation and camelCase identifiers with byte-level BPE splits.
+    Implements TokenCounter with APPROXIMATED source (calibrated synthetic profile).
+    """
+
+    _PATTERN = re.compile(
+        r"[A-Z]?[a-z]+|[A-Z]+(?=[A-Z][a-z]|\b)|[0-9]+|[^\w\s]|\s+",
+    )
+
+    def __init__(self, model_name: str = "qwen-profile") -> None:
+        self._model_name = model_name
+
+    def count(self, text: str) -> int:
+        if not text:
+            return 0
+        matches = self._PATTERN.findall(text)
+        total = 0
+        for token in matches:
+            if not token.strip():
+                total += max(1, len(token) // 5)
+            elif len(token) > 8:
+                total += (len(token) + 4) // 5
+            else:
+                total += 1
+        return total
+
+    def count_batch(self, texts: list[str]) -> list[int]:
+        return [self.count(t) for t in texts]
+
+    @property
+    def encoding_name(self) -> str:
+        return f"profile-{self._model_name}"
+
+    @property
+    def measurement_source(self) -> TokenMeasurementSource:
+        return TokenMeasurementSource.APPROXIMATED
+
+
+class DeterministicWordTokenCounter:
+    """Offline word-token approximation fallback.
+
+    Explicitly labeled with APPROXIMATED source.
+    """
+
+    _TOKEN_PATTERN = re.compile(
+        r"[A-Za-z0-9]+(?:(?:[+#._-]+)[A-Za-z0-9]+)*|[^\w\s]"
+    )
+
+    def count(self, text: str) -> int:
+        if not text:
+            return 0
+        return len(self._TOKEN_PATTERN.findall(text))
+
+    def count_batch(self, texts: list[str]) -> list[int]:
+        return [self.count(text) for text in texts]
+
+    @property
+    def encoding_name(self) -> str:
+        return "deterministic-word-approximation"
+
+    @property
+    def measurement_source(self) -> TokenMeasurementSource:
+        return TokenMeasurementSource.APPROXIMATED
+
+
+def get_token_counter_for_model(
+    model_id: str,
+    tokenizer_family: str | None = None,
+) -> TokenCounter:
+    """Select the appropriate token counter for a model ID or family."""
+    fam = (tokenizer_family or "").lower()
+    mid = model_id.lower()
+
+    # FakeProvider's default is an offline simulation, not a real BPE tokenizer.
+    if fam in {"deterministic", "approximation"} or mid == "fake-default":
+        return DeterministicWordTokenCounter()
+
+    if "claude" in fam or "claude" in mid or "anthropic" in fam:
+        return ClaudeProfileTokenCounter(model_name=model_id)
+
+    if "qwen" in fam or "qwen" in mid:
+        return QwenProfileTokenCounter(model_name=model_id)
+
+    if "cl100k" in fam or "o200k" in fam or "gpt" in mid or "openai" in fam:
+        encoding = "o200k_base" if "o200k" in fam or "omni" in mid or "gpt-4o" in mid else "cl100k_base"
+        try:
+            return TiktokenCounter(encoding_name=encoding)
+        except Exception:
+            return TiktokenCounter("cl100k_base")
+
+    # Default fallback: try tiktoken cl100k_base
+    try:
+        return TiktokenCounter("cl100k_base")
+    except Exception:
+        return DeterministicWordTokenCounter()
+
+
+def recount_cross_model(
+    text: str,
+    target_models_or_families: list[str],
+) -> dict[str, tuple[int, TokenMeasurementSource]]:
+    """Recount the SAME compiled context text across distinct model profiles.
+
+    Returns a dict mapping model_or_family to (token_count, measurement_source).
+    """
+    results: dict[str, tuple[int, TokenMeasurementSource]] = {}
+    for target in target_models_or_families:
+        counter = get_token_counter_for_model(target, target)
+        count = counter.count(text)
+        results[target] = (count, counter.measurement_source)
+    return results

contextos/storage/__init__.py

@@ -0,0 +1 @@
+"""Storage package for ContextOS."""

contextos/storage/connector_repo.py

@@ -0,0 +1,67 @@
+"""Persistence for connector cursors and source identity, never source bodies."""
+from __future__ import annotations
+import json
+from datetime import datetime, timezone
+from uuid import UUID
+import aiosqlite
+from contextos.connectors.models import ConnectorItem, ConnectorSyncState
+
+def _now() -> str: return datetime.now(timezone.utc).isoformat()
+
+class SqliteConnectorRepository:
+    def __init__(self, connection: aiosqlite.Connection) -> None: self._db = connection
+    async def state(self, connector_id: str) -> ConnectorSyncState | None:
+        row = await (await self._db.execute("SELECT * FROM connector_state WHERE connector_id=?", (connector_id,))).fetchone()
+        if row is None: return None
+        return ConnectorSyncState(connector_id=row["connector_id"], connector_type=row["connector_type"], cursor=row["cursor"], enabled=bool(row["enabled"]), status=row["status"], error_code=row["error_code"], last_success_at=row["last_success_at"], last_attempt_at=row["last_attempt_at"])
+    async def save_state(self, state: ConnectorSyncState) -> None:
+        await self._db.execute("INSERT INTO connector_state(connector_id,connector_type,cursor,last_success_at,last_attempt_at,enabled,status,error_code) VALUES(?,?,?,?,?,?,?,?) ON CONFLICT(connector_id) DO UPDATE SET cursor=excluded.cursor,last_success_at=excluded.last_success_at,last_attempt_at=excluded.last_attempt_at,enabled=excluded.enabled,status=excluded.status,error_code=excluded.error_code", (state.connector_id,state.connector_type,state.cursor,state.last_success_at.isoformat() if state.last_success_at else None,state.last_attempt_at.isoformat() if state.last_attempt_at else None,int(state.enabled),state.status,state.error_code)); await self._db.commit()
+    async def item_is_current(self, connector_id: str, item: ConnectorItem, content_hash: str) -> bool:
+        row = await (await self._db.execute("SELECT revision,content_hash,deleted FROM connector_items WHERE connector_id=? AND external_id=?", (connector_id,item.external_id))).fetchone()
+        return row is not None and not row["deleted"] and row["revision"] == item.revision and row["content_hash"] == content_hash
+    async def save_item(self, connector_id: str, item: ConnectorItem, content_hash: str, memory_ids: list[UUID], deleted: bool = False) -> None:
+        is_deleted = int(deleted or item.deleted)
+        await self._db.execute("INSERT INTO connector_items(connector_id,external_id,revision,content_hash,source_uri,memory_ids,last_seen_at,deleted) VALUES(?,?,?,?,?,?,?,?) ON CONFLICT(connector_id,external_id) DO UPDATE SET revision=excluded.revision,content_hash=excluded.content_hash,source_uri=excluded.source_uri,memory_ids=excluded.memory_ids,last_seen_at=excluded.last_seen_at,deleted=excluded.deleted", (connector_id,item.external_id,item.revision,content_hash,item.source_uri,json.dumps([str(value) for value in memory_ids]),_now(),is_deleted)); await self._db.commit()
+
+    async def get_item_memory_ids(self, connector_id: str, external_id: str) -> list[UUID]:
+        row = await (await self._db.execute("SELECT memory_ids FROM connector_items WHERE connector_id=? AND external_id=?", (connector_id, external_id))).fetchone()
+        if row is None or not row["memory_ids"]:
+            return []
+        try:
+            raw = json.loads(row["memory_ids"])
+            return [UUID(val) for val in raw]
+        except Exception:
+            return []
+
+    async def count_active_references(self, memory_id: UUID) -> int:
+        target = str(memory_id)
+        cursor = await self._db.execute("SELECT memory_ids FROM connector_items WHERE deleted = 0")
+        rows = await cursor.fetchall()
+        count = 0
+        for row in rows:
+            if not row["memory_ids"]:
+                continue
+            try:
+                raw = json.loads(row["memory_ids"])
+                if target in raw:
+                    count += 1
+            except Exception:
+                pass
+        return count
+
+    async def list_states(self) -> list[ConnectorSyncState]:
+        cursor = await self._db.execute("SELECT * FROM connector_state ORDER BY connector_id")
+        rows = await cursor.fetchall()
+        return [
+            ConnectorSyncState(
+                connector_id=row["connector_id"],
+                connector_type=row["connector_type"],
+                cursor=row["cursor"],
+                enabled=bool(row["enabled"]),
+                status=row["status"],
+                error_code=row["error_code"],
+                last_success_at=row["last_success_at"],
+                last_attempt_at=row["last_attempt_at"],
+            )
+            for row in rows
+        ]