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

contextos/storage/telemetry_repo.py

@@ -0,0 +1,481 @@
+"""SQLite persistence for model invocation telemetry."""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from datetime import datetime
+from typing import Any
+from uuid import UUID
+
+import aiosqlite
+
+from contextos.core.enums import ModelFinishReason, RoutingPolicy, TokenMeasurementSource
+from contextos.core.models import ModelInvocationTelemetry, TelemetrySummary
+
+logger = logging.getLogger(__name__)
+
+_SECRET_PATTERNS = [
+    re.compile(r"Bearer\s+[A-Za-z0-9_\-\.~+/]+=*", re.IGNORECASE),
+    re.compile(r"sk-[A-Za-z0-9_\-]{8,}", re.IGNORECASE),
+    re.compile(r"gh[pousr]-[A-Za-z0-9_]{16,}", re.IGNORECASE),
+    re.compile(r"(?:api[-_]?key|secret|password|token)\s*[:=]\s*['\"]?[A-Za-z0-9_\-\.~+/]+['\"]?", re.IGNORECASE),
+]
+
+_SENSITIVE_KEY_SUBSTRINGS = {
+    "api_key",
+    "apikey",
+    "secret",
+    "password",
+    "token",
+    "auth",
+    "authorization",
+    "raw_prompt",
+    "raw_response",
+    "cookie",
+    "credential",
+}
+
+
+def sanitize_telemetry_metadata(val: Any) -> Any:
+    """Recursively sanitize metadata dicts/lists to strip secrets and raw payloads."""
+    if isinstance(val, dict):
+        cleaned: dict[str, Any] = {}
+        for k, v in val.items():
+            key_lower = str(k).lower()
+            if any(s in key_lower for s in _SENSITIVE_KEY_SUBSTRINGS):
+                cleaned[str(k)] = "[REDACTED]"
+            else:
+                cleaned[str(k)] = sanitize_telemetry_metadata(v)
+        return cleaned
+    elif isinstance(val, list):
+        return [sanitize_telemetry_metadata(item) for item in val]
+    elif isinstance(val, str):
+        sanitized_str = val
+        for pat in _SECRET_PATTERNS:
+            sanitized_str = pat.sub("[REDACTED_SECRET]", sanitized_str)
+        return sanitized_str
+    return val
+
+
+
+class SqliteTelemetryRepository:
+    """Stores and queries model invocation metrics in SQLite."""
+
+    def __init__(self, conn: aiosqlite.Connection) -> None:
+        self._conn = conn
+
+    async def record(self, telemetry: ModelInvocationTelemetry) -> None:
+        """Persist a single model invocation telemetry record."""
+        row_id = str(telemetry.invocation_id)
+        safe_meta = sanitize_telemetry_metadata(telemetry.metadata)
+        metadata_json = json.dumps(safe_meta)
+        timestamp_str = telemetry.timestamp.isoformat()
+
+        sql = """
+        INSERT INTO model_invocations (
+            id,
+            invocation_id,
+            session_id,
+            provider_id,
+            model_id,
+            is_local,
+            timestamp,
+            candidate_context_tokens,
+            retrieved_context_tokens,
+            optimized_context_tokens,
+            compiled_context_tokens,
+            prompt_tokens_before_context,
+            final_input_tokens,
+            provider_input_tokens,
+            provider_output_tokens,
+            provider_total_tokens,
+            token_measurement_source,
+            context_tokens_avoided,
+            reduction_ratio,
+            lexical_candidate_count,
+            dense_candidate_count,
+            hybrid_candidate_count,
+            graph_expanded_count,
+            temporal_filtered_count,
+            selected_memory_count,
+            compiled_fact_count,
+            retrieval_ms,
+            optimization_ms,
+            compilation_ms,
+            routing_ms,
+            token_counting_ms,
+            provider_latency_ms,
+            end_to_end_ms,
+            routing_policy,
+            routing_reason,
+            selected_provider,
+            selected_model,
+            fallback_used,
+            fallback_reason,
+            finish_reason,
+            status,
+            error_code,
+            metadata,
+            context_token_measurement_source,
+            context_tokenizer
+        ) VALUES (
+            ?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
+            ?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
+            ?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
+            ?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
+            ?, ?, ?, ?, ?
+        )
+        """
+        params = (
+            row_id,
+            str(telemetry.invocation_id),
+            telemetry.session_id,
+            telemetry.provider_id,
+            telemetry.model_id,
+            1 if telemetry.is_local else 0,
+            timestamp_str,
+            telemetry.candidate_context_tokens,
+            telemetry.retrieved_context_tokens,
+            telemetry.optimized_context_tokens,
+            telemetry.compiled_context_tokens,
+            telemetry.prompt_tokens_before_context,
+            telemetry.final_input_tokens,
+            telemetry.provider_input_tokens,
+            telemetry.provider_output_tokens,
+            telemetry.provider_total_tokens,
+            telemetry.token_measurement_source.value,
+            telemetry.context_tokens_avoided,
+            telemetry.reduction_ratio,
+            telemetry.lexical_candidate_count,
+            telemetry.dense_candidate_count,
+            telemetry.hybrid_candidate_count,
+            telemetry.graph_expanded_count,
+            telemetry.temporal_filtered_count,
+            telemetry.selected_memory_count,
+            telemetry.compiled_fact_count,
+            telemetry.retrieval_ms,
+            telemetry.optimization_ms,
+            telemetry.compilation_ms,
+            telemetry.routing_ms,
+            telemetry.token_counting_ms,
+            telemetry.provider_latency_ms,
+            telemetry.end_to_end_ms,
+            telemetry.routing_policy.value,
+            telemetry.routing_reason,
+            telemetry.selected_provider,
+            telemetry.selected_model,
+            1 if telemetry.fallback_used else 0,
+            telemetry.fallback_reason,
+            telemetry.finish_reason.value,
+            telemetry.status,
+            telemetry.error_code,
+            metadata_json,
+            telemetry.context_token_measurement_source.value if telemetry.context_token_measurement_source else None,
+            telemetry.context_tokenizer,
+        )
+        await self._conn.execute(sql, params)
+        await self._conn.commit()
+
+    async def get(self, invocation_id: UUID) -> ModelInvocationTelemetry | None:
+        """Retrieve a telemetry record by invocation UUID."""
+        sql = "SELECT * FROM model_invocations WHERE invocation_id = ?"
+        async with self._conn.execute(sql, (str(invocation_id),)) as cursor:
+            row = await cursor.fetchone()
+            if row is None:
+                return None
+            return self._row_to_model(row)
+
+    async def list_recent(
+        self, limit: int = 50, model_id: str | None = None,
+        provider_id: str | None = None, start: datetime | None = None,
+    ) -> list[ModelInvocationTelemetry]:
+        """List recent invocations in descending chronological order."""
+        sql = "SELECT * FROM model_invocations WHERE 1=1"
+        params: list[Any] = []
+        if model_id is not None:
+            sql += " AND model_id = ?"
+            params.append(model_id)
+        if provider_id is not None:
+            sql += " AND provider_id = ?"
+            params.append(provider_id)
+        if start is not None:
+            sql += " AND timestamp >= ?"
+            params.append(start.isoformat())
+        sql += " ORDER BY timestamp DESC LIMIT ?"
+        params.append(max(1, min(limit, 100)))
+        async with self._conn.execute(sql, params) as cursor:
+            rows = await cursor.fetchall()
+            return [self._row_to_model(r) for r in rows]
+
+    async def provider_model_breakdown(
+        self, start: datetime | None = None, provider_id: str | None = None,
+        model_id: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """Group successful context counts by provider, model and tokenizer basis."""
+        conditions = ["1=1"]
+        params: list[Any] = []
+        if start is not None:
+            conditions.append("timestamp >= ?")
+            params.append(start.isoformat())
+        if provider_id is not None:
+            conditions.append("provider_id = ?")
+            params.append(provider_id)
+        if model_id is not None:
+            conditions.append("model_id = ?")
+            params.append(model_id)
+        sql = f"""
+        SELECT provider_id, model_id, is_local,
+               COALESCE(context_token_measurement_source, 'unknown') AS basis,
+               COALESCE(context_tokenizer, 'unknown') AS tokenizer,
+               SUM(CASE WHEN status = 'success' THEN 1 ELSE 0 END) AS successful,
+               SUM(CASE WHEN status != 'success' THEN 1 ELSE 0 END) AS errors,
+               COALESCE(SUM(CASE WHEN status = 'success' THEN candidate_context_tokens END), 0),
+               COALESCE(SUM(CASE WHEN status = 'success' THEN compiled_context_tokens END), 0),
+               COALESCE(SUM(CASE WHEN status = 'success' THEN context_tokens_avoided END), 0),
+               COALESCE(SUM(CASE WHEN status = 'success' THEN final_input_tokens END), 0),
+               COALESCE(SUM(CASE WHEN status = 'success' THEN provider_input_tokens END), 0),
+               COALESCE(SUM(CASE WHEN status = 'success' THEN provider_output_tokens END), 0),
+               COALESCE(SUM(graph_expanded_count), 0),
+               COALESCE(SUM(selected_memory_count), 0),
+               COALESCE(AVG(retrieval_ms), 0),
+               COALESCE(AVG(compilation_ms), 0),
+               COALESCE(AVG(provider_latency_ms), 0)
+        FROM model_invocations WHERE {' AND '.join(conditions)}
+        GROUP BY provider_id, model_id, is_local, basis, tokenizer
+        ORDER BY successful DESC, provider_id, model_id LIMIT 50
+        """
+        async with self._conn.execute(sql, params) as cursor:
+            rows = await cursor.fetchall()
+        return [{
+            "provider": row[0], "model": row[1], "local": bool(row[2]),
+            "context_measurement_source": row[3], "context_tokenizer": row[4],
+            "invocations": row[5], "errors": row[6],
+            "candidate_context_tokens": row[7], "compiled_context_tokens": row[8],
+            "context_tokens_avoided": row[9], "preflight_input_tokens": row[10],
+            "provider_input_tokens": row[11], "provider_output_tokens": row[12],
+            "graph_expanded_count": row[13], "selected_memory_count": row[14],
+            "average_retrieval_ms": row[15], "average_compilation_ms": row[16],
+            "average_provider_ms": row[17],
+            "weighted_reduction_ratio": row[9] / row[7] if row[7] else None,
+        } for row in rows]
+
+    async def count(self) -> int:
+        """Count total recorded invocations."""
+        sql = "SELECT COUNT(*) FROM model_invocations"
+        async with self._conn.execute(sql) as cursor:
+            row = await cursor.fetchone()
+            return row[0] if row else 0
+
+    async def context_measurement_bases(
+        self, model_id: str | None = None, provider_id: str | None = None,
+        start: datetime | None = None,
+    ) -> list[dict[str, str]]:
+        """Distinct provenance bases; unknown legacy rows are preserved as unknown."""
+        sql = """SELECT DISTINCT COALESCE(context_token_measurement_source, 'unknown'),
+                        COALESCE(context_tokenizer, 'unknown') FROM model_invocations WHERE status = 'success'"""
+        params: list[str] = []
+        if model_id is not None:
+            sql += " AND model_id = ?"
+            params.append(model_id)
+        if provider_id is not None:
+            sql += " AND provider_id = ?"
+            params.append(provider_id)
+        if start is not None:
+            sql += " AND timestamp >= ?"
+            params.append(start.isoformat())
+        async with self._conn.execute(sql, params) as cursor:
+            rows = await cursor.fetchall()
+        return [{"source": row[0], "tokenizer": row[1]} for row in rows]
+
+    async def summary(
+        self,
+        start: datetime | None = None,
+        end: datetime | None = None,
+        provider_id: str | None = None,
+        model_id: str | None = None,
+        success_only: bool = False,
+    ) -> TelemetrySummary:
+        """Compute aggregated token usage, avoidance, and latency statistics."""
+        conditions: list[str] = []
+        params: list[Any] = []
+
+        if start is not None:
+            conditions.append("timestamp >= ?")
+            params.append(start.isoformat())
+        if end is not None:
+            conditions.append("timestamp <= ?")
+            params.append(end.isoformat())
+        if provider_id is not None:
+            conditions.append("provider_id = ?")
+            params.append(provider_id)
+        if model_id is not None:
+            conditions.append("model_id = ?")
+            params.append(model_id)
+        if success_only:
+            conditions.append("status = 'success'")
+
+        where_clause = f" WHERE {' AND '.join(conditions)}" if conditions else ""
+
+        agg_sql = f"""
+        SELECT
+            COUNT(*),
+            COALESCE(SUM(provider_input_tokens), 0),
+            COALESCE(SUM(provider_output_tokens), 0),
+            COALESCE(SUM(context_tokens_avoided), 0),
+            COALESCE(AVG(reduction_ratio), 0.0),
+            COALESCE(AVG(provider_latency_ms), 0.0),
+            COALESCE(SUM(CASE WHEN is_local = 1 THEN 1 ELSE 0 END), 0),
+            COALESCE(SUM(CASE WHEN is_local = 0 THEN 1 ELSE 0 END), 0),
+            COALESCE(SUM(candidate_context_tokens), 0)
+        FROM model_invocations{where_clause}
+        """
+
+        async with self._conn.execute(agg_sql, params) as cursor:
+            row = await cursor.fetchone()
+
+        total_invocations = row[0] if row else 0
+        total_input_tokens = row[1] if row else 0
+        total_output_tokens = row[2] if row else 0
+        total_tokens_avoided = row[3] if row else 0
+        average_reduction_ratio = float(row[4]) if row else 0.0
+        average_provider_latency_ms = float(row[5]) if row else 0.0
+        local_invocations = row[6] if row else 0
+        remote_invocations = row[7] if row else 0
+        total_candidate_tokens = row[8] if row else 0
+        weighted_reduction_ratio = (
+            float(total_tokens_avoided / total_candidate_tokens)
+            if total_candidate_tokens > 0
+            else 0.0
+        )
+
+        # By provider breakdown
+        by_provider: dict[str, Any] = {}
+        prov_sql = f"""
+        SELECT
+            provider_id,
+            COUNT(*),
+            COALESCE(SUM(provider_input_tokens), 0),
+            COALESCE(SUM(provider_output_tokens), 0),
+            COALESCE(SUM(context_tokens_avoided), 0),
+            COALESCE(AVG(reduction_ratio), 0.0),
+            COALESCE(AVG(provider_latency_ms), 0.0),
+            COALESCE(SUM(candidate_context_tokens), 0)
+        FROM model_invocations{where_clause}
+        GROUP BY provider_id
+        """
+        async with self._conn.execute(prov_sql, params) as cursor:
+            async for p_row in cursor:
+                p_avoided = p_row[4]
+                p_cand = p_row[7]
+                by_provider[p_row[0]] = {
+                    "invocations": p_row[1],
+                    "input_tokens": p_row[2],
+                    "output_tokens": p_row[3],
+                    "tokens_avoided": p_avoided,
+                    "average_reduction_ratio": float(p_row[5]),
+                    "weighted_reduction_ratio": float(p_avoided / p_cand) if p_cand > 0 else 0.0,
+                    "average_latency_ms": float(p_row[6]),
+                }
+
+        # By model breakdown
+        by_model: dict[str, Any] = {}
+        model_sql = f"""
+        SELECT
+            model_id,
+            COUNT(*),
+            COALESCE(SUM(provider_input_tokens), 0),
+            COALESCE(SUM(provider_output_tokens), 0),
+            COALESCE(SUM(context_tokens_avoided), 0),
+            COALESCE(AVG(reduction_ratio), 0.0),
+            COALESCE(AVG(provider_latency_ms), 0.0),
+            COALESCE(SUM(candidate_context_tokens), 0)
+        FROM model_invocations{where_clause}
+        GROUP BY model_id
+        """
+        async with self._conn.execute(model_sql, params) as cursor:
+            async for m_row in cursor:
+                m_avoided = m_row[4]
+                m_cand = m_row[7]
+                by_model[m_row[0]] = {
+                    "invocations": m_row[1],
+                    "input_tokens": m_row[2],
+                    "output_tokens": m_row[3],
+                    "tokens_avoided": m_avoided,
+                    "average_reduction_ratio": float(m_row[5]),
+                    "weighted_reduction_ratio": float(m_avoided / m_cand) if m_cand > 0 else 0.0,
+                    "average_latency_ms": float(m_row[6]),
+                }
+
+        return TelemetrySummary(
+            total_invocations=total_invocations,
+            total_input_tokens=total_input_tokens,
+            total_output_tokens=total_output_tokens,
+            total_tokens_avoided=total_tokens_avoided,
+            average_reduction_ratio=average_reduction_ratio,
+            weighted_reduction_ratio=weighted_reduction_ratio,
+            average_provider_latency_ms=average_provider_latency_ms,
+            local_invocations=local_invocations,
+            remote_invocations=remote_invocations,
+            by_provider=by_provider,
+            by_model=by_model,
+        )
+
+    def _row_to_model(self, row: Any) -> ModelInvocationTelemetry:
+        metadata_raw = row["metadata"] if "metadata" in row.keys() else "{}"
+        try:
+            metadata = json.loads(metadata_raw) if isinstance(metadata_raw, str) else metadata_raw
+        except Exception:
+            metadata = {}
+
+        return ModelInvocationTelemetry(
+            invocation_id=UUID(row["invocation_id"]),
+            session_id=row["session_id"],
+            provider_id=row["provider_id"],
+            model_id=row["model_id"],
+            is_local=bool(row["is_local"]),
+            timestamp=datetime.fromisoformat(row["timestamp"]),
+            candidate_context_tokens=row["candidate_context_tokens"],
+            retrieved_context_tokens=row["retrieved_context_tokens"],
+            optimized_context_tokens=row["optimized_context_tokens"],
+            compiled_context_tokens=row["compiled_context_tokens"],
+            prompt_tokens_before_context=row["prompt_tokens_before_context"],
+            preflight_input_tokens=row["final_input_tokens"],
+            final_input_tokens=row["final_input_tokens"],
+            provider_input_tokens=row["provider_input_tokens"],
+            provider_output_tokens=row["provider_output_tokens"],
+            provider_total_tokens=row["provider_total_tokens"],
+            token_measurement_source=TokenMeasurementSource(row["token_measurement_source"]),
+            context_token_measurement_source=(
+                TokenMeasurementSource(row["context_token_measurement_source"])
+                if row["context_token_measurement_source"] else None
+            ),
+            context_tokenizer=row["context_tokenizer"],
+            context_tokens_avoided=row["context_tokens_avoided"],
+            reduction_ratio=float(row["reduction_ratio"]),
+            lexical_candidate_count=row["lexical_candidate_count"],
+            dense_candidate_count=row["dense_candidate_count"],
+            hybrid_candidate_count=row["hybrid_candidate_count"],
+            graph_expanded_count=row["graph_expanded_count"],
+            temporal_filtered_count=row["temporal_filtered_count"],
+            selected_memory_count=row["selected_memory_count"],
+            compiled_fact_count=row["compiled_fact_count"],
+            retrieval_ms=float(row["retrieval_ms"]),
+            optimization_ms=float(row["optimization_ms"]),
+            compilation_ms=float(row["compilation_ms"]),
+            routing_ms=float(row["routing_ms"]),
+            token_counting_ms=float(row["token_counting_ms"]),
+            provider_latency_ms=float(row["provider_latency_ms"]),
+            end_to_end_ms=float(row["end_to_end_ms"]),
+            routing_policy=RoutingPolicy(row["routing_policy"]),
+            routing_reason=row["routing_reason"],
+            selected_provider=row["selected_provider"],
+            selected_model=row["selected_model"],
+            fallback_used=bool(row["fallback_used"]),
+            fallback_reason=row["fallback_reason"],
+            finish_reason=ModelFinishReason(row["finish_reason"]),
+            status=row["status"],
+            error_code=row["error_code"],
+            metadata=metadata if isinstance(metadata, dict) else {},
+        )

contextos/storage/vector/__init__.py

@@ -0,0 +1 @@
+"""Vector storage sub-package for ContextOS."""

contextos/storage/vector/in_memory.py

@@ -0,0 +1,162 @@
+"""In-memory vector store using numpy for ContextOS.
+
+Phase 1 implementation: simple brute-force cosine similarity search.
+No ANN index — at Phase 1 scale (< 10K vectors), brute force is fast enough
+and avoids external library complexity.
+
+This will be replaced by sqlite-vec or LanceDB when we need ANN performance.
+It satisfies the VectorStore protocol and passes the same contract tests.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import numpy as np
+
+from contextos.core.models import VectorResult
+
+logger = logging.getLogger(__name__)
+
+
+class InMemoryVectorStore:
+    """Brute-force cosine similarity vector store.
+
+    Implements the VectorStore protocol.
+
+    Suitable for Phase 1 (< 10K vectors). O(n) search, but n is small.
+    """
+
+    def __init__(self, dimension: int = 384) -> None:
+        self._dimension = dimension
+        self._ids: list[str] = []
+        self._vectors: np.ndarray = np.empty((0, dimension), dtype=np.float32)
+        self._metadata: dict[str, dict[str, Any]] = {}
+
+    def contains(self, doc_id: str | Any) -> bool:
+        """Check whether a document ID is present in the vector store."""
+        return str(doc_id) in self._metadata
+
+    def __contains__(self, doc_id: str | Any) -> bool:
+        return str(doc_id) in self._metadata
+
+    async def add(
+        self,
+        ids: list[str],
+        vectors: list[list[float]],
+        metadata: list[dict[str, Any]],
+    ) -> None:
+        """Add vectors to the store."""
+        if not ids:
+            return
+        if len(ids) != len(vectors) or len(ids) != len(metadata):
+            raise ValueError("ids, vectors, and metadata must have equal lengths")
+
+        new_vectors = np.array(vectors, dtype=np.float32)
+        if new_vectors.ndim == 1:
+            new_vectors = new_vectors.reshape(1, -1)
+        if new_vectors.ndim != 2 or new_vectors.shape[1] != self._dimension:
+            raise ValueError(
+                f"Expected vectors with dimension {self._dimension}, got shape {new_vectors.shape}"
+            )
+        if not np.isfinite(new_vectors).all():
+            raise ValueError("Vectors must contain only finite values")
+
+        # Normalize for cosine similarity
+        norms = np.linalg.norm(new_vectors, axis=1, keepdims=True)
+        norms[norms == 0] = 1.0  # Avoid division by zero
+        new_vectors = new_vectors / norms
+
+        for i, doc_id in enumerate(ids):
+            if doc_id in self._metadata:
+                # Update: replace existing
+                idx = self._ids.index(doc_id)
+                self._vectors[idx] = new_vectors[i]
+                self._metadata[doc_id] = metadata[i] if i < len(metadata) else {}
+            else:
+                # Add new
+                self._ids.append(doc_id)
+                self._vectors = np.vstack([self._vectors, new_vectors[i:i + 1]])
+                self._metadata[doc_id] = metadata[i] if i < len(metadata) else {}
+
+    async def search(
+        self,
+        vector: list[float],
+        top_k: int = 20,
+        filters: dict[str, Any] | None = None,
+    ) -> list[VectorResult]:
+        """Search by cosine similarity."""
+        if len(self._ids) == 0:
+            return []
+
+        query = np.array(vector, dtype=np.float32)
+        if query.ndim != 1 or query.shape[0] != self._dimension:
+            raise ValueError(
+                f"Expected query dimension {self._dimension}, got shape {query.shape}"
+            )
+        if not np.isfinite(query).all():
+            raise ValueError("Query vector must contain only finite values")
+        norm = np.linalg.norm(query)
+        if norm == 0:
+            return []
+        query = query / norm
+
+        # Cosine similarity (vectors are pre-normalized)
+        similarities = self._vectors @ query
+
+        # Build results
+        indices = sorted(
+            range(len(self._ids)),
+            key=lambda index: (-float(similarities[index]), self._ids[index]),
+        )
+        results: list[VectorResult] = []
+
+        for idx in indices:
+            if len(results) >= top_k:
+                break
+
+            doc_id = self._ids[idx]
+            score = float(similarities[idx])
+
+            if score <= 0:
+                continue
+
+            # Apply filters
+            if filters:
+                meta = self._metadata.get(doc_id, {})
+                if not all(meta.get(k) == v for k, v in filters.items()):
+                    continue
+
+            results.append(VectorResult(
+                id=doc_id,
+                score=score,
+                metadata=self._metadata.get(doc_id, {}),
+            ))
+
+        return results
+
+    async def delete(self, ids: list[str]) -> None:
+        """Remove vectors from the store."""
+        for doc_id in ids:
+            if doc_id in self._metadata:
+                idx = self._ids.index(doc_id)
+                self._ids.pop(idx)
+                self._vectors = np.delete(self._vectors, idx, axis=0)
+                del self._metadata[doc_id]
+
+    async def count(self) -> int:
+        return len(self._ids)
+
+    async def rebuild(
+        self,
+        ids: list[str],
+        vectors: list[list[float]],
+        metadata: list[dict[str, Any]],
+    ) -> None:
+        """Replace the complete index after validating the new corpus."""
+        replacement = InMemoryVectorStore(self._dimension)
+        await replacement.add(ids, vectors, metadata)
+        self._ids = replacement._ids
+        self._vectors = replacement._vectors
+        self._metadata = replacement._metadata