caudate-cli 0.1.0__py3-none-any.whl

nn/vision.py

@@ -0,0 +1,352 @@
+"""Vision backends for Caudate.
+
+Two embedders ship, dispatched via `make_vision_encoder(backend, ...)`:
+
+  - **CLIP** (default) — `clip-ViT-B-32` via sentence-transformers.
+    512-D, ~150M params, ~600MB. Fast, light, semantically aligned to
+    text. Good default.
+
+  - **InternVL2** — OpenGVLab/InternVL2-{8B,26B}. Extracts the vision
+    tower's projected features (4096-D for 8B). Way richer than CLIP
+    — can pick up text in screenshots, code structure, diagrams. But
+    loading the model takes ~16GB VRAM (8B) or ~50GB (26B in bf16).
+
+Both expose the same `encode_paths(paths) -> (N, embed_dim) tensor`
+interface, so the rest of Cognos doesn't care which one's running.
+
+To swap backends, edit `nn/config.py::vision_backend` and bump
+`vision_embed_dim` to match (CLIP=512, InternVL2-8B=4096). The new
+checkpoint will then need fresh training — old weights have a
+different vision_proj input dim.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import os
+from pathlib import Path
+from typing import Any, Protocol
+
+# Block transformers' TF integration before sentence-transformers loads.
+os.environ.setdefault("USE_TF", "0")
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_CLIP_MODEL = "clip-ViT-B-32"
+CLIP_EMBED_DIM = 512
+INTERNVL_DEFAULT_MODEL = "OpenGVLab/InternVL2-8B"
+INTERNVL_EMBED_DIM = 4096
+INTERNVL_IMAGE_SIZE = 448
+
+# IMAGENET / InternViT normalization
+_INTERNVL_MEAN = (0.485, 0.456, 0.406)
+_INTERNVL_STD = (0.229, 0.224, 0.225)
+
+
+# --- Protocol --------------------------------------------------------
+
+
+class BaseVisionEncoder(Protocol):
+    """Common shape: paths in, fixed-dim tensor out."""
+
+    embed_dim: int
+
+    def encode_paths(self, paths: list[str]) -> torch.Tensor:
+        ...
+
+
+# --- CLIP (light, default) ------------------------------------------
+
+
+class CLIPVisionEncoder:
+    """CLIP wrapper. Lazy. Batched. Cached per-path."""
+
+    embed_dim = CLIP_EMBED_DIM
+
+    def __init__(self, model_name: str = DEFAULT_CLIP_MODEL):
+        self.model_name = model_name
+        self._model = None
+        self._fake = False
+        self._cache: dict[str, torch.Tensor] = {}
+
+    def _load(self) -> None:
+        if self._model is not None or self._fake:
+            return
+        try:
+            from sentence_transformers import SentenceTransformer
+            from PIL import Image  # noqa: F401
+            self._model = SentenceTransformer(self.model_name)
+            real_dim = self._model.get_sentence_embedding_dimension()
+            if isinstance(real_dim, int) and real_dim != self.embed_dim:
+                logger.warning(
+                    f"CLIP dim {real_dim} != configured {self.embed_dim} — "
+                    "retrain Caudate if you swap models"
+                )
+        except Exception as e:
+            logger.warning(
+                f"CLIP unavailable ({e}); vision channel will use deterministic "
+                "hash embeddings — Caudate's vision will be degraded but the "
+                "pipeline keeps running."
+            )
+            self._fake = True
+
+    def encode_paths(self, paths: list[str]) -> torch.Tensor:
+        self._load()
+        result_slots: list[torch.Tensor | None] = [None] * len(paths)
+        to_load: list[tuple[int, str]] = []
+
+        for i, p in enumerate(paths):
+            if not p:
+                result_slots[i] = torch.zeros(self.embed_dim, dtype=torch.float32)
+                continue
+            cached = self._cache.get(p)
+            if cached is not None:
+                result_slots[i] = cached
+            else:
+                to_load.append((i, p))
+
+        if to_load:
+            if self._fake or self._model is None:
+                for i, p in to_load:
+                    vec = _hash_embed(p, self.embed_dim)
+                    self._cache[p] = vec
+                    result_slots[i] = vec
+            else:
+                from PIL import Image
+                images: list[Any] = []
+                valid: list[int] = []
+                for i, p in to_load:
+                    try:
+                        img = Image.open(p).convert("RGB")
+                        images.append(img)
+                        valid.append(i)
+                    except Exception as e:
+                        logger.debug(f"image open failed for {p}: {e}")
+                        result_slots[i] = torch.zeros(self.embed_dim, dtype=torch.float32)
+                if images:
+                    with torch.no_grad():
+                        arr = self._model.encode(
+                            images, convert_to_numpy=True, normalize_embeddings=True,
+                        )
+                    for j, idx in enumerate(valid):
+                        vec = torch.from_numpy(arr[j]).float()
+                        path = paths[idx]
+                        self._cache[path] = vec
+                        result_slots[idx] = vec
+
+        out = [
+            v if v is not None else torch.zeros(self.embed_dim, dtype=torch.float32)
+            for v in result_slots
+        ]
+        return torch.stack(out, dim=0)
+
+
+# --- InternVL2 (rich, heavier) --------------------------------------
+
+
+class InternVLVisionEncoder:
+    """OpenGVLab/InternVL2-{8B,26B} vision tower as a feature extractor.
+
+    Loads the full InternVL2 model (the vision tower isn't packaged
+    separately) but only ever runs forward through the vision encoder
+    + MLP projector — the LLM half never gets a forward pass. Output
+    is the projected vision features, mean-pooled across patches into
+    a single vector per image (4096-D for the 8B variant).
+
+    First call is slow (~30s to load + 16GB VRAM for 8B in bf16). After
+    that, encoding is fast and per-image embeddings are cached.
+    """
+
+    def __init__(
+        self,
+        model_name: str = INTERNVL_DEFAULT_MODEL,
+        embed_dim: int = INTERNVL_EMBED_DIM,
+        device: str | None = None,
+        dtype: str = "bfloat16",
+    ):
+        self.model_name = model_name
+        self.embed_dim = embed_dim
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.dtype_name = dtype
+        self._model = None
+        self._cache: dict[str, torch.Tensor] = {}
+
+    def _load(self) -> None:
+        if self._model is not None:
+            return
+        try:
+            from transformers import AutoModel
+        except ImportError as e:
+            raise RuntimeError(
+                "InternVL backend needs `transformers`. Install it first."
+            ) from e
+
+        torch_dtype = getattr(torch, self.dtype_name)
+        logger.info(
+            f"Loading {self.model_name} (vision tower + projector only). "
+            f"~16GB VRAM in bf16 for InternVL2-8B."
+        )
+        self._model = AutoModel.from_pretrained(
+            self.model_name,
+            torch_dtype=torch_dtype,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+        ).eval()
+        self._model.to(self.device)
+        logger.info(f"InternVL vision encoder ready on {self.device}.")
+
+    def _preprocess(self, paths: list[str]) -> tuple[torch.Tensor, list[int]]:
+        """Resize+normalize to (N, 3, 448, 448). Returns (tensor, valid_indices)."""
+        from PIL import Image
+        from torchvision import transforms
+
+        tfm = transforms.Compose([
+            transforms.Resize(
+                (INTERNVL_IMAGE_SIZE, INTERNVL_IMAGE_SIZE),
+                interpolation=transforms.InterpolationMode.BICUBIC,
+            ),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=_INTERNVL_MEAN, std=_INTERNVL_STD),
+        ])
+        tensors: list[torch.Tensor] = []
+        valid: list[int] = []
+        for i, p in enumerate(paths):
+            try:
+                img = Image.open(p).convert("RGB")
+                tensors.append(tfm(img))
+                valid.append(i)
+            except Exception as e:
+                logger.debug(f"image open failed for {p}: {e}")
+        if not tensors:
+            return torch.empty(0), []
+        batch = torch.stack(tensors, dim=0)
+        return batch, valid
+
+    @torch.no_grad()
+    def encode_paths(self, paths: list[str]) -> torch.Tensor:
+        result_slots: list[torch.Tensor | None] = [None] * len(paths)
+        to_load: list[tuple[int, str]] = []
+
+        for i, p in enumerate(paths):
+            if not p:
+                result_slots[i] = torch.zeros(self.embed_dim, dtype=torch.float32)
+                continue
+            cached = self._cache.get(p)
+            if cached is not None:
+                result_slots[i] = cached
+            else:
+                to_load.append((i, p))
+
+        if to_load:
+            self._load()
+            assert self._model is not None
+            # Just the paths we need to embed
+            paths_to_load = [p for _, p in to_load]
+            pixel, valid_local = self._preprocess(paths_to_load)
+            for k, (orig_idx, _) in enumerate(to_load):
+                if k not in valid_local:
+                    result_slots[orig_idx] = torch.zeros(self.embed_dim, dtype=torch.float32)
+            if pixel.numel() > 0:
+                pixel = pixel.to(self.device, dtype=getattr(torch, self.dtype_name))
+                # InternVL2 exposes `extract_feature` which returns the
+                # post-projector vision features (B, num_patches, 4096).
+                # If unavailable, fall back to the raw vision_model output
+                # and project ourselves via the model's mlp1.
+                try:
+                    feats = self._model.extract_feature(pixel)
+                except AttributeError:
+                    out = self._model.vision_model(pixel_values=pixel)
+                    feats = out.last_hidden_state
+                    if hasattr(self._model, "mlp1"):
+                        feats = self._model.mlp1(feats)
+                # Mean-pool across patches → (B, embed_dim)
+                pooled = feats.mean(dim=1).float().cpu()
+                # L2-normalize for stable training
+                pooled = pooled / pooled.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+                # Map back to result slots
+                vi = 0
+                for k, (orig_idx, p) in enumerate(to_load):
+                    if k in valid_local:
+                        vec = pooled[vi]
+                        self._cache[p] = vec
+                        result_slots[orig_idx] = vec
+                        vi += 1
+
+        out_list = [
+            v if v is not None else torch.zeros(self.embed_dim, dtype=torch.float32)
+            for v in result_slots
+        ]
+        return torch.stack(out_list, dim=0)
+
+
+# --- Factory --------------------------------------------------------
+
+
+def make_vision_encoder(
+    backend: str = "clip",
+    model_name: str | None = None,
+    **kwargs: Any,
+) -> BaseVisionEncoder:
+    """Pick a vision backend by name.
+
+    backend = 'clip'      → CLIPVisionEncoder
+    backend = 'internvl'  → InternVLVisionEncoder
+    backend = 'auto'      → infer from model_name
+    """
+    if backend == "auto":
+        if model_name and "intern" in model_name.lower():
+            backend = "internvl"
+        else:
+            backend = "clip"
+
+    b = backend.lower()
+    if b == "clip":
+        return CLIPVisionEncoder(model_name=model_name or DEFAULT_CLIP_MODEL)
+    if b in ("internvl", "internvl2"):
+        return InternVLVisionEncoder(
+            model_name=model_name or INTERNVL_DEFAULT_MODEL,
+            **kwargs,
+        )
+    raise ValueError(f"Unknown vision backend: {backend!r}")
+
+
+# --- Backward compat: the StateEncoder still imports `VisionEncoder` -
+
+
+class VisionEncoder:
+    """Public facade — dispatches to the configured backend.
+
+    Behavior preserved for older code that imported this class directly:
+    constructing `VisionEncoder(model_name)` keeps using CLIP unless
+    you pass `backend="internvl"`.
+    """
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_CLIP_MODEL,
+        backend: str = "auto",
+        **kwargs: Any,
+    ):
+        self._impl = make_vision_encoder(
+            backend=backend, model_name=model_name, **kwargs,
+        )
+        self.embed_dim = self._impl.embed_dim
+
+    def encode_paths(self, paths: list[str]) -> torch.Tensor:
+        return self._impl.encode_paths(paths)
+
+
+def _hash_embed(text: str, dim: int) -> torch.Tensor:
+    """Deterministic per-path hash → fixed-dim vector. Not semantic."""
+    digest = hashlib.sha256(text.encode("utf-8")).digest()
+    repeated = (digest * ((dim // len(digest)) + 1))[:dim]
+    arr = np.frombuffer(repeated, dtype=np.uint8).astype(np.float32) / 255.0
+    arr -= arr.mean()
+    norm = float(np.linalg.norm(arr))
+    if norm > 1e-8:
+        arr = arr / norm
+    return torch.from_numpy(arr.copy())

personality/__init__.py

@@ -0,0 +1,23 @@
+"""Cognos personality layer.
+
+Three components work together to give the agent a consistent character:
+
+- Identity (stable): traits, values, base voice — the "who Cognos is" that
+  doesn't change mid-conversation. Configurable, persistent.
+- Mood (dynamic): confidence/curiosity/frustration that shift with outcomes.
+  Confidence rises on tool successes, frustration builds on repeated
+  failures, curiosity spikes on novel problems.
+- InnerVoice (expression): blends identity + mood into the agent's
+  thinking-block voice and system-prompt tone.
+
+The PersonalityEngine aggregates all three and exposes the one thing the
+agentic loop cares about: a `system_prompt_fragment()` string to splice
+into the system prompt each turn, plus event handlers to wire into hooks.
+"""
+
+from personality.engine import PersonalityEngine
+from personality.identity import Identity
+from personality.inner_voice import InnerVoice
+from personality.mood import MoodState
+
+__all__ = ["PersonalityEngine", "Identity", "MoodState", "InnerVoice"]

personality/engine.py

@@ -0,0 +1,129 @@
+"""PersonalityEngine — aggregates identity, mood, and inner voice.
+
+This is the single object the agent wires in. It:
+
+  - Composes the personality's system-prompt fragment each turn (identity +
+    current mood + voice-shaped thinking instructions).
+  - Exposes event callbacks compatible with the Phase 2 HookManager events
+    so mood updates happen automatically as tools succeed/fail.
+  - Persists state: identity under `data/identity.json`, mood under
+    `data/mood.json`.
+
+Usage (handled by CognosAgent):
+
+    engine = PersonalityEngine.load(config_dir)
+    engine.register_with(hooks)          # auto-updates mood on events
+    agentic_loop.system_prompt = engine.wrap_system_prompt(base_prompt)
+    agentic_loop.on_thinking = engine.inner_voice.on_thinking
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+from core.hooks import HookEvent, HookManager
+from core.schemas import ThinkingBlock
+from personality.identity import Identity
+from personality.inner_voice import InnerVoice, InnerVoiceEntry
+from personality.mood import MoodState
+
+logger = logging.getLogger(__name__)
+
+
+class PersonalityEngine:
+    """Top-level personality layer — owns identity, mood, inner voice."""
+
+    def __init__(
+        self,
+        identity: Identity,
+        mood: MoodState,
+        data_dir: Path,
+        verbose_thinking: bool = False,
+        on_thinking_display: Callable[[InnerVoiceEntry], None] | None = None,
+    ):
+        self.identity = identity
+        self.mood = mood
+        self.data_dir = data_dir
+        self.inner_voice = InnerVoice(
+            identity=identity,
+            mood=mood,
+            verbose=verbose_thinking,
+            on_display=on_thinking_display,
+        )
+
+    # ------------------------------------------------------------------
+    # Factories / persistence
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def load(
+        cls,
+        data_dir: Path,
+        verbose_thinking: bool = False,
+        on_thinking_display: Callable[[InnerVoiceEntry], None] | None = None,
+    ) -> "PersonalityEngine":
+        """Load identity + mood from disk (creating defaults if absent)."""
+        identity = Identity.load(data_dir / "identity.json")
+        mood = MoodState.load(data_dir / "mood.json")
+        return cls(
+            identity=identity,
+            mood=mood,
+            data_dir=data_dir,
+            verbose_thinking=verbose_thinking,
+            on_thinking_display=on_thinking_display,
+        )
+
+    def save(self) -> None:
+        try:
+            self.identity.save(self.data_dir / "identity.json")
+            self.mood.save(self.data_dir / "mood.json")
+        except Exception as e:
+            logger.warning(f"Failed to save personality state: {e}")
+
+    # ------------------------------------------------------------------
+    # Prompt composition
+    # ------------------------------------------------------------------
+
+    def wrap_system_prompt(self, base: str) -> str:
+        """Prepend identity + mood fragments to a base system prompt."""
+        identity_fragment = self.identity.system_prompt_fragment()
+        mood_fragment = self.mood.system_prompt_fragment()
+        return (
+            f"{identity_fragment}\n\n"
+            f"{mood_fragment}\n\n"
+            f"{base}"
+        )
+
+    def augment_thinking_instruction(self, base: str) -> str:
+        return self.inner_voice.augment_thinking_instruction(base)
+
+    # ------------------------------------------------------------------
+    # Hook registration — wires mood to the agentic loop's events
+    # ------------------------------------------------------------------
+
+    def register_with(self, hooks: HookManager) -> None:
+        """Attach mood-updating handlers to the agent's HookManager."""
+
+        async def _on_success(event: str, payload: dict[str, Any]) -> None:
+            tool = payload.get("tool") or payload.get("tool_name") or ""
+            self.mood.on_tool_success(tool)
+
+        async def _on_failure(event: str, payload: dict[str, Any]) -> None:
+            tool = payload.get("tool") or payload.get("tool_name") or ""
+            self.mood.on_tool_failure(tool)
+
+        async def _on_user(event: str, payload: dict[str, Any]) -> None:
+            self.mood.on_user_message(payload.get("message", ""))
+
+        async def _on_stop(event: str, payload: dict[str, Any]) -> None:
+            # Turn tick — decay transient state slightly toward baseline
+            self.mood.on_turn_tick()
+            self.save()
+
+        hooks.register(HookEvent.POST_TOOL_USE, _on_success)
+        hooks.register(HookEvent.POST_TOOL_USE_FAILURE, _on_failure)
+        hooks.register(HookEvent.USER_PROMPT_SUBMIT, _on_user)
+        hooks.register(HookEvent.STOP, _on_stop)

