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:
parent
d26c56c2f0
commit
a13719f754
9 changed files with 5644 additions and 0 deletions
220
src/mnemosyne/context_assembler.py
Normal file
220
src/mnemosyne/context_assembler.py
Normal 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
116
src/mnemosyne/embedder.py
Normal 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
|
||||
547
src/mnemosyne/object_store.py
Normal file
547
src/mnemosyne/object_store.py
Normal 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()
|
||||
568
src/mnemosyne/pgvector_backend.py
Normal file
568
src/mnemosyne/pgvector_backend.py
Normal 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
927
src/mnemosyne/segmenter.py
Normal 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",
|
||||
}
|
||||
)
|
||||
408
tests/test_context_assembler.py
Normal file
408
tests/test_context_assembler.py
Normal 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
996
tests/test_object_store.py
Normal 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
|
||||
882
tests/test_pgvector_backend.py
Normal file
882
tests/test_pgvector_backend.py
Normal 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
980
tests/test_segmenter.py
Normal 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)
|
||||
Loading…
Add table
Add a link
Reference in a new issue