diff --git a/src/mnemosyne/admission.py b/src/mnemosyne/admission.py new file mode 100644 index 0000000..f046c3b --- /dev/null +++ b/src/mnemosyne/admission.py @@ -0,0 +1,146 @@ +"""Admission control for incoming semantic objects. + +Scores incoming content on type/novelty/utility/recency and rejects +objects below a threshold to prevent the ObjectStore from filling +with low-value content. + +Phase 4d of the Mnemosyne context manager. +See ARCHITECTURE.md §5.3 for the admission control specification. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class AdmissionScore: + """Admission score breakdown. + + Each component is 0-1, with total being a weighted combination. + """ + + type_score: float # 0-1, based on object type value + novelty_score: float # 0-1, 1.0 if no duplicate exists + utility_score: float # 0-1, based on content quality signals + size_score: float # 0-1, penalize very small or very large objects + total: float # weighted combination + + +# Type value weights (higher = more valuable to retain in context) +TYPE_WEIGHTS: dict[str, float] = { + "design_decision": 0.9, + "debugging_session": 0.8, + "file_context": 0.7, + "plan": 0.7, + "error_context": 0.6, + "tool_result": 0.5, + "external_reference": 0.4, + "conversation_phase": 0.3, +} + +# Objects scoring below this threshold are rejected +DEFAULT_THRESHOLD = 0.35 + + +class AdmissionController: + """Gate-keeps incoming semantic objects by scoring content value. + + Scores each incoming object on four axes: + - type_score: inherent value of the object type + - novelty_score: penalizes duplicates heavily + - utility_score: heuristic quality signals (code, entities, length) + - size_score: penalizes very small or very large objects + + Objects below the threshold are rejected and never stored. + """ + + def __init__(self, threshold: float = DEFAULT_THRESHOLD) -> None: + self.threshold = threshold + self.stats: dict[str, int] = {"admitted": 0, "rejected": 0} + + def score( + self, + content: str, + object_type: str, + has_duplicate: bool, + key_entities: list[str] | None = None, + ) -> AdmissionScore: + """Compute admission score for a candidate object. + + Args: + content: Full text content of the object. + object_type: Semantic object type (e.g. 'file_context'). + has_duplicate: Whether a duplicate source_key already exists. + key_entities: Extracted entities (file paths, function names, etc.). + + Returns: + AdmissionScore with component breakdown and weighted total. + """ + if key_entities is None: + key_entities = [] + + # Type score: known types get their weight, unknown defaults to 0.3 + type_score = TYPE_WEIGHTS.get(object_type, 0.3) + + # Novelty: duplicates are heavily penalized but not zero + novelty_score = 0.2 if has_duplicate else 1.0 + + # Utility: count code blocks, file paths, function names, length + utility = 0.0 + if "```" in content: + utility += 0.3 + if "def " in content or "function " in content: + utility += 0.2 + if "/" in content and "." in content: # file paths heuristic + utility += 0.1 + if len(key_entities) > 2: + utility += 0.2 + if len(content) > 200: + utility += 0.2 + utility_score = min(1.0, utility) + + # Size: penalize very small (< 100 chars) or very large (> 50k chars) + length = len(content) + if length < 100: + size_score = length / 100 if length > 0 else 0.0 + elif length > 50000: + size_score = max(0.3, 1.0 - (length - 50000) / 100000) + else: + size_score = 1.0 + + total = type_score * 0.3 + novelty_score * 0.3 + utility_score * 0.25 + size_score * 0.15 + + return AdmissionScore( + type_score=type_score, + novelty_score=novelty_score, + utility_score=utility_score, + size_score=size_score, + total=total, + ) + + def should_admit( + self, + content: str, + object_type: str, + has_duplicate: bool, + key_entities: list[str] | None = None, + ) -> tuple[bool, AdmissionScore]: + """Decide whether to admit an object into the ObjectStore. + + Args: + content: Full text content of the object. + object_type: Semantic object type. + has_duplicate: Whether a duplicate source_key already exists. + key_entities: Extracted entities. + + Returns: + (admitted, score) tuple. admitted is True if score >= threshold. + """ + score = self.score(content, object_type, has_duplicate, key_entities) + admitted = score.total >= self.threshold + if admitted: + self.stats["admitted"] += 1 + else: + self.stats["rejected"] += 1 + return admitted, score diff --git a/src/mnemosyne/entropy.py b/src/mnemosyne/entropy.py new file mode 100644 index 0000000..08a192d --- /dev/null +++ b/src/mnemosyne/entropy.py @@ -0,0 +1,156 @@ +"""Entropy-gated faulting — detect when model might need evicted content. + +Analyzes assistant responses for hedging language, uncertainty markers, +and references to evicted entities. When uncertainty is high and the +response references evicted content, proactively triggers micro-fault +retrieval so the content is available on the next turn. + +Phase 4e of the Mnemosyne context manager. +See ARCHITECTURE.md §7.4 for the entropy-gated faulting specification. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field + + +@dataclass +class EntropySignal: + """Entropy analysis of an assistant response. + + Higher score means the model is more likely uncertain and may + need evicted content to produce a better response. + """ + + has_hedging: bool # "I think", "probably", "might be" + has_uncertainty: bool # "I'm not sure", "unclear", "not certain" + has_hallucination_risk: bool # fabricated file paths, wrong function names + references_evicted: bool # mentions content that was evicted + referenced_entities: list[str] # entities from evicted content found in response + score: float # 0-1, higher = more likely needs evicted content + + +# Hedging language patterns — model is guessing rather than knowing +HEDGING_PATTERNS: list[re.Pattern[str]] = [ + re.compile(r"\bI think\b", re.IGNORECASE), + re.compile(r"\bprobably\b", re.IGNORECASE), + re.compile(r"\bmight be\b", re.IGNORECASE), + re.compile(r"\bpossibly\b", re.IGNORECASE), + re.compile(r"\bif I recall\b", re.IGNORECASE), + re.compile(r"\bI believe\b", re.IGNORECASE), +] + +# Explicit uncertainty — model knows it doesn't know +UNCERTAINTY_PATTERNS: list[re.Pattern[str]] = [ + re.compile(r"\bI'm not sure\b", re.IGNORECASE), + re.compile(r"\bunclear\b", re.IGNORECASE), + re.compile(r"\bnot certain\b", re.IGNORECASE), + re.compile(r"\bI don't remember\b", re.IGNORECASE), + re.compile(r"\bI can't recall\b", re.IGNORECASE), +] + +# Threshold above which we trigger proactive faulting +DEFAULT_FAULT_THRESHOLD = 0.6 + + +class EntropyDetector: + """Detects model uncertainty from response text and evicted entity references. + + Used in the post-response path to decide whether to proactively + fault evicted content back in for the next turn. + """ + + def __init__(self, fault_threshold: float = DEFAULT_FAULT_THRESHOLD) -> None: + self.fault_threshold = fault_threshold + self.stats: dict[str, int] = {"analyzed": 0, "faults_triggered": 0} + + def analyze_response( + self, + response_text: str, + evicted_entities: list[str] | None = None, + ) -> EntropySignal: + """Analyze an assistant response for uncertainty signals. + + Args: + response_text: The assistant's response text. + evicted_entities: Key entities from evicted (L4) objects. + If the response references these, it may be hallucinating + or working from degraded context. + + Returns: + EntropySignal with component flags and composite score. + """ + if evicted_entities is None: + evicted_entities = [] + + self.stats["analyzed"] += 1 + + # Check hedging patterns + hedging_count = sum(1 for p in HEDGING_PATTERNS if p.search(response_text)) + has_hedging = hedging_count > 0 + + # Check uncertainty patterns + uncertainty_count = sum(1 for p in UNCERTAINTY_PATTERNS if p.search(response_text)) + has_uncertainty = uncertainty_count > 0 + + # Check for references to evicted entities + referenced: list[str] = [] + for entity in evicted_entities: + if not entity: + continue + # Match entity as a word or path component in the response + # Use re.escape to handle file paths with special chars + pattern = re.escape(entity) + if re.search(pattern, response_text): + referenced.append(entity) + + references_evicted = len(referenced) > 0 + + # Hallucination risk: hedging + referencing evicted content + has_hallucination_risk = has_hedging and references_evicted + + # Composite score: weighted combination + score = 0.0 + if has_hedging: + score += min(0.3, hedging_count * 0.1) + if has_uncertainty: + score += min(0.3, uncertainty_count * 0.15) + if references_evicted: + score += min(0.4, len(referenced) * 0.15) + if has_hallucination_risk: + score += 0.1 # bonus for combined signals + + score = min(1.0, score) + + return EntropySignal( + has_hedging=has_hedging, + has_uncertainty=has_uncertainty, + has_hallucination_risk=has_hallucination_risk, + references_evicted=references_evicted, + referenced_entities=referenced, + score=score, + ) + + def should_fault(self, signal: EntropySignal) -> list[str]: + """Decide which entities should be faulted back in. + + Returns a list of entity names/paths that should be retrieved + from evicted storage. Only triggers if the entropy score exceeds + the fault threshold AND the response references evicted content. + + Args: + signal: EntropySignal from analyze_response(). + + Returns: + List of entity identifiers to fault back in. Empty if no + faulting is needed. + """ + if signal.score < self.fault_threshold: + return [] + + if not signal.references_evicted: + return [] + + self.stats["faults_triggered"] += 1 + return list(signal.referenced_entities) diff --git a/src/mnemosyne/hierarchy.py b/src/mnemosyne/hierarchy.py new file mode 100644 index 0000000..daa5e4c --- /dev/null +++ b/src/mnemosyne/hierarchy.py @@ -0,0 +1,722 @@ +"""Hierarchical segmentation for long sessions (Strategy B). + +Extends the rule-based Segmenter (Strategy A, turns 1-50) with +embedding-based hierarchical clustering for turns 50+. Implements +the xMemory-style 4-level hierarchy: + + Level 0: Individual messages/tool results (existing StoredObjects) + Level 1: Episodes — groups of related objects from coherence clustering + Level 2: Semantics — abstract themes spanning multiple episodes + Level 3: Themes — top-level categories for the entire session + +Key advantages over flat retrieval: + - Top-down search prevents redundant retrieval + - Hierarchical gating reduces search space at each level + - Periodic re-clustering adapts to evolving session structure + - Incremental add_object avoids full rebuild on every turn + +Uses numpy for all vector operations. No external clustering libraries. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass, field +from typing import Protocol, runtime_checkable + +import numpy as np + + +# ── Embedder protocol (re-exported for convenience) ───────────────────── + + +@runtime_checkable +class Embedder(Protocol): + """Protocol for embedding providers (mirrors object_store.Embedder).""" + + def embed(self, text: str) -> list[float]: ... + + def embed_batch(self, texts: list[str]) -> list[list[float]]: ... + + +# ── Vector math utilities ─────────────────────────────────────────────── + + +def _cosine_similarity(a: list[float], b: list[float]) -> float: + """Compute cosine similarity between two vectors. + + Both vectors are expected to be L2-normalized (from the embedder), + so this is just a dot product. Falls back to full formula for safety. + """ + if not a or not b: + return 0.0 + va = np.asarray(a, dtype=np.float64) + vb = np.asarray(b, dtype=np.float64) + dot = float(np.dot(va, vb)) + norm_a = float(np.linalg.norm(va)) + norm_b = float(np.linalg.norm(vb)) + if norm_a == 0.0 or norm_b == 0.0: + return 0.0 + return dot / (norm_a * norm_b) + + +def _cosine_similarity_matrix(embeddings: np.ndarray) -> np.ndarray: + """Compute pairwise cosine similarity matrix for a set of embeddings. + + Args: + embeddings: (N, D) array of embedding vectors. + + Returns: + (N, N) similarity matrix where entry [i,j] is cosine_sim(i, j). + """ + if embeddings.shape[0] == 0: + return np.empty((0, 0), dtype=np.float64) + norms = np.linalg.norm(embeddings, axis=1, keepdims=True) + # Avoid division by zero + norms = np.where(norms == 0, 1.0, norms) + normalized = embeddings / norms + return normalized @ normalized.T + + +def _centroid(embeddings: list[list[float]]) -> list[float]: + """Compute the centroid (mean) of a list of embedding vectors. + + Returns L2-normalized centroid so cosine similarity works correctly. + """ + if not embeddings: + return [] + arr = np.asarray(embeddings, dtype=np.float64) + mean = arr.mean(axis=0) + norm = float(np.linalg.norm(mean)) + if norm > 0: + mean = mean / norm + return mean.tolist() + + +# ── Agglomerative clustering ─────────────────────────────────────────── + + +def _agglomerative_cluster( + embeddings: list[list[float]], + threshold: float, +) -> list[list[int]]: + """Single-linkage agglomerative clustering using cosine similarity. + + Merges items whose similarity exceeds `threshold` into clusters. + Uses a union-find approach for efficiency. + + Args: + embeddings: List of embedding vectors (one per item). + threshold: Minimum cosine similarity to merge two items. + + Returns: + List of clusters, where each cluster is a list of item indices. + """ + n = len(embeddings) + if n == 0: + return [] + if n == 1: + return [[0]] + + # Compute pairwise similarity matrix + emb_array = np.asarray(embeddings, dtype=np.float64) + sim_matrix = _cosine_similarity_matrix(emb_array) + + # Union-Find + parent = list(range(n)) + + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] # path compression + x = parent[x] + return x + + def union(x: int, y: int) -> None: + rx, ry = find(x), find(y) + if rx != ry: + parent[rx] = ry + + # Merge pairs above threshold + for i in range(n): + for j in range(i + 1, n): + if sim_matrix[i, j] >= threshold: + union(i, j) + + # Collect clusters + clusters: dict[int, list[int]] = {} + for i in range(n): + root = find(i) + if root not in clusters: + clusters[root] = [] + clusters[root].append(i) + + return list(clusters.values()) + + +# ── Dataclasses ───────────────────────────────────────────────────────── + + +# Import StoredObject for type annotations +from mnemosyne.object_store import StoredObject + + +@dataclass +class Episode: + """Level 1: A group of related StoredObjects forming a coherent episode. + + Episodes are created by clustering objects whose embeddings have + cosine similarity > 0.6. Each episode has a centroid embedding + (mean of its object embeddings) and an auto-generated summary. + """ + + id: str + objects: list[StoredObject] = field(default_factory=list) + embedding: list[float] = field(default_factory=list) + summary: str = "" + turn_range: tuple[int, int] = (0, 0) + object_types: set[str] = field(default_factory=set) + + def __hash__(self) -> int: + return hash(self.id) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Episode): + return NotImplemented + return self.id == other.id + + +@dataclass +class SemanticTheme: + """Level 2: An abstract theme spanning multiple episodes. + + Semantic themes are created by clustering episodes whose centroid + embeddings have cosine similarity > 0.4. Each theme has a centroid + embedding and an auto-generated label. + """ + + id: str + episodes: list[Episode] = field(default_factory=list) + embedding: list[float] = field(default_factory=list) + label: str = "" + + def __hash__(self) -> int: + return hash(self.id) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, SemanticTheme): + return NotImplemented + return self.id == other.id + + +@dataclass +class Theme: + """Level 3: A top-level category for the entire session. + + Themes are created by clustering semantic themes whose centroid + embeddings have cosine similarity > 0.25. + """ + + id: str + semantic_themes: list[SemanticTheme] = field(default_factory=list) + embedding: list[float] = field(default_factory=list) + label: str = "" + + def __hash__(self) -> int: + return hash(self.id) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Theme): + return NotImplemented + return self.id == other.id + + +# ── Helper functions ──────────────────────────────────────────────────── + + +def _generate_episode_summary(objects: list[StoredObject]) -> str: + """Auto-generate a summary for an episode from its objects' stubs.""" + if not objects: + return "(empty episode)" + stubs = [obj.stub for obj in objects if obj.stub] + if not stubs: + # Fall back to first line of content + stubs = [obj.content_full.split("\n", 1)[0][:100] for obj in objects[:3]] + # Take up to 5 stubs + preview = "; ".join(stubs[:5]) + if len(stubs) > 5: + preview += f" (+{len(stubs) - 5} more)" + return preview + + +def _generate_theme_label(episodes: list[Episode]) -> str: + """Auto-generate a label for a semantic theme from its episodes.""" + if not episodes: + return "(empty theme)" + # Collect unique object types across episodes + all_types: set[str] = set() + for ep in episodes: + all_types.update(ep.object_types) + type_str = ", ".join(sorted(all_types)[:4]) + return f"Theme [{type_str}] ({len(episodes)} episodes)" + + +def _generate_top_theme_label(semantic_themes: list[SemanticTheme]) -> str: + """Auto-generate a label for a top-level theme.""" + if not semantic_themes: + return "(empty top theme)" + labels = [st.label for st in semantic_themes[:3]] + return f"Top: {'; '.join(labels)}" + + +def _make_episode(objects: list[StoredObject]) -> Episode: + """Create an Episode from a list of StoredObjects.""" + embeddings = [obj.embedding for obj in objects if obj.embedding] + centroid = _centroid(embeddings) if embeddings else [] + + turns = [] + for obj in objects: + if obj.source_turn_start is not None: + turns.append(obj.source_turn_start) + if obj.source_turn_end is not None: + turns.append(obj.source_turn_end) + + turn_range = (min(turns), max(turns)) if turns else (0, 0) + obj_types = {obj.object_type for obj in objects} + + return Episode( + id=uuid.uuid4().hex[:16], + objects=list(objects), + embedding=centroid, + summary=_generate_episode_summary(objects), + turn_range=turn_range, + object_types=obj_types, + ) + + +def _make_semantic_theme(episodes: list[Episode]) -> SemanticTheme: + """Create a SemanticTheme from a list of Episodes.""" + embeddings = [ep.embedding for ep in episodes if ep.embedding] + centroid = _centroid(embeddings) if embeddings else [] + + return SemanticTheme( + id=uuid.uuid4().hex[:16], + episodes=list(episodes), + embedding=centroid, + label=_generate_theme_label(episodes), + ) + + +def _make_theme(semantic_themes: list[SemanticTheme]) -> Theme: + """Create a Theme from a list of SemanticThemes.""" + embeddings = [st.embedding for st in semantic_themes if st.embedding] + centroid = _centroid(embeddings) if embeddings else [] + + return Theme( + id=uuid.uuid4().hex[:16], + semantic_themes=list(semantic_themes), + embedding=centroid, + label=_generate_top_theme_label(semantic_themes), + ) + + +# ── SessionHierarchy ──────────────────────────────────────────────────── + + +# Clustering thresholds +EPISODE_SIMILARITY_THRESHOLD = 0.6 +SEMANTIC_SIMILARITY_THRESHOLD = 0.4 +THEME_SIMILARITY_THRESHOLD = 0.25 + +# Retrieval gating thresholds +THEME_RETRIEVAL_THRESHOLD = 0.3 +EPISODE_RETRIEVAL_THRESHOLD = 0.4 + +# Re-clustering interval +RECLUSTER_INTERVAL = 20 + + +class SessionHierarchy: + """Builds and maintains a 4-level hierarchy over StoredObjects. + + The hierarchy enables efficient top-down retrieval for long sessions + (50+ turns) where flat search becomes noisy. Objects are grouped into + episodes (Level 1), episodes into semantic themes (Level 2), and + semantic themes into top-level themes (Level 3). + + Usage: + hierarchy = SessionHierarchy(embedder) + + # Incremental: add objects one at a time + hierarchy.add_object(stored_obj) + + # Full rebuild from scratch + hierarchy.rebuild(all_objects) + + # Top-down retrieval + results = hierarchy.retrieve(query_embedding, limit=10) + + # Periodic maintenance + hierarchy.maintenance(turn=60) + """ + + def __init__(self, embedder: Embedder) -> None: + """Initialize the hierarchy with an embedder for computing similarities. + + Args: + embedder: An Embedder instance for computing text embeddings. + """ + self._embedder = embedder + self._episodes: list[Episode] = [] + self._semantic_themes: list[SemanticTheme] = [] + self._themes: list[Theme] = [] + self._all_objects: list[StoredObject] = [] + self._object_ids: set[str] = set() # Track seen object IDs + self._last_rebuild_turn: int = 0 + self._enabled: bool = True + self._last_goal_hash: str | None = None + + # ── Properties ────────────────────────────────────────────────── + + @property + def enabled(self) -> bool: + """Whether the hierarchy is active.""" + return self._enabled + + @enabled.setter + def enabled(self, value: bool) -> None: + self._enabled = value + + @property + def object_count(self) -> int: + """Total number of objects in the hierarchy.""" + return len(self._all_objects) + + @property + def episode_count(self) -> int: + """Total number of episodes.""" + return len(self._episodes) + + @property + def theme_count(self) -> int: + """Total number of top-level themes.""" + return len(self._themes) + + # ── Public API ────────────────────────────────────────────────── + + def add_object(self, obj: StoredObject) -> Episode | None: + """Assign an object to the nearest episode by embedding similarity. + + If no episode has similarity > EPISODE_SIMILARITY_THRESHOLD (0.6), + creates a new episode for this object. + + Args: + obj: The StoredObject to add. + + Returns: + The Episode the object was assigned to, or None if the + hierarchy is disabled or the object has no embedding. + """ + if not self._enabled: + return None + + if obj.id in self._object_ids: + return None # Already tracked + + self._all_objects.append(obj) + self._object_ids.add(obj.id) + + if not obj.embedding: + # Object has no embedding — create a singleton episode + episode = _make_episode([obj]) + self._episodes.append(episode) + self._rebuild_upper_levels() + return episode + + # Find the nearest episode + best_episode: Episode | None = None + best_sim = -1.0 + + for episode in self._episodes: + if not episode.embedding: + continue + sim = _cosine_similarity(obj.embedding, episode.embedding) + if sim > best_sim: + best_sim = sim + best_episode = episode + + if best_episode is not None and best_sim >= EPISODE_SIMILARITY_THRESHOLD: + # Add to existing episode and update centroid + best_episode.objects.append(obj) + embeddings = [o.embedding for o in best_episode.objects if o.embedding] + best_episode.embedding = _centroid(embeddings) if embeddings else [] + best_episode.summary = _generate_episode_summary(best_episode.objects) + best_episode.object_types.add(obj.object_type) + + # Update turn range + turns = [] + for o in best_episode.objects: + if o.source_turn_start is not None: + turns.append(o.source_turn_start) + if o.source_turn_end is not None: + turns.append(o.source_turn_end) + if turns: + best_episode.turn_range = (min(turns), max(turns)) + + return best_episode + + # No matching episode — create a new one + episode = _make_episode([obj]) + self._episodes.append(episode) + self._rebuild_upper_levels() + return episode + + def rebuild(self, objects: list[StoredObject]) -> None: + """Full re-cluster from scratch. + + Clears the existing hierarchy and rebuilds all levels: + 1. Cluster objects into episodes (similarity > 0.6) + 2. Cluster episodes into semantic themes (similarity > 0.4) + 3. Cluster semantic themes into top-level themes (similarity > 0.25) + + Args: + objects: All StoredObjects to include in the hierarchy. + """ + self._all_objects = list(objects) + self._object_ids = {obj.id for obj in objects} + self._episodes = [] + self._semantic_themes = [] + self._themes = [] + + if not objects: + return + + # Filter objects with embeddings for clustering + objects_with_embeddings = [obj for obj in objects if obj.embedding] + objects_without_embeddings = [obj for obj in objects if not obj.embedding] + + # Level 1: Cluster objects into episodes + if objects_with_embeddings: + embeddings = [obj.embedding for obj in objects_with_embeddings] + clusters = _agglomerative_cluster(embeddings, EPISODE_SIMILARITY_THRESHOLD) + + for cluster_indices in clusters: + cluster_objects = [objects_with_embeddings[i] for i in cluster_indices] + episode = _make_episode(cluster_objects) + self._episodes.append(episode) + + # Create singleton episodes for objects without embeddings + for obj in objects_without_embeddings: + episode = _make_episode([obj]) + self._episodes.append(episode) + + # Level 2 & 3: Build upper levels + self._rebuild_upper_levels() + + def _rebuild_upper_levels(self) -> None: + """Rebuild semantic themes and top-level themes from current episodes.""" + # Level 2: Cluster episodes into semantic themes + self._semantic_themes = [] + episodes_with_embeddings = [ep for ep in self._episodes if ep.embedding] + episodes_without_embeddings = [ep for ep in self._episodes if not ep.embedding] + + if episodes_with_embeddings: + ep_embeddings = [ep.embedding for ep in episodes_with_embeddings] + ep_clusters = _agglomerative_cluster(ep_embeddings, SEMANTIC_SIMILARITY_THRESHOLD) + + for cluster_indices in ep_clusters: + cluster_episodes = [episodes_with_embeddings[i] for i in cluster_indices] + theme = _make_semantic_theme(cluster_episodes) + self._semantic_themes.append(theme) + + for ep in episodes_without_embeddings: + theme = _make_semantic_theme([ep]) + self._semantic_themes.append(theme) + + # Level 3: Cluster semantic themes into top-level themes + self._themes = [] + themes_with_embeddings = [st for st in self._semantic_themes if st.embedding] + themes_without_embeddings = [st for st in self._semantic_themes if not st.embedding] + + if themes_with_embeddings: + st_embeddings = [st.embedding for st in themes_with_embeddings] + st_clusters = _agglomerative_cluster(st_embeddings, THEME_SIMILARITY_THRESHOLD) + + for cluster_indices in st_clusters: + cluster_themes = [themes_with_embeddings[i] for i in cluster_indices] + top_theme = _make_theme(cluster_themes) + self._themes.append(top_theme) + + for st in themes_without_embeddings: + top_theme = _make_theme([st]) + self._themes.append(top_theme) + + def retrieve( + self, + query_embedding: list[float], + limit: int = 10, + ) -> list[StoredObject]: + """Top-down hierarchical retrieval. + + Finds matching themes → expands to semantics → expands to episodes + → returns objects, with similarity gating at each level. This + prevents redundant retrieval (key xMemory advantage). + + Algorithm: + 1. Compute query similarity to all theme embeddings + 2. Expand themes with similarity > 0.3 to their semantic themes + 3. Expand semantic themes to their episodes + 4. Expand episodes with similarity > 0.4 to their objects + 5. Return objects sorted by direct similarity to query, up to limit + + If no themes pass the gate, falls back to flat search over all objects. + + Args: + query_embedding: The query embedding vector. + limit: Maximum number of objects to return. + + Returns: + List of StoredObjects sorted by similarity to query. + """ + if not self._enabled or not self._themes: + return self._flat_search(query_embedding, limit) + + if not query_embedding: + return [] + + # Step 1: Find matching top-level themes + candidate_semantic_themes: list[SemanticTheme] = [] + for theme in self._themes: + if not theme.embedding: + # Include themes without embeddings as fallback + candidate_semantic_themes.extend(theme.semantic_themes) + continue + sim = _cosine_similarity(query_embedding, theme.embedding) + if sim >= THEME_RETRIEVAL_THRESHOLD: + candidate_semantic_themes.extend(theme.semantic_themes) + + # Fallback: if no themes matched, use all semantic themes + if not candidate_semantic_themes: + candidate_semantic_themes = list(self._semantic_themes) + + # Step 2: Expand semantic themes to episodes + candidate_episodes: list[Episode] = [] + for st in candidate_semantic_themes: + candidate_episodes.extend(st.episodes) + + # Step 3: Gate episodes by similarity + candidate_objects: list[StoredObject] = [] + for episode in candidate_episodes: + if not episode.embedding: + candidate_objects.extend(episode.objects) + continue + sim = _cosine_similarity(query_embedding, episode.embedding) + if sim >= EPISODE_RETRIEVAL_THRESHOLD: + candidate_objects.extend(episode.objects) + + # Fallback: if no episodes matched, expand all + if not candidate_objects: + for episode in candidate_episodes: + candidate_objects.extend(episode.objects) + + # Step 4: Score and sort by direct similarity to query + scored: list[tuple[StoredObject, float]] = [] + seen_ids: set[str] = set() + for obj in candidate_objects: + if obj.id in seen_ids: + continue + seen_ids.add(obj.id) + if obj.embedding: + sim = _cosine_similarity(query_embedding, obj.embedding) + else: + sim = 0.0 + scored.append((obj, sim)) + + scored.sort(key=lambda x: x[1], reverse=True) + return [obj for obj, _sim in scored[:limit]] + + def _flat_search( + self, + query_embedding: list[float], + limit: int, + ) -> list[StoredObject]: + """Fallback flat search over all objects (no hierarchy).""" + if not query_embedding: + return [] + + scored: list[tuple[StoredObject, float]] = [] + for obj in self._all_objects: + if obj.embedding: + sim = _cosine_similarity(query_embedding, obj.embedding) + else: + sim = 0.0 + scored.append((obj, sim)) + + scored.sort(key=lambda x: x[1], reverse=True) + return [obj for obj, _sim in scored[:limit]] + + def get_episodes(self) -> list[Episode]: + """Return all episodes in the hierarchy.""" + return list(self._episodes) + + def get_themes(self) -> list[SemanticTheme]: + """Return all semantic themes (Level 2) in the hierarchy.""" + return list(self._semantic_themes) + + def get_top_themes(self) -> list[Theme]: + """Return all top-level themes (Level 3) in the hierarchy.""" + return list(self._themes) + + def maintenance(self, turn: int, goal_hash: str | None = None) -> bool: + """Periodic re-clustering maintenance. + + Triggers a full rebuild if: + - turn is a multiple of RECLUSTER_INTERVAL (20), or + - goal_hash changed since last maintenance (topic shift) + + Args: + turn: Current conversation turn number. + goal_hash: Optional hash of current goal for change detection. + + Returns: + True if a rebuild was triggered, False otherwise. + """ + if not self._enabled: + return False + + should_rebuild = False + + # Check periodic interval + if turn > 0 and turn % RECLUSTER_INTERVAL == 0: + if turn != self._last_rebuild_turn: + should_rebuild = True + + # Check goal change + if goal_hash is not None and goal_hash != self._last_goal_hash: + should_rebuild = True + self._last_goal_hash = goal_hash + + if should_rebuild: + self.rebuild(self._all_objects) + self._last_rebuild_turn = turn + return True + + return False + + def summary(self) -> dict[str, object]: + """Return a summary of the hierarchy state for debugging/telemetry.""" + total_objects = len(self._all_objects) + total_episodes = len(self._episodes) + total_semantic_themes = len(self._semantic_themes) + total_themes = len(self._themes) + + avg_episode_size = total_objects / total_episodes if total_episodes > 0 else 0.0 + + return { + "enabled": self._enabled, + "objects": total_objects, + "episodes": total_episodes, + "semantic_themes": total_semantic_themes, + "themes": total_themes, + "avg_episode_size": round(avg_episode_size, 1), + "last_rebuild_turn": self._last_rebuild_turn, + } diff --git a/src/mnemosyne/phantom.py b/src/mnemosyne/phantom.py new file mode 100644 index 0000000..0c0131c --- /dev/null +++ b/src/mnemosyne/phantom.py @@ -0,0 +1,975 @@ +"""Phantom tools — side-channel communication between proxy and model. + +The proxy injects phantom tools into the request. The model can call +them to communicate memory management hints. The proxy intercepts +these calls from the SSE stream before the framework sees them. + +The framework never knows. The model and the proxy have a private +channel. + +Tools: + memory_fault(paths): Model requests evicted content back. + The proxy resolves from PageStore, no file system round trip. + +Note: memory_release is handled via inline tags +(see tags.py) or by the framework's native memory_release tool. +It no longer needs phantom tool infrastructure. +""" + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass, field + + +class CleanupTagFilter: + """Strips ... tags from SSE text streams + and executes the cleanup operations inline. + + Tags can span multiple text_delta chunks. The filter buffers pending + text when it sees a partial tag opening, accumulates the tag body, + executes operations when the closing tag arrives, and re-emits clean + text outside tags. + + Pass block_store and page_store to execute operations as tags are + stripped. Without them, the filter only strips (no execution). + """ + + _OPEN = "" + _CLOSE = "" + + def __init__(self, block_store=None, page_store=None) -> None: + self._inside = False + self._pending = "" # text that might be the start of a tag + self._tag_body: list[str] = [] # accumulates content inside tag + self._bs = block_store + self._ps = page_store + self.executed_ops: list[str] = [] # log of executed operations + + def _execute_tag_body(self) -> None: + """Parse and execute cleanup operations from the accumulated tag body.""" + from mnemosyne.tags import parse_cleanup_tags + + body = "".join(self._tag_body) + self._tag_body = [] + ops = parse_cleanup_tags(f"{body}") + if ops.empty: + return + if self._bs is not None: + for block_id in ops.drops: + if self._bs.drop(block_id): + self.executed_ops.append(f"dropped {block_id}") + for block_id, summary in ops.summaries: + if self._bs.summarize(block_id, summary): + self.executed_ops.append(f"summarized {block_id}") + for block_id in ops.anchors: + if self._bs.anchor(block_id): + self.executed_ops.append(f"anchored {block_id}") + if self._ps is not None and ops.releases: + for path in ops.releases: + self._ps.mark_released(path) + self.executed_ops.append(f"released {len(ops.releases)} path(s)") + + def filter(self, text: str) -> str: + """Filter cleanup tags from a text chunk. Returns clean text.""" + result = [] + + i = 0 + while i < len(text): + if self._inside: + # Look for closing tag + close_pos = text.find(self._CLOSE, i) + if close_pos >= 0: + # Accumulate body up to the closing tag, then execute + self._tag_body.append(text[i:close_pos]) + self._execute_tag_body() + self._inside = False + i = close_pos + len(self._CLOSE) + # Skip optional newline after closing tag + if i < len(text) and text[i] == "\n": + i += 1 + else: + # Entire remaining text is inside tag — accumulate + self._tag_body.append(text[i:]) + break + + elif self._pending: + # We have buffered text that might be a partial open tag + combined = self._pending + text[i:] + if self._OPEN in combined: + # Tag completed — emit everything before it + tag_pos = combined.find(self._OPEN) + result.append(combined[:tag_pos]) + self._pending = "" + self._inside = True + i = tag_pos + len(self._OPEN) - len(self._pending) + # Recalculate: skip past the open tag in the original text + consumed_from_text = len(combined) - len(text) + i + i = max(0, consumed_from_text) + continue + elif self._OPEN.startswith(combined): + # Still a partial match — keep buffering + self._pending = combined + break + else: + # False alarm — emit pending and restart + result.append(self._pending) + self._pending = "" + # Don't advance i — reprocess from same position + + else: + # Normal state — look for tag opening + open_pos = text.find(self._OPEN, i) + partial_start = self._find_partial_open(text, i) + + if open_pos >= 0 and (partial_start < 0 or open_pos <= partial_start): + result.append(text[i:open_pos]) + self._inside = True + i = open_pos + len(self._OPEN) + elif partial_start >= 0: + # Partial tag at end of chunk — buffer it + result.append(text[i:partial_start]) + self._pending = text[partial_start:] + break + else: + result.append(text[i:]) + break + + return "".join(result) + + def _find_partial_open(self, text: str, start: int) -> int: + """Find position of a partial tag at end of text.""" + tag = self._OPEN + for length in range(len(tag) - 1, 0, -1): + if text.endswith(tag[:length], start): + pos = len(text) - length + if pos >= start: + return pos + return -1 + + def flush(self) -> str: + """Return any buffered text that turned out not to be a tag.""" + pending = self._pending + self._pending = "" + # If we're still inside a tag at flush, that's a malformed tag — discard + self._inside = False + return pending + + +PHANTOM_TOOL_NAMES = frozenset( + {"yuyay", "recall", "memory_fault", "qunqay", "tiqsiy", "memory_query"} +) + +PHANTOM_TOOL_DEFINITIONS = [ + { + "name": "yuyay", + "description": ( + "Remember — restore evicted content by tensor handle. " + "When you see '[tensor:xxxx — description]' markers in " + "your context, the original content has been evicted to " + "save space. Call this with the tensor handle(s) to get " + "the content back. Faster and cheaper than re-reading " + "files. You can also pass file paths or tool_use_ids " + "for backward compatibility." + ), + "input_schema": { + "type": "object", + "properties": { + "handles": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Tensor handles (e.g. 'a3f2b901'), file paths, " + "or tool_use_ids to restore from eviction cache" + ), + } + }, + "required": ["handles"], + }, + }, + { + "name": "qunqay", + "description": ( + "Release — mark content for eviction from the working set. " + "Use this to proactively free context space by releasing " + "tensors, file reads, or tool results you no longer need. " + "The content remains in the backing store and can be " + "restored later with yuyay. Use this when: (1) you rewrote " + "a file and the old read is stale, (2) a tool result has " + "been fully consumed, (3) you want to free space for " + "content you value more. Every token you release saves " + "O(n^2) compute on every subsequent turn." + ), + "input_schema": { + "type": "object", + "properties": { + "handles": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Tensor handles, file paths, or tool_use_ids " + "to release from the working set" + ), + }, + "reason": { + "type": "string", + "description": ("Why this content is being released (logged for audit trail)"), + }, + }, + "required": ["handles"], + }, + }, + { + "name": "tiqsiy", + "description": ( + "Compact — structurally compress older conversation turns. " + "Replaces a range of older turns with a single user/assistant " + "summary pair you provide. Original turns are archived to " + "the backing store. Use this when conversation history is " + "consuming context but you've captured the important " + "conclusions. This reduces both content tokens and structural " + "overhead (role labels, turn boundaries). The compaction is " + "applied on the next turn." + ), + "input_schema": { + "type": "object", + "properties": { + "older_than": { + "type": "integer", + "description": ( + "Compact messages more than this many turns " + "from the current turn. Messages newer than " + "this are preserved." + ), + }, + "summary": { + "type": "string", + "description": ( + "Your summary of the compacted conversation " + "segment. Should capture conclusions, decisions, " + "and any context needed for future reference. " + "This becomes the content of the replacement " + "assistant message." + ), + }, + }, + "required": ["older_than", "summary"], + }, + }, + { + "name": "memory_query", + "description": ( + "Ask a question about evicted content without restoring it. " + "When you need a specific fact from content that was compressed " + "or evicted (shown as stubs or summaries), use this instead of " + "yuyay. The proxy will search the backing store and return a " + "targeted answer (50-200 tokens) — much cheaper than restoring " + "the full content. Use yuyay only when you need to edit or " + "deeply reference the full content." + ), + "input_schema": { + "type": "object", + "properties": { + "question": { + "type": "string", + "description": "The specific question to answer from evicted content", + }, + "scope": { + "type": "string", + "description": "Optional hint about what type of content to search (e.g. 'auth files', 'error logs')", + }, + "max_tokens": { + "type": "integer", + "description": "Maximum tokens for the answer (default 200)", + }, + }, + "required": ["question"], + }, + }, +] + + +@dataclass +class PhantomCall: + """A captured phantom tool call from the model's response.""" + + name: str + tool_use_id: str + input: dict = field(default_factory=dict) + + +def extract_phantom_calls(messages: list[dict]) -> list[PhantomCall]: + """Extract unresolved phantom tool calls from the conversation messages. + + Scans assistant messages for tool_use blocks whose name is in + PHANTOM_TOOL_NAMES, then checks if a corresponding tool_result + exists in a subsequent user message. Returns only calls that + either have no result or have a pending placeholder result. + """ + # Collect all tool_result IDs already in messages + resolved_ids: set[str] = set() + for msg in messages: + if msg.get("role") != "user": + continue + content = msg.get("content", []) + if isinstance(content, str): + continue + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_result": + tool_use_id = block.get("tool_use_id", "") + result_text = block.get("content", "") + if isinstance(result_text, str) and "[memory_query:pending]" in result_text: + continue # Pending placeholder — treat as unresolved + resolved_ids.add(tool_use_id) + + # Scan assistant messages for phantom tool calls + calls: list[PhantomCall] = [] + for msg in messages: + if msg.get("role") != "assistant": + continue + content = msg.get("content", []) + if isinstance(content, str): + continue + if isinstance(content, list): + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") != "tool_use": + continue + name = block.get("name", "") + if name not in PHANTOM_TOOL_NAMES: + continue + tool_use_id = block.get("id", "") + if tool_use_id in resolved_ids: + continue + calls.append( + PhantomCall( + name=name, + tool_use_id=tool_use_id, + input=block.get("input", {}), + ) + ) + return calls + + +def inject_tools(body: dict) -> set[str]: + """Add phantom tool definitions to the request's tools array. + + Returns a set of tool names that the framework already provides + (observe-only). Tools NOT in this set are fully intercepted by + the gateway — their events are suppressed from the stream and + handled via continuation. + + Returns empty set when no phantom tools are framework-provided + (all are intercepted). + """ + tools = body.get("tools", []) + existing_names = {t.get("name") for t in tools} + observe_only = PHANTOM_TOOL_NAMES & existing_names + # Inject definitions for tools the framework doesn't provide + for defn in PHANTOM_TOOL_DEFINITIONS: + if defn["name"] not in existing_names: + tools.append(defn) + body["tools"] = tools + return observe_only + + +def inject_phantom_results( + messages: list[dict], + phantom_calls: list[PhantomCall], + page_store, + observe_only: bool | set[str] = False, + context_assembler=None, + session_id: str = "", +) -> list[dict]: + """Inject tool_result messages for phantom calls from the previous turn. + + The model called phantom tools, but the framework never saw them. + We need to inject the results so the model's next turn sees a + coherent conversation — it called a tool, it got a result. + + For memory_fault calls, the result includes the restored content. + + observe_only can be: + - True: skip all injection (framework handled everything) + - False: inject all (gateway intercepted everything) + - set[str]: skip injection for tools in set (framework handled + those); inject for tools NOT in set (gateway intercepted those) + """ + if not phantom_calls: + return messages + if observe_only is True: + return messages + + # Filter to only intercepted calls (not framework-handled) + if isinstance(observe_only, set) and observe_only: + phantom_calls = [c for c in phantom_calls if c.name not in observe_only] + if not phantom_calls: + return messages + + # Find the last assistant message — that's where the phantom calls were + last_assistant_idx = None + for i in range(len(messages) - 1, -1, -1): + if messages[i].get("role") == "assistant": + last_assistant_idx = i + break + + if last_assistant_idx is None: + return messages + + # Re-insert the phantom tool_use blocks into the assistant message + assistant_msg = messages[last_assistant_idx] + content = assistant_msg.get("content", []) + if isinstance(content, str): + content = [{"type": "text", "text": content}] + + existing_tool_ids = { + b.get("id") for b in content if isinstance(b, dict) and b.get("type") == "tool_use" + } + for call in phantom_calls: + if call.tool_use_id in existing_tool_ids: + continue # already present — don't duplicate + content.append( + { + "type": "tool_use", + "id": call.tool_use_id, + "name": call.name, + "input": call.input, + } + ) + assistant_msg["content"] = content + + # Build tool_result messages for each phantom call + results = [] + for call in phantom_calls: + result_content = _handle_phantom_call( + call, + page_store, + context_assembler=context_assembler, + session_id=session_id, + ) + results.append( + { + "type": "tool_result", + "tool_use_id": call.tool_use_id, + "content": result_content, + } + ) + + # Find where to inject — after the assistant message's tool results + # The next user message (if any) should contain the results + inject_idx = last_assistant_idx + 1 + if inject_idx < len(messages) and messages[inject_idx].get("role") == "user": + # Append to existing user message, skipping duplicates + user_msg = messages[inject_idx] + user_content = user_msg.get("content", []) + if isinstance(user_content, str): + user_content = [{"type": "text", "text": user_content}] + existing_result_ids = { + b.get("tool_use_id") + for b in user_content + if isinstance(b, dict) and b.get("type") == "tool_result" + } + for r in results: + if r["tool_use_id"] not in existing_result_ids: + user_content.append(r) + user_msg["content"] = user_content + else: + # Insert a new user message with the results + messages.insert(inject_idx, {"role": "user", "content": results}) + + return messages + + +def _run_async_phantom(coro): + """Run an async coroutine from sync phantom handler context.""" + import asyncio + import concurrent.futures + + try: + asyncio.get_running_loop() + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(asyncio.run, coro) + return future.result(timeout=10) + except RuntimeError: + return asyncio.run(coro) + + +def _handle_phantom_call( + call: PhantomCall, + page_store, + block_store=None, + context_assembler=None, + session_id: str = "", +) -> str: + """Execute a phantom tool call and return the result text.""" + if call.name == "qunqay": + identifiers = call.input.get("handles", []) + reason = call.input.get("reason", "model requested release") + released = [] + not_found = [] + for identifier in identifiers: + found = False + # Try page store (file reads, tool results) + if page_store is not None: + if page_store.mark_released(identifier): + found = True + # Also try eviction index (file paths) + elif identifier in getattr(page_store, "_eviction_index", {}): + page_store.mark_released(identifier) + found = True + # Try block store (conversation blocks) + if not found and block_store is not None: + if block_store.drop(identifier): + found = True + if found: + released.append(identifier) + else: + not_found.append(identifier) + parts = [] + if released: + parts.append(f"Released {len(released)} tensor(s): {', '.join(released)}") + if not_found: + parts.append(f"Not found: {', '.join(not_found)}") + if reason: + parts.append(f"Reason: {reason}") + return " | ".join(parts) if parts else "Nothing to release." + + if call.name in ("yuyay", "recall", "memory_fault"): + # Accept both "handles" (new) and "paths" (legacy) parameter names + identifiers = call.input.get("handles", call.input.get("paths", [])) + restored = [] + not_found = [] + for identifier in identifiers: + entry = None + + # 1. Try tensor handle (unified addressing) + if page_store is not None: + entry = page_store.resolve_tensor(identifier) + if entry is None and block_store is not None: + content = block_store.restore(identifier) + if content is not None: + restored.append( + { + "label": f"tensor:{identifier}", + "content": content, + "size": len(content), + } + ) + continue + + # 2. Try file path (Read results) + if entry is None and page_store is not None: + entry = page_store._eviction_index.get(identifier) + + # 3. Try tool_use_id (legacy) + if entry is None and page_store is not None: + entry = page_store.pages.get(identifier) + + if entry is not None: + label = f"tensor:{identifier}" + if entry.tool_name == "Read": + path = entry.tool_input.get("file_path", identifier) + label = f"tensor:{identifier} ({path})" + elif entry.tool_name: + label = f"tensor:{identifier} ({entry.tool_name})" + restored.append( + { + "label": label, + "content": entry.original_content, + "size": entry.original_size, + } + ) + else: + not_found.append(identifier) + + parts = [] + for r in restored: + parts.append(f"--- {r['label']} ({r['size']} bytes) ---\n{r['content']}") + if not_found: + parts.append(f"Not in cache (use Read or re-run): {', '.join(not_found)}") + return "\n\n".join(parts) if parts else "No handles requested." + + if call.name == "memory_query": + question = call.input.get("question", "") + scope = call.input.get("scope") + max_tokens = call.input.get("max_tokens", 200) + + # Resolve via ContextAssembler if available + if context_assembler is not None and session_id: + try: + result = _run_async_phantom( + context_assembler.handle_micro_fault( + session_id=session_id, + question=question, + scope=scope, + max_tokens=max_tokens, + ) + ) + sources_str = ", ".join(result.sources) if result.sources else "none" + return ( + f"{result.answer}\n\n" + f"[memory_query resolved: {result.answer_tokens} tokens, " + f"saved ~{result.avoided_tokens} tokens, " + f"sources: {sources_str}, " + f"{result.latency_ms:.0f}ms]" + ) + except Exception as exc: + import sys + + print( + f" [memory_query] resolution failed: {exc}", + file=sys.stderr, + ) + # Fall through to pending placeholder + + # Fallback: return pending placeholder when ContextAssembler unavailable + return ( + f"[memory_query:pending] question={question!r} scope={scope!r} max_tokens={max_tokens}" + ) + + if call.name == "tiqsiy": + # Structural compaction is deferred — we store the request + # in the session and apply it on the next inbound request. + # The handler in proxy.py reads pending_compaction from session. + older_than = call.input.get("older_than", 20) + summary = call.input.get("summary", "") + if not summary: + return "Error: summary is required for compaction." + # Return confirmation — actual compaction happens next turn. + return ( + f"Compaction scheduled: messages older than {older_than} turns " + f"will be replaced with your summary on the next turn. " + f"Summary length: {len(summary)} chars." + ) + + return f"Unknown phantom tool: {call.name}" + + +@dataclass +class CompactionResult: + """Result of a structural compaction operation.""" + + messages_removed: int = 0 + messages_after: int = 0 + chars_removed: int = 0 + archived: list[dict] = field(default_factory=list) + + +def apply_compaction( + messages: list[dict], + older_than: int, + summary: str, +) -> CompactionResult: + """Replace older conversation turns with a summary pair. + + Finds messages older than `older_than` turns from the end, + replaces them with a single user/assistant summary pair. + Returns the archived messages for storage. + + Preserves the first message (often contains system context + from the framework) and maintains valid alternating structure. + """ + result = CompactionResult() + total = len(messages) + + if total < 4: # need at least 2 turns to compact + return result + + # Count turns (pairs of user/assistant messages) + # Each turn is roughly 2 messages + preserve_count = older_than * 2 + if preserve_count >= total: + return result # nothing old enough to compact + + # Split: archive the old, keep the recent + # Always preserve the first message (framework context) + compact_end = total - preserve_count + if compact_end <= 1: + return result + + archived = messages[1:compact_end] # skip first message + preserved_first = messages[0] + preserved_recent = messages[compact_end:] + + # Calculate what we're removing + archived_chars = sum(len(json.dumps(m).encode("utf-8")) for m in archived) + + # Build the summary pair — a single user/assistant exchange + # that replaces the archived conversation segment + summary_pair = [ + { + "role": "user", + "content": ( + f"[Compacted: {len(archived)} messages from earlier " + f"in this conversation were structurally compressed " + f"into this summary by the model.]" + ), + }, + { + "role": "assistant", + "content": summary, + }, + ] + + # Ensure valid alternation after first message + # The first preserved message sets the expected next role + new_messages = [preserved_first] + summary_pair + preserved_recent + + # Validate alternation — fix if needed + _fix_alternation(new_messages) + + result.messages_removed = len(archived) + result.messages_after = len(new_messages) + result.chars_removed = archived_chars + result.archived = archived + + # Replace in-place + messages.clear() + messages.extend(new_messages) + + return result + + +def _fix_alternation(messages: list[dict]) -> None: + """Ensure strict user/assistant alternation. + + If two consecutive messages have the same role, merge the + second into the first. + """ + i = 1 + while i < len(messages): + if messages[i].get("role") == messages[i - 1].get("role"): + # Merge: append content of messages[i] to messages[i-1] + prev = messages[i - 1] + curr = messages[i] + prev_content = prev.get("content", "") + curr_content = curr.get("content", "") + if isinstance(prev_content, str) and isinstance(curr_content, str): + prev["content"] = prev_content + "\n\n" + curr_content + elif isinstance(prev_content, list) and isinstance(curr_content, list): + prev["content"] = prev_content + curr_content + elif isinstance(prev_content, str): + prev["content"] = [{"type": "text", "text": prev_content}] + ( + curr_content + if isinstance(curr_content, list) + else [{"type": "text", "text": str(curr_content)}] + ) + else: + prev["content"] = prev_content + [{"type": "text", "text": str(curr_content)}] + messages.pop(i) + else: + i += 1 + + +def filtered_stream( + byte_iter, + collected_chunks, + phantom_calls_out, + observe_only: bool | set[str] = False, + block_store=None, + page_store=None, + session_id: str = "", + continuation_needed: list | None = None, +): + """Filter phantom tool events from an SSE byte stream. + + Yields filtered bytes to the framework. Phantom tool_use events + are suppressed. Completed phantom calls are appended to + phantom_calls_out. + + When ALL tool_use blocks in a response are phantom, the stop_reason + is rewritten from "tool_use" to "end_turn" so the framework doesn't + enter tool execution mode for tools it never saw. + + observe_only can be: + - False: all phantom tools are intercepted (suppress from stream) + - True: all phantom tools are observed only (legacy compat) + - set[str]: tools in the set are observed; others are intercepted + + collected_chunks receives ALL bytes (including phantom) for logging. + """ + # Normalize observe_only to a set of tool names + if observe_only is True: + _observe_set = PHANTOM_TOOL_NAMES # all are observe-only + elif observe_only is False: + _observe_set: set[str] = set() # none are observe-only + else: + _observe_set = observe_only # per-tool set + + phantom_block_indices: set[int] = set() + real_tool_indices: set[int] = set() + # Track which phantom blocks are intercept-only (not observe-only) + intercept_block_indices: set[int] = set() + phantom_building: dict[int, dict] = {} + cleanup_filter = CleanupTagFilter(block_store=block_store, page_store=page_store) + if continuation_needed is None: + continuation_needed = [] + buffer = b"" + + for chunk in byte_iter: + collected_chunks.append(chunk) + buffer += chunk + + # Process complete SSE events (delimited by \n\n) + while b"\n\n" in buffer: + event_bytes, buffer = buffer.split(b"\n\n", 1) + event_text = event_bytes.decode("utf-8", errors="replace") + + suppress = False + data_str = None + + for line in event_text.split("\n"): + if line.startswith("data: "): + data_str = line[6:] + + # Suppress [DONE] when continuation is pending — + # the continuation will provide its own [DONE] + if data_str == "[DONE]" and continuation_needed: + continue + + if data_str and data_str != "[DONE]": + try: + event_data = json.loads(data_str) + + # Strip cleanup tags from text_delta events before + # the framework sees them. Tags can span chunks so + # CleanupTagFilter maintains state across events. + # Only strip if we have any intercept-mode tools active. + if ( + _observe_set != PHANTOM_TOOL_NAMES + and event_data.get("type") == "content_block_delta" + ): + delta = event_data.get("delta", {}) + if delta.get("type") == "text_delta": + raw_text = delta.get("text", "") + clean_text = cleanup_filter.filter(raw_text) + if clean_text != raw_text: + delta["text"] = clean_text + event_data["delta"] = delta + data_str = json.dumps(event_data) + event_bytes = f"data: {data_str}".encode("utf-8") + + suppress = _classify_event( + event_data, + phantom_block_indices, + real_tool_indices, + phantom_building, + phantom_calls_out, + intercept_indices=intercept_block_indices, + observe_set=_observe_set, + ) + # When all content blocks were intercepted phantom + # tools (no real tools, no observe-only phantom tools), + # the gateway needs to auto-continue. + if ( + intercept_block_indices + and event_data.get("type") == "message_delta" + and not real_tool_indices + # Only continue if ALL phantom blocks are intercepted + and phantom_block_indices == intercept_block_indices + ): + delta = event_data.get("delta", {}) + if delta.get("stop_reason") == "tool_use": + # Signal continuation needed + continuation_needed.append(True) + # Suppress this event — continuation will + # provide its own stop events + suppress = True + except json.JSONDecodeError: + pass + + # Also suppress message_stop when continuation is pending + if not suppress and continuation_needed and data_str and data_str != "[DONE]": + try: + evt = json.loads(data_str) + if evt.get("type") == "message_stop": + suppress = True + except json.JSONDecodeError: + pass + + if not suppress: + yield event_bytes + b"\n\n" + + # Flush remaining buffer + if buffer: + yield buffer + + # Flush cleanup filter — discard any partial/unclosed tag + cleanup_filter.flush() + if cleanup_filter.executed_ops: + import sys + + print( + f" [{session_id}] CLEANUP (stream): {'; '.join(cleanup_filter.executed_ops)}", + file=sys.stderr, + ) + + +def _classify_event( + event_data: dict, + phantom_indices: set[int], + real_tool_indices: set[int], + building: dict[int, dict], + completed: list[PhantomCall], + intercept_indices: set[int] | None = None, + observe_set: set[str] | None = None, +) -> bool: + """Returns True if this event should be suppressed. + + When observe_set is provided, only tools NOT in the set are + suppressed (intercept mode). Tools in the set are captured + but not suppressed (observe mode). + """ + if intercept_indices is None: + intercept_indices = set() + if observe_set is None: + observe_set = set() + + event_type = event_data.get("type", "") + index = event_data.get("index") + + if event_type == "content_block_start": + cb = event_data.get("content_block", {}) + if cb.get("type") == "tool_use": + tool_name = cb.get("name", "") + if tool_name in PHANTOM_TOOL_NAMES: + phantom_indices.add(index) + building[index] = { + "name": tool_name, + "id": cb.get("id", ""), + "input_json": "", + } + if tool_name not in observe_set: + intercept_indices.add(index) + return True # suppress — we're intercepting + return False # observe — let it through + else: + # Real tool — track so we know stop_reason is legitimate + real_tool_indices.add(index) + + elif event_type == "content_block_delta" and index in phantom_indices: + delta = event_data.get("delta", {}) + if delta.get("type") == "input_json_delta": + building[index]["input_json"] += delta.get("partial_json", "") + return index in intercept_indices + + elif event_type == "content_block_stop" and index in phantom_indices: + if index in building: + rec = building.pop(index) + try: + parsed_input = json.loads(rec["input_json"]) + except json.JSONDecodeError: + parsed_input = {} + completed.append( + PhantomCall( + name=rec["name"], + tool_use_id=rec["id"], + input=parsed_input, + ) + ) + return index in intercept_indices + + return False diff --git a/tests/test_admission.py b/tests/test_admission.py new file mode 100644 index 0000000..8c7ef09 --- /dev/null +++ b/tests/test_admission.py @@ -0,0 +1,156 @@ +"""Tests for the admission control module. + +Tests the AdmissionController scoring logic across all four axes: +type_score, novelty_score, utility_score, size_score, and the +threshold-based admit/reject decision. +""" + +from __future__ import annotations + +from mnemosyne.admission import AdmissionController, AdmissionScore, DEFAULT_THRESHOLD + + +# ── Type score ─────────────────────────────────────────────────────────── + + +def test_admission_score_design_decision(): + """design_decision has the highest type weight (0.9).""" + ctrl = AdmissionController() + score = ctrl.score("x" * 300, "design_decision", has_duplicate=False) + assert score.type_score == 0.9 + + +def test_admission_score_conversation_phase(): + """conversation_phase has the lowest type weight (0.3).""" + ctrl = AdmissionController() + score = ctrl.score("x" * 300, "conversation_phase", has_duplicate=False) + assert score.type_score == 0.3 + + +def test_admission_unknown_type_default(): + """Unknown object types default to 0.3.""" + ctrl = AdmissionController() + score = ctrl.score("x" * 300, "totally_unknown_type", has_duplicate=False) + assert score.type_score == 0.3 + + +# ── Novelty score ──────────────────────────────────────────────────────── + + +def test_admission_duplicate_penalized(): + """Duplicates get novelty_score=0.2 instead of 1.0.""" + ctrl = AdmissionController() + score_unique = ctrl.score("x" * 300, "file_context", has_duplicate=False) + score_dup = ctrl.score("x" * 300, "file_context", has_duplicate=True) + assert score_unique.novelty_score == 1.0 + assert score_dup.novelty_score == 0.2 + assert score_dup.total < score_unique.total + + +# ── Utility score ──────────────────────────────────────────────────────── + + +def test_admission_utility_code_blocks(): + """Content with code blocks gets utility boost.""" + ctrl = AdmissionController() + content_with_code = "Here is code:\n```python\ndef foo(): pass\n```\n" + "x" * 200 + content_plain = "Just some plain text without code blocks. " + "x" * 200 + score_code = ctrl.score(content_with_code, "file_context", has_duplicate=False) + score_plain = ctrl.score(content_plain, "file_context", has_duplicate=False) + assert score_code.utility_score > score_plain.utility_score + + +def test_admission_utility_entities(): + """Content with many key_entities gets utility boost.""" + ctrl = AdmissionController() + entities = ["foo.py", "bar.py", "baz.py"] + score_with = ctrl.score("x" * 300, "file_context", has_duplicate=False, key_entities=entities) + score_without = ctrl.score("x" * 300, "file_context", has_duplicate=False, key_entities=[]) + assert score_with.utility_score > score_without.utility_score + + +# ── Size score ─────────────────────────────────────────────────────────── + + +def test_admission_size_small_penalized(): + """Very small content (< 100 chars) gets penalized size_score.""" + ctrl = AdmissionController() + score = ctrl.score("tiny", "file_context", has_duplicate=False) + assert score.size_score < 1.0 + assert score.size_score == len("tiny") / 100 + + +def test_admission_size_large_penalized(): + """Very large content (> 50k chars) gets penalized size_score.""" + ctrl = AdmissionController() + score = ctrl.score("x" * 80000, "file_context", has_duplicate=False) + assert score.size_score < 1.0 + assert score.size_score >= 0.3 + + +def test_admission_size_normal(): + """Normal-sized content (100-50k chars) gets size_score=1.0.""" + ctrl = AdmissionController() + score = ctrl.score("x" * 500, "file_context", has_duplicate=False) + assert score.size_score == 1.0 + + +# ── Threshold decisions ────────────────────────────────────────────────── + + +def test_admission_threshold_admit(): + """High-value content is admitted.""" + ctrl = AdmissionController() + # design_decision (0.9 type) + unique (1.0 novelty) + code + entities + length + content = "```python\ndef important(): pass\n```\n" + "x" * 300 + admitted, score = ctrl.should_admit( + content, + "design_decision", + has_duplicate=False, + key_entities=["foo.py", "bar.py", "baz.py"], + ) + assert admitted is True + assert score.total >= DEFAULT_THRESHOLD + + +def test_admission_threshold_reject(): + """Low-value duplicate conversation_phase is rejected.""" + ctrl = AdmissionController() + # conversation_phase (0.3 type) + duplicate (0.2 novelty) + tiny content + admitted, score = ctrl.should_admit("hi", "conversation_phase", has_duplicate=True) + assert admitted is False + assert score.total < DEFAULT_THRESHOLD + + +# ── Stats tracking ─────────────────────────────────────────────────────── + + +def test_admission_stats_tracking(): + """Stats correctly count admitted and rejected objects.""" + ctrl = AdmissionController() + assert ctrl.stats == {"admitted": 0, "rejected": 0} + + # Admit a high-value object + ctrl.should_admit( + "```python\ndef foo(): pass\n```\n" + "x" * 300, + "design_decision", + has_duplicate=False, + key_entities=["a.py", "b.py", "c.py"], + ) + assert ctrl.stats["admitted"] == 1 + + # Reject a low-value object + ctrl.should_admit("hi", "conversation_phase", has_duplicate=True) + assert ctrl.stats["rejected"] == 1 + assert ctrl.stats["admitted"] == 1 + + +# ── Edge cases ─────────────────────────────────────────────────────────── + + +def test_admission_empty_content(): + """Empty content gets size_score=0 and low utility.""" + ctrl = AdmissionController() + score = ctrl.score("", "file_context", has_duplicate=False) + assert score.size_score == 0.0 + assert score.utility_score == 0.0 diff --git a/tests/test_entropy.py b/tests/test_entropy.py new file mode 100644 index 0000000..77d9c51 --- /dev/null +++ b/tests/test_entropy.py @@ -0,0 +1,182 @@ +"""Tests for the entropy-gated faulting module. + +Tests the EntropyDetector's ability to detect hedging, uncertainty, +evicted entity references, and the should_fault decision logic. +""" + +from __future__ import annotations + +from mnemosyne.entropy import ( + DEFAULT_FAULT_THRESHOLD, + EntropyDetector, + EntropySignal, +) + + +# ── Hedging detection ──────────────────────────────────────────────────── + + +def test_entropy_hedging_detected(): + """Hedging language like 'I think' and 'probably' is detected.""" + detector = EntropyDetector() + signal = detector.analyze_response( + "I think the function is probably in utils.py", + evicted_entities=[], + ) + assert signal.has_hedging is True + assert signal.score > 0 + + +def test_entropy_multiple_patterns(): + """Multiple hedging patterns increase the score.""" + detector = EntropyDetector() + signal_one = detector.analyze_response( + "I think it might work.", + evicted_entities=[], + ) + signal_many = detector.analyze_response( + "I think it probably might be in the file, if I recall correctly. I believe so.", + evicted_entities=[], + ) + assert signal_many.score >= signal_one.score + + +# ── Uncertainty detection ──────────────────────────────────────────────── + + +def test_entropy_uncertainty_detected(): + """Explicit uncertainty like 'I'm not sure' is detected.""" + detector = EntropyDetector() + signal = detector.analyze_response( + "I'm not sure about the exact implementation. It's unclear to me.", + evicted_entities=[], + ) + assert signal.has_uncertainty is True + assert signal.score > 0 + + +# ── No signals ─────────────────────────────────────────────────────────── + + +def test_entropy_no_signals(): + """Clean, confident response has no entropy signals.""" + detector = EntropyDetector() + signal = detector.analyze_response( + "The function `calculate_total` is defined in src/utils.py at line 42. " + "It takes two arguments: price and quantity.", + evicted_entities=[], + ) + assert signal.has_hedging is False + assert signal.has_uncertainty is False + assert signal.references_evicted is False + assert signal.has_hallucination_risk is False + assert signal.score == 0.0 + + +# ── Evicted entity references ─────────────────────────────────────────── + + +def test_entropy_evicted_references(): + """References to evicted entities are detected.""" + detector = EntropyDetector() + signal = detector.analyze_response( + "The config is in settings.py and the handler is in views.py", + evicted_entities=["settings.py", "views.py", "models.py"], + ) + assert signal.references_evicted is True + assert "settings.py" in signal.referenced_entities + assert "views.py" in signal.referenced_entities + assert "models.py" not in signal.referenced_entities + + +def test_entropy_hallucination_risk(): + """Hedging + evicted references = hallucination risk.""" + detector = EntropyDetector() + signal = detector.analyze_response( + "I think the config is probably in settings.py somewhere", + evicted_entities=["settings.py"], + ) + assert signal.has_hedging is True + assert signal.references_evicted is True + assert signal.has_hallucination_risk is True + # Hallucination risk adds a bonus to the score + assert signal.score > 0.3 + + +# ── Composite score ────────────────────────────────────────────────────── + + +def test_entropy_composite_score(): + """Score combines hedging, uncertainty, and evicted references.""" + detector = EntropyDetector() + # All signals present + signal = detector.analyze_response( + "I'm not sure, but I think the code is probably in utils.py", + evicted_entities=["utils.py"], + ) + assert signal.has_hedging is True + assert signal.has_uncertainty is True + assert signal.references_evicted is True + assert signal.score > 0.5 # Should be high with all signals + + +# ── should_fault decisions ─────────────────────────────────────────────── + + +def test_entropy_should_fault_above_threshold(): + """Faulting triggers when score >= threshold AND references evicted.""" + detector = EntropyDetector(fault_threshold=0.4) + signal = detector.analyze_response( + "I think the code is probably in utils.py, if I recall correctly. I believe so.", + evicted_entities=["utils.py"], + ) + # Score should be above 0.4 with hedging + evicted ref + hallucination bonus + entities = detector.should_fault(signal) + assert len(entities) > 0 + assert "utils.py" in entities + + +def test_entropy_should_fault_below_threshold(): + """No faulting when score is below threshold.""" + detector = EntropyDetector(fault_threshold=0.99) # Very high threshold + signal = detector.analyze_response( + "I think it might be there.", + evicted_entities=[], + ) + entities = detector.should_fault(signal) + assert entities == [] + + +def test_entropy_should_fault_no_evicted_refs(): + """No faulting when there are no evicted entity references, even with high hedging.""" + detector = EntropyDetector(fault_threshold=0.1) # Very low threshold + signal = detector.analyze_response( + "I think it probably might be somewhere, if I recall. I believe so.", + evicted_entities=[], + ) + # Even though hedging score is high, no evicted refs → no fault + entities = detector.should_fault(signal) + assert entities == [] + + +# ── Stats tracking ─────────────────────────────────────────────────────── + + +def test_entropy_stats_tracking(): + """Stats correctly count analyzed responses and triggered faults.""" + detector = EntropyDetector(fault_threshold=0.3) + assert detector.stats == {"analyzed": 0, "faults_triggered": 0} + + # Analyze a clean response + detector.analyze_response("The function is in utils.py.", evicted_entities=[]) + assert detector.stats["analyzed"] == 1 + assert detector.stats["faults_triggered"] == 0 + + # Analyze and fault a hedging response with evicted refs + signal = detector.analyze_response( + "I think the code is probably in utils.py, if I recall correctly. I believe so.", + evicted_entities=["utils.py"], + ) + assert detector.stats["analyzed"] == 2 + detector.should_fault(signal) + assert detector.stats["faults_triggered"] == 1 diff --git a/tests/test_goal_aware.py b/tests/test_goal_aware.py new file mode 100644 index 0000000..1ce53b9 --- /dev/null +++ b/tests/test_goal_aware.py @@ -0,0 +1,130 @@ +"""Tests for goal-aware retrieval integration in the gateway. + +Tests the Session attributes for goal tracking, cosine similarity-based +topic shift detection, and graceful degradation when helper_llm is absent. +""" + +from __future__ import annotations + +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from mnemosyne.gateway import Session +from mnemosyne.helper_llm import GoalClassification +from mnemosyne.object_store import DummyEmbedder, _cosine_similarity + + +# ── Fixtures ───────────────────────────────────────────────────────────── + + +@pytest.fixture +def tmp_log_dir(): + with TemporaryDirectory() as d: + yield Path(d) + + +@pytest.fixture +def session(tmp_log_dir): + return Session("goal01", tmp_log_dir) + + +# ── Session attributes ─────────────────────────────────────────────────── + + +class TestGoalSessionAttributes: + """Verify Session.__init__ creates goal-tracking attributes.""" + + def test_goal_session_attributes_exist(self, session): + """Session has _last_user_embedding and _current_goal attributes.""" + assert hasattr(session, "_last_user_embedding") + assert session._last_user_embedding is None + assert hasattr(session, "_current_goal") + assert session._current_goal is None + assert hasattr(session, "entropy_detector") + assert session.entropy_detector is not None + + +# ── Cosine similarity topic shift ──────────────────────────────────────── + + +class TestCosineTopicShift: + """Test topic shift detection via cosine similarity.""" + + def test_goal_cosine_similarity_topic_shift(self): + """Different topics produce low cosine similarity (< 0.5).""" + embedder = DummyEmbedder() + # Two very different texts should produce different embeddings + emb_a = embedder.embed("How do I configure the database connection pool?") + emb_b = embedder.embed("What color should the login button be?") + sim = _cosine_similarity(emb_a, emb_b) + # DummyEmbedder uses hash-based vectors, so different texts + # produce essentially random vectors with low expected similarity + # For 384-dim random unit vectors, expected |cos| ≈ 0.05 + assert sim < 0.5 + + def test_goal_cosine_similarity_same_topic(self): + """Identical text produces cosine similarity of 1.0.""" + embedder = DummyEmbedder() + text = "How do I configure the database connection pool?" + emb_a = embedder.embed(text) + emb_b = embedder.embed(text) + sim = _cosine_similarity(emb_a, emb_b) + assert sim > 0.99 # Same text → same hash → same embedding + + def test_goal_first_message_always_classifies(self, session): + """First message (no prior embedding) should always trigger classification.""" + # _last_user_embedding is None → goal_changed should be True + assert session._last_user_embedding is None + # This is the logic from gateway._preprocess step 1c: + # if session._last_user_embedding is not None: check sim + # else: goal_changed = True + goal_changed = session._last_user_embedding is None + assert goal_changed is True + + def test_goal_classification_fallback_no_helper(self, session): + """Goal detection is skipped gracefully when helper_llm is None. + + The gateway code wraps goal detection in try/except and checks + `if goal_changed and helper_llm is not None`. When helper_llm + is None, no classification occurs and _current_goal stays None. + """ + # Simulate the gateway logic: helper_llm is None + helper_llm = None + goal_changed = True # First message + + # This mirrors the gateway code path + if goal_changed and helper_llm is not None: + # Would call helper_llm.classify_goal(...) + session._current_goal = GoalClassification(goal="test") + + # Goal should remain None since helper_llm is None + assert session._current_goal is None + + def test_goal_embedding_updated_after_message(self, session): + """_last_user_embedding is updated after processing a message.""" + embedder = DummyEmbedder() + user_text = "Tell me about the authentication system" + current_embedding = embedder.embed(user_text) + + # Simulate the gateway update + session._last_user_embedding = current_embedding + + assert session._last_user_embedding is not None + assert len(session._last_user_embedding) == 384 + + def test_goal_classification_stored_on_session(self, session): + """GoalClassification is stored on session._current_goal.""" + goal = GoalClassification( + goal="Implement authentication", + relevant_types=["file_context", "design_decision"], + relevant_tags=["auth", "security"], + predicted_needs=["auth.py", "middleware.py"], + ) + session._current_goal = goal + + assert session._current_goal.goal == "Implement authentication" + assert "file_context" in session._current_goal.relevant_types + assert "auth" in session._current_goal.relevant_tags diff --git a/tests/test_hierarchy.py b/tests/test_hierarchy.py new file mode 100644 index 0000000..bad2bf9 --- /dev/null +++ b/tests/test_hierarchy.py @@ -0,0 +1,893 @@ +"""Tests for the hierarchical segmentation system (Strategy B). + +Covers: + - Episode creation and clustering + - SemanticTheme creation + - Theme creation + - SessionHierarchy: add_object, rebuild, retrieve, maintenance + - Edge cases: empty, single object, identical embeddings, no embeddings + - Vector math utilities + - Incremental vs full rebuild consistency +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from mnemosyne.hierarchy import ( + EPISODE_RETRIEVAL_THRESHOLD, + EPISODE_SIMILARITY_THRESHOLD, + Episode, + RECLUSTER_INTERVAL, + SEMANTIC_SIMILARITY_THRESHOLD, + SemanticTheme, + SessionHierarchy, + THEME_RETRIEVAL_THRESHOLD, + THEME_SIMILARITY_THRESHOLD, + Theme, + _agglomerative_cluster, + _centroid, + _cosine_similarity, + _cosine_similarity_matrix, + _generate_episode_summary, + _generate_theme_label, + _make_episode, + _make_semantic_theme, + _make_theme, +) +from mnemosyne.object_store import DummyEmbedder, StoredObject, _estimate_tokens + + +# ── Helpers ────────────────────────────────────────────────────────────── + + +def _make_stored_object( + content: str = "test content", + *, + session_id: str = "sess-1", + object_type: str = "file_context", + source_tool: str | None = "Read", + source_key: str | None = None, + stub: str | None = None, + embedding: list[float] | None = None, + object_id: str | None = None, + turn: int | None = None, +) -> StoredObject: + """Create a StoredObject with sensible defaults for testing.""" + return StoredObject( + id=object_id or f"obj-{hash(content) % 100000:05d}", + session_id=session_id, + object_type=object_type, + source_tool=source_tool, + source_key=source_key, + content_full=content, + summary_detailed=None, + summary_compact=None, + stub=stub or f"{object_type}: {content[:40]}", + tokens_l0=_estimate_tokens(content), + tokens_l3=_estimate_tokens(stub or f"{object_type}: test"), + embedding=embedding or [], + created_at="2025-01-01T00:00:00+00:00", + last_accessed="2025-01-01T00:00:00+00:00", + source_turn_start=turn, + source_turn_end=turn, + ) + + +def _make_unit_vector(dim: int = 384, seed: int = 42) -> list[float]: + """Create a deterministic unit vector.""" + rng = np.random.default_rng(seed) + vec = rng.standard_normal(dim) + vec = vec / np.linalg.norm(vec) + return vec.tolist() + + +def _make_similar_vector(base: list[float], noise: float = 0.05, seed: int = 99) -> list[float]: + """Create a vector similar to `base` by adding small noise.""" + rng = np.random.default_rng(seed) + arr = np.asarray(base, dtype=np.float64) + perturbation = rng.standard_normal(len(base)) * noise + result = arr + perturbation + result = result / np.linalg.norm(result) + return result.tolist() + + +def _make_orthogonal_vector(base: list[float], seed: int = 77) -> list[float]: + """Create a vector roughly orthogonal to `base`.""" + rng = np.random.default_rng(seed) + random_vec = rng.standard_normal(len(base)) + arr = np.asarray(base, dtype=np.float64) + # Gram-Schmidt: subtract projection onto base + proj = np.dot(random_vec, arr) / np.dot(arr, arr) * arr + ortho = random_vec - proj + norm = np.linalg.norm(ortho) + if norm > 0: + ortho = ortho / norm + return ortho.tolist() + + +# ── Vector math tests ──────────────────────────────────────────────────── + + +class TestCosineSimlarity: + def test_identical_vectors(self): + v = _make_unit_vector(seed=1) + assert _cosine_similarity(v, v) == pytest.approx(1.0, abs=1e-6) + + def test_orthogonal_vectors(self): + v1 = _make_unit_vector(seed=1) + v2 = _make_orthogonal_vector(v1, seed=2) + assert abs(_cosine_similarity(v1, v2)) < 0.05 + + def test_opposite_vectors(self): + v = _make_unit_vector(seed=1) + neg_v = [-x for x in v] + assert _cosine_similarity(v, neg_v) == pytest.approx(-1.0, abs=1e-6) + + def test_empty_vectors(self): + assert _cosine_similarity([], []) == 0.0 + assert _cosine_similarity([1.0, 0.0], []) == 0.0 + assert _cosine_similarity([], [1.0, 0.0]) == 0.0 + + def test_zero_vector(self): + v = _make_unit_vector(seed=1) + zero = [0.0] * len(v) + assert _cosine_similarity(v, zero) == 0.0 + + def test_similar_vectors_high_similarity(self): + v1 = _make_unit_vector(seed=1) + v2 = _make_similar_vector(v1, noise=0.01, seed=2) + sim = _cosine_similarity(v1, v2) + assert sim > 0.9 + + +class TestCosineSimlarityMatrix: + def test_single_vector(self): + v = _make_unit_vector(seed=1) + mat = _cosine_similarity_matrix(np.array([v])) + assert mat.shape == (1, 1) + assert mat[0, 0] == pytest.approx(1.0, abs=1e-6) + + def test_identity_diagonal(self): + vecs = [_make_unit_vector(seed=i) for i in range(5)] + mat = _cosine_similarity_matrix(np.array(vecs)) + for i in range(5): + assert mat[i, i] == pytest.approx(1.0, abs=1e-6) + + def test_symmetry(self): + vecs = [_make_unit_vector(seed=i) for i in range(3)] + mat = _cosine_similarity_matrix(np.array(vecs)) + for i in range(3): + for j in range(3): + assert mat[i, j] == pytest.approx(mat[j, i], abs=1e-10) + + def test_empty(self): + mat = _cosine_similarity_matrix(np.empty((0, 384))) + assert mat.shape == (0, 0) + + +class TestCentroid: + def test_single_vector(self): + v = _make_unit_vector(seed=1) + c = _centroid([v]) + # Centroid of a single unit vector is itself + assert _cosine_similarity(v, c) == pytest.approx(1.0, abs=1e-6) + + def test_identical_vectors(self): + v = _make_unit_vector(seed=1) + c = _centroid([v, v, v]) + assert _cosine_similarity(v, c) == pytest.approx(1.0, abs=1e-6) + + def test_empty(self): + assert _centroid([]) == [] + + def test_centroid_is_normalized(self): + vecs = [_make_unit_vector(seed=i) for i in range(5)] + c = _centroid(vecs) + norm = float(np.linalg.norm(c)) + assert norm == pytest.approx(1.0, abs=1e-6) + + +# ── Agglomerative clustering tests ────────────────────────────────────── + + +class TestAgglomerativeClustering: + def test_empty(self): + assert _agglomerative_cluster([], 0.5) == [] + + def test_single_item(self): + v = _make_unit_vector(seed=1) + clusters = _agglomerative_cluster([v], 0.5) + assert len(clusters) == 1 + assert clusters[0] == [0] + + def test_identical_items_merge(self): + v = _make_unit_vector(seed=1) + clusters = _agglomerative_cluster([v, v, v], 0.5) + assert len(clusters) == 1 + assert sorted(clusters[0]) == [0, 1, 2] + + def test_dissimilar_items_separate(self): + v1 = _make_unit_vector(seed=1) + v2 = _make_orthogonal_vector(v1, seed=2) + clusters = _agglomerative_cluster([v1, v2], 0.5) + assert len(clusters) == 2 + + def test_similar_items_merge(self): + v1 = _make_unit_vector(seed=1) + v2 = _make_similar_vector(v1, noise=0.02, seed=2) + clusters = _agglomerative_cluster([v1, v2], 0.5) + assert len(clusters) == 1 + + def test_mixed_clusters(self): + """Two groups of similar vectors should form two clusters.""" + v1 = _make_unit_vector(seed=1) + v1b = _make_similar_vector(v1, noise=0.02, seed=10) + v2 = _make_orthogonal_vector(v1, seed=2) + v2b = _make_similar_vector(v2, noise=0.02, seed=20) + clusters = _agglomerative_cluster([v1, v1b, v2, v2b], 0.5) + assert len(clusters) == 2 + + def test_threshold_boundary(self): + """Items with high similarity should merge at a reasonable threshold.""" + v = _make_unit_vector(seed=1) + # Identical vectors have similarity ~1.0 (floating point) + clusters = _agglomerative_cluster([v, v], 0.99) + assert len(clusters) == 1 + + +# ── Episode dataclass tests ───────────────────────────────────────────── + + +class TestEpisode: + def test_creation(self): + ep = Episode(id="ep-1") + assert ep.id == "ep-1" + assert ep.objects == [] + assert ep.embedding == [] + assert ep.summary == "" + assert ep.turn_range == (0, 0) + assert ep.object_types == set() + + def test_hash_and_equality(self): + ep1 = Episode(id="ep-1") + ep2 = Episode(id="ep-1") + ep3 = Episode(id="ep-2") + assert ep1 == ep2 + assert ep1 != ep3 + assert hash(ep1) == hash(ep2) + + def test_make_episode(self): + embedder = DummyEmbedder() + obj1 = _make_stored_object( + "content A", object_id="a", turn=5, embedding=embedder.embed("content A") + ) + obj2 = _make_stored_object( + "content B", object_id="b", turn=8, embedding=embedder.embed("content B") + ) + ep = _make_episode([obj1, obj2]) + assert len(ep.objects) == 2 + assert ep.turn_range == (5, 8) + assert "file_context" in ep.object_types + assert ep.embedding # Should have a centroid + assert ep.summary # Should have auto-generated summary + + +class TestSemanticTheme: + def test_creation(self): + st = SemanticTheme(id="st-1") + assert st.id == "st-1" + assert st.episodes == [] + assert st.embedding == [] + assert st.label == "" + + def test_hash_and_equality(self): + st1 = SemanticTheme(id="st-1") + st2 = SemanticTheme(id="st-1") + st3 = SemanticTheme(id="st-2") + assert st1 == st2 + assert st1 != st3 + + +class TestTheme: + def test_creation(self): + t = Theme(id="t-1") + assert t.id == "t-1" + assert t.semantic_themes == [] + + def test_hash_and_equality(self): + t1 = Theme(id="t-1") + t2 = Theme(id="t-1") + t3 = Theme(id="t-2") + assert t1 == t2 + assert t1 != t3 + + +# ── Summary generation tests ──────────────────────────────────────────── + + +class TestSummaryGeneration: + def test_episode_summary_from_stubs(self): + obj = _make_stored_object("hello world", stub="file_context: hello") + summary = _generate_episode_summary([obj]) + assert "file_context: hello" in summary + + def test_episode_summary_empty(self): + summary = _generate_episode_summary([]) + assert "empty" in summary.lower() + + def test_episode_summary_many_objects(self): + objs = [ + _make_stored_object(f"content {i}", object_id=f"o{i}", stub=f"stub {i}") + for i in range(10) + ] + summary = _generate_episode_summary(objs) + assert "+5 more" in summary + + def test_theme_label_from_episodes(self): + ep = Episode(id="ep-1", object_types={"file_context", "tool_result"}) + label = _generate_theme_label([ep]) + assert "file_context" in label or "tool_result" in label + assert "1 episodes" in label + + def test_theme_label_empty(self): + label = _generate_theme_label([]) + assert "empty" in label.lower() + + +# ── SessionHierarchy tests ────────────────────────────────────────────── + + +class TestSessionHierarchyInit: + def test_init(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + assert h.enabled is True + assert h.object_count == 0 + assert h.episode_count == 0 + assert h.theme_count == 0 + + def test_disable(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + h.enabled = False + assert h.enabled is False + + def test_summary_empty(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + s = h.summary() + assert s["objects"] == 0 + assert s["episodes"] == 0 + assert s["themes"] == 0 + + +class TestSessionHierarchyAddObject: + def test_add_single_object(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + obj = _make_stored_object( + "test content", object_id="o1", embedding=embedder.embed("test content") + ) + ep = h.add_object(obj) + assert ep is not None + assert h.object_count == 1 + assert h.episode_count == 1 + + def test_add_duplicate_ignored(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + obj = _make_stored_object("test", object_id="o1", embedding=embedder.embed("test")) + h.add_object(obj) + result = h.add_object(obj) + assert result is None + assert h.object_count == 1 + + def test_add_similar_objects_same_episode(self): + """Objects with very similar embeddings should join the same episode.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + + # Use the same text to get identical embeddings + emb = embedder.embed("shared topic content") + obj1 = _make_stored_object("shared topic content A", object_id="o1", embedding=emb) + obj2 = _make_stored_object( + "shared topic content B", object_id="o2", embedding=emb + ) # Same embedding + + ep1 = h.add_object(obj1) + ep2 = h.add_object(obj2) + assert ep1 is not None + assert ep2 is not None + assert ep1.id == ep2.id # Same episode + assert h.episode_count == 1 + + def test_add_dissimilar_objects_different_episodes(self): + """Objects with very different embeddings should go to different episodes.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + + # DummyEmbedder produces different vectors for different text + obj1 = _make_stored_object( + "python programming language", + object_id="o1", + embedding=embedder.embed("python programming language"), + ) + obj2 = _make_stored_object( + "cooking recipes for dinner", + object_id="o2", + embedding=embedder.embed("cooking recipes for dinner"), + ) + + h.add_object(obj1) + h.add_object(obj2) + # DummyEmbedder is hash-based, so different texts → different vectors + # They should be in different episodes (similarity < 0.6) + assert h.episode_count == 2 + + def test_add_object_without_embedding(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + obj = _make_stored_object("no embedding", object_id="o1", embedding=[]) + ep = h.add_object(obj) + assert ep is not None + assert h.episode_count == 1 + + def test_add_object_disabled(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + h.enabled = False + obj = _make_stored_object("test", object_id="o1", embedding=embedder.embed("test")) + result = h.add_object(obj) + assert result is None + assert h.object_count == 0 + + def test_add_updates_episode_metadata(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("shared content") + obj1 = _make_stored_object( + "shared content A", object_id="o1", object_type="file_context", embedding=emb, turn=5 + ) + obj2 = _make_stored_object( + "shared content B", object_id="o2", object_type="tool_result", embedding=emb, turn=10 + ) + h.add_object(obj1) + ep = h.add_object(obj2) + assert ep is not None + assert "file_context" in ep.object_types + assert "tool_result" in ep.object_types + assert ep.turn_range == (5, 10) + + +class TestSessionHierarchyRebuild: + def test_rebuild_empty(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + h.rebuild([]) + assert h.object_count == 0 + assert h.episode_count == 0 + + def test_rebuild_single_object(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + obj = _make_stored_object("test", object_id="o1", embedding=embedder.embed("test")) + h.rebuild([obj]) + assert h.object_count == 1 + assert h.episode_count == 1 + assert len(h.get_themes()) >= 1 + assert len(h.get_top_themes()) >= 1 + + def test_rebuild_clusters_similar_objects(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("shared topic") + objects = [ + _make_stored_object(f"shared topic {i}", object_id=f"o{i}", embedding=emb) + for i in range(5) + ] + h.rebuild(objects) + assert h.object_count == 5 + # All identical embeddings → 1 episode + assert h.episode_count == 1 + + def test_rebuild_separates_dissimilar_objects(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + objects = [ + _make_stored_object( + f"unique topic number {i} with distinct content", + object_id=f"o{i}", + embedding=embedder.embed(f"unique topic number {i} with distinct content"), + ) + for i in range(10) + ] + h.rebuild(objects) + assert h.object_count == 10 + # DummyEmbedder with different texts → different vectors → multiple episodes + assert h.episode_count > 1 + + def test_rebuild_handles_mixed_embeddings(self): + """Objects with and without embeddings should both be included.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + obj_with = _make_stored_object( + "with embedding", object_id="o1", embedding=embedder.embed("with embedding") + ) + obj_without = _make_stored_object("no embedding", object_id="o2", embedding=[]) + h.rebuild([obj_with, obj_without]) + assert h.object_count == 2 + assert h.episode_count == 2 # Each in its own episode + + def test_rebuild_clears_previous_state(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + obj1 = _make_stored_object("first", object_id="o1", embedding=embedder.embed("first")) + h.rebuild([obj1]) + assert h.object_count == 1 + + obj2 = _make_stored_object("second", object_id="o2", embedding=embedder.embed("second")) + h.rebuild([obj2]) + assert h.object_count == 1 # Only the new object + assert h.episode_count == 1 + + def test_rebuild_creates_all_levels(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + objects = [ + _make_stored_object( + f"content {i}", object_id=f"o{i}", embedding=embedder.embed(f"content {i}") + ) + for i in range(20) + ] + h.rebuild(objects) + assert h.object_count == 20 + assert h.episode_count > 0 + assert len(h.get_themes()) > 0 + assert len(h.get_top_themes()) > 0 + + +class TestSessionHierarchyRetrieve: + def test_retrieve_empty_hierarchy(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + query = embedder.embed("test query") + results = h.retrieve(query, limit=5) + assert results == [] + + def test_retrieve_returns_relevant_objects(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("python programming") + obj = _make_stored_object("python programming guide", object_id="o1", embedding=emb) + h.rebuild([obj]) + results = h.retrieve(emb, limit=5) + assert len(results) == 1 + assert results[0].id == "o1" + + def test_retrieve_respects_limit(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("shared") + objects = [ + _make_stored_object(f"shared content {i}", object_id=f"o{i}", embedding=emb) + for i in range(20) + ] + h.rebuild(objects) + results = h.retrieve(emb, limit=5) + assert len(results) <= 5 + + def test_retrieve_empty_query(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + obj = _make_stored_object("test", object_id="o1", embedding=embedder.embed("test")) + h.rebuild([obj]) + results = h.retrieve([], limit=5) + assert results == [] + + def test_retrieve_disabled_falls_back_to_flat(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test content", object_id="o1", embedding=emb) + h.rebuild([obj]) + h.enabled = False + results = h.retrieve(emb, limit=5) + # Falls back to flat search + assert len(results) == 1 + + def test_retrieve_sorted_by_similarity(self): + """Results should be sorted by similarity to query (highest first).""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + + query_emb = embedder.embed("target query") + # Create objects with varying similarity to query + obj_close = _make_stored_object( + "target query exact", object_id="close", embedding=query_emb + ) + obj_far = _make_stored_object( + "completely unrelated xyz", + object_id="far", + embedding=embedder.embed("completely unrelated xyz"), + ) + + h.rebuild([obj_close, obj_far]) + results = h.retrieve(query_emb, limit=10) + assert len(results) >= 1 + # The close object should be first + assert results[0].id == "close" + + def test_retrieve_no_duplicates(self): + """Retrieve should not return duplicate objects.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("shared") + objects = [ + _make_stored_object(f"shared {i}", object_id=f"o{i}", embedding=emb) for i in range(5) + ] + h.rebuild(objects) + results = h.retrieve(emb, limit=10) + ids = [r.id for r in results] + assert len(ids) == len(set(ids)) + + def test_flat_search_fallback(self): + """When hierarchy has no themes, should fall back to flat search.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test", object_id="o1", embedding=emb) + # Add object but clear themes to force fallback + h._all_objects = [obj] + h._object_ids = {obj.id} + h._themes = [] + results = h.retrieve(emb, limit=5) + assert len(results) == 1 + + +class TestSessionHierarchyMaintenance: + def test_maintenance_at_interval(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test", object_id="o1", embedding=emb) + h.add_object(obj) + + # Turn 20 should trigger rebuild + rebuilt = h.maintenance(turn=RECLUSTER_INTERVAL) + assert rebuilt is True + + def test_maintenance_not_at_interval(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test", object_id="o1", embedding=emb) + h.add_object(obj) + + rebuilt = h.maintenance(turn=15) + assert rebuilt is False + + def test_maintenance_on_goal_change(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test", object_id="o1", embedding=emb) + h.add_object(obj) + + rebuilt = h.maintenance(turn=5, goal_hash="goal-1") + assert rebuilt is True # First goal → always rebuild + + rebuilt = h.maintenance(turn=6, goal_hash="goal-1") + assert rebuilt is False # Same goal, not at interval + + rebuilt = h.maintenance(turn=7, goal_hash="goal-2") + assert rebuilt is True # Goal changed + + def test_maintenance_disabled(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + h.enabled = False + rebuilt = h.maintenance(turn=RECLUSTER_INTERVAL) + assert rebuilt is False + + def test_maintenance_no_double_rebuild(self): + """Same turn should not trigger rebuild twice.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test", object_id="o1", embedding=emb) + h.add_object(obj) + + rebuilt1 = h.maintenance(turn=RECLUSTER_INTERVAL) + assert rebuilt1 is True + rebuilt2 = h.maintenance(turn=RECLUSTER_INTERVAL) + assert rebuilt2 is False + + def test_maintenance_at_multiple_intervals(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test", object_id="o1", embedding=emb) + h.add_object(obj) + + assert h.maintenance(turn=20) is True + assert h.maintenance(turn=40) is True + assert h.maintenance(turn=60) is True + + def test_maintenance_turn_zero_no_rebuild(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + rebuilt = h.maintenance(turn=0) + assert rebuilt is False + + +class TestSessionHierarchyGetters: + def test_get_episodes(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test", object_id="o1", embedding=emb) + h.add_object(obj) + episodes = h.get_episodes() + assert len(episodes) == 1 + assert episodes[0].objects[0].id == "o1" + + def test_get_themes(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test", object_id="o1", embedding=emb) + h.add_object(obj) + themes = h.get_themes() + assert len(themes) >= 1 + + def test_get_top_themes(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("test") + obj = _make_stored_object("test", object_id="o1", embedding=emb) + h.add_object(obj) + top_themes = h.get_top_themes() + assert len(top_themes) >= 1 + + +class TestSessionHierarchySummary: + def test_summary_after_rebuild(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + objects = [ + _make_stored_object( + f"content {i}", object_id=f"o{i}", embedding=embedder.embed(f"content {i}") + ) + for i in range(10) + ] + h.rebuild(objects) + s = h.summary() + assert s["objects"] == 10 + assert s["episodes"] > 0 + assert s["themes"] > 0 + assert s["enabled"] is True + assert isinstance(s["avg_episode_size"], float) + + +# ── Edge case tests ───────────────────────────────────────────────────── + + +class TestEdgeCases: + def test_all_identical_embeddings(self): + """All objects with identical embeddings → single episode.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + emb = embedder.embed("identical") + objects = [ + _make_stored_object(f"identical {i}", object_id=f"o{i}", embedding=emb) + for i in range(10) + ] + h.rebuild(objects) + assert h.episode_count == 1 + assert h.object_count == 10 + + def test_all_objects_no_embeddings(self): + """Objects without embeddings each get their own episode.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + objects = [ + _make_stored_object(f"no emb {i}", object_id=f"o{i}", embedding=[]) for i in range(5) + ] + h.rebuild(objects) + assert h.episode_count == 5 + + def test_large_number_of_objects(self): + """Hierarchy should handle 100+ objects without error.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + objects = [ + _make_stored_object( + f"object number {i} with unique content", + object_id=f"o{i}", + embedding=embedder.embed(f"object number {i} with unique content"), + ) + for i in range(100) + ] + h.rebuild(objects) + assert h.object_count == 100 + assert h.episode_count > 0 + assert len(h.get_top_themes()) > 0 + + def test_retrieve_with_many_objects(self): + """Retrieval should work efficiently with many objects.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + objects = [ + _make_stored_object( + f"content {i}", object_id=f"o{i}", embedding=embedder.embed(f"content {i}") + ) + for i in range(50) + ] + h.rebuild(objects) + query = embedder.embed("content 0") + results = h.retrieve(query, limit=5) + assert len(results) <= 5 + # The exact match should be in results + result_ids = [r.id for r in results] + assert "o0" in result_ids + + def test_incremental_then_rebuild_consistency(self): + """Incremental adds followed by rebuild should produce valid hierarchy.""" + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + + objects = [] + for i in range(10): + obj = _make_stored_object( + f"content {i}", object_id=f"o{i}", embedding=embedder.embed(f"content {i}") + ) + objects.append(obj) + h.add_object(obj) + + # Now rebuild from scratch + h.rebuild(objects) + assert h.object_count == 10 + assert h.episode_count > 0 + + def test_rebuild_with_single_object_no_crash(self): + embedder = DummyEmbedder() + h = SessionHierarchy(embedder) + obj = _make_stored_object("solo", object_id="o1", embedding=embedder.embed("solo")) + h.rebuild([obj]) + results = h.retrieve(embedder.embed("solo"), limit=5) + assert len(results) == 1 + + def test_episode_not_equal_to_non_episode(self): + ep = Episode(id="ep-1") + assert ep != "not an episode" + + def test_semantic_theme_not_equal_to_non_theme(self): + st = SemanticTheme(id="st-1") + assert st != 42 + + def test_theme_not_equal_to_non_theme(self): + t = Theme(id="t-1") + assert t != [] + + +# ── Constants tests ───────────────────────────────────────────────────── + + +class TestConstants: + def test_episode_threshold(self): + assert EPISODE_SIMILARITY_THRESHOLD == 0.6 + + def test_semantic_threshold(self): + assert SEMANTIC_SIMILARITY_THRESHOLD == 0.4 + + def test_theme_threshold(self): + assert THEME_SIMILARITY_THRESHOLD == 0.25 + + def test_retrieval_thresholds(self): + assert THEME_RETRIEVAL_THRESHOLD == 0.3 + assert EPISODE_RETRIEVAL_THRESHOLD == 0.4 + + def test_recluster_interval(self): + assert RECLUSTER_INTERVAL == 20 diff --git a/tests/test_phantom_memory_query.py b/tests/test_phantom_memory_query.py new file mode 100644 index 0000000..8452236 --- /dev/null +++ b/tests/test_phantom_memory_query.py @@ -0,0 +1,180 @@ +"""Tests for the memory_query phantom tool addition to phantom.py.""" + +from __future__ import annotations + +import copy + +from mnemosyne.phantom import ( + PHANTOM_TOOL_DEFINITIONS, + PHANTOM_TOOL_NAMES, + PhantomCall, + _handle_phantom_call, + inject_tools, +) + + +# ── memory_query in PHANTOM_TOOL_NAMES ─────────────────────── + + +def test_memory_query_in_phantom_tool_names(): + """memory_query must be a recognized phantom tool name.""" + assert "memory_query" in PHANTOM_TOOL_NAMES + + +def test_phantom_tool_names_is_frozenset(): + """PHANTOM_TOOL_NAMES must remain a frozenset (immutable).""" + assert isinstance(PHANTOM_TOOL_NAMES, frozenset) + + +# ── memory_query tool definition schema ────────────────────── + + +def _get_memory_query_def() -> dict: + """Helper: find the memory_query definition from PHANTOM_TOOL_DEFINITIONS.""" + for defn in PHANTOM_TOOL_DEFINITIONS: + if defn["name"] == "memory_query": + return defn + raise AssertionError("memory_query not found in PHANTOM_TOOL_DEFINITIONS") + + +def test_memory_query_definition_exists(): + """memory_query must have a definition in PHANTOM_TOOL_DEFINITIONS.""" + defn = _get_memory_query_def() + assert defn["name"] == "memory_query" + + +def test_memory_query_has_description(): + """memory_query definition must have a non-empty description.""" + defn = _get_memory_query_def() + assert isinstance(defn["description"], str) + assert len(defn["description"]) > 20 + + +def test_memory_query_schema_properties(): + """memory_query input_schema must have question, scope, max_tokens.""" + defn = _get_memory_query_def() + schema = defn["input_schema"] + assert schema["type"] == "object" + props = schema["properties"] + assert "question" in props + assert props["question"]["type"] == "string" + assert "scope" in props + assert props["scope"]["type"] == "string" + assert "max_tokens" in props + assert props["max_tokens"]["type"] == "integer" + + +def test_memory_query_required_fields(): + """Only 'question' should be required for memory_query.""" + defn = _get_memory_query_def() + required = defn["input_schema"]["required"] + assert required == ["question"] + + +# ── inject_tools adds memory_query ─────────────────────────── + + +def test_inject_tools_adds_memory_query(): + """inject_tools should add memory_query to the tools array.""" + body = {"tools": [], "messages": []} + inject_tools(body) + tool_names = {t["name"] for t in body["tools"]} + assert "memory_query" in tool_names + + +def test_inject_tools_does_not_duplicate_memory_query(): + """inject_tools should not duplicate memory_query if already present.""" + existing_def = { + "name": "memory_query", + "description": "already here", + "input_schema": {"type": "object", "properties": {}}, + } + body = {"tools": [existing_def], "messages": []} + inject_tools(body) + mq_count = sum(1 for t in body["tools"] if t["name"] == "memory_query") + assert mq_count == 1 + + +def test_inject_tools_returns_memory_query_as_observe_only_when_framework_provides(): + """If framework already has memory_query, it should be in observe_only set.""" + existing_def = { + "name": "memory_query", + "description": "framework provided", + "input_schema": {"type": "object", "properties": {}}, + } + body = {"tools": [existing_def], "messages": []} + observe_only = inject_tools(body) + assert "memory_query" in observe_only + + +# ── _handle_phantom_call returns pending placeholder ───────── + + +def test_handle_phantom_call_memory_query_returns_pending(): + """memory_query handler should return a pending placeholder string.""" + call = PhantomCall( + name="memory_query", + tool_use_id="toolu_test123", + input={"question": "What auth library is used?"}, + ) + result = _handle_phantom_call(call, page_store=None) + assert "[memory_query:pending]" in result + assert "What auth library is used?" in result + + +def test_handle_phantom_call_memory_query_includes_scope(): + """memory_query handler should include scope in the placeholder.""" + call = PhantomCall( + name="memory_query", + tool_use_id="toolu_test456", + input={ + "question": "What port does the server run on?", + "scope": "config files", + }, + ) + result = _handle_phantom_call(call, page_store=None) + assert "[memory_query:pending]" in result + assert "config files" in result + + +def test_handle_phantom_call_memory_query_includes_max_tokens(): + """memory_query handler should include max_tokens in the placeholder.""" + call = PhantomCall( + name="memory_query", + tool_use_id="toolu_test789", + input={ + "question": "What is the DB schema?", + "max_tokens": 100, + }, + ) + result = _handle_phantom_call(call, page_store=None) + assert "[memory_query:pending]" in result + assert "100" in result + + +def test_handle_phantom_call_memory_query_defaults(): + """memory_query handler should use defaults for optional fields.""" + call = PhantomCall( + name="memory_query", + tool_use_id="toolu_defaults", + input={"question": "test question"}, + ) + result = _handle_phantom_call(call, page_store=None) + assert "[memory_query:pending]" in result + # Default scope is None, default max_tokens is 200 + assert "None" in result + assert "200" in result + + +# ── Existing tools still work ──────────────────────────────── + + +def test_existing_phantom_tools_still_present(): + """All original phantom tool names must still be in PHANTOM_TOOL_NAMES.""" + for name in ("yuyay", "recall", "memory_fault", "qunqay", "tiqsiy"): + assert name in PHANTOM_TOOL_NAMES, f"{name} missing from PHANTOM_TOOL_NAMES" + + +def test_existing_tool_definitions_count(): + """PHANTOM_TOOL_DEFINITIONS should now have 4 tools (3 original + memory_query).""" + assert len(PHANTOM_TOOL_DEFINITIONS) == 4