feat: add object store with semantic segmentation

Object-addressed memory: segment messages into semantic objects,
embed with sentence-transformers, store in pgvector-backed store,
and reassemble context via goal-aware retrieval.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-opencode)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
Joey Yakimowich-Payne 2026-03-13 11:41:04 -06:00
commit a13719f754
9 changed files with 5644 additions and 0 deletions

View file

@ -0,0 +1,220 @@
"""Context assembly and micro-fault orchestration.
Orchestrates the full micro-fault flow for the memory_query phantom tool:
1. Embed the question via ObjectStore's embedder
2. Search ObjectStore for relevant objects (including evicted ones)
3. Send question + retrieved full content to HelperLLM
4. Return a targeted answer (50-200 tokens)
Also provides context assembly building the context window from
semantic objects at their current fidelity levels.
See ARCHITECTURE.md §7.1-7.3 for the full specification.
"""
from __future__ import annotations
import time
from dataclasses import dataclass, field
from mnemosyne.fidelity import FidelityLevel, FidelityManager
from mnemosyne.helper_llm import HelperLLM
from mnemosyne.object_store import ObjectStore, StoredObject
def _estimate_tokens(text: str | None) -> int:
"""Estimate token count using ~4 chars per token heuristic."""
if text is None:
return 0
return max(1, len(text) // 4)
# ── Data classes ─────────────────────────────────────────────
@dataclass
class MicroFaultResult:
"""Result of a memory_query micro-fault resolution."""
answer: str
sources: list[str] # IDs of objects consulted
answer_tokens: int
avoided_tokens: int # tokens saved vs full restore
latency_ms: float
@dataclass
class ContextWindow:
"""Snapshot of the assembled context window."""
objects: list[tuple[str, int, str]] # (object_id, fidelity, content)
total_tokens: int
pressure_zone: str
# ── ContextAssembler ─────────────────────────────────────────
class ContextAssembler:
"""Orchestrates context assembly, micro-faults, and goal-aware retrieval."""
def __init__(
self,
object_store: ObjectStore,
helper_llm: HelperLLM | None = None,
) -> None:
self.store = object_store
self.helper = helper_llm
async def handle_micro_fault(
self,
session_id: str,
question: str,
scope: str | None = None,
max_tokens: int = 200,
) -> MicroFaultResult:
"""Handle a memory_query micro-fault.
Flow:
1. Search ObjectStore for relevant objects (including evicted ones)
using hybrid semantic search
2. If HelperLLM is available, send question + retrieved full content
to get a targeted answer
3. If HelperLLM is unavailable, fall back to returning the best
available summaries from the top-k objects
Records micro-fault access on each consulted object.
Args:
session_id: The session to search within.
question: The specific question to answer.
scope: Optional hint to narrow search (e.g. 'auth files').
max_tokens: Maximum tokens for the answer.
Returns:
MicroFaultResult with answer, sources, and token savings.
"""
start = time.monotonic()
# Build search query — incorporate scope hint if provided
search_query = question
if scope:
search_query = f"{scope}: {question}"
# Search for relevant objects (top 5)
search_results = await self.store.semantic_search(session_id, search_query, limit=5)
if not search_results:
elapsed = (time.monotonic() - start) * 1000
return MicroFaultResult(
answer="No relevant content found in the backing store for this question.",
sources=[],
answer_tokens=_estimate_tokens(
"No relevant content found in the backing store for this question."
),
avoided_tokens=0,
latency_ms=elapsed,
)
# Collect full content from matched objects and record micro-fault access
source_ids: list[str] = []
full_contents: list[str] = []
total_source_tokens = 0
for obj, _score in search_results:
source_ids.append(obj.id)
full_contents.append(obj.content_full)
total_source_tokens += obj.tokens_l0
# Record micro-fault access on each consulted object
await self.store.record_fault(obj.id, is_micro=True)
# Generate answer
if self.helper is not None:
answer = await self.helper.answer_micro_fault(
question=question,
relevant_contents=full_contents,
max_tokens=max_tokens,
)
else:
# Fallback: return best available summaries
summaries: list[str] = []
for obj, _score in search_results:
# Prefer the most compressed available summary
if obj.summary_compact:
summaries.append(f"[{obj.id}] {obj.summary_compact}")
elif obj.summary_detailed:
summaries.append(f"[{obj.id}] {obj.summary_detailed}")
else:
# Use stub as last resort
summaries.append(f"[{obj.id}] {obj.stub}")
answer = f"(HelperLLM unavailable — returning summaries)\n" + "\n---\n".join(summaries)
answer_tokens = _estimate_tokens(answer)
avoided_tokens = max(0, total_source_tokens - answer_tokens)
elapsed = (time.monotonic() - start) * 1000
return MicroFaultResult(
answer=answer,
sources=source_ids,
answer_tokens=answer_tokens,
avoided_tokens=avoided_tokens,
latency_ms=elapsed,
)
async def assemble_context(
self,
session_id: str,
fidelity_manager: FidelityManager,
current_turn: int,
) -> list[dict]:
"""Assemble the context window from objects at their current fidelity levels.
Returns a list of context blocks (dicts with role/content) that
represent the managed context. Objects at L0 include full content,
L1/L2 include summaries, L3 include stubs, L4 are excluded.
Args:
session_id: The session to assemble context for.
fidelity_manager: The fidelity manager tracking object states.
current_turn: Current conversation turn number.
Returns:
List of context block dicts with 'object_id', 'fidelity',
'content', and 'tokens' keys.
"""
# Get all non-evicted objects from the store
objects = await self.store.get_session_objects(session_id, include_evicted=False)
context_blocks: list[dict] = []
for obj in objects:
# Use the fidelity manager's view if available, else use stored fidelity
fm_obj = fidelity_manager.get_object(obj.id)
if fm_obj is not None:
fidelity = int(fm_obj.current_fidelity)
else:
fidelity = obj.current_fidelity
# Skip evicted objects
if fidelity >= FidelityLevel.L4:
continue
# Get content at the current fidelity level
content = obj.content_at(fidelity)
if content is None:
# Fall back to stub if the expected level has no content
content = obj.stub
fidelity = FidelityLevel.L3
tokens = _estimate_tokens(content)
context_blocks.append(
{
"object_id": obj.id,
"fidelity": fidelity,
"content": content,
"tokens": tokens,
}
)
return context_blocks

116
src/mnemosyne/embedder.py Normal file
View file

@ -0,0 +1,116 @@
"""Embedding providers for semantic search.
Implements the Embedder protocol from object_store.py using
sentence-transformers (all-MiniLM-L6-v2) for production use.
Falls back gracefully if sentence-transformers is not installed.
The default model produces 384-dimensional embeddings in ~20ms
per text on CPU, suitable for real-time use in the proxy path.
"""
from __future__ import annotations
import logging
from functools import lru_cache
logger = logging.getLogger(__name__)
# Embedding dimension — must match the model and pgvector schema.
EMBEDDING_DIM = 384
# Default model — small, fast, good enough for semantic similarity.
DEFAULT_MODEL = "all-MiniLM-L6-v2"
class SentenceTransformerEmbedder:
"""Production embedder using sentence-transformers.
Lazily loads the model on first embed() call to avoid slowing
proxy startup. Thread-safe after initialization (the model is
read-only once loaded).
Truncates input to max_chars to avoid OOM on huge tool results.
"""
def __init__(
self,
model_name: str = DEFAULT_MODEL,
max_chars: int = 8192,
) -> None:
self._model_name = model_name
self._max_chars = max_chars
self._model = None
def _load_model(self):
"""Lazy-load the sentence-transformers model."""
if self._model is not None:
return
try:
from sentence_transformers import SentenceTransformer
self._model = SentenceTransformer(self._model_name)
logger.info("Loaded embedding model: %s", self._model_name)
except ImportError:
logger.warning(
"sentence-transformers not installed — embeddings disabled. "
"Install with: pip install sentence-transformers"
)
raise
except Exception:
logger.exception("Failed to load embedding model %s", self._model_name)
raise
def embed(self, text: str) -> list[float]:
"""Embed a single text string into a 384-dim vector.
Truncates to max_chars before encoding. Returns a
unit-normalized vector (L2 norm = 1.0).
"""
self._load_model()
assert self._model is not None
truncated = text[: self._max_chars] if len(text) > self._max_chars else text
# encode returns numpy array of shape (dim,) for single input
vec = self._model.encode(truncated, normalize_embeddings=True)
return vec.tolist()
def embed_batch(self, texts: list[str]) -> list[list[float]]:
"""Embed multiple texts in a single batch call.
More efficient than calling embed() in a loop the model
batches the forward pass.
"""
self._load_model()
assert self._model is not None
truncated = [t[: self._max_chars] if len(t) > self._max_chars else t for t in texts]
# encode returns numpy array of shape (n, dim)
vecs = self._model.encode(truncated, normalize_embeddings=True)
return [v.tolist() for v in vecs]
@lru_cache(maxsize=1)
def get_embedder(model_name: str = DEFAULT_MODEL) -> SentenceTransformerEmbedder:
"""Get or create a cached embedder instance.
Cached by model_name so we only load one model per process.
Call this from gateway startup to share the embedder across sessions.
"""
return SentenceTransformerEmbedder(model_name=model_name)
def try_get_embedder(model_name: str = DEFAULT_MODEL) -> SentenceTransformerEmbedder | None:
"""Try to create an embedder, returning None if sentence-transformers isn't available.
Use this for graceful degradation the proxy works without embeddings,
just without semantic search capability.
"""
try:
embedder = SentenceTransformerEmbedder(model_name=model_name)
# Test that the model can actually load
embedder._load_model()
return embedder
except (ImportError, Exception) as exc:
logger.warning("Embedder unavailable: %s — semantic search disabled", exc)
return None

View file

@ -0,0 +1,547 @@
"""Backing store for semantic objects.
Provides CRUD operations, embedding generation, and semantic search
for the Mnemosyne context manager. Designed around an abstract backend
so we can swap in-memory (Phase 1-2) for PostgreSQL+pgvector (Phase 3+)
without changing the facade.
Key types:
StoredObject dataclass matching the semantic_objects schema
Embedder protocol for pluggable embedding providers
DummyEmbedder deterministic hash-based embedder for testing
ObjectStoreBackend abstract base for storage backends
InMemoryBackend dict-backed implementation with cosine search
ObjectStore facade that wraps backend + embedder
"""
from __future__ import annotations
import asyncio
import hashlib
import math
import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Protocol, runtime_checkable
import numpy as np
# ── StoredObject ─────────────────────────────────────────────
def _estimate_tokens(text: str | None) -> int:
"""Estimate token count using ~4 chars per token heuristic."""
if text is None:
return 0
return max(1, len(text) // 4)
@dataclass
class StoredObject:
"""A semantic object persisted in the backing store.
Maps 1:1 to the semantic_objects table in SCHEMA.md.
Multi-fidelity content at L0-L3, with declared losses,
embedding vector, and access/fault counters.
"""
id: str
session_id: str
object_type: str # conversation_phase, file_context, etc.
source_tool: str | None # Read, Bash, Grep, etc.
source_key: str | None # dedup key (file path for reads)
# Multi-fidelity content
content_full: str # L0: always preserved
summary_detailed: str | None # L1
summary_compact: str | None # L2
stub: str # L3: always present
# Declared losses
losses_l1: list[str] = field(default_factory=list)
losses_l2: list[str] = field(default_factory=list)
can_answer_l1: list[str] = field(default_factory=list)
can_answer_l2: list[str] = field(default_factory=list)
fault_when: list[str] = field(default_factory=list)
# Metadata
key_entities: list[str] = field(default_factory=list)
tags: list[str] = field(default_factory=list)
current_fidelity: int = 0 # 0-4
pinned: bool = False
# Token counts
tokens_l0: int = 0
tokens_l1: int | None = None
tokens_l2: int | None = None
tokens_l3: int = 0
# Source range
source_turn_start: int | None = None
source_turn_end: int | None = None
# Embedding (384-dim float list, or empty if not computed yet)
embedding: list[float] = field(default_factory=list)
# Timestamps
created_at: str = "" # ISO format
last_accessed: str = "" # ISO format
access_count: int = 0
fault_count: int = 0
micro_fault_count: int = 0
def content_at(self, level: int) -> str | None:
"""Return content at the given fidelity level."""
if level == 0:
return self.content_full
if level == 1:
return self.summary_detailed
if level == 2:
return self.summary_compact
if level == 3:
return self.stub
return None # L4: evicted
def tokens_at(self, level: int) -> int:
"""Return token count at the given fidelity level."""
if level == 0:
return self.tokens_l0
if level == 1:
return self.tokens_l1 if self.tokens_l1 is not None else 0
if level == 2:
return self.tokens_l2 if self.tokens_l2 is not None else 0
if level == 3:
return self.tokens_l3
return 0 # L4: evicted
@property
def current_tokens(self) -> int:
"""Token count at the object's current fidelity level."""
return self.tokens_at(self.current_fidelity)
# ── Embedder protocol ────────────────────────────────────────
@runtime_checkable
class Embedder(Protocol):
"""Protocol for embedding providers.
Implementations must produce 384-dimensional float vectors.
The protocol is intentionally synchronous embedding is CPU-bound
and typically fast enough to not need async.
"""
def embed(self, text: str) -> list[float]: ...
def embed_batch(self, texts: list[str]) -> list[list[float]]: ...
class DummyEmbedder:
"""Deterministic hash-based embedder for testing.
Produces 384-dim vectors seeded by a hash of the input text,
so identical inputs always produce identical embeddings.
Not semantically meaningful use only for testing search plumbing.
"""
dim: int = 384
def embed(self, text: str) -> list[float]:
seed = int(hashlib.sha256(text.encode()).hexdigest(), 16) % (2**32)
rng = np.random.default_rng(seed)
vec = rng.standard_normal(self.dim)
# L2-normalize so cosine similarity is just dot product
norm = float(np.linalg.norm(vec))
if norm > 0:
vec = vec / norm
return vec.tolist()
def embed_batch(self, texts: list[str]) -> list[list[float]]:
return [self.embed(t) for t in texts]
# ── ObjectStoreBackend (abstract) ────────────────────────────
class ObjectStoreBackend(ABC):
"""Abstract base for semantic object storage.
All methods are async to support both in-memory and database
backends uniformly. Session-scoped queries filter by session_id
to enforce isolation.
"""
@abstractmethod
async def store(self, obj: StoredObject) -> None:
"""Persist a StoredObject. Overwrites if id already exists."""
@abstractmethod
async def get(self, object_id: str) -> StoredObject | None:
"""Retrieve a single object by ID, or None if not found."""
@abstractmethod
async def get_by_session(self, session_id: str, fidelity_max: int = 4) -> list[StoredObject]:
"""Return all objects for a session with current_fidelity <= fidelity_max.
Default fidelity_max=4 includes evicted objects. Use fidelity_max=3
to get only objects still in context.
"""
@abstractmethod
async def update_fidelity(
self,
object_id: str,
new_fidelity: int,
summary: str | None = None,
losses: list[str] | None = None,
) -> None:
"""Update an object's fidelity level, optionally setting a new summary."""
@abstractmethod
async def search_by_embedding(
self,
session_id: str,
query_embedding: list[float],
limit: int = 5,
) -> list[tuple[StoredObject, float]]:
"""Find objects by cosine similarity to query_embedding.
Returns (object, similarity_score) pairs sorted by descending similarity.
Only searches within the given session.
"""
@abstractmethod
async def search_by_text(
self, session_id: str, query: str, limit: int = 5
) -> list[StoredObject]:
"""Find objects by substring match on content_full, stub, and key_entities.
Simple text search not full-text ranking. Returns up to `limit` results.
"""
@abstractmethod
async def delete_session(self, session_id: str) -> int:
"""Delete all objects for a session. Returns count of deleted objects."""
@abstractmethod
async def get_by_source_key(self, session_id: str, source_key: str) -> StoredObject | None:
"""Find the most recent object with this source_key in the session.
Used for deduplication (e.g., re-reading the same file).
"""
# ── InMemoryBackend ──────────────────────────────────────────
def _cosine_similarity(a: list[float], b: list[float]) -> float:
"""Compute cosine similarity between two vectors using numpy."""
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 or norm_b == 0:
return 0.0
return dot / (norm_a * norm_b)
class InMemoryBackend(ObjectStoreBackend):
"""Dict-backed in-memory implementation for Phase 1-2.
Thread-safe via asyncio.Lock. Indexes objects by both id and
session_id for efficient session-scoped queries.
"""
def __init__(self) -> None:
self._objects: dict[str, StoredObject] = {}
self._session_index: dict[str, list[str]] = {} # session_id → [object_id]
self._lock = asyncio.Lock()
async def store(self, obj: StoredObject) -> None:
async with self._lock:
is_new = obj.id not in self._objects
self._objects[obj.id] = obj
if is_new:
if obj.session_id not in self._session_index:
self._session_index[obj.session_id] = []
self._session_index[obj.session_id].append(obj.id)
async def get(self, object_id: str) -> StoredObject | None:
return self._objects.get(object_id)
async def get_by_session(self, session_id: str, fidelity_max: int = 4) -> list[StoredObject]:
ids = self._session_index.get(session_id, [])
return [
self._objects[oid]
for oid in ids
if oid in self._objects and self._objects[oid].current_fidelity <= fidelity_max
]
async def update_fidelity(
self,
object_id: str,
new_fidelity: int,
summary: str | None = None,
losses: list[str] | None = None,
) -> None:
async with self._lock:
obj = self._objects.get(object_id)
if obj is None:
return
old_fidelity = obj.current_fidelity
obj.current_fidelity = new_fidelity
# Attach summary at the appropriate level
if summary is not None:
if new_fidelity == 1:
obj.summary_detailed = summary
obj.tokens_l1 = _estimate_tokens(summary)
elif new_fidelity == 2:
obj.summary_compact = summary
obj.tokens_l2 = _estimate_tokens(summary)
if losses is not None:
if new_fidelity == 1:
obj.losses_l1 = losses
elif new_fidelity == 2:
obj.losses_l2 = losses
async def search_by_embedding(
self,
session_id: str,
query_embedding: list[float],
limit: int = 5,
) -> list[tuple[StoredObject, float]]:
ids = self._session_index.get(session_id, [])
scored: list[tuple[StoredObject, float]] = []
for oid in ids:
obj = self._objects.get(oid)
if obj is None or not obj.embedding:
continue
sim = _cosine_similarity(query_embedding, obj.embedding)
scored.append((obj, sim))
scored.sort(key=lambda x: x[1], reverse=True)
return scored[:limit]
async def search_by_text(
self, session_id: str, query: str, limit: int = 5
) -> list[StoredObject]:
ids = self._session_index.get(session_id, [])
query_lower = query.lower()
results: list[StoredObject] = []
for oid in ids:
obj = self._objects.get(oid)
if obj is None:
continue
# Search across content_full, stub, and key_entities
searchable = obj.content_full.lower()
searchable += " " + obj.stub.lower()
searchable += " " + " ".join(e.lower() for e in obj.key_entities)
if query_lower in searchable:
results.append(obj)
if len(results) >= limit:
break
return results
async def delete_session(self, session_id: str) -> int:
async with self._lock:
ids = self._session_index.pop(session_id, [])
count = 0
for oid in ids:
if oid in self._objects:
del self._objects[oid]
count += 1
return count
async def get_by_source_key(self, session_id: str, source_key: str) -> StoredObject | None:
ids = self._session_index.get(session_id, [])
# Return the most recently created object with this source_key
best: StoredObject | None = None
for oid in ids:
obj = self._objects.get(oid)
if obj is not None and obj.source_key == source_key:
if best is None or obj.created_at > best.created_at:
best = obj
return best
# ── ObjectStore facade ───────────────────────────────────────
class ObjectStore:
"""High-level facade over a storage backend + embedder.
Handles object creation (ID generation, timestamps, token estimation,
embedding), access tracking, fault recording, and hybrid semantic search.
"""
def __init__(
self,
backend: ObjectStoreBackend,
embedder: Embedder | None = None,
) -> None:
self._backend = backend
self._embedder = embedder
async def store_object(
self,
session_id: str,
content: str,
object_type: str,
*,
source_tool: str | None = None,
source_key: str | None = None,
stub: str | None = None,
tags: list[str] | None = None,
key_entities: list[str] | None = None,
turn: int | None = None,
) -> StoredObject:
"""Create and store a new semantic object.
Auto-generates: ID, timestamps, token estimates, embedding, stub.
"""
now = datetime.now(timezone.utc).isoformat()
object_id = uuid.uuid4().hex[:16]
# Auto-generate stub if not provided
if stub is None:
preview = content[:80].replace("\n", " ")
stub = f"{object_type}: {preview}..."
# Compute embedding
embedding: list[float] = []
if self._embedder is not None:
embedding = self._embedder.embed(content)
obj = StoredObject(
id=object_id,
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,
tags=tags or [],
key_entities=key_entities or [],
current_fidelity=0,
pinned=False,
tokens_l0=_estimate_tokens(content),
tokens_l1=None,
tokens_l2=None,
tokens_l3=_estimate_tokens(stub),
source_turn_start=turn,
source_turn_end=turn,
embedding=embedding,
created_at=now,
last_accessed=now,
access_count=0,
fault_count=0,
micro_fault_count=0,
)
await self._backend.store(obj)
return obj
async def get(self, object_id: str) -> StoredObject | None:
"""Retrieve a single object by ID."""
return await self._backend.get(object_id)
async def get_session_objects(
self, session_id: str, *, include_evicted: bool = False
) -> list[StoredObject]:
"""Return all objects for a session.
By default excludes evicted (fidelity=4) objects. Pass
include_evicted=True to get everything.
"""
fidelity_max = 4 if include_evicted else 3
return await self._backend.get_by_session(session_id, fidelity_max=fidelity_max)
async def update_fidelity(
self,
object_id: str,
new_fidelity: int,
summary: str | None = None,
losses: list[str] | None = None,
) -> None:
"""Update an object's fidelity level."""
await self._backend.update_fidelity(object_id, new_fidelity, summary, losses)
async def semantic_search(
self,
session_id: str,
query: str,
limit: int = 5,
) -> list[tuple[StoredObject, float]]:
"""Hybrid search: embed query, search by embedding + text.
Combines vector similarity (weight 0.7) with text match (weight 0.3)
following the hybrid search pattern from SCHEMA.md §2.3.
Returns (object, combined_score) pairs sorted by descending score.
"""
# Embedding search
embedding_results: list[tuple[StoredObject, float]] = []
if self._embedder is not None:
query_embedding = self._embedder.embed(query)
embedding_results = await self._backend.search_by_embedding(
session_id, query_embedding, limit=limit * 4
)
# Text search
text_results = await self._backend.search_by_text(session_id, query, limit=limit * 4)
text_ids = {obj.id for obj in text_results}
# Merge: vector score * 0.7 + text match * 0.3
scored: dict[str, tuple[StoredObject, float]] = {}
for obj, vec_score in embedding_results:
text_bonus = 0.3 if obj.id in text_ids else 0.0
combined = vec_score * 0.7 + text_bonus
scored[obj.id] = (obj, combined)
# Add text-only results not already in embedding results
for obj in text_results:
if obj.id not in scored:
scored[obj.id] = (obj, 0.3) # text match only
results = sorted(scored.values(), key=lambda x: x[1], reverse=True)
return results[:limit]
async def find_duplicate(self, session_id: str, source_key: str) -> StoredObject | None:
"""Check if an object with this source_key already exists in the session."""
return await self._backend.get_by_source_key(session_id, source_key)
async def record_access(self, object_id: str) -> None:
"""Record that an object was accessed (updates timestamp and counter)."""
obj = await self._backend.get(object_id)
if obj is None:
return
obj.last_accessed = datetime.now(timezone.utc).isoformat()
obj.access_count += 1
async def record_fault(self, object_id: str, *, is_micro: bool = False) -> None:
"""Record a fault (full restore) or micro-fault on an object."""
obj = await self._backend.get(object_id)
if obj is None:
return
if is_micro:
obj.micro_fault_count += 1
else:
obj.fault_count += 1
obj.last_accessed = datetime.now(timezone.utc).isoformat()

View file

@ -0,0 +1,568 @@
"""PostgreSQL + pgvector backend for the Mnemosyne object store.
Implements ObjectStoreBackend using asyncpg for async connection pooling
and pgvector for embedding similarity search. Maps the gateway's string
session IDs (external_id) to the internal UUID-based sessions table.
Requires PostgreSQL 16+ with the pgvector extension and the schema from
sql/init.sql already applied.
"""
from __future__ import annotations
import json
import uuid
from datetime import datetime, timezone
from typing import Any
import asyncpg
from pgvector.asyncpg import register_vector
from mnemosyne.object_store import ObjectStoreBackend, StoredObject
class PgVectorBackend(ObjectStoreBackend):
"""PostgreSQL + pgvector implementation of ObjectStoreBackend.
Uses asyncpg connection pool with pgvector codec registration.
Transparently maps gateway string session IDs (external_id) to
internal UUID session IDs via the sessions table.
"""
def __init__(
self,
*,
host: str = "localhost",
port: int = 5433,
database: str = "mnemosyne",
user: str = "mnemosyne",
password: str = "mnemosyne_dev",
min_connections: int = 2,
max_connections: int = 10,
) -> None:
self._host = host
self._port = port
self._database = database
self._user = user
self._password = password
self._min_connections = min_connections
self._max_connections = max_connections
self._pool: asyncpg.Pool | None = None
# Cache: external_id → internal UUID to avoid repeated lookups
self._session_cache: dict[str, uuid.UUID] = {}
async def connect(self) -> None:
"""Create the connection pool and register the pgvector codec."""
self._pool = await asyncpg.create_pool(
host=self._host,
port=self._port,
database=self._database,
user=self._user,
password=self._password,
min_size=self._min_connections,
max_size=self._max_connections,
init=_init_connection,
)
async def close(self) -> None:
"""Close the connection pool."""
if self._pool is not None:
await self._pool.close()
self._pool = None
self._session_cache.clear()
async def health_check(self) -> bool:
"""Check database connectivity."""
if self._pool is None:
return False
try:
async with self._pool.acquire() as conn:
await conn.fetchval("SELECT 1")
return True
except Exception:
return False
# ── Session management ───────────────────────────────────────
async def ensure_session(self, external_id: str, model: str = "unknown") -> uuid.UUID:
"""Upsert a session by external_id and return its internal UUID.
Uses an INSERT ... ON CONFLICT to create the session if it doesn't
exist, or update last_active_at if it does.
"""
cached = self._session_cache.get(external_id)
if cached is not None:
return cached
pool = self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
INSERT INTO sessions (external_id, model)
VALUES ($1, $2)
ON CONFLICT (external_id) DO UPDATE
SET last_active_at = now()
RETURNING id
""",
external_id,
model,
)
assert row is not None
internal_id = row["id"]
self._session_cache[external_id] = internal_id
return internal_id
async def _resolve_session_id(self, external_id: str) -> uuid.UUID | None:
"""Look up the internal UUID for an external session ID.
Returns None if the session doesn't exist.
"""
cached = self._session_cache.get(external_id)
if cached is not None:
return cached
pool = self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"SELECT id FROM sessions WHERE external_id = $1",
external_id,
)
if row is None:
return None
internal_id = row["id"]
self._session_cache[external_id] = internal_id
return internal_id
async def _get_or_create_session(self, external_id: str) -> uuid.UUID:
"""Resolve session ID, creating the session if needed."""
resolved = await self._resolve_session_id(external_id)
if resolved is not None:
return resolved
return await self.ensure_session(external_id)
# ── ObjectStoreBackend interface ─────────────────────────────
async def store(self, obj: StoredObject) -> None:
"""Persist a StoredObject. Upserts by object ID."""
pool = self._get_pool()
session_uuid = await self._get_or_create_session(obj.session_id)
# Convert embedding to numpy array for pgvector
import numpy as np
embedding = (
np.array(obj.embedding, dtype=np.float32)
if obj.embedding
else np.zeros(384, dtype=np.float32)
)
async with pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO semantic_objects (
id, session_id,
object_type, source_tool, source_key,
content_full, summary_detailed, summary_compact, stub,
losses_l1, losses_l2, can_answer_l1, can_answer_l2, fault_when,
key_entities, tags,
current_fidelity, pinned,
tokens_l0, tokens_l1, tokens_l2, tokens_l3,
source_turn_start, source_turn_end,
embedding,
created_at, last_accessed,
access_count, fault_count, micro_fault_count
) VALUES (
$1, $2,
$3, $4, $5,
$6, $7, $8, $9,
$10, $11, $12, $13, $14,
$15, $16,
$17, $18,
$19, $20, $21, $22,
$23, $24,
$25,
$26, $27,
$28, $29, $30
)
ON CONFLICT (id) DO UPDATE SET
object_type = EXCLUDED.object_type,
source_tool = EXCLUDED.source_tool,
source_key = EXCLUDED.source_key,
content_full = EXCLUDED.content_full,
summary_detailed = EXCLUDED.summary_detailed,
summary_compact = EXCLUDED.summary_compact,
stub = EXCLUDED.stub,
losses_l1 = EXCLUDED.losses_l1,
losses_l2 = EXCLUDED.losses_l2,
can_answer_l1 = EXCLUDED.can_answer_l1,
can_answer_l2 = EXCLUDED.can_answer_l2,
fault_when = EXCLUDED.fault_when,
key_entities = EXCLUDED.key_entities,
tags = EXCLUDED.tags,
current_fidelity = EXCLUDED.current_fidelity,
pinned = EXCLUDED.pinned,
tokens_l0 = EXCLUDED.tokens_l0,
tokens_l1 = EXCLUDED.tokens_l1,
tokens_l2 = EXCLUDED.tokens_l2,
tokens_l3 = EXCLUDED.tokens_l3,
source_turn_start = EXCLUDED.source_turn_start,
source_turn_end = EXCLUDED.source_turn_end,
embedding = EXCLUDED.embedding,
last_accessed = EXCLUDED.last_accessed,
access_count = EXCLUDED.access_count,
fault_count = EXCLUDED.fault_count,
micro_fault_count = EXCLUDED.micro_fault_count
""",
uuid.UUID(obj.id) if not isinstance(obj.id, uuid.UUID) else obj.id,
session_uuid,
obj.object_type,
obj.source_tool,
obj.source_key,
obj.content_full,
obj.summary_detailed,
obj.summary_compact,
obj.stub,
json.dumps(obj.losses_l1),
json.dumps(obj.losses_l2),
json.dumps(obj.can_answer_l1),
json.dumps(obj.can_answer_l2),
json.dumps(obj.fault_when),
json.dumps(obj.key_entities),
obj.tags,
obj.current_fidelity,
obj.pinned,
obj.tokens_l0,
obj.tokens_l1,
obj.tokens_l2,
obj.tokens_l3,
obj.source_turn_start,
obj.source_turn_end,
embedding,
_parse_timestamp(obj.created_at),
_parse_timestamp(obj.last_accessed),
obj.access_count,
obj.fault_count,
obj.micro_fault_count,
)
async def get(self, object_id: str) -> StoredObject | None:
"""Retrieve a single object by ID."""
pool = self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
SELECT so.*, s.external_id AS session_external_id
FROM semantic_objects so
JOIN sessions s ON s.id = so.session_id
WHERE so.id = $1
""",
uuid.UUID(object_id) if not isinstance(object_id, uuid.UUID) else object_id,
)
if row is None:
return None
return _row_to_stored_object(row)
async def get_by_session(self, session_id: str, fidelity_max: int = 4) -> list[StoredObject]:
"""Return all objects for a session with current_fidelity <= fidelity_max."""
session_uuid = await self._resolve_session_id(session_id)
if session_uuid is None:
return []
pool = self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT so.*, $3::text AS session_external_id
FROM semantic_objects so
WHERE so.session_id = $1
AND so.current_fidelity <= $2
ORDER BY
so.pinned DESC,
so.current_fidelity ASC,
so.last_accessed DESC
""",
session_uuid,
fidelity_max,
session_id,
)
return [_row_to_stored_object(row) for row in rows]
async def update_fidelity(
self,
object_id: str,
new_fidelity: int,
summary: str | None = None,
losses: list[str] | None = None,
) -> None:
"""Update an object's fidelity level, optionally setting a new summary."""
pool = self._get_pool()
obj_uuid = uuid.UUID(object_id) if not isinstance(object_id, uuid.UUID) else object_id
# Build dynamic SET clause based on what's provided
set_parts = ["current_fidelity = $2"]
params: list[Any] = [obj_uuid, new_fidelity]
param_idx = 3
if summary is not None:
if new_fidelity == 1:
set_parts.append(f"summary_detailed = ${param_idx}")
params.append(summary)
param_idx += 1
set_parts.append(f"tokens_l1 = ${param_idx}")
params.append(_estimate_tokens(summary))
param_idx += 1
elif new_fidelity == 2:
set_parts.append(f"summary_compact = ${param_idx}")
params.append(summary)
param_idx += 1
set_parts.append(f"tokens_l2 = ${param_idx}")
params.append(_estimate_tokens(summary))
param_idx += 1
if losses is not None:
if new_fidelity == 1:
set_parts.append(f"losses_l1 = ${param_idx}")
params.append(json.dumps(losses))
param_idx += 1
elif new_fidelity == 2:
set_parts.append(f"losses_l2 = ${param_idx}")
params.append(json.dumps(losses))
param_idx += 1
query = f"UPDATE semantic_objects SET {', '.join(set_parts)} WHERE id = $1"
async with pool.acquire() as conn:
await conn.execute(query, *params)
async def search_by_embedding(
self,
session_id: str,
query_embedding: list[float],
limit: int = 5,
) -> list[tuple[StoredObject, float]]:
"""Find objects by cosine similarity to query_embedding."""
session_uuid = await self._resolve_session_id(session_id)
if session_uuid is None:
return []
import numpy as np
embedding = np.array(query_embedding, dtype=np.float32)
pool = self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT so.*, $4::text AS session_external_id,
1 - (so.embedding <=> $2) AS similarity
FROM semantic_objects so
WHERE so.session_id = $1
ORDER BY so.embedding <=> $2
LIMIT $3
""",
session_uuid,
embedding,
limit,
session_id,
)
results: list[tuple[StoredObject, float]] = []
for row in rows:
obj = _row_to_stored_object(row)
similarity = float(row["similarity"])
results.append((obj, similarity))
return results
async def search_by_text(
self, session_id: str, query: str, limit: int = 5
) -> list[StoredObject]:
"""Find objects by PostgreSQL full-text search on content_full."""
session_uuid = await self._resolve_session_id(session_id)
if session_uuid is None:
return []
pool = self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT so.*, $4::text AS session_external_id
FROM semantic_objects so
WHERE so.session_id = $1
AND to_tsvector('english', so.content_full) @@
plainto_tsquery('english', $2)
ORDER BY ts_rank(
to_tsvector('english', so.content_full),
plainto_tsquery('english', $2)
) DESC
LIMIT $3
""",
session_uuid,
query,
limit,
session_id,
)
return [_row_to_stored_object(row) for row in rows]
async def delete_session(self, session_id: str) -> int:
"""Delete all objects for a session. Returns count of deleted objects."""
session_uuid = await self._resolve_session_id(session_id)
if session_uuid is None:
return 0
pool = self._get_pool()
async with pool.acquire() as conn:
result = await conn.execute(
"DELETE FROM semantic_objects WHERE session_id = $1",
session_uuid,
)
# asyncpg returns "DELETE N" where N is the count
count = int(result.split()[-1])
return count
async def get_by_source_key(self, session_id: str, source_key: str) -> StoredObject | None:
"""Find the most recent object with this source_key in the session."""
session_uuid = await self._resolve_session_id(session_id)
if session_uuid is None:
return None
pool = self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
SELECT so.*, $3::text AS session_external_id
FROM semantic_objects so
WHERE so.session_id = $1
AND so.source_key = $2
ORDER BY so.created_at DESC
LIMIT 1
""",
session_uuid,
source_key,
session_id,
)
if row is None:
return None
return _row_to_stored_object(row)
# ── Internal helpers ─────────────────────────────────────────
def _get_pool(self) -> asyncpg.Pool:
"""Return the connection pool, raising if not connected."""
if self._pool is None:
raise RuntimeError("PgVectorBackend is not connected. Call connect() first.")
return self._pool
# ── Module-level helpers ─────────────────────────────────────────
async def _init_connection(conn: asyncpg.Connection) -> None:
"""Register pgvector codec on each new connection."""
await register_vector(conn)
def _estimate_tokens(text: str | None) -> int:
"""Estimate token count using ~4 chars per token heuristic."""
if text is None:
return 0
return max(1, len(text) // 4)
def _parse_timestamp(ts: str) -> datetime:
"""Parse an ISO timestamp string to a timezone-aware datetime.
Falls back to UTC now if the string is empty or unparseable.
"""
if not ts:
return datetime.now(timezone.utc)
try:
dt = datetime.fromisoformat(ts)
if dt.tzinfo is None:
dt = dt.replace(tzinfo=timezone.utc)
return dt
except (ValueError, TypeError):
return datetime.now(timezone.utc)
def _row_to_stored_object(row: asyncpg.Record) -> StoredObject:
"""Convert a database row to a StoredObject dataclass."""
# Parse JSONB fields
losses_l1 = _parse_jsonb(row["losses_l1"])
losses_l2 = _parse_jsonb(row["losses_l2"])
can_answer_l1 = _parse_jsonb(row["can_answer_l1"])
can_answer_l2 = _parse_jsonb(row["can_answer_l2"])
fault_when = _parse_jsonb(row["fault_when"])
key_entities = _parse_jsonb(row["key_entities"])
# Convert tags from PostgreSQL TEXT[] to list[str]
tags = list(row["tags"]) if row["tags"] else []
# Convert embedding from pgvector to list[float]
embedding_raw = row["embedding"]
if embedding_raw is not None:
import numpy as np
embedding = np.array(embedding_raw, dtype=np.float32).tolist()
else:
embedding = []
# Use external_id for session_id (gateway uses string IDs)
session_id = row["session_external_id"]
# Convert UUID to string for the object ID
object_id = str(row["id"])
# Convert timestamps to ISO format strings
created_at = row["created_at"].isoformat() if row["created_at"] else ""
last_accessed = row["last_accessed"].isoformat() if row["last_accessed"] else ""
return StoredObject(
id=object_id,
session_id=session_id,
object_type=row["object_type"],
source_tool=row["source_tool"],
source_key=row["source_key"],
content_full=row["content_full"],
summary_detailed=row["summary_detailed"],
summary_compact=row["summary_compact"],
stub=row["stub"],
losses_l1=losses_l1,
losses_l2=losses_l2,
can_answer_l1=can_answer_l1,
can_answer_l2=can_answer_l2,
fault_when=fault_when,
key_entities=key_entities,
tags=tags,
current_fidelity=row["current_fidelity"],
pinned=row["pinned"],
tokens_l0=row["tokens_l0"],
tokens_l1=row["tokens_l1"],
tokens_l2=row["tokens_l2"],
tokens_l3=row["tokens_l3"],
source_turn_start=row["source_turn_start"],
source_turn_end=row["source_turn_end"],
embedding=embedding,
created_at=created_at,
last_accessed=last_accessed,
access_count=row["access_count"],
fault_count=row["fault_count"],
micro_fault_count=row["micro_fault_count"],
)
def _parse_jsonb(value: Any) -> list[str]:
"""Parse a JSONB value to a list of strings.
asyncpg automatically deserializes JSONB to Python objects,
so this handles both pre-parsed lists and raw JSON strings.
"""
if value is None:
return []
if isinstance(value, list):
return [str(v) for v in value]
if isinstance(value, str):
try:
parsed = json.loads(value)
if isinstance(parsed, list):
return [str(v) for v in parsed]
except (json.JSONDecodeError, TypeError):
pass
return []

927
src/mnemosyne/segmenter.py Normal file
View file

@ -0,0 +1,927 @@
"""Semantic object segmenter for conversation message streams.
Splits Anthropic Messages API conversation arrays into typed semantic
objects using structural signals (tool boundaries, role transitions,
content patterns). This is Strategy A fast, rule-based segmentation
for turns 1-50. Strategy B (embedding-based) is Phase 6.
Object types per ARCHITECTURE.md §5.1:
conversation_phase, design_decision, debugging_session, file_context,
tool_result, plan, error_context, external_reference
"""
from __future__ import annotations
import re
from dataclasses import dataclass, field
# ---------------------------------------------------------------------------
# Object types (mirrors fidelity.VALID_OBJECT_TYPES)
# ---------------------------------------------------------------------------
VALID_OBJECT_TYPES = frozenset(
{
"conversation_phase",
"design_decision",
"debugging_session",
"file_context",
"tool_result",
"plan",
"error_context",
"external_reference",
}
)
# ---------------------------------------------------------------------------
# Tool name → object type mapping
# ---------------------------------------------------------------------------
_READ_TOOLS = frozenset({"Read", "read", "ReadFile", "read_file", "cat"})
_BASH_TOOLS = frozenset({"Bash", "bash", "shell", "execute", "run"})
_GREP_TOOLS = frozenset({"Grep", "grep", "search", "Search", "rg", "ripgrep"})
_WRITE_TOOLS = frozenset({"Write", "write", "Edit", "edit", "WriteFile", "EditFile"})
_EXTERNAL_TOOLS = frozenset({"WebFetch", "web_fetch", "fetch", "Browse", "browse"})
# ---------------------------------------------------------------------------
# Classification patterns
# ---------------------------------------------------------------------------
_ERROR_PATTERNS = [
re.compile(r"(?i)\b(?:error|exception|traceback|stack\s*trace|panic|fatal)\b"),
re.compile(r"(?i)\b(?:TypeError|ValueError|KeyError|AttributeError|RuntimeError)\b"),
re.compile(r"(?i)\b(?:FAILED|FAIL|ERR|E\d{4})\b"),
re.compile(r"File \"[^\"]+\", line \d+"),
re.compile(r"at [\w.]+\([\w./]+:\d+:\d+\)"), # JS stack trace
]
_PLAN_PATTERNS = [
re.compile(r"(?i)\b(?:step\s+\d|phase\s+\d)\b"),
re.compile(r"(?m)^\s*\d+\.\s+\S"), # numbered list
re.compile(r"(?i)\b(?:TODO|plan|implementation plan|roadmap)\b"),
re.compile(r"(?i)\b(?:first|second|third|finally|next),?\s+(?:we|I|let)"),
]
_DECISION_PATTERNS = [
re.compile(r"(?i)\bI (?:decided|chose|picked|selected|went with)\b"),
re.compile(r"(?i)\b(?:because|the reason|rationale|trade-?off)\b"),
re.compile(r"(?i)\b(?:chose .+ over|instead of|rather than)\b"),
re.compile(r"(?i)\b(?:design decision|architectural choice)\b"),
]
_DEBUG_PATTERNS = [
re.compile(r"(?i)\b(?:investigating|debugging|diagnosing|root cause)\b"),
re.compile(r"(?i)\b(?:fixed by|the fix|the issue was|the bug was)\b"),
re.compile(r"(?i)\b(?:hypothesis|reproduce|bisect|narrowed down)\b"),
re.compile(r"(?i)\b(?:turns out|found the problem|the culprit)\b"),
]
# ---------------------------------------------------------------------------
# Entity extraction patterns
# ---------------------------------------------------------------------------
_FILE_PATH_RE = re.compile(
r"(?:^|[\s\"'`(,])("
r"(?:[a-zA-Z_][\w.-]*/)+[\w.-]+\.[\w]+" # relative: src/foo/bar.py
r"|/(?:[\w.-]+/)+[\w.-]+\.[\w]+" # absolute: /path/to/file.py
r")"
)
_FUNCTION_RE = re.compile(
r"(?:"
r"(?:def|function|fn|func|async\s+function|async\s+fn)\s+([\w]+)" # def foo
r"|"
r"\b([\w]+)\s*\(" # foo(
r")"
)
_IMPORT_RE = re.compile(
r"(?:"
r"(?:from|import)\s+([\w.]+)" # Python: from X import / import X
r"|"
r"require\(['\"]([^'\"]+)['\"]\)" # JS: require('X')
r"|"
r"from\s+['\"]([^'\"]+)['\"]" # JS: from 'X'
r")"
)
# Short user messages that should merge with the next assistant response
_SHORT_MSG_RE = re.compile(
r"^(?:ok|okay|yes|yeah|yep|no|nope|sure|continue|go ahead|"
r"do it|proceed|thanks|thank you|lgtm|looks good|great|perfect|"
r"got it|right|correct|exactly|please|next|done)[\s.!?]*$",
re.IGNORECASE,
)
def _estimate_tokens(text: str) -> int:
"""Estimate token count using ~4 chars per token heuristic."""
return max(1, len(text) // 4)
# ---------------------------------------------------------------------------
# SegmentedObject
# ---------------------------------------------------------------------------
@dataclass
class SegmentedObject:
"""A semantic object extracted from the conversation stream.
Produced by the Segmenter. Each object represents a coherent unit
of conversation content a file read, a tool result, a planning
phase, a debugging session, etc.
"""
content: str # The full text content
object_type: str # One of VALID_OBJECT_TYPES
source_tool: str | None = None # Tool that produced this (Read, Bash, etc.)
source_key: str | None = None # Dedup key (file path for Read results)
stub: str = "" # Auto-generated one-line stub
turn_start: int = 0 # First turn index
turn_end: int = 0 # Last turn index
token_estimate: int = 0 # Estimated tokens (len/4)
key_entities: list[str] = field(default_factory=list)
tags: list[str] = field(default_factory=list)
# ---------------------------------------------------------------------------
# Internal intermediate representation
# ---------------------------------------------------------------------------
@dataclass
class _RawSegment:
"""Intermediate segment before merging and classification."""
content: str
object_type: str
source_tool: str | None = None
source_key: str | None = None
turn_index: int = 0
role: str = "" # "user" or "assistant"
is_short_user: bool = False # Short user message to merge
# ---------------------------------------------------------------------------
# Segmenter
# ---------------------------------------------------------------------------
class Segmenter:
"""Splits conversation messages into typed semantic objects.
Uses structural signals (tool boundaries, role transitions, content size)
to identify semantic boundaries and classify object types.
This implements Strategy A (structural segmentation) from
ARCHITECTURE.md §5.2 fast, rule-based, suitable for turns 1-50.
"""
def __init__(
self,
min_object_tokens: int = 125,
max_object_tokens: int = 8000,
) -> None:
"""
Args:
min_object_tokens: Minimum tokens for an object (below this,
merge with adjacent).
max_object_tokens: Maximum tokens before splitting.
"""
self.min_object_tokens = min_object_tokens
self.max_object_tokens = max_object_tokens
# -------------------------------------------------------------------
# Public API
# -------------------------------------------------------------------
def segment_messages(
self,
messages: list[dict],
start_turn: int = 0,
) -> list[SegmentedObject]:
"""Process a list of Anthropic-format messages and return semantic objects.
Each message is a dict with "role" and "content" (str or list of blocks).
Content blocks can be: text, tool_use, tool_result, image.
"""
raw = self._extract_raw_segments(messages, start_turn)
merged = self._merge_segments(raw)
return self._finalize(merged)
def segment_incremental(
self,
new_messages: list[dict],
existing_objects: list[SegmentedObject],
start_turn: int = 0,
) -> list[SegmentedObject]:
"""Process only new messages, extending existing objects if appropriate.
If the last existing object and the first new segment are compatible
(same type, same turn range), they get merged. Otherwise the new
segments are appended.
"""
new_raw = self._extract_raw_segments(new_messages, start_turn)
new_merged = self._merge_segments(new_raw)
new_objects = self._finalize(new_merged)
if not new_objects:
return list(existing_objects)
if not existing_objects:
return new_objects
# Try to merge the last existing object with the first new one
last = existing_objects[-1]
first_new = new_objects[0]
if self._can_merge_objects(last, first_new):
combined = self._merge_two_objects(last, first_new)
return list(existing_objects[:-1]) + [combined] + new_objects[1:]
return list(existing_objects) + new_objects
# -------------------------------------------------------------------
# Phase 1: Extract raw segments from messages
# -------------------------------------------------------------------
def _extract_raw_segments(
self,
messages: list[dict],
start_turn: int,
) -> list[_RawSegment]:
"""Walk messages and produce raw segments."""
segments: list[_RawSegment] = []
turn = start_turn
# Track tool_use IDs to tool names for matching with tool_results
tool_use_map: dict[str, str] = {}
# Track tool_use IDs to their input for source_key extraction
tool_input_map: dict[str, dict] = {}
for msg in messages:
role = msg.get("role", "")
content = msg.get("content", "")
if role == "assistant":
segs = self._extract_assistant(content, turn, tool_use_map, tool_input_map)
segments.extend(segs)
elif role == "user":
segs = self._extract_user(content, turn, tool_use_map, tool_input_map)
segments.extend(segs)
turn += 1
return segments
def _extract_assistant(
self,
content: str | list,
turn: int,
tool_use_map: dict[str, str],
tool_input_map: dict[str, dict],
) -> list[_RawSegment]:
"""Extract segments from an assistant message."""
if isinstance(content, str):
if not content.strip():
return []
obj_type = self._classify_text(content)
return [
_RawSegment(
content=content,
object_type=obj_type,
turn_index=turn,
role="assistant",
)
]
if not isinstance(content, list):
return []
segments: list[_RawSegment] = []
text_buffer: list[str] = []
for block in content:
if not isinstance(block, dict):
continue
block_type = block.get("type", "")
if block_type == "text":
text = block.get("text", "")
if text.strip():
text_buffer.append(text)
elif block_type == "tool_use":
# Record the tool_use for later matching with tool_result
tool_id = block.get("id", "")
tool_name = block.get("name", "")
tool_input = block.get("input", {})
if tool_id and tool_name:
tool_use_map[tool_id] = tool_name
tool_input_map[tool_id] = tool_input if isinstance(tool_input, dict) else {}
# Don't create a standalone object for tool_use — it merges
# with preceding text or following result
# Flush any accumulated text
if text_buffer:
combined = "\n".join(text_buffer)
obj_type = self._classify_text(combined)
segments.append(
_RawSegment(
content=combined,
object_type=obj_type,
turn_index=turn,
role="assistant",
)
)
return segments
def _extract_user(
self,
content: str | list,
turn: int,
tool_use_map: dict[str, str],
tool_input_map: dict[str, dict],
) -> list[_RawSegment]:
"""Extract segments from a user message.
User messages can contain:
- Plain text (str or text blocks)
- tool_result blocks (responses to assistant's tool_use)
"""
if isinstance(content, str):
if not content.strip():
return []
is_short = bool(_SHORT_MSG_RE.match(content.strip()))
return [
_RawSegment(
content=content,
object_type="conversation_phase",
turn_index=turn,
role="user",
is_short_user=is_short,
)
]
if not isinstance(content, list):
return []
segments: list[_RawSegment] = []
text_buffer: list[str] = []
for block in content:
if not isinstance(block, dict):
continue
block_type = block.get("type", "")
if block_type == "text":
text = block.get("text", "")
if text.strip():
text_buffer.append(text)
elif block_type == "tool_result":
# Flush text buffer first
if text_buffer:
combined = "\n".join(text_buffer)
is_short = bool(_SHORT_MSG_RE.match(combined.strip()))
segments.append(
_RawSegment(
content=combined,
object_type="conversation_phase",
turn_index=turn,
role="user",
is_short_user=is_short,
)
)
text_buffer = []
# Process tool result
seg = self._process_tool_result(block, turn, tool_use_map, tool_input_map)
if seg is not None:
segments.append(seg)
# Flush remaining text
if text_buffer:
combined = "\n".join(text_buffer)
is_short = bool(_SHORT_MSG_RE.match(combined.strip()))
segments.append(
_RawSegment(
content=combined,
object_type="conversation_phase",
turn_index=turn,
role="user",
is_short_user=is_short,
)
)
return segments
def _process_tool_result(
self,
block: dict,
turn: int,
tool_use_map: dict[str, str],
tool_input_map: dict[str, dict],
) -> _RawSegment | None:
"""Process a tool_result block into a raw segment."""
tool_use_id = block.get("tool_use_id", "")
tool_name = tool_use_map.get(tool_use_id, "")
tool_input = tool_input_map.get(tool_use_id, {})
# Extract content from tool_result
result_content = block.get("content", "")
if isinstance(result_content, list):
# tool_result content can be a list of blocks
parts = []
for part in result_content:
if isinstance(part, dict) and part.get("type") == "text":
parts.append(part.get("text", ""))
elif isinstance(part, str):
parts.append(part)
result_content = "\n".join(parts)
elif not isinstance(result_content, str):
result_content = str(result_content)
if not result_content.strip():
return None
# Classify by tool name
obj_type, source_key = self._classify_tool_result(tool_name, tool_input, result_content)
return _RawSegment(
content=result_content,
object_type=obj_type,
source_tool=tool_name or None,
source_key=source_key,
turn_index=turn,
role="user", # tool_results come in user messages
)
def _classify_tool_result(
self,
tool_name: str,
tool_input: dict,
content: str,
) -> tuple[str, str | None]:
"""Classify a tool result into object type and source key.
Returns:
(object_type, source_key) tuple.
"""
source_key: str | None = None
if tool_name in _READ_TOOLS:
# Extract file path from tool input
source_key = (
tool_input.get("file_path") or tool_input.get("filePath") or tool_input.get("path")
)
return "file_context", source_key
if tool_name in _EXTERNAL_TOOLS:
source_key = tool_input.get("url") or tool_input.get("uri")
return "external_reference", source_key
# All other tools → tool_result
# But check content for error patterns
if self._has_error_content(content):
return "error_context", source_key
return "tool_result", source_key
# -------------------------------------------------------------------
# Phase 2: Merge segments
# -------------------------------------------------------------------
def _merge_segments(self, segments: list[_RawSegment]) -> list[_RawSegment]:
"""Merge adjacent segments according to rules.
1. Short user messages merge with the next segment.
2. Adjacent conversation_phase segments in the same turn pair merge.
3. Segments below min_object_tokens merge with adjacent compatible.
"""
if not segments:
return []
# Pass 1: Merge short user messages with the next segment
merged: list[_RawSegment] = []
i = 0
while i < len(segments):
seg = segments[i]
if seg.is_short_user and i + 1 < len(segments):
# Merge with next segment
next_seg = segments[i + 1]
combined_content = seg.content + "\n" + next_seg.content
merged.append(
_RawSegment(
content=combined_content,
object_type=next_seg.object_type,
source_tool=next_seg.source_tool,
source_key=next_seg.source_key,
turn_index=seg.turn_index,
role=next_seg.role,
)
)
i += 2
else:
merged.append(seg)
i += 1
# Pass 2: Merge adjacent conversation_phase segments in same turn pair
merged2: list[_RawSegment] = []
for seg in merged:
if (
merged2
and merged2[-1].object_type == "conversation_phase"
and seg.object_type == "conversation_phase"
and abs(seg.turn_index - merged2[-1].turn_index) <= 1
):
prev = merged2[-1]
prev.content = prev.content + "\n" + seg.content
# Keep the earlier turn_index
else:
merged2.append(seg)
# Pass 3: Merge undersized segments with adjacent compatible
merged3: list[_RawSegment] = []
for seg in merged2:
tokens = _estimate_tokens(seg.content)
if (
tokens < self.min_object_tokens
and merged3
and self._compatible_types(merged3[-1].object_type, seg.object_type)
):
prev = merged3[-1]
prev.content = prev.content + "\n" + seg.content
# Prefer the more specific type
if seg.object_type != "conversation_phase":
prev.object_type = seg.object_type
if seg.source_tool and not prev.source_tool:
prev.source_tool = seg.source_tool
if seg.source_key and not prev.source_key:
prev.source_key = seg.source_key
else:
merged3.append(seg)
return merged3
def _compatible_types(self, type_a: str, type_b: str) -> bool:
"""Check if two object types are compatible for merging."""
if type_a == type_b:
return True
# conversation_phase is compatible with most types
if "conversation_phase" in (type_a, type_b):
return True
# error_context and debugging_session are compatible
if {type_a, type_b} == {"error_context", "debugging_session"}:
return True
return False
# -------------------------------------------------------------------
# Phase 3: Finalize into SegmentedObjects
# -------------------------------------------------------------------
def _finalize(self, segments: list[_RawSegment]) -> list[SegmentedObject]:
"""Convert raw segments into final SegmentedObjects.
Handles splitting oversized segments and generating metadata.
"""
objects: list[SegmentedObject] = []
for seg in segments:
tokens = _estimate_tokens(seg.content)
if tokens > self.max_object_tokens:
# Split oversized segments
parts = self._split_content(seg.content, self.max_object_tokens)
for i, part in enumerate(parts):
obj = self._make_object(
content=part,
object_type=seg.object_type,
source_tool=seg.source_tool,
source_key=seg.source_key,
turn_index=seg.turn_index,
)
objects.append(obj)
else:
obj = self._make_object(
content=seg.content,
object_type=seg.object_type,
source_tool=seg.source_tool,
source_key=seg.source_key,
turn_index=seg.turn_index,
)
objects.append(obj)
return objects
def _make_object(
self,
content: str,
object_type: str,
source_tool: str | None,
source_key: str | None,
turn_index: int,
) -> SegmentedObject:
"""Create a SegmentedObject with auto-generated metadata."""
entities = self._extract_entities(content)
tags = self._generate_tags(object_type, source_tool, content)
stub = self._generate_stub(object_type, content)
return SegmentedObject(
content=content,
object_type=object_type,
source_tool=source_tool,
source_key=source_key,
stub=stub,
turn_start=turn_index,
turn_end=turn_index,
token_estimate=_estimate_tokens(content),
key_entities=entities,
tags=tags,
)
def _split_content(self, content: str, max_tokens: int) -> list[str]:
"""Split content into chunks of at most max_tokens.
Tries to split on paragraph boundaries first, then line boundaries.
"""
max_chars = max_tokens * 4 # Reverse the token estimate
if len(content) <= max_chars:
return [content]
parts: list[str] = []
# Try splitting on double newlines (paragraphs)
paragraphs = content.split("\n\n")
current: list[str] = []
current_len = 0
for para in paragraphs:
para_len = len(para) + 2 # +2 for the \n\n
if current_len + para_len > max_chars and current:
parts.append("\n\n".join(current))
current = [para]
current_len = para_len
else:
current.append(para)
current_len += para_len
if current:
parts.append("\n\n".join(current))
# If any part is still too large, split on lines
final_parts: list[str] = []
for part in parts:
if len(part) > max_chars:
lines = part.split("\n")
if len(lines) <= 1:
# No newlines — hard split by character count
for i in range(0, len(part), max_chars):
final_parts.append(part[i : i + max_chars])
else:
chunk: list[str] = []
chunk_len = 0
for line in lines:
line_len = len(line) + 1
if chunk_len + line_len > max_chars and chunk:
final_parts.append("\n".join(chunk))
chunk = [line]
chunk_len = line_len
else:
chunk.append(line)
chunk_len += line_len
if chunk:
final_parts.append("\n".join(chunk))
else:
final_parts.append(part)
return final_parts if final_parts else [content]
# -------------------------------------------------------------------
# Classification helpers
# -------------------------------------------------------------------
def _classify_text(self, text: str) -> str:
"""Classify assistant text content into an object type.
Priority order: error > debug > plan > decision > conversation_phase
"""
if self._has_error_content(text):
return "error_context"
if self._matches_patterns(text, _DEBUG_PATTERNS):
return "debugging_session"
if self._matches_patterns(text, _PLAN_PATTERNS):
return "plan"
if self._matches_patterns(text, _DECISION_PATTERNS):
return "design_decision"
return "conversation_phase"
def _has_error_content(self, text: str) -> bool:
"""Check if text contains error/stack trace patterns."""
# Require at least 2 error pattern matches for confidence
matches = sum(1 for p in _ERROR_PATTERNS if p.search(text))
return matches >= 2
def _matches_patterns(self, text: str, patterns: list[re.Pattern]) -> bool:
"""Check if text matches at least one pattern from the list."""
return any(p.search(text) for p in patterns)
# -------------------------------------------------------------------
# Entity extraction
# -------------------------------------------------------------------
def _extract_entities(self, content: str) -> list[str]:
"""Extract key entities from content (file paths, functions, packages)."""
entities: list[str] = []
seen: set[str] = set()
# File paths
for m in _FILE_PATH_RE.finditer(content):
path = m.group(1)
if path not in seen and len(path) > 3:
seen.add(path)
entities.append(path)
# Function names (only from def/function declarations, not all calls)
for m in _FUNCTION_RE.finditer(content):
name = m.group(1) or m.group(2)
if name and name not in seen and len(name) > 2:
# Filter out common keywords and builtins
if name not in _COMMON_WORDS:
seen.add(name)
entities.append(name)
# Import/require names
for m in _IMPORT_RE.finditer(content):
name = m.group(1) or m.group(2) or m.group(3)
if name and name not in seen:
seen.add(name)
entities.append(name)
return entities[:20] # Cap at 20 entities
# -------------------------------------------------------------------
# Tag generation
# -------------------------------------------------------------------
def _generate_tags(
self,
object_type: str,
source_tool: str | None,
content: str,
) -> list[str]:
"""Auto-generate tags based on type, tool, and content."""
tags: list[str] = [object_type]
if source_tool:
tags.append(source_tool.lower())
# File extension tags
ext_matches = _FILE_EXT_RE.findall(content)
seen_exts: set[str] = set()
for ext in ext_matches:
dotted = f".{ext.lower()}"
if dotted not in seen_exts and dotted in _KNOWN_EXTENSIONS:
seen_exts.add(dotted)
tags.append(dotted)
return tags
# -------------------------------------------------------------------
# Stub generation
# -------------------------------------------------------------------
def _generate_stub(self, object_type: str, content: str) -> str:
"""Generate a one-line stub: [{type}: {first_line_truncated}]"""
first_line = content.split("\n", 1)[0].strip()
if len(first_line) > 100:
first_line = first_line[:97] + "..."
return f"[{object_type}: {first_line}]"
# -------------------------------------------------------------------
# Incremental merge helpers
# -------------------------------------------------------------------
def _can_merge_objects(self, a: SegmentedObject, b: SegmentedObject) -> bool:
"""Check if two SegmentedObjects can be merged."""
if not self._compatible_types(a.object_type, b.object_type):
return False
# Only merge if they're adjacent turns
if b.turn_start - a.turn_end > 1:
return False
# Don't merge file_context or tool_result objects
if a.object_type in ("file_context", "tool_result"):
return False
if b.object_type in ("file_context", "tool_result"):
return False
# Don't merge if combined would be too large
combined_tokens = a.token_estimate + b.token_estimate
if combined_tokens > self.max_object_tokens:
return False
return True
def _merge_two_objects(self, a: SegmentedObject, b: SegmentedObject) -> SegmentedObject:
"""Merge two SegmentedObjects into one."""
combined_content = a.content + "\n" + b.content
# Prefer the more specific type
obj_type = a.object_type
if b.object_type != "conversation_phase":
obj_type = b.object_type
entities = list(dict.fromkeys(a.key_entities + b.key_entities))[:20]
tags = list(dict.fromkeys(a.tags + b.tags))
return SegmentedObject(
content=combined_content,
object_type=obj_type,
source_tool=a.source_tool or b.source_tool,
source_key=a.source_key or b.source_key,
stub=self._generate_stub(obj_type, combined_content),
turn_start=min(a.turn_start, b.turn_start),
turn_end=max(a.turn_end, b.turn_end),
token_estimate=_estimate_tokens(combined_content),
key_entities=entities,
tags=tags,
)
# ---------------------------------------------------------------------------
# Module-level constants (used by entity extraction / tag generation)
# ---------------------------------------------------------------------------
_COMMON_WORDS = frozenset(
{
"if",
"for",
"in",
"is",
"it",
"to",
"do",
"or",
"and",
"not",
"the",
"this",
"that",
"with",
"from",
"self",
"None",
"True",
"False",
"return",
"class",
"import",
"print",
"len",
"str",
"int",
"list",
"dict",
"set",
"type",
"def",
"var",
"let",
"const",
"new",
"get",
"has",
"map",
"any",
"all",
}
)
_FILE_EXT_RE = re.compile(
r"\.(ts|tsx|js|jsx|py|rs|go|rb|java|cpp|c|h|css|html|md|json|yaml|yml|toml|sql)\b"
)
_KNOWN_EXTENSIONS = frozenset(
{
".ts",
".tsx",
".js",
".jsx",
".py",
".rs",
".go",
".rb",
".java",
".cpp",
".c",
".h",
".css",
".html",
".md",
".json",
".yaml",
".yml",
".toml",
".sql",
}
)

View file

@ -0,0 +1,408 @@
"""Tests for the ContextAssembler micro-fault and context assembly module."""
from __future__ import annotations
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from mnemosyne.context_assembler import (
ContextAssembler,
ContextWindow,
MicroFaultResult,
_estimate_tokens,
)
from mnemosyne.fidelity import FidelityLevel, FidelityManager, make_object
from mnemosyne.object_store import (
DummyEmbedder,
InMemoryBackend,
ObjectStore,
StoredObject,
)
# ── Fixtures ─────────────────────────────────────────────────
@pytest.fixture
def object_store():
"""Create an ObjectStore with InMemoryBackend and DummyEmbedder."""
backend = InMemoryBackend()
embedder = DummyEmbedder()
return ObjectStore(backend=backend, embedder=embedder)
@pytest.fixture
def mock_helper_llm():
"""Create a mock HelperLLM with answer_micro_fault returning a canned answer."""
helper = AsyncMock()
helper.answer_micro_fault = AsyncMock(
return_value="The auth library used is passport.js with JWT strategy."
)
return helper
@pytest.fixture
def assembler(object_store, mock_helper_llm):
"""Create a ContextAssembler with real ObjectStore and mock HelperLLM."""
return ContextAssembler(object_store=object_store, helper_llm=mock_helper_llm)
@pytest.fixture
def assembler_no_helper(object_store):
"""Create a ContextAssembler without HelperLLM (fallback mode)."""
return ContextAssembler(object_store=object_store, helper_llm=None)
async def _seed_objects(store: ObjectStore, session_id: str = "sess_test") -> list[StoredObject]:
"""Seed the store with a few test objects and return them."""
obj1 = await store.store_object(
session_id=session_id,
content="Authentication uses passport.js with JWT strategy. "
"The secret key is loaded from AUTH_SECRET env var. "
"Token expiry is set to 24 hours.",
object_type="file_context",
source_tool="Read",
source_key="src/auth/middleware.ts",
stub="[file_context | auth middleware | passport.js JWT]",
key_entities=["passport.js", "JWT", "AUTH_SECRET"],
)
obj2 = await store.store_object(
session_id=session_id,
content="Database schema has users, posts, and comments tables. "
"Users table has id, email, password_hash, created_at columns. "
"Posts reference users via author_id foreign key.",
object_type="file_context",
source_tool="Read",
source_key="schema.sql",
stub="[file_context | DB schema | users, posts, comments]",
key_entities=["users", "posts", "comments", "schema.sql"],
)
obj3 = await store.store_object(
session_id=session_id,
content="Server runs on port 3000 by default. "
"Configuration loaded from .env file. "
"CORS enabled for localhost:5173 in development.",
object_type="file_context",
source_tool="Read",
source_key="src/config.ts",
stub="[file_context | server config | port 3000, CORS]",
key_entities=["port 3000", "CORS", ".env"],
)
return [obj1, obj2, obj3]
# ── MicroFaultResult dataclass ───────────────────────────────
def test_micro_fault_result_fields():
"""MicroFaultResult should store all expected fields."""
result = MicroFaultResult(
answer="The answer is 42.",
sources=["obj_abc", "obj_def"],
answer_tokens=5,
avoided_tokens=500,
latency_ms=12.5,
)
assert result.answer == "The answer is 42."
assert result.sources == ["obj_abc", "obj_def"]
assert result.answer_tokens == 5
assert result.avoided_tokens == 500
assert result.latency_ms == 12.5
def test_micro_fault_result_token_savings():
"""avoided_tokens should represent tokens saved vs full restore."""
result = MicroFaultResult(
answer="short answer",
sources=["obj1"],
answer_tokens=3,
avoided_tokens=997,
latency_ms=10.0,
)
# The savings ratio: avoided / (avoided + answer_tokens)
savings_ratio = result.avoided_tokens / (result.avoided_tokens + result.answer_tokens)
assert savings_ratio > 0.99 # 99%+ savings
# ── ContextWindow dataclass ──────────────────────────────────
def test_context_window_fields():
"""ContextWindow should store objects, total_tokens, and pressure_zone."""
window = ContextWindow(
objects=[("obj1", 0, "full content"), ("obj2", 2, "compact summary")],
total_tokens=150,
pressure_zone="NORMAL",
)
assert len(window.objects) == 2
assert window.total_tokens == 150
assert window.pressure_zone == "NORMAL"
# ── handle_micro_fault with HelperLLM ────────────────────────
@pytest.mark.asyncio
async def test_micro_fault_with_helper(assembler, object_store, mock_helper_llm):
"""Micro-fault with HelperLLM should search, call helper, return answer."""
await _seed_objects(object_store, "sess_test")
result = await assembler.handle_micro_fault(
session_id="sess_test",
question="What auth library is used?",
)
assert isinstance(result, MicroFaultResult)
assert result.answer == "The auth library used is passport.js with JWT strategy."
assert len(result.sources) > 0
assert result.answer_tokens > 0
assert result.avoided_tokens > 0
assert result.latency_ms >= 0
# Verify HelperLLM was called
mock_helper_llm.answer_micro_fault.assert_called_once()
call_kwargs = mock_helper_llm.answer_micro_fault.call_args
assert call_kwargs.kwargs["question"] == "What auth library is used?"
assert len(call_kwargs.kwargs["relevant_contents"]) > 0
@pytest.mark.asyncio
async def test_micro_fault_with_scope(assembler, object_store, mock_helper_llm):
"""Micro-fault with scope hint should incorporate it into the search query."""
await _seed_objects(object_store, "sess_test")
result = await assembler.handle_micro_fault(
session_id="sess_test",
question="What port?",
scope="config files",
)
assert isinstance(result, MicroFaultResult)
assert len(result.sources) > 0
mock_helper_llm.answer_micro_fault.assert_called_once()
@pytest.mark.asyncio
async def test_micro_fault_with_max_tokens(assembler, object_store, mock_helper_llm):
"""Micro-fault should pass max_tokens to the HelperLLM."""
await _seed_objects(object_store, "sess_test")
await assembler.handle_micro_fault(
session_id="sess_test",
question="What is the DB schema?",
max_tokens=100,
)
call_kwargs = mock_helper_llm.answer_micro_fault.call_args
assert call_kwargs.kwargs["max_tokens"] == 100
# ── handle_micro_fault fallback (no HelperLLM) ──────────────
@pytest.mark.asyncio
async def test_micro_fault_fallback_no_helper(assembler_no_helper, object_store):
"""Without HelperLLM, micro-fault should return summaries as fallback."""
objects = await _seed_objects(object_store, "sess_test")
# Set a summary on one object so fallback has something to show
await object_store.update_fidelity(objects[0].id, 2, summary="Auth uses passport.js JWT")
result = await assembler_no_helper.handle_micro_fault(
session_id="sess_test",
question="What auth library is used?",
)
assert isinstance(result, MicroFaultResult)
assert "HelperLLM unavailable" in result.answer
assert len(result.sources) > 0
assert result.answer_tokens > 0
@pytest.mark.asyncio
async def test_micro_fault_fallback_uses_stub_when_no_summary(assembler_no_helper, object_store):
"""Fallback should use stub when no summaries are available."""
await _seed_objects(object_store, "sess_test")
result = await assembler_no_helper.handle_micro_fault(
session_id="sess_test",
question="What auth library is used?",
)
assert isinstance(result, MicroFaultResult)
assert "HelperLLM unavailable" in result.answer
# Should contain stub content since no summaries exist
assert len(result.sources) > 0
# ── handle_micro_fault with no search results ────────────────
@pytest.mark.asyncio
async def test_micro_fault_no_results(assembler):
"""Micro-fault with no matching objects should return a 'not found' message."""
# Don't seed any objects — empty store
result = await assembler.handle_micro_fault(
session_id="sess_empty",
question="What is the meaning of life?",
)
assert isinstance(result, MicroFaultResult)
assert "No relevant content found" in result.answer
assert result.sources == []
assert result.avoided_tokens == 0
# ── Micro-fault records access on consulted objects ──────────
@pytest.mark.asyncio
async def test_micro_fault_records_access(assembler, object_store):
"""Micro-fault should call record_fault(is_micro=True) on consulted objects."""
objects = await _seed_objects(object_store, "sess_test")
result = await assembler.handle_micro_fault(
session_id="sess_test",
question="What auth library is used?",
)
# At least one source should have been consulted
assert len(result.sources) > 0
# Check that micro_fault_count was incremented on consulted objects
for source_id in result.sources:
obj = await object_store.get(source_id)
assert obj is not None
assert obj.micro_fault_count > 0
# ── assemble_context ─────────────────────────────────────────
@pytest.mark.asyncio
async def test_assemble_context_all_l0(assembler, object_store):
"""assemble_context with all L0 objects should include full content."""
objects = await _seed_objects(object_store, "sess_test")
fm = FidelityManager(window_size=200_000)
for obj in objects:
fm_obj = make_object(
object_type=obj.object_type,
content_full=obj.content_full,
stub=obj.stub,
)
fm_obj.id = obj.id # match IDs
fm.register_object(fm_obj)
blocks = await assembler.assemble_context(
session_id="sess_test",
fidelity_manager=fm,
current_turn=1,
)
assert len(blocks) == 3
for block in blocks:
assert block["fidelity"] == 0 # L0
assert len(block["content"]) > 50 # full content
assert block["tokens"] > 0
@pytest.mark.asyncio
async def test_assemble_context_mixed_fidelity(assembler, object_store):
"""assemble_context should respect per-object fidelity levels."""
objects = await _seed_objects(object_store, "sess_test")
# Set different fidelity levels in the store
await object_store.update_fidelity(
objects[0].id,
0, # L0: full
)
await object_store.update_fidelity(
objects[1].id,
2,
summary="DB has users, posts, comments tables", # L2: compact
)
await object_store.update_fidelity(
objects[2].id,
3, # L3: stub
)
fm = FidelityManager(window_size=200_000)
for obj in objects:
fm_obj = make_object(
object_type=obj.object_type,
content_full=obj.content_full,
stub=obj.stub,
)
fm_obj.id = obj.id
# Set fidelity to match what we set in the store
stored = await object_store.get(obj.id)
fm_obj.current_fidelity = FidelityLevel(stored.current_fidelity)
if stored.summary_compact:
fm_obj.summary_compact = stored.summary_compact
fm.register_object(fm_obj)
blocks = await assembler.assemble_context(
session_id="sess_test",
fidelity_manager=fm,
current_turn=5,
)
assert len(blocks) == 3
fidelities = {b["object_id"]: b["fidelity"] for b in blocks}
assert fidelities[objects[0].id] == 0 # L0
assert fidelities[objects[1].id] == 2 # L2
assert fidelities[objects[2].id] == 3 # L3
@pytest.mark.asyncio
async def test_assemble_context_excludes_evicted(assembler, object_store):
"""assemble_context should exclude L4 (evicted) objects."""
objects = await _seed_objects(object_store, "sess_test")
# Evict one object
await object_store.update_fidelity(objects[1].id, 4)
fm = FidelityManager(window_size=200_000)
for obj in objects:
fm_obj = make_object(
object_type=obj.object_type,
content_full=obj.content_full,
stub=obj.stub,
)
fm_obj.id = obj.id
stored = await object_store.get(obj.id)
fm_obj.current_fidelity = FidelityLevel(min(stored.current_fidelity, 4))
fm.register_object(fm_obj)
blocks = await assembler.assemble_context(
session_id="sess_test",
fidelity_manager=fm,
current_turn=5,
)
# Only 2 objects should be in context (one was evicted)
block_ids = {b["object_id"] for b in blocks}
assert objects[1].id not in block_ids
assert len(blocks) == 2
# ── _estimate_tokens helper ──────────────────────────────────
def test_estimate_tokens_none():
"""_estimate_tokens(None) should return 0."""
assert _estimate_tokens(None) == 0
def test_estimate_tokens_empty():
"""_estimate_tokens('') should return 1 (minimum)."""
assert _estimate_tokens("") == 1
def test_estimate_tokens_normal():
"""_estimate_tokens should use ~4 chars per token heuristic."""
text = "a" * 400
assert _estimate_tokens(text) == 100

996
tests/test_object_store.py Normal file
View file

@ -0,0 +1,996 @@
"""Tests for the semantic object backing store."""
from __future__ import annotations
import asyncio
import pytest
from mnemosyne.object_store import (
DummyEmbedder,
InMemoryBackend,
ObjectStore,
ObjectStoreBackend,
StoredObject,
_cosine_similarity,
_estimate_tokens,
)
# ── Helpers ──────────────────────────────────────────────────
def _make_stored_object(
session_id: str = "sess-1",
content: str = "x" * 400,
*,
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,
) -> StoredObject:
"""Create a StoredObject with sensible defaults for testing."""
return StoredObject(
id=object_id or f"obj-{id(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}: test object",
tokens_l0=_estimate_tokens(content),
tokens_l3=_estimate_tokens(stub or f"{object_type}: test object"),
embedding=embedding or [],
created_at="2025-01-01T00:00:00+00:00",
last_accessed="2025-01-01T00:00:00+00:00",
)
# ── StoredObject dataclass ───────────────────────────────────
class TestStoredObject:
def test_required_fields(self):
obj = _make_stored_object()
assert obj.id
assert obj.session_id == "sess-1"
assert obj.object_type == "file_context"
assert obj.content_full == "x" * 400
assert obj.current_fidelity == 0
assert obj.pinned is False
assert obj.fault_count == 0
assert obj.micro_fault_count == 0
def test_default_lists_are_empty(self):
obj = _make_stored_object()
assert obj.losses_l1 == []
assert obj.losses_l2 == []
assert obj.can_answer_l1 == []
assert obj.can_answer_l2 == []
assert obj.fault_when == []
assert obj.key_entities == []
assert obj.tags == []
def test_content_at_levels(self):
obj = _make_stored_object(content="full content")
obj.summary_detailed = "detailed"
obj.summary_compact = "compact"
obj.stub = "stub line"
assert obj.content_at(0) == "full content"
assert obj.content_at(1) == "detailed"
assert obj.content_at(2) == "compact"
assert obj.content_at(3) == "stub line"
assert obj.content_at(4) is None
def test_tokens_at_levels(self):
obj = _make_stored_object(content="a" * 400)
obj.tokens_l0 = 100
obj.tokens_l1 = 30
obj.tokens_l2 = 5
obj.tokens_l3 = 3
assert obj.tokens_at(0) == 100
assert obj.tokens_at(1) == 30
assert obj.tokens_at(2) == 5
assert obj.tokens_at(3) == 3
assert obj.tokens_at(4) == 0
def test_tokens_at_none_levels(self):
obj = _make_stored_object()
obj.tokens_l1 = None
obj.tokens_l2 = None
assert obj.tokens_at(1) == 0
assert obj.tokens_at(2) == 0
def test_current_tokens_tracks_fidelity(self):
obj = _make_stored_object(content="a" * 400)
obj.tokens_l0 = 100
obj.tokens_l1 = 30
assert obj.current_tokens == 100 # L0
obj.current_fidelity = 1
assert obj.current_tokens == 30 # L1
def test_embedding_default_empty(self):
obj = _make_stored_object()
assert obj.embedding == []
# ── DummyEmbedder ────────────────────────────────────────────
class TestDummyEmbedder:
def test_embed_returns_384_dim(self):
emb = DummyEmbedder()
vec = emb.embed("hello world")
assert len(vec) == 384
def test_embed_deterministic(self):
emb = DummyEmbedder()
v1 = emb.embed("test input")
v2 = emb.embed("test input")
assert v1 == v2
def test_embed_different_inputs_differ(self):
emb = DummyEmbedder()
v1 = emb.embed("input A")
v2 = emb.embed("input B")
assert v1 != v2
def test_embed_is_normalized(self):
import numpy as np
emb = DummyEmbedder()
vec = emb.embed("normalize me")
norm = float(np.linalg.norm(vec))
assert abs(norm - 1.0) < 1e-6
def test_embed_batch(self):
emb = DummyEmbedder()
texts = ["alpha", "beta", "gamma"]
vecs = emb.embed_batch(texts)
assert len(vecs) == 3
assert all(len(v) == 384 for v in vecs)
def test_embed_batch_matches_single(self):
emb = DummyEmbedder()
texts = ["one", "two"]
batch = emb.embed_batch(texts)
singles = [emb.embed(t) for t in texts]
assert batch == singles
# ── Cosine similarity ────────────────────────────────────────
class TestCosineSimilarity:
def test_identical_vectors(self):
v = [1.0, 0.0, 0.0]
assert abs(_cosine_similarity(v, v) - 1.0) < 1e-9
def test_orthogonal_vectors(self):
a = [1.0, 0.0, 0.0]
b = [0.0, 1.0, 0.0]
assert abs(_cosine_similarity(a, b)) < 1e-9
def test_opposite_vectors(self):
a = [1.0, 0.0]
b = [-1.0, 0.0]
assert abs(_cosine_similarity(a, b) - (-1.0)) < 1e-9
def test_empty_vectors(self):
assert _cosine_similarity([], []) == 0.0
assert _cosine_similarity([1.0], []) == 0.0
def test_zero_vector(self):
assert _cosine_similarity([0.0, 0.0], [1.0, 0.0]) == 0.0
# ── InMemoryBackend CRUD ─────────────────────────────────────
class TestInMemoryBackendCRUD:
async def test_store_and_get(self):
backend = InMemoryBackend()
obj = _make_stored_object(object_id="obj-001")
await backend.store(obj)
result = await backend.get("obj-001")
assert result is obj
async def test_get_nonexistent(self):
backend = InMemoryBackend()
assert await backend.get("nonexistent") is None
async def test_store_overwrites(self):
backend = InMemoryBackend()
obj1 = _make_stored_object(object_id="obj-001", content="original")
obj2 = _make_stored_object(object_id="obj-001", content="updated")
await backend.store(obj1)
await backend.store(obj2)
result = await backend.get("obj-001")
assert result is not None
assert result.content_full == "updated"
async def test_get_by_session(self):
backend = InMemoryBackend()
obj1 = _make_stored_object(session_id="s1", object_id="o1")
obj2 = _make_stored_object(session_id="s1", object_id="o2")
obj3 = _make_stored_object(session_id="s2", object_id="o3")
await backend.store(obj1)
await backend.store(obj2)
await backend.store(obj3)
s1_objs = await backend.get_by_session("s1")
assert len(s1_objs) == 2
assert {o.id for o in s1_objs} == {"o1", "o2"}
async def test_get_by_session_filters_fidelity(self):
backend = InMemoryBackend()
obj1 = _make_stored_object(session_id="s1", object_id="o1")
obj2 = _make_stored_object(session_id="s1", object_id="o2")
obj2.current_fidelity = 4 # evicted
await backend.store(obj1)
await backend.store(obj2)
# Default: include evicted
all_objs = await backend.get_by_session("s1", fidelity_max=4)
assert len(all_objs) == 2
# Exclude evicted
active_objs = await backend.get_by_session("s1", fidelity_max=3)
assert len(active_objs) == 1
assert active_objs[0].id == "o1"
async def test_get_by_session_empty(self):
backend = InMemoryBackend()
assert await backend.get_by_session("nonexistent") == []
async def test_delete_session(self):
backend = InMemoryBackend()
obj1 = _make_stored_object(session_id="s1", object_id="o1")
obj2 = _make_stored_object(session_id="s1", object_id="o2")
obj3 = _make_stored_object(session_id="s2", object_id="o3")
await backend.store(obj1)
await backend.store(obj2)
await backend.store(obj3)
count = await backend.delete_session("s1")
assert count == 2
assert await backend.get("o1") is None
assert await backend.get("o2") is None
assert await backend.get("o3") is not None
async def test_delete_session_nonexistent(self):
backend = InMemoryBackend()
count = await backend.delete_session("nonexistent")
assert count == 0
# ── InMemoryBackend fidelity updates ─────────────────────────
class TestInMemoryBackendFidelity:
async def test_update_fidelity_basic(self):
backend = InMemoryBackend()
obj = _make_stored_object(object_id="o1")
await backend.store(obj)
await backend.update_fidelity("o1", 2)
result = await backend.get("o1")
assert result is not None
assert result.current_fidelity == 2
async def test_update_fidelity_with_summary_l1(self):
backend = InMemoryBackend()
obj = _make_stored_object(object_id="o1")
await backend.store(obj)
await backend.update_fidelity(
"o1", 1, summary="detailed summary", losses=["exact line numbers"]
)
result = await backend.get("o1")
assert result is not None
assert result.current_fidelity == 1
assert result.summary_detailed == "detailed summary"
assert result.losses_l1 == ["exact line numbers"]
assert result.tokens_l1 is not None
assert result.tokens_l1 > 0
async def test_update_fidelity_with_summary_l2(self):
backend = InMemoryBackend()
obj = _make_stored_object(object_id="o1")
await backend.store(obj)
await backend.update_fidelity(
"o1", 2, summary="compact summary", losses=["function bodies"]
)
result = await backend.get("o1")
assert result is not None
assert result.summary_compact == "compact summary"
assert result.losses_l2 == ["function bodies"]
async def test_update_fidelity_nonexistent(self):
backend = InMemoryBackend()
# Should not raise
await backend.update_fidelity("nonexistent", 1)
# ── InMemoryBackend source_key dedup ─────────────────────────
class TestInMemoryBackendSourceKey:
async def test_get_by_source_key(self):
backend = InMemoryBackend()
obj = _make_stored_object(session_id="s1", object_id="o1", source_key="src/auth.py")
await backend.store(obj)
result = await backend.get_by_source_key("s1", "src/auth.py")
assert result is not None
assert result.id == "o1"
async def test_get_by_source_key_not_found(self):
backend = InMemoryBackend()
assert await backend.get_by_source_key("s1", "nonexistent.py") is None
async def test_get_by_source_key_session_isolation(self):
backend = InMemoryBackend()
obj = _make_stored_object(session_id="s1", object_id="o1", source_key="src/auth.py")
await backend.store(obj)
# Different session should not find it
assert await backend.get_by_source_key("s2", "src/auth.py") is None
async def test_get_by_source_key_returns_latest(self):
backend = InMemoryBackend()
obj_old = _make_stored_object(session_id="s1", object_id="o1", source_key="src/auth.py")
obj_old.created_at = "2025-01-01T00:00:00+00:00"
obj_new = _make_stored_object(session_id="s1", object_id="o2", source_key="src/auth.py")
obj_new.created_at = "2025-01-02T00:00:00+00:00"
await backend.store(obj_old)
await backend.store(obj_new)
result = await backend.get_by_source_key("s1", "src/auth.py")
assert result is not None
assert result.id == "o2"
# ── InMemoryBackend embedding search ─────────────────────────
class TestInMemoryBackendEmbeddingSearch:
async def test_search_by_embedding_basic(self):
emb = DummyEmbedder()
backend = InMemoryBackend()
obj = _make_stored_object(
session_id="s1",
object_id="o1",
content="authentication middleware",
embedding=emb.embed("authentication middleware"),
)
await backend.store(obj)
query_vec = emb.embed("authentication middleware")
results = await backend.search_by_embedding("s1", query_vec, limit=5)
assert len(results) == 1
assert results[0][0].id == "o1"
# Same text → same embedding → similarity ≈ 1.0
assert results[0][1] > 0.99
async def test_search_by_embedding_ranking(self):
emb = DummyEmbedder()
backend = InMemoryBackend()
# Store objects with different content
for i, content in enumerate(["auth login", "database schema", "test runner"]):
obj = _make_stored_object(
session_id="s1",
object_id=f"o{i}",
content=content,
embedding=emb.embed(content),
)
await backend.store(obj)
# Search for exact match
query_vec = emb.embed("auth login")
results = await backend.search_by_embedding("s1", query_vec, limit=3)
assert len(results) == 3
# Exact match should be first with highest similarity
assert results[0][0].id == "o0"
assert results[0][1] > results[1][1]
async def test_search_by_embedding_session_isolation(self):
emb = DummyEmbedder()
backend = InMemoryBackend()
obj_s1 = _make_stored_object(
session_id="s1",
object_id="o1",
content="hello",
embedding=emb.embed("hello"),
)
obj_s2 = _make_stored_object(
session_id="s2",
object_id="o2",
content="hello",
embedding=emb.embed("hello"),
)
await backend.store(obj_s1)
await backend.store(obj_s2)
results = await backend.search_by_embedding("s1", emb.embed("hello"), limit=10)
assert len(results) == 1
assert results[0][0].id == "o1"
async def test_search_by_embedding_respects_limit(self):
emb = DummyEmbedder()
backend = InMemoryBackend()
for i in range(10):
obj = _make_stored_object(
session_id="s1",
object_id=f"o{i}",
content=f"content {i}",
embedding=emb.embed(f"content {i}"),
)
await backend.store(obj)
results = await backend.search_by_embedding("s1", emb.embed("content 0"), limit=3)
assert len(results) == 3
async def test_search_by_embedding_skips_no_embedding(self):
backend = InMemoryBackend()
obj = _make_stored_object(session_id="s1", object_id="o1", embedding=[])
await backend.store(obj)
results = await backend.search_by_embedding("s1", [0.1] * 384, limit=5)
assert len(results) == 0
# ── InMemoryBackend text search ──────────────────────────────
class TestInMemoryBackendTextSearch:
async def test_search_by_text_content(self):
backend = InMemoryBackend()
obj = _make_stored_object(
session_id="s1",
object_id="o1",
content="authentication middleware for JWT tokens",
)
await backend.store(obj)
results = await backend.search_by_text("s1", "JWT")
assert len(results) == 1
assert results[0].id == "o1"
async def test_search_by_text_case_insensitive(self):
backend = InMemoryBackend()
obj = _make_stored_object(session_id="s1", object_id="o1", content="Authentication")
await backend.store(obj)
results = await backend.search_by_text("s1", "authentication")
assert len(results) == 1
async def test_search_by_text_in_stub(self):
backend = InMemoryBackend()
obj = _make_stored_object(
session_id="s1",
object_id="o1",
content="some content",
stub="file_context: auth middleware",
)
await backend.store(obj)
results = await backend.search_by_text("s1", "auth middleware")
assert len(results) == 1
async def test_search_by_text_in_key_entities(self):
backend = InMemoryBackend()
obj = _make_stored_object(session_id="s1", object_id="o1", content="code")
obj.key_entities = ["AuthService", "JWTValidator"]
await backend.store(obj)
results = await backend.search_by_text("s1", "JWTValidator")
assert len(results) == 1
async def test_search_by_text_no_match(self):
backend = InMemoryBackend()
obj = _make_stored_object(session_id="s1", object_id="o1", content="hello")
await backend.store(obj)
results = await backend.search_by_text("s1", "nonexistent")
assert len(results) == 0
async def test_search_by_text_session_isolation(self):
backend = InMemoryBackend()
obj = _make_stored_object(session_id="s1", object_id="o1", content="shared keyword")
await backend.store(obj)
results = await backend.search_by_text("s2", "shared keyword")
assert len(results) == 0
async def test_search_by_text_respects_limit(self):
backend = InMemoryBackend()
for i in range(10):
obj = _make_stored_object(
session_id="s1",
object_id=f"o{i}",
content=f"common keyword item {i}",
)
await backend.store(obj)
results = await backend.search_by_text("s1", "common keyword", limit=3)
assert len(results) == 3
# ── Session isolation ────────────────────────────────────────
class TestSessionIsolation:
async def test_objects_isolated_by_session(self):
backend = InMemoryBackend()
obj_s1 = _make_stored_object(session_id="s1", object_id="o1")
obj_s2 = _make_stored_object(session_id="s2", object_id="o2")
await backend.store(obj_s1)
await backend.store(obj_s2)
s1_objs = await backend.get_by_session("s1")
s2_objs = await backend.get_by_session("s2")
assert len(s1_objs) == 1
assert s1_objs[0].id == "o1"
assert len(s2_objs) == 1
assert s2_objs[0].id == "o2"
async def test_delete_session_does_not_affect_other(self):
backend = InMemoryBackend()
obj_s1 = _make_stored_object(session_id="s1", object_id="o1")
obj_s2 = _make_stored_object(session_id="s2", object_id="o2")
await backend.store(obj_s1)
await backend.store(obj_s2)
await backend.delete_session("s1")
assert await backend.get("o1") is None
assert await backend.get("o2") is not None
assert len(await backend.get_by_session("s2")) == 1
async def test_source_key_scoped_to_session(self):
backend = InMemoryBackend()
obj_s1 = _make_stored_object(session_id="s1", object_id="o1", source_key="file.py")
obj_s2 = _make_stored_object(session_id="s2", object_id="o2", source_key="file.py")
await backend.store(obj_s1)
await backend.store(obj_s2)
result = await backend.get_by_source_key("s1", "file.py")
assert result is not None
assert result.id == "o1"
# ── ObjectStore facade ───────────────────────────────────────
class TestObjectStoreFacade:
async def test_store_object_creates_with_defaults(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1",
content="def authenticate(token): ...",
object_type="file_context",
source_tool="Read",
)
assert obj.id
assert len(obj.id) == 16
assert obj.session_id == "s1"
assert obj.object_type == "file_context"
assert obj.content_full == "def authenticate(token): ..."
assert obj.source_tool == "Read"
assert obj.current_fidelity == 0
assert obj.created_at
assert obj.last_accessed
assert obj.tokens_l0 > 0
assert len(obj.embedding) == 384
async def test_store_object_auto_generates_stub(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1",
content="some content here",
object_type="tool_result",
)
assert "tool_result:" in obj.stub
async def test_store_object_custom_stub(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1",
content="content",
object_type="file_context",
stub="Read src/auth.py (150 lines)",
)
assert obj.stub == "Read src/auth.py (150 lines)"
async def test_store_object_with_tags_and_entities(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1",
content="content",
object_type="file_context",
tags=["auth", "middleware"],
key_entities=["AuthService"],
)
assert obj.tags == ["auth", "middleware"]
assert obj.key_entities == ["AuthService"]
async def test_store_object_with_turn(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1",
content="content",
object_type="file_context",
turn=7,
)
assert obj.source_turn_start == 7
assert obj.source_turn_end == 7
async def test_store_object_no_embedder(self):
store = ObjectStore(InMemoryBackend(), embedder=None)
obj = await store.store_object(
session_id="s1",
content="content",
object_type="file_context",
)
assert obj.embedding == []
async def test_get_retrieves_stored(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1", content="content", object_type="file_context"
)
result = await store.get(obj.id)
assert result is obj
async def test_get_nonexistent(self):
store = ObjectStore(InMemoryBackend())
assert await store.get("nonexistent") is None
async def test_get_session_objects_excludes_evicted(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj1 = await store.store_object(
session_id="s1", content="active", object_type="file_context"
)
obj2 = await store.store_object(
session_id="s1", content="evicted", object_type="file_context"
)
await store.update_fidelity(obj2.id, 4)
active = await store.get_session_objects("s1")
assert len(active) == 1
assert active[0].id == obj1.id
async def test_get_session_objects_includes_evicted(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
await store.store_object(session_id="s1", content="active", object_type="file_context")
obj2 = await store.store_object(
session_id="s1", content="evicted", object_type="file_context"
)
await store.update_fidelity(obj2.id, 4)
all_objs = await store.get_session_objects("s1", include_evicted=True)
assert len(all_objs) == 2
async def test_update_fidelity(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1", content="content", object_type="file_context"
)
await store.update_fidelity(obj.id, 1, summary="summary", losses=["details"])
result = await store.get(obj.id)
assert result is not None
assert result.current_fidelity == 1
assert result.summary_detailed == "summary"
assert result.losses_l1 == ["details"]
# ── ObjectStore semantic search ──────────────────────────────
class TestObjectStoreSemanticSearch:
async def test_semantic_search_finds_exact_match(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
await store.store_object(
session_id="s1",
content="authentication middleware for JWT",
object_type="file_context",
)
await store.store_object(
session_id="s1",
content="database migration script",
object_type="file_context",
)
results = await store.semantic_search("s1", "authentication middleware for JWT")
assert len(results) >= 1
# Exact content match should rank first
assert results[0][0].content_full == "authentication middleware for JWT"
async def test_semantic_search_hybrid_boost(self):
"""Text match should boost ranking via hybrid scoring."""
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
await store.store_object(
session_id="s1",
content="the quick brown fox jumps over the lazy dog",
object_type="file_context",
)
await store.store_object(
session_id="s1",
content="unrelated content about databases",
object_type="file_context",
)
results = await store.semantic_search("s1", "quick brown fox")
assert len(results) >= 1
# Text match should help the fox content rank higher
assert "fox" in results[0][0].content_full
async def test_semantic_search_session_isolation(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
await store.store_object(
session_id="s1", content="session one content", object_type="file_context"
)
await store.store_object(
session_id="s2", content="session two content", object_type="file_context"
)
results = await store.semantic_search("s1", "session one content")
assert all(r[0].session_id == "s1" for r in results)
async def test_semantic_search_respects_limit(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
for i in range(10):
await store.store_object(
session_id="s1",
content=f"content item {i}",
object_type="file_context",
)
results = await store.semantic_search("s1", "content item", limit=3)
assert len(results) <= 3
async def test_semantic_search_no_embedder(self):
"""Without embedder, search falls back to text-only."""
store = ObjectStore(InMemoryBackend(), embedder=None)
await store.store_object(
session_id="s1",
content="findable keyword here",
object_type="file_context",
)
results = await store.semantic_search("s1", "findable keyword")
assert len(results) == 1
assert results[0][1] == 0.3 # text-only score
async def test_semantic_search_empty_session(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
results = await store.semantic_search("empty", "anything")
assert results == []
# ── ObjectStore deduplication ────────────────────────────────
class TestObjectStoreDedup:
async def test_find_duplicate_exists(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
await store.store_object(
session_id="s1",
content="file content",
object_type="file_context",
source_key="src/auth.py",
)
dup = await store.find_duplicate("s1", "src/auth.py")
assert dup is not None
assert dup.source_key == "src/auth.py"
async def test_find_duplicate_not_found(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
dup = await store.find_duplicate("s1", "nonexistent.py")
assert dup is None
async def test_find_duplicate_session_scoped(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
await store.store_object(
session_id="s1",
content="content",
object_type="file_context",
source_key="src/auth.py",
)
dup = await store.find_duplicate("s2", "src/auth.py")
assert dup is None
# ── ObjectStore access/fault tracking ────────────────────────
class TestObjectStoreTracking:
async def test_record_access(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1", content="content", object_type="file_context"
)
original_accessed = obj.last_accessed
await store.record_access(obj.id)
assert obj.access_count == 1
assert obj.last_accessed >= original_accessed
async def test_record_access_increments(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1", content="content", object_type="file_context"
)
await store.record_access(obj.id)
await store.record_access(obj.id)
await store.record_access(obj.id)
assert obj.access_count == 3
async def test_record_access_nonexistent(self):
store = ObjectStore(InMemoryBackend())
# Should not raise
await store.record_access("nonexistent")
async def test_record_fault(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1", content="content", object_type="file_context"
)
await store.record_fault(obj.id)
assert obj.fault_count == 1
assert obj.micro_fault_count == 0
async def test_record_micro_fault(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1", content="content", object_type="file_context"
)
await store.record_fault(obj.id, is_micro=True)
assert obj.fault_count == 0
assert obj.micro_fault_count == 1
async def test_record_fault_nonexistent(self):
store = ObjectStore(InMemoryBackend())
# Should not raise
await store.record_fault("nonexistent")
async def test_mixed_faults(self):
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
obj = await store.store_object(
session_id="s1", content="content", object_type="file_context"
)
await store.record_fault(obj.id)
await store.record_fault(obj.id, is_micro=True)
await store.record_fault(obj.id)
await store.record_fault(obj.id, is_micro=True)
await store.record_fault(obj.id, is_micro=True)
assert obj.fault_count == 2
assert obj.micro_fault_count == 3
# ── Integration: full lifecycle ──────────────────────────────
class TestIntegration:
async def test_store_search_degrade_cycle(self):
"""Full lifecycle: store → search → degrade → search again."""
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
# Store several objects
obj1 = await store.store_object(
session_id="s1",
content="authentication middleware handles JWT validation",
object_type="file_context",
source_tool="Read",
source_key="src/auth/middleware.ts",
)
obj2 = await store.store_object(
session_id="s1",
content="database migration adds users table with email column",
object_type="file_context",
source_tool="Read",
source_key="migrations/001_users.sql",
)
# Search finds relevant object
results = await store.semantic_search("s1", "JWT validation")
assert len(results) >= 1
# Degrade first object
await store.update_fidelity(
obj1.id,
1,
summary="Auth middleware: validates JWT tokens",
losses=["exact error handling code"],
)
degraded = await store.get(obj1.id)
assert degraded is not None
assert degraded.current_fidelity == 1
assert degraded.summary_detailed is not None
# Search still works after degradation
results2 = await store.semantic_search("s1", "JWT validation")
assert len(results2) >= 1
async def test_dedup_workflow(self):
"""Dedup: check for existing → store if new."""
store = ObjectStore(InMemoryBackend(), DummyEmbedder())
# First read
dup = await store.find_duplicate("s1", "src/auth.py")
assert dup is None
obj = await store.store_object(
session_id="s1",
content="original content",
object_type="file_context",
source_key="src/auth.py",
)
# Second read — duplicate found
dup = await store.find_duplicate("s1", "src/auth.py")
assert dup is not None
assert dup.id == obj.id
async def test_multi_session_lifecycle(self):
"""Multiple sessions operate independently."""
backend = InMemoryBackend()
store = ObjectStore(backend, DummyEmbedder())
await store.store_object(
session_id="s1", content="session 1 auth code", object_type="file_context"
)
await store.store_object(
session_id="s2", content="session 2 db code", object_type="file_context"
)
s1_objs = await store.get_session_objects("s1")
s2_objs = await store.get_session_objects("s2")
assert len(s1_objs) == 1
assert len(s2_objs) == 1
# Delete session 1
count = await backend.delete_session("s1")
assert count == 1
# Session 2 unaffected
s2_objs = await store.get_session_objects("s2")
assert len(s2_objs) == 1
# Session 1 empty
s1_objs = await store.get_session_objects("s1")
assert len(s1_objs) == 0

View file

@ -0,0 +1,882 @@
"""Tests for the PostgreSQL + pgvector backend.
Unit tests use mocked asyncpg connections to verify SQL generation and
data conversion without requiring a database. Integration tests require
a running PostgreSQL instance (Docker on port 5433) and are skipped
automatically when unavailable.
"""
from __future__ import annotations
import asyncio
import json
import os
import socket
import uuid
from datetime import datetime, timezone
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import numpy as np
import pytest
from mnemosyne.object_store import StoredObject, _estimate_tokens
from mnemosyne.pgvector_backend import (
PgVectorBackend,
_parse_jsonb,
_parse_timestamp,
_row_to_stored_object,
)
class _MockPool:
"""A mock asyncpg pool that properly supports async context manager on acquire()."""
def __init__(self, conn: AsyncMock):
self._conn = conn
def acquire(self):
return _MockAcquire(self._conn)
class _MockAcquire:
"""Async context manager returned by pool.acquire()."""
def __init__(self, conn: AsyncMock):
self._conn = conn
async def __aenter__(self):
return self._conn
async def __aexit__(self, *args):
return False
def _make_mock_pool(conn: AsyncMock) -> _MockPool:
"""Create a mock pool with proper async context manager support."""
return _MockPool(conn)
# ── Helpers ──────────────────────────────────────────────────
def _make_stored_object(
session_id: str = "test-session",
content: str = "Test content for pgvector backend",
*,
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,
) -> StoredObject:
"""Create a StoredObject with sensible defaults for testing."""
oid = object_id or uuid.uuid4().hex
return StoredObject(
id=oid,
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}: test object",
tokens_l0=_estimate_tokens(content),
tokens_l3=_estimate_tokens(stub or f"{object_type}: test object"),
embedding=embedding or [0.1] * 384,
created_at="2025-01-01T00:00:00+00:00",
last_accessed="2025-01-01T00:00:00+00:00",
)
def _make_mock_row(
obj: StoredObject | None = None,
*,
session_external_id: str = "test-session",
similarity: float | None = None,
) -> MagicMock:
"""Create a mock asyncpg.Record from a StoredObject."""
if obj is None:
obj = _make_stored_object()
row = MagicMock()
row_data: dict[str, Any] = {
"id": uuid.UUID(obj.id) if len(obj.id) == 32 else uuid.uuid4(),
"session_id": uuid.uuid4(),
"session_external_id": session_external_id,
"object_type": obj.object_type,
"source_tool": obj.source_tool,
"source_key": obj.source_key,
"content_full": obj.content_full,
"summary_detailed": obj.summary_detailed,
"summary_compact": obj.summary_compact,
"stub": obj.stub,
"losses_l1": obj.losses_l1,
"losses_l2": obj.losses_l2,
"can_answer_l1": obj.can_answer_l1,
"can_answer_l2": obj.can_answer_l2,
"fault_when": obj.fault_when,
"key_entities": obj.key_entities,
"tags": obj.tags,
"current_fidelity": obj.current_fidelity,
"pinned": obj.pinned,
"tokens_l0": obj.tokens_l0,
"tokens_l1": obj.tokens_l1,
"tokens_l2": obj.tokens_l2,
"tokens_l3": obj.tokens_l3,
"source_turn_start": obj.source_turn_start,
"source_turn_end": obj.source_turn_end,
"embedding": np.array(obj.embedding, dtype=np.float32) if obj.embedding else None,
"created_at": datetime(2025, 1, 1, tzinfo=timezone.utc),
"last_accessed": datetime(2025, 1, 1, tzinfo=timezone.utc),
"access_count": obj.access_count,
"fault_count": obj.fault_count,
"micro_fault_count": obj.micro_fault_count,
}
if similarity is not None:
row_data["similarity"] = similarity
row.__getitem__ = lambda self, key: row_data[key]
return row
def _pg_available() -> bool:
"""Check if PostgreSQL is reachable on port 5433."""
try:
with socket.create_connection(("localhost", 5433), timeout=1):
return True
except (OSError, ConnectionRefusedError):
return False
# ── Unit Tests: Data Conversion ──────────────────────────────
class TestParseJsonb:
def test_none_returns_empty_list(self):
assert _parse_jsonb(None) == []
def test_list_passthrough(self):
assert _parse_jsonb(["a", "b", "c"]) == ["a", "b", "c"]
def test_list_converts_to_strings(self):
assert _parse_jsonb([1, 2, 3]) == ["1", "2", "3"]
def test_json_string(self):
assert _parse_jsonb('["x", "y"]') == ["x", "y"]
def test_invalid_json_string(self):
assert _parse_jsonb("not json") == []
def test_empty_list(self):
assert _parse_jsonb([]) == []
class TestParseTimestamp:
def test_iso_format(self):
dt = _parse_timestamp("2025-01-01T00:00:00+00:00")
assert dt.year == 2025
assert dt.tzinfo is not None
def test_empty_string_returns_now(self):
dt = _parse_timestamp("")
assert dt.tzinfo is not None
# Should be close to now
diff = abs((datetime.now(timezone.utc) - dt).total_seconds())
assert diff < 5
def test_naive_timestamp_gets_utc(self):
dt = _parse_timestamp("2025-06-15T12:00:00")
assert dt.tzinfo == timezone.utc
def test_invalid_returns_now(self):
dt = _parse_timestamp("not-a-date")
assert dt.tzinfo is not None
class TestRowToStoredObject:
def test_basic_conversion(self):
obj = _make_stored_object()
row = _make_mock_row(obj)
result = _row_to_stored_object(row)
assert result.object_type == "file_context"
assert result.content_full == obj.content_full
assert result.stub == obj.stub
assert result.session_id == "test-session"
assert result.current_fidelity == 0
assert result.pinned is False
def test_embedding_conversion(self):
obj = _make_stored_object(embedding=[0.5] * 384)
row = _make_mock_row(obj)
result = _row_to_stored_object(row)
assert len(result.embedding) == 384
assert abs(result.embedding[0] - 0.5) < 1e-6
def test_jsonb_fields_parsed(self):
obj = _make_stored_object()
obj.losses_l1 = ["detail_a", "detail_b"]
obj.key_entities = ["src/main.py", "Config"]
row = _make_mock_row(obj)
result = _row_to_stored_object(row)
assert result.losses_l1 == ["detail_a", "detail_b"]
assert result.key_entities == ["src/main.py", "Config"]
def test_tags_conversion(self):
obj = _make_stored_object()
obj.tags = ["auth", "middleware"]
row = _make_mock_row(obj)
result = _row_to_stored_object(row)
assert result.tags == ["auth", "middleware"]
def test_null_embedding(self):
obj = _make_stored_object(embedding=[])
row_data: dict[str, Any] = {
"id": uuid.uuid4(),
"session_id": uuid.uuid4(),
"session_external_id": "test-session",
"object_type": obj.object_type,
"source_tool": obj.source_tool,
"source_key": obj.source_key,
"content_full": obj.content_full,
"summary_detailed": obj.summary_detailed,
"summary_compact": obj.summary_compact,
"stub": obj.stub,
"losses_l1": [],
"losses_l2": [],
"can_answer_l1": [],
"can_answer_l2": [],
"fault_when": [],
"key_entities": [],
"tags": [],
"current_fidelity": 0,
"pinned": False,
"tokens_l0": obj.tokens_l0,
"tokens_l1": None,
"tokens_l2": None,
"tokens_l3": obj.tokens_l3,
"source_turn_start": None,
"source_turn_end": None,
"embedding": None,
"created_at": datetime(2025, 1, 1, tzinfo=timezone.utc),
"last_accessed": datetime(2025, 1, 1, tzinfo=timezone.utc),
"access_count": 0,
"fault_count": 0,
"micro_fault_count": 0,
}
row = MagicMock()
row.__getitem__ = lambda self, key: row_data[key]
result = _row_to_stored_object(row)
assert result.embedding == []
# ── Unit Tests: Backend Methods (Mocked DB) ─────────────────
class TestPgVectorBackendInit:
def test_default_config(self):
backend = PgVectorBackend()
assert backend._host == "localhost"
assert backend._port == 5433
assert backend._database == "mnemosyne"
assert backend._pool is None
def test_custom_config(self):
backend = PgVectorBackend(
host="db.example.com",
port=5432,
database="mydb",
user="myuser",
password="secret",
min_connections=5,
max_connections=20,
)
assert backend._host == "db.example.com"
assert backend._port == 5432
assert backend._min_connections == 5
assert backend._max_connections == 20
def test_get_pool_raises_when_not_connected(self):
backend = PgVectorBackend()
with pytest.raises(RuntimeError, match="not connected"):
backend._get_pool()
class TestPgVectorBackendStore:
"""Test store() with mocked pool."""
async def test_store_calls_execute(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.execute = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value={"id": uuid.uuid4()})
backend._pool = _make_mock_pool(mock_conn)
obj = _make_stored_object(object_id=uuid.uuid4().hex)
await backend.store(obj)
# Should have called execute for the INSERT (store) and fetchrow for session upsert
assert mock_conn.execute.called or mock_conn.fetchrow.called
async def test_store_creates_session_if_needed(self):
backend = PgVectorBackend()
session_uuid = uuid.uuid4()
mock_conn = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value={"id": session_uuid})
mock_conn.execute = AsyncMock()
backend._pool = _make_mock_pool(mock_conn)
obj = _make_stored_object(object_id=uuid.uuid4().hex)
await backend.store(obj)
# Session should be cached after creation
assert "test-session" in backend._session_cache
class TestPgVectorBackendGet:
"""Test get() with mocked pool."""
async def test_get_returns_none_when_not_found(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value=None)
backend._pool = _make_mock_pool(mock_conn)
result = await backend.get(uuid.uuid4().hex)
assert result is None
async def test_get_returns_stored_object(self):
backend = PgVectorBackend()
obj = _make_stored_object(object_id=uuid.uuid4().hex)
mock_row = _make_mock_row(obj)
mock_conn = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value=mock_row)
backend._pool = _make_mock_pool(mock_conn)
result = await backend.get(obj.id)
assert result is not None
assert result.content_full == obj.content_full
class TestPgVectorBackendGetBySession:
"""Test get_by_session() with mocked pool."""
async def test_returns_empty_for_unknown_session(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value=None)
backend._pool = _make_mock_pool(mock_conn)
result = await backend.get_by_session("nonexistent")
assert result == []
class TestPgVectorBackendUpdateFidelity:
"""Test update_fidelity() with mocked pool."""
async def test_update_fidelity_basic(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.execute = AsyncMock()
backend._pool = _make_mock_pool(mock_conn)
await backend.update_fidelity(uuid.uuid4().hex, 1)
assert mock_conn.execute.called
async def test_update_fidelity_with_summary_l1(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.execute = AsyncMock()
backend._pool = _make_mock_pool(mock_conn)
await backend.update_fidelity(
uuid.uuid4().hex, 1, summary="A summary", losses=["lost detail"]
)
# Verify the SQL includes summary_detailed and losses_l1
call_args = mock_conn.execute.call_args
query = call_args[0][0]
assert "summary_detailed" in query
assert "losses_l1" in query
async def test_update_fidelity_with_summary_l2(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.execute = AsyncMock()
backend._pool = _make_mock_pool(mock_conn)
await backend.update_fidelity(uuid.uuid4().hex, 2, summary="Compact", losses=["more lost"])
call_args = mock_conn.execute.call_args
query = call_args[0][0]
assert "summary_compact" in query
assert "losses_l2" in query
class TestPgVectorBackendSearchByEmbedding:
"""Test search_by_embedding() with mocked pool."""
async def test_returns_empty_for_unknown_session(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value=None)
backend._pool = _make_mock_pool(mock_conn)
result = await backend.search_by_embedding("nonexistent", [0.1] * 384)
assert result == []
class TestPgVectorBackendSearchByText:
"""Test search_by_text() with mocked pool."""
async def test_returns_empty_for_unknown_session(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value=None)
backend._pool = _make_mock_pool(mock_conn)
result = await backend.search_by_text("nonexistent", "test query")
assert result == []
class TestPgVectorBackendDeleteSession:
"""Test delete_session() with mocked pool."""
async def test_returns_zero_for_unknown_session(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value=None)
backend._pool = _make_mock_pool(mock_conn)
result = await backend.delete_session("nonexistent")
assert result == 0
class TestPgVectorBackendGetBySourceKey:
"""Test get_by_source_key() with mocked pool."""
async def test_returns_none_for_unknown_session(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value=None)
backend._pool = _make_mock_pool(mock_conn)
result = await backend.get_by_source_key("nonexistent", "src/main.py")
assert result is None
class TestPgVectorBackendHealthCheck:
"""Test health_check()."""
async def test_returns_false_when_not_connected(self):
backend = PgVectorBackend()
assert await backend.health_check() is False
async def test_returns_true_with_healthy_pool(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.fetchval = AsyncMock(return_value=1)
backend._pool = _make_mock_pool(mock_conn)
assert await backend.health_check() is True
async def test_returns_false_on_error(self):
backend = PgVectorBackend()
mock_conn = AsyncMock()
mock_conn.fetchval = AsyncMock(side_effect=Exception("connection lost"))
backend._pool = _make_mock_pool(mock_conn)
assert await backend.health_check() is False
class TestPgVectorBackendSessionCache:
"""Test session ID caching behavior."""
async def test_session_cache_populated_on_ensure(self):
backend = PgVectorBackend()
session_uuid = uuid.uuid4()
mock_conn = AsyncMock()
mock_conn.fetchrow = AsyncMock(return_value={"id": session_uuid})
backend._pool = _make_mock_pool(mock_conn)
result = await backend.ensure_session("my-session", "claude-opus-4")
assert result == session_uuid
assert backend._session_cache["my-session"] == session_uuid
async def test_session_cache_avoids_db_lookup(self):
backend = PgVectorBackend()
session_uuid = uuid.uuid4()
backend._session_cache["cached-session"] = session_uuid
# Pool shouldn't be needed since cache is populated
backend._pool = _make_mock_pool(AsyncMock())
result = await backend._resolve_session_id("cached-session")
assert result == session_uuid
async def test_close_clears_cache(self):
backend = PgVectorBackend()
backend._session_cache["test"] = uuid.uuid4()
mock_pool = AsyncMock()
mock_pool.close = AsyncMock()
backend._pool = mock_pool
await backend.close()
assert len(backend._session_cache) == 0
assert backend._pool is None
# ── Integration Tests (require Docker PostgreSQL) ────────────
_skip_no_pg = pytest.mark.skipif(
not _pg_available(),
reason="PostgreSQL not available on localhost:5433 (run: docker compose up -d)",
)
@_skip_no_pg
class TestPgVectorBackendIntegration:
"""Integration tests against a real PostgreSQL instance.
These tests require Docker PostgreSQL running on port 5433 with
the schema from sql/init.sql applied. Start with:
cd ~/Projects/contextmanager && docker compose up -d
"""
@pytest.fixture
async def backend(self):
"""Create a connected backend and clean up after test."""
b = PgVectorBackend(
host="localhost",
port=5433,
database="mnemosyne",
user="mnemosyne",
password="mnemosyne_dev",
min_connections=1,
max_connections=3,
)
await b.connect()
# Create a unique session for this test
test_session = f"integration-test-{uuid.uuid4().hex[:8]}"
yield b, test_session
# Cleanup: delete test session data
try:
session_uuid = await b._resolve_session_id(test_session)
if session_uuid is not None:
pool = b._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"DELETE FROM semantic_objects WHERE session_id = $1",
session_uuid,
)
await conn.execute(
"DELETE FROM sessions WHERE id = $1",
session_uuid,
)
except Exception:
pass
await b.close()
async def test_health_check(self, backend):
b, _ = backend
assert await b.health_check() is True
async def test_ensure_session(self, backend):
b, test_session = backend
session_uuid = await b.ensure_session(test_session, "claude-opus-4")
assert isinstance(session_uuid, uuid.UUID)
# Second call should return same UUID
session_uuid2 = await b.ensure_session(test_session, "claude-opus-4")
assert session_uuid == session_uuid2
async def test_store_and_get(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
obj_id = uuid.uuid4().hex
obj = _make_stored_object(
session_id=test_session,
content="Integration test content",
object_id=obj_id,
)
await b.store(obj)
retrieved = await b.get(obj_id)
assert retrieved is not None
assert retrieved.content_full == "Integration test content"
assert retrieved.session_id == test_session
assert retrieved.object_type == "file_context"
async def test_store_upsert(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
obj_id = uuid.uuid4().hex
obj = _make_stored_object(
session_id=test_session,
content="Original content",
object_id=obj_id,
)
await b.store(obj)
# Update the same object
obj.content_full = "Updated content"
await b.store(obj)
retrieved = await b.get(obj_id)
assert retrieved is not None
assert retrieved.content_full == "Updated content"
async def test_get_by_session(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
# Store 3 objects
for i in range(3):
obj = _make_stored_object(
session_id=test_session,
content=f"Content {i}",
object_id=uuid.uuid4().hex,
)
await b.store(obj)
results = await b.get_by_session(test_session)
assert len(results) == 3
async def test_get_by_session_fidelity_filter(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
# Store object at fidelity 0
obj0 = _make_stored_object(
session_id=test_session,
content="Fidelity 0",
object_id=uuid.uuid4().hex,
)
await b.store(obj0)
# Store object at fidelity 4 (evicted)
obj4 = _make_stored_object(
session_id=test_session,
content="Fidelity 4",
object_id=uuid.uuid4().hex,
)
obj4.current_fidelity = 4
await b.store(obj4)
# Default: include all
all_results = await b.get_by_session(test_session, fidelity_max=4)
assert len(all_results) == 2
# Exclude evicted
active_results = await b.get_by_session(test_session, fidelity_max=3)
assert len(active_results) == 1
assert active_results[0].content_full == "Fidelity 0"
async def test_update_fidelity(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
obj_id = uuid.uuid4().hex
obj = _make_stored_object(
session_id=test_session,
content="Will be degraded",
object_id=obj_id,
)
await b.store(obj)
await b.update_fidelity(
obj_id,
1,
summary="Detailed summary",
losses=["lost some detail"],
)
retrieved = await b.get(obj_id)
assert retrieved is not None
assert retrieved.current_fidelity == 1
assert retrieved.summary_detailed == "Detailed summary"
async def test_search_by_embedding(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
# Store objects with different embeddings
rng = np.random.default_rng(42)
for i in range(5):
vec = rng.standard_normal(384).astype(np.float32)
vec = vec / np.linalg.norm(vec)
obj = _make_stored_object(
session_id=test_session,
content=f"Embedding test {i}",
object_id=uuid.uuid4().hex,
embedding=vec.tolist(),
)
await b.store(obj)
# Search with the first object's embedding
query_vec = rng.standard_normal(384).astype(np.float32)
query_vec = query_vec / np.linalg.norm(query_vec)
results = await b.search_by_embedding(test_session, query_vec.tolist(), limit=3)
assert len(results) <= 3
# Results should be (StoredObject, float) tuples
for obj, score in results:
assert isinstance(obj, StoredObject)
assert isinstance(score, float)
assert -1.0 <= score <= 1.0
# Scores should be in descending order
scores = [s for _, s in results]
assert scores == sorted(scores, reverse=True)
async def test_search_by_text(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
obj = _make_stored_object(
session_id=test_session,
content="The authentication middleware validates JWT tokens and checks expiration dates",
object_id=uuid.uuid4().hex,
)
await b.store(obj)
obj2 = _make_stored_object(
session_id=test_session,
content="Database connection pooling configuration for PostgreSQL",
object_id=uuid.uuid4().hex,
)
await b.store(obj2)
# Search for auth-related content
results = await b.search_by_text(test_session, "authentication JWT tokens")
assert len(results) >= 1
assert any("authentication" in r.content_full for r in results)
async def test_delete_session(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
# Store some objects
for i in range(3):
obj = _make_stored_object(
session_id=test_session,
content=f"Delete test {i}",
object_id=uuid.uuid4().hex,
)
await b.store(obj)
count = await b.delete_session(test_session)
assert count == 3
# Verify they're gone
results = await b.get_by_session(test_session)
assert len(results) == 0
async def test_get_by_source_key(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
obj = _make_stored_object(
session_id=test_session,
content="File content",
source_key="src/main.py",
object_id=uuid.uuid4().hex,
)
await b.store(obj)
result = await b.get_by_source_key(test_session, "src/main.py")
assert result is not None
assert result.source_key == "src/main.py"
# Non-existent key
result2 = await b.get_by_source_key(test_session, "nonexistent.py")
assert result2 is None
async def test_get_by_source_key_returns_most_recent(self, backend):
b, test_session = backend
await b.ensure_session(test_session)
# Store two objects with same source_key
obj1 = _make_stored_object(
session_id=test_session,
content="Old version",
source_key="src/config.py",
object_id=uuid.uuid4().hex,
)
obj1.created_at = "2025-01-01T00:00:00+00:00"
await b.store(obj1)
obj2 = _make_stored_object(
session_id=test_session,
content="New version",
source_key="src/config.py",
object_id=uuid.uuid4().hex,
)
obj2.created_at = "2025-06-01T00:00:00+00:00"
await b.store(obj2)
result = await b.get_by_source_key(test_session, "src/config.py")
assert result is not None
assert result.content_full == "New version"
async def test_session_isolation(self, backend):
b, test_session = backend
other_session = f"other-{uuid.uuid4().hex[:8]}"
await b.ensure_session(test_session)
await b.ensure_session(other_session)
# Store in test_session
obj = _make_stored_object(
session_id=test_session,
content="Session A content",
object_id=uuid.uuid4().hex,
)
await b.store(obj)
# Store in other_session
obj2 = _make_stored_object(
session_id=other_session,
content="Session B content",
object_id=uuid.uuid4().hex,
)
await b.store(obj2)
# Each session should only see its own objects
results_a = await b.get_by_session(test_session)
results_b = await b.get_by_session(other_session)
assert len(results_a) == 1
assert results_a[0].content_full == "Session A content"
assert len(results_b) == 1
assert results_b[0].content_full == "Session B content"
# Cleanup other session
try:
session_uuid = await b._resolve_session_id(other_session)
if session_uuid:
pool = b._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"DELETE FROM semantic_objects WHERE session_id = $1",
session_uuid,
)
await conn.execute(
"DELETE FROM sessions WHERE id = $1",
session_uuid,
)
except Exception:
pass

980
tests/test_segmenter.py Normal file
View file

@ -0,0 +1,980 @@
"""Tests for the semantic object segmenter.
Covers: tool result segmentation, text classification, user message handling,
entity extraction, stub/tag generation, merging, incremental segmentation,
realistic multi-turn payloads, and edge cases.
"""
from __future__ import annotations
import pytest
from mnemosyne.segmenter import (
VALID_OBJECT_TYPES,
Segmenter,
SegmentedObject,
_estimate_tokens,
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def seg() -> Segmenter:
"""Default segmenter with standard thresholds."""
return Segmenter()
@pytest.fixture
def small_seg() -> Segmenter:
"""Segmenter with low min_object_tokens for testing merging."""
return Segmenter(min_object_tokens=10, max_object_tokens=500)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _msg(role: str, content: str | list) -> dict:
"""Shorthand for creating a message dict."""
return {"role": role, "content": content}
def _tool_use_block(tool_id: str, name: str, input_: dict | None = None) -> dict:
return {
"type": "tool_use",
"id": tool_id,
"name": name,
"input": input_ or {},
}
def _tool_result_block(tool_use_id: str, content: str) -> dict:
return {
"type": "tool_result",
"tool_use_id": tool_use_id,
"content": content,
}
def _text_block(text: str) -> dict:
return {"type": "text", "text": text}
def _long_text(base: str = "x", tokens: int = 200) -> str:
"""Generate text of approximately `tokens` estimated tokens."""
return base * (tokens * 4)
# ===========================================================================
# 1. Tool result segmentation
# ===========================================================================
class TestToolResultSegmentation:
"""Test that tool results are classified by tool name."""
def test_read_tool_produces_file_context(self, seg: Segmenter):
messages = [
_msg(
"assistant",
[
_text_block("Let me read the file."),
_tool_use_block("t1", "Read", {"file_path": "src/main.py"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("def main():\n pass\n")),
],
),
]
objects = seg.segment_messages(messages)
file_objs = [o for o in objects if o.object_type == "file_context"]
assert len(file_objs) >= 1
assert file_objs[0].source_tool == "Read"
assert file_objs[0].source_key == "src/main.py"
def test_bash_tool_produces_tool_result(self, seg: Segmenter):
messages = [
_msg(
"assistant",
[
_text_block("Running tests."),
_tool_use_block("t1", "Bash", {"command": "pytest"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("All 42 tests passed.\n")),
],
),
]
objects = seg.segment_messages(messages)
tool_objs = [o for o in objects if o.object_type == "tool_result"]
assert len(tool_objs) >= 1
assert tool_objs[0].source_tool == "Bash"
def test_grep_tool_produces_tool_result(self, seg: Segmenter):
messages = [
_msg(
"assistant",
[
_text_block("Searching for usage."),
_tool_use_block("t1", "Grep", {"pattern": "handleAuth"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("src/auth.ts:15: handleAuth()\n")),
],
),
]
objects = seg.segment_messages(messages)
tool_objs = [o for o in objects if o.object_type == "tool_result"]
assert len(tool_objs) >= 1
assert tool_objs[0].source_tool == "Grep"
def test_write_tool_produces_tool_result(self, seg: Segmenter):
messages = [
_msg(
"assistant",
[
_text_block("Writing the file."),
_tool_use_block("t1", "Write", {"file_path": "out.py"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("File written successfully.\n")),
],
),
]
objects = seg.segment_messages(messages)
tool_objs = [o for o in objects if o.object_type == "tool_result"]
assert len(tool_objs) >= 1
def test_unknown_tool_produces_tool_result(self, seg: Segmenter):
messages = [
_msg(
"assistant",
[
_text_block("Using custom tool."),
_tool_use_block("t1", "CustomTool", {}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("Custom output here.\n")),
],
),
]
objects = seg.segment_messages(messages)
tool_objs = [o for o in objects if o.object_type == "tool_result"]
assert len(tool_objs) >= 1
assert tool_objs[0].source_tool == "CustomTool"
def test_web_fetch_produces_external_reference(self, seg: Segmenter):
messages = [
_msg(
"assistant",
[
_text_block("Fetching docs."),
_tool_use_block("t1", "WebFetch", {"url": "https://docs.example.com"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("# API Documentation\n")),
],
),
]
objects = seg.segment_messages(messages)
ext_objs = [o for o in objects if o.object_type == "external_reference"]
assert len(ext_objs) >= 1
assert ext_objs[0].source_key == "https://docs.example.com"
def test_read_with_filePath_key(self, seg: Segmenter):
"""Test that filePath (camelCase) is also recognized."""
messages = [
_msg(
"assistant",
[
_tool_use_block("t1", "Read", {"filePath": "/etc/config.toml"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("key = 'value'\n")),
],
),
]
objects = seg.segment_messages(messages)
file_objs = [o for o in objects if o.object_type == "file_context"]
assert len(file_objs) >= 1
assert file_objs[0].source_key == "/etc/config.toml"
# ===========================================================================
# 2. Text classification
# ===========================================================================
class TestTextClassification:
"""Test assistant text classification into object types."""
def test_error_text_classified(self, seg: Segmenter):
error_text = _long_text(
"Traceback (most recent call last):\n"
' File "src/main.py", line 42, in run\n'
" raise ValueError('bad input')\n"
"ValueError: bad input\n"
)
messages = [_msg("assistant", error_text)]
objects = seg.segment_messages(messages)
assert any(o.object_type == "error_context" for o in objects)
def test_plan_text_classified(self, seg: Segmenter):
plan_text = _long_text(
"Here's the implementation plan:\n"
"1. Create the database schema\n"
"2. Implement the API endpoints\n"
"3. Write integration tests\n"
"4. Deploy to staging\n"
)
messages = [_msg("assistant", plan_text)]
objects = seg.segment_messages(messages)
assert any(o.object_type == "plan" for o in objects)
def test_decision_text_classified(self, seg: Segmenter):
decision_text = _long_text(
"I decided to use JWT tokens for authentication because "
"they're stateless and work well with our microservice architecture. "
"I chose JWT over session cookies because we need cross-domain support.\n"
)
messages = [_msg("assistant", decision_text)]
objects = seg.segment_messages(messages)
assert any(o.object_type == "design_decision" for o in objects)
def test_debug_text_classified(self, seg: Segmenter):
debug_text = _long_text(
"I'm investigating the race condition. After debugging, "
"I found the root cause: the mutex wasn't being held during "
"the token refresh. Fixed by adding a lock around the critical section.\n"
)
messages = [_msg("assistant", debug_text)]
objects = seg.segment_messages(messages)
assert any(o.object_type == "debugging_session" for o in objects)
def test_default_conversation_phase(self, seg: Segmenter):
text = _long_text(
"Sure, I'll help you with that. Let me take a look at the code "
"and see what we can improve here.\n"
)
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
assert any(o.object_type == "conversation_phase" for o in objects)
def test_error_needs_multiple_signals(self, seg: Segmenter):
"""A single 'error' word shouldn't trigger error_context."""
text = _long_text("There might be an error somewhere in the logic.\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
# Should NOT be error_context with just one weak signal
assert all(o.object_type != "error_context" for o in objects)
# ===========================================================================
# 3. User message handling
# ===========================================================================
class TestUserMessageHandling:
"""Test user message boundary and merge behavior."""
def test_short_user_message_merges(self, seg: Segmenter):
"""Short messages like 'ok' should merge with next segment."""
messages = [
_msg("user", "ok"),
_msg("assistant", _long_text("I'll proceed with the implementation.\n")),
]
objects = seg.segment_messages(messages)
# Should produce a single merged object, not two
assert len(objects) <= 2
def test_normal_user_message_creates_boundary(self, seg: Segmenter):
long_user = _long_text(
"Can you refactor the authentication module to use OAuth2 "
"instead of the current basic auth approach?\n"
)
messages = [
_msg("user", long_user),
_msg("assistant", _long_text("I'll refactor the auth module.\n")),
]
objects = seg.segment_messages(messages)
assert len(objects) >= 1
def test_short_messages_recognized(self, seg: Segmenter):
"""Various short messages should be recognized."""
short_msgs = ["ok", "yes", "continue", "go ahead", "thanks", "lgtm", "done"]
for short in short_msgs:
messages = [
_msg("user", short),
_msg("assistant", _long_text("Continuing...\n")),
]
objects = seg.segment_messages(messages)
# Should merge — at most 1 object
assert len(objects) <= 2, f"'{short}' was not merged"
def test_user_message_with_tool_results(self, seg: Segmenter):
"""User messages can contain both text and tool_result blocks."""
messages = [
_msg(
"assistant",
[
_text_block("Let me read the file."),
_tool_use_block("t1", "Read", {"file_path": "src/app.py"}),
],
),
_msg(
"user",
[
_tool_result_block(
"t1", _long_text("import flask\napp = flask.Flask(__name__)\n")
),
],
),
]
objects = seg.segment_messages(messages)
assert any(o.object_type == "file_context" for o in objects)
# ===========================================================================
# 4. Entity extraction
# ===========================================================================
class TestEntityExtraction:
"""Test extraction of file paths, function names, and packages."""
def test_file_paths_extracted(self, seg: Segmenter):
text = _long_text(
"I read src/auth/middleware.ts and tests/auth.test.ts. "
"The main logic is in src/core/handler.py.\n"
)
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
entities = objects[0].key_entities
assert any("middleware.ts" in e for e in entities)
assert any("handler.py" in e for e in entities)
def test_function_names_extracted(self, seg: Segmenter):
text = _long_text("def handleAuth(request):\n return authenticate(request.token)\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
entities = objects[0].key_entities
assert any("handleAuth" in e for e in entities)
def test_import_names_extracted(self, seg: Segmenter):
text = _long_text("import flask\nfrom sqlalchemy import Column\nimport os\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
entities = objects[0].key_entities
assert any("flask" in e for e in entities)
assert any("sqlalchemy" in e for e in entities)
def test_js_require_extracted(self, seg: Segmenter):
text = _long_text(
"const express = require('express');\nconst jwt = require('jsonwebtoken');\n"
)
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
entities = objects[0].key_entities
assert any("express" in e for e in entities)
assert any("jsonwebtoken" in e for e in entities)
def test_entity_cap_at_20(self, seg: Segmenter):
"""Entities should be capped at 20."""
# Generate text with many unique file paths
paths = [f"src/module{i}/file{i}.py" for i in range(30)]
text = _long_text(" ".join(paths) + "\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
assert len(objects[0].key_entities) <= 20
# ===========================================================================
# 5. Stub generation
# ===========================================================================
class TestStubGeneration:
"""Test auto-generated stubs."""
def test_stub_format(self, seg: Segmenter):
text = _long_text("This is the first line of content.\nSecond line here.\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
stub = objects[0].stub
assert stub.startswith("[")
assert stub.endswith("]")
assert ":" in stub
def test_stub_contains_type(self, seg: Segmenter):
text = _long_text("Some conversation content here.\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
assert objects[0].object_type in objects[0].stub
def test_stub_truncates_long_first_line(self, seg: Segmenter):
long_line = "A" * 200 + "\nSecond line."
text = _long_text(long_line)
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
stub = objects[0].stub
# Type prefix + ": " + content, content part should be <= 100 chars
content_part = stub.split(": ", 1)[1].rstrip("]")
assert len(content_part) <= 103 # 100 + "..."
def test_stub_uses_first_line_only(self, seg: Segmenter):
text = _long_text("First line here.\nSecond line should not appear.\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
assert "Second line" not in objects[0].stub
# ===========================================================================
# 6. Tag generation
# ===========================================================================
class TestTagGeneration:
"""Test auto-generated tags."""
def test_tags_include_object_type(self, seg: Segmenter):
text = _long_text("Some content.\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
assert objects[0].object_type in objects[0].tags
def test_tags_include_source_tool(self, seg: Segmenter):
messages = [
_msg(
"assistant",
[
_tool_use_block("t1", "Read", {"file_path": "x.py"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("content of x.py\n")),
],
),
]
objects = seg.segment_messages(messages)
file_objs = [o for o in objects if o.source_tool == "Read"]
assert len(file_objs) >= 1
assert "read" in file_objs[0].tags
def test_tags_include_file_extensions(self, seg: Segmenter):
text = _long_text("Modified src/auth.ts and src/handler.py to fix the issue.\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
tags = objects[0].tags
assert ".ts" in tags
assert ".py" in tags
def test_tags_no_duplicate_extensions(self, seg: Segmenter):
text = _long_text("Read src/a.py and src/b.py and src/c.py.\n")
messages = [_msg("assistant", text)]
objects = seg.segment_messages(messages)
py_count = sum(1 for t in objects[0].tags if t == ".py")
assert py_count == 1
# ===========================================================================
# 7. Minimum object size merging
# ===========================================================================
class TestMinObjectSizeMerging:
"""Test that undersized objects get merged."""
def test_small_segments_merge(self, seg: Segmenter):
"""Very small adjacent segments should merge."""
messages = [
_msg("assistant", "Hi."),
_msg("user", "Hello."),
_msg("assistant", "How can I help?"),
]
objects = seg.segment_messages(messages)
# These are all tiny — should merge into fewer objects
assert len(objects) <= 2
def test_large_segments_stay_separate(self, seg: Segmenter):
"""Segments above min_object_tokens stay separate."""
messages = [
_msg("assistant", _long_text("First large block of content.\n")),
_msg("user", _long_text("Second large block of content.\n")),
]
objects = seg.segment_messages(messages)
assert len(objects) >= 1 # At least one object
def test_incompatible_types_dont_merge(self, small_seg: Segmenter):
"""file_context and plan shouldn't merge even if small."""
messages = [
_msg(
"assistant",
[
_tool_use_block("t1", "Read", {"file_path": "a.py"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", "x = 1"),
],
),
_msg(
"assistant",
[
_tool_use_block("t2", "Read", {"file_path": "b.py"}),
],
),
_msg(
"user",
[
_tool_result_block("t2", "y = 2"),
],
),
]
objects = small_seg.segment_messages(messages)
file_objs = [o for o in objects if o.object_type == "file_context"]
# Even if small, file_context objects from different files should
# remain separate (they have different source_keys)
# But they might merge if compatible — the key thing is they exist
assert len(file_objs) >= 1
# ===========================================================================
# 8. Incremental segmentation
# ===========================================================================
class TestIncrementalSegmentation:
"""Test segment_incremental for extending existing objects."""
def test_incremental_appends_new(self, seg: Segmenter):
existing = [
SegmentedObject(
content=_long_text("Previous conversation.\n"),
object_type="conversation_phase",
turn_start=0,
turn_end=1,
token_estimate=200,
)
]
new_messages = [
_msg(
"assistant",
[
_tool_use_block("t1", "Read", {"file_path": "new.py"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("new file content\n")),
],
),
]
result = seg.segment_incremental(new_messages, existing, start_turn=2)
assert len(result) > len(existing)
assert result[0] is existing[0] # First object unchanged
def test_incremental_merges_compatible(self, seg: Segmenter):
existing = [
SegmentedObject(
content=_long_text("Starting the discussion.\n"),
object_type="conversation_phase",
turn_start=0,
turn_end=0,
token_estimate=200,
)
]
new_messages = [
_msg("assistant", _long_text("Continuing the discussion.\n")),
]
result = seg.segment_incremental(new_messages, existing, start_turn=1)
# Should merge the conversation_phase objects
assert len(result) >= 1
def test_incremental_empty_new(self, seg: Segmenter):
existing = [
SegmentedObject(
content="test",
object_type="conversation_phase",
turn_start=0,
turn_end=0,
token_estimate=1,
)
]
result = seg.segment_incremental([], existing)
assert len(result) == 1
def test_incremental_empty_existing(self, seg: Segmenter):
new_messages = [
_msg("assistant", _long_text("Hello world.\n")),
]
result = seg.segment_incremental(new_messages, [])
assert len(result) >= 1
def test_incremental_no_merge_across_types(self, seg: Segmenter):
"""file_context shouldn't merge with conversation_phase."""
existing = [
SegmentedObject(
content=_long_text("file content here\n"),
object_type="file_context",
source_tool="Read",
turn_start=0,
turn_end=0,
token_estimate=200,
)
]
new_messages = [
_msg("assistant", _long_text("Now let me explain what I found.\n")),
]
result = seg.segment_incremental(new_messages, existing, start_turn=1)
# Should NOT merge file_context with conversation_phase
assert len(result) >= 2
# ===========================================================================
# 9. Realistic multi-turn conversation payloads
# ===========================================================================
class TestRealisticPayloads:
"""Test with realistic multi-turn conversation structures."""
def test_typical_coding_session(self, seg: Segmenter):
"""Simulate: user asks → assistant reads file → assistant explains."""
messages = [
_msg(
"user",
_long_text(
"Can you look at the auth middleware and tell me "
"how it handles token refresh?\n"
),
),
_msg(
"assistant",
[
_text_block("I'll read the auth middleware file."),
_tool_use_block("t1", "Read", {"file_path": "src/auth/middleware.ts"}),
],
),
_msg(
"user",
[
_tool_result_block(
"t1",
_long_text(
"import jwt from 'jsonwebtoken';\n"
"export function handleAuth(req, res, next) {\n"
" const token = req.headers.authorization;\n"
" // ... token validation logic\n"
"}\n"
),
),
],
),
_msg(
"assistant",
_long_text(
"The auth middleware in src/auth/middleware.ts handles "
"token refresh by checking the JWT expiry and issuing "
"a new token if within the refresh window.\n"
),
),
]
objects = seg.segment_messages(messages)
types = {o.object_type for o in objects}
assert "file_context" in types
assert len(objects) >= 2
def test_debugging_flow(self, seg: Segmenter):
"""Simulate: error → investigation → fix."""
messages = [
_msg("user", _long_text("The tests are failing with this error.\n")),
_msg(
"assistant",
_long_text(
"I'm investigating the test failure. Let me look at the "
"root cause. The error seems to be a TypeError in the "
"authentication module.\n"
),
),
_msg(
"assistant",
[
_text_block("Let me run the failing test."),
_tool_use_block("t1", "Bash", {"command": "pytest tests/test_auth.py -v"}),
],
),
_msg(
"user",
[
_tool_result_block(
"t1",
_long_text(
"FAILED tests/test_auth.py::test_refresh - TypeError: "
"'NoneType' object is not subscriptable\n"
"Traceback (most recent call last):\n"
' File "tests/test_auth.py", line 42\n'
" token = response['access_token']\n"
"TypeError: 'NoneType' object is not subscriptable\n"
),
),
],
),
_msg(
"assistant",
_long_text(
"Found the problem! The root cause is that the refresh "
"endpoint returns None when the token is expired. "
"Fixed by adding a null check before accessing the response.\n"
),
),
]
objects = seg.segment_messages(messages)
types = {o.object_type for o in objects}
# Should have debugging and/or error objects
assert types & {"debugging_session", "error_context", "tool_result"}
def test_planning_session(self, seg: Segmenter):
"""Simulate: user asks for plan → assistant creates plan."""
messages = [
_msg(
"user",
_long_text(
"I need to add OAuth2 support. Can you create an implementation plan?\n"
),
),
_msg(
"assistant",
_long_text(
"Here's the implementation plan for OAuth2:\n"
"1. Install the oauth2 library\n"
"2. Create the OAuth2 provider configuration\n"
"3. Implement the authorization flow\n"
"4. Add callback handling\n"
"5. Write integration tests\n"
"Step 1 involves adding the dependency to package.json.\n"
),
),
]
objects = seg.segment_messages(messages)
assert any(o.object_type == "plan" for o in objects)
def test_multi_tool_turn(self, seg: Segmenter):
"""Assistant uses multiple tools in one turn."""
messages = [
_msg(
"assistant",
[
_text_block("Let me check both files."),
_tool_use_block("t1", "Read", {"file_path": "src/a.py"}),
_tool_use_block("t2", "Read", {"file_path": "src/b.py"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("# File A content\nclass A:\n pass\n")),
_tool_result_block("t2", _long_text("# File B content\nclass B:\n pass\n")),
],
),
]
objects = seg.segment_messages(messages)
file_objs = [o for o in objects if o.object_type == "file_context"]
assert len(file_objs) >= 2
def test_long_conversation_produces_multiple_objects(self, seg: Segmenter):
"""A long conversation should produce multiple objects."""
messages = []
for i in range(10):
messages.append(_msg("user", _long_text(f"Question {i} about the codebase.\n")))
messages.append(_msg("assistant", _long_text(f"Answer {i} with details.\n")))
objects = seg.segment_messages(messages)
assert len(objects) >= 3 # Should have multiple objects
# ===========================================================================
# 10. Empty and edge cases
# ===========================================================================
class TestEdgeCases:
"""Test empty inputs, malformed messages, and boundary conditions."""
def test_empty_messages(self, seg: Segmenter):
assert seg.segment_messages([]) == []
def test_empty_content_string(self, seg: Segmenter):
messages = [_msg("assistant", "")]
assert seg.segment_messages(messages) == []
def test_empty_content_list(self, seg: Segmenter):
messages = [_msg("assistant", [])]
assert seg.segment_messages(messages) == []
def test_whitespace_only_content(self, seg: Segmenter):
messages = [_msg("assistant", " \n\t ")]
assert seg.segment_messages(messages) == []
def test_missing_role(self, seg: Segmenter):
"""Messages without role should be handled gracefully."""
messages = [{"content": "no role here"}]
# Should not crash
result = seg.segment_messages(messages)
assert isinstance(result, list)
def test_non_dict_content_blocks(self, seg: Segmenter):
"""Content list with non-dict items should be handled."""
messages = [_msg("assistant", ["not a dict", 42, None])]
result = seg.segment_messages(messages)
assert isinstance(result, list)
def test_tool_result_with_list_content(self, seg: Segmenter):
"""tool_result content can be a list of text blocks."""
messages = [
_msg(
"assistant",
[
_tool_use_block("t1", "Read", {"file_path": "x.py"}),
],
),
_msg(
"user",
[
{
"type": "tool_result",
"tool_use_id": "t1",
"content": [{"type": "text", "text": _long_text("file content\n")}],
},
],
),
]
objects = seg.segment_messages(messages)
file_objs = [o for o in objects if o.object_type == "file_context"]
assert len(file_objs) >= 1
def test_token_estimate_accuracy(self, seg: Segmenter):
"""Token estimate should be approximately len/4."""
text = "a" * 400
assert _estimate_tokens(text) == 100
def test_token_estimate_minimum(self):
"""Token estimate should be at least 1."""
assert _estimate_tokens("") == 1
assert _estimate_tokens("a") == 1
def test_all_object_types_valid(self, seg: Segmenter):
"""All produced object types should be in VALID_OBJECT_TYPES."""
messages = [
_msg("user", _long_text("Question.\n")),
_msg(
"assistant",
_long_text(
'Traceback (most recent call last):\n File "x.py", line 1\nValueError: bad\n'
),
),
_msg("assistant", _long_text("Here's the plan:\n1. Do this\n2. Do that\n")),
_msg(
"assistant",
[
_tool_use_block("t1", "Read", {"file_path": "f.py"}),
],
),
_msg(
"user",
[
_tool_result_block("t1", _long_text("content\n")),
],
),
]
objects = seg.segment_messages(messages)
for obj in objects:
assert obj.object_type in VALID_OBJECT_TYPES, f"Invalid type: {obj.object_type}"
def test_oversized_content_gets_split(self, seg: Segmenter):
"""Content exceeding max_object_tokens should be split."""
small_seg = Segmenter(min_object_tokens=10, max_object_tokens=100)
huge_text = "x" * 2000 # ~500 tokens, well over 100
messages = [_msg("assistant", huge_text)]
objects = small_seg.segment_messages(messages)
assert len(objects) >= 2
for obj in objects:
assert obj.token_estimate <= 200 # Some slack for splitting
def test_segmented_object_defaults(self):
"""SegmentedObject should have sensible defaults."""
obj = SegmentedObject(content="test", object_type="conversation_phase")
assert obj.source_tool is None
assert obj.source_key is None
assert obj.stub == ""
assert obj.turn_start == 0
assert obj.turn_end == 0
assert obj.token_estimate == 0
assert obj.key_entities == []
assert obj.tags == []
def test_tool_result_with_error_content(self, seg: Segmenter):
"""Tool result containing errors should be classified as error_context."""
messages = [
_msg(
"assistant",
[
_tool_use_block("t1", "Bash", {"command": "npm test"}),
],
),
_msg(
"user",
[
_tool_result_block(
"t1",
_long_text(
"FAILED test_auth.py\n"
"TypeError: Cannot read property 'token' of undefined\n"
" at handleAuth (src/auth.js:42:15)\n"
" at processTicksAndRejections (internal/process/task_queues.js:95:5)\n"
),
),
],
),
]
objects = seg.segment_messages(messages)
assert any(o.object_type == "error_context" for o in objects)
def test_system_role_ignored(self, seg: Segmenter):
"""System messages should be ignored (not user or assistant)."""
messages = [
{"role": "system", "content": "You are a helpful assistant."},
_msg("assistant", _long_text("Hello!\n")),
]
objects = seg.segment_messages(messages)
# Should only have objects from the assistant message
assert all("helpful assistant" not in o.content for o in objects)