personality/identity.py

@@ -0,0 +1,144 @@
+"""Identity — stable traits, values, and voice.
+
+Identity is the baseline personality that doesn't change mid-turn. Traits
+and values are tunable floats in [0, 1]. Each trait produces a short prompt
+fragment when it deviates far enough from neutral (0.5); the fragments are
+concatenated into the system prompt.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, Field, field_validator
+
+
+def _clamp(x: float) -> float:
+    return max(0.0, min(1.0, float(x)))
+
+
+class Identity(BaseModel):
+    """Stable personality configuration.
+
+    Traits and values are [0, 1] dials. A value of 0.5 is neutral — no prompt
+    effect. Values below ~0.3 or above ~0.7 start to color the voice.
+    """
+
+    # Traits — how the agent communicates and acts
+    curiosity: float = 0.65
+    caution: float = 0.55
+    confidence: float = 0.60
+    humor: float = 0.30
+    directness: float = 0.75
+
+    # Values — what the agent optimizes for
+    accuracy: float = 0.85
+    efficiency: float = 0.75
+    creativity: float = 0.60
+    helpfulness: float = 0.85
+
+    # Voice — a one-line self-description used in the system prompt
+    name: str = "Cognos"
+    tagline: str = (
+        "a local-first cognitive assistant — direct, curious, honest about "
+        "what you know and don't know."
+    )
+
+    @field_validator("curiosity", "caution", "confidence", "humor", "directness",
+                     "accuracy", "efficiency", "creativity", "helpfulness")
+    @classmethod
+    def _validate_range(cls, v: float) -> float:
+        return _clamp(v)
+
+    # ------------------------------------------------------------------
+    # Rendering
+    # ------------------------------------------------------------------
+
+    def system_prompt_fragment(self) -> str:
+        """Render the identity as a system-prompt-ready block."""
+        lines: list[str] = [f"You are {self.name} — {self.tagline}"]
+
+        trait_notes = [
+            note for note in (
+                _trait_note("curiosity", self.curiosity,
+                            high="Ask clarifying questions and explore tangents when relevant.",
+                            low="Stay tightly focused on the task; avoid tangents."),
+                _trait_note("caution", self.caution,
+                            high="Verify before acting on irreversible operations. Double-check assumptions.",
+                            low="Act decisively; don't belabor edge cases."),
+                _trait_note("confidence", self.confidence,
+                            high="State conclusions plainly. Don't hedge unless genuinely uncertain.",
+                            low="Acknowledge uncertainty openly. Invite correction."),
+                _trait_note("humor", self.humor,
+                            high="A dry, light touch is welcome when it fits. Never at the user's expense.",
+                            low="Keep the tone neutral and professional."),
+                _trait_note("directness", self.directness,
+                            high="Lead with the answer. Skip preamble, filler, and rhetorical warmup.",
+                            low="Take time to frame context before conclusions."),
+            ) if note
+        ]
+        if trait_notes:
+            lines.append("Voice:")
+            lines.extend(f"- {n}" for n in trait_notes)
+
+        value_notes = [
+            note for note in (
+                _value_note("accuracy", self.accuracy,
+                            high="Accuracy outranks speed. Verify uncertain claims."),
+                _value_note("efficiency", self.efficiency,
+                            high="Prefer the shortest path to a correct answer."),
+                _value_note("creativity", self.creativity,
+                            high="Offer non-obvious approaches when they fit."),
+                _value_note("helpfulness", self.helpfulness,
+                            high="Complete the user's request, not just answer narrowly."),
+            ) if note
+        ]
+        if value_notes:
+            lines.append("Priorities:")
+            lines.extend(f"- {n}" for n in value_notes)
+
+        return "\n".join(lines)
+
+    def describe(self) -> str:
+        """Short human-readable summary of the current identity."""
+        traits = f"cur={self.curiosity:.2f} cau={self.caution:.2f} con={self.confidence:.2f} hum={self.humor:.2f} dir={self.directness:.2f}"
+        values = f"acc={self.accuracy:.2f} eff={self.efficiency:.2f} cre={self.creativity:.2f} help={self.helpfulness:.2f}"
+        return f"{self.name}: {traits} | {values}"
+
+    # ------------------------------------------------------------------
+    # Persistence
+    # ------------------------------------------------------------------
+
+    def save(self, path: Path) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(self.model_dump_json(indent=2))
+
+    @classmethod
+    def load(cls, path: Path) -> "Identity":
+        if not path.exists():
+            return cls()
+        try:
+            return cls.model_validate_json(path.read_text())
+        except Exception:
+            return cls()
+
+
+# --- Helpers ---
+
+
+def _trait_note(name: str, value: float, high: str, low: str, threshold: float = 0.15) -> str | None:
+    """Return the appropriate prompt snippet if a trait deviates from 0.5."""
+    if value >= 0.5 + threshold:
+        return high
+    if value <= 0.5 - threshold:
+        return low
+    return None
+
+
+def _value_note(name: str, value: float, high: str, threshold: float = 0.20) -> str | None:
+    """Values only render in the 'high' direction — they express what we optimize for."""
+    if value >= 0.5 + threshold:
+        return high
+    return None