tribalmemory 0.1.0__py3-none-any.whl

tribalmemory/services/memory.py

@@ -0,0 +1,275 @@
+"""Tribal Memory Service - Main API for agents."""
+
+import os
+from datetime import datetime
+from typing import Optional
+import uuid
+
+from ..interfaces import (
+    IMemoryService,
+    IEmbeddingService,
+    IVectorStore,
+    MemoryEntry,
+    MemorySource,
+    RecallResult,
+    StoreResult,
+)
+from .deduplication import SemanticDeduplicationService
+
+
+class TribalMemoryService(IMemoryService):
+    """Production tribal memory service.
+    
+    Usage:
+        service = TribalMemoryService(
+            instance_id="clawdio-1",
+            embedding_service=embedding_service,
+            vector_store=vector_store
+        )
+        
+        await service.remember("Joe prefers TypeScript")
+        results = await service.recall("What language for Wally?")
+    """
+    
+    def __init__(
+        self,
+        instance_id: str,
+        embedding_service: IEmbeddingService,
+        vector_store: IVectorStore,
+        dedup_exact_threshold: float = 0.98,
+        dedup_near_threshold: float = 0.90,
+        auto_reject_duplicates: bool = True,
+    ):
+        self.instance_id = instance_id
+        self.embedding_service = embedding_service
+        self.vector_store = vector_store
+        self.auto_reject_duplicates = auto_reject_duplicates
+        
+        self.dedup_service = SemanticDeduplicationService(
+            vector_store=vector_store,
+            embedding_service=embedding_service,
+            exact_threshold=dedup_exact_threshold,
+            near_threshold=dedup_near_threshold,
+        )
+    
+    async def remember(
+        self,
+        content: str,
+        source_type: MemorySource = MemorySource.AUTO_CAPTURE,
+        context: Optional[str] = None,
+        tags: Optional[list[str]] = None,
+        skip_dedup: bool = False,
+    ) -> StoreResult:
+        """Store a new memory."""
+        if not content or not content.strip():
+            return StoreResult(success=False, error="TribalMemory: Empty content not allowed")
+        
+        content = content.strip()
+        
+        try:
+            embedding = await self.embedding_service.embed(content)
+        except Exception as e:
+            return StoreResult(success=False, error=f"Embedding generation failed: {e}")
+        
+        if not skip_dedup and self.auto_reject_duplicates:
+            is_dup, dup_id = await self.dedup_service.is_duplicate(content, embedding)
+            if is_dup:
+                return StoreResult(success=False, duplicate_of=dup_id)
+        
+        entry = MemoryEntry(
+            id=str(uuid.uuid4()),
+            content=content,
+            embedding=embedding,
+            source_instance=self.instance_id,
+            source_type=source_type,
+            created_at=datetime.utcnow(),
+            updated_at=datetime.utcnow(),
+            tags=tags or [],
+            context=context,
+            confidence=1.0,
+        )
+        
+        return await self.vector_store.store(entry)
+    
+    async def recall(
+        self,
+        query: str,
+        limit: int = 5,
+        min_relevance: float = 0.7,
+        tags: Optional[list[str]] = None,
+    ) -> list[RecallResult]:
+        """Recall relevant memories.
+        
+        Args:
+            query: Natural language query
+            limit: Maximum results
+            min_relevance: Minimum similarity score
+            tags: Filter by tags (e.g., ["work", "preferences"])
+        """
+        try:
+            query_embedding = await self.embedding_service.embed(query)
+        except Exception:
+            return []
+        
+        filters = {"tags": tags} if tags else None
+        
+        results = await self.vector_store.recall(
+            query_embedding,
+            limit=limit,
+            min_similarity=min_relevance,
+            filters=filters,
+        )
+        
+        return self._filter_superseded(results)
+    
+    async def correct(
+        self,
+        original_id: str,
+        corrected_content: str,
+        context: Optional[str] = None,
+    ) -> StoreResult:
+        """Store a correction to an existing memory."""
+        original = await self.vector_store.get(original_id)
+        if not original:
+            return StoreResult(success=False, error=f"Original memory {original_id} not found")
+        
+        try:
+            embedding = await self.embedding_service.embed(corrected_content)
+        except Exception as e:
+            return StoreResult(success=False, error=f"Embedding generation failed: {e}")
+        
+        entry = MemoryEntry(
+            id=str(uuid.uuid4()),
+            content=corrected_content,
+            embedding=embedding,
+            source_instance=self.instance_id,
+            source_type=MemorySource.CORRECTION,
+            created_at=datetime.utcnow(),
+            updated_at=datetime.utcnow(),
+            tags=original.tags,
+            context=context or f"Correction of memory {original_id}",
+            confidence=1.0,
+            supersedes=original_id,
+            related_to=[original_id],
+        )
+        
+        return await self.vector_store.store(entry)
+    
+    async def forget(self, memory_id: str) -> bool:
+        """Forget (soft delete) a memory."""
+        return await self.vector_store.delete(memory_id)
+    
+    async def get(self, memory_id: str) -> Optional[MemoryEntry]:
+        """Get a memory by ID with full provenance."""
+        return await self.vector_store.get(memory_id)
+    
+    async def get_stats(self) -> dict:
+        """Get memory statistics.
+        
+        Note: Stats are computed over up to 10,000 most recent memories.
+        For systems with >10k memories, consider using count() with filters.
+        """
+        all_memories = await self.vector_store.list(limit=10000)
+        
+        by_source: dict[str, int] = {}
+        by_instance: dict[str, int] = {}
+        by_tag: dict[str, int] = {}
+        
+        for m in all_memories:
+            source = m.source_type.value
+            by_source[source] = by_source.get(source, 0) + 1
+            
+            instance = m.source_instance
+            by_instance[instance] = by_instance.get(instance, 0) + 1
+            
+            for tag in m.tags:
+                by_tag[tag] = by_tag.get(tag, 0) + 1
+        
+        corrections = sum(1 for m in all_memories if m.supersedes)
+        
+        return {
+            "total_memories": len(all_memories),
+            "by_source_type": by_source,
+            "by_tag": by_tag,
+            "by_instance": by_instance,
+            "corrections": corrections,
+        }
+
+    @staticmethod
+    def _filter_superseded(results: list[RecallResult]) -> list[RecallResult]:
+        """Remove memories that are superseded by corrections in the result set."""
+        superseded_ids = {
+            r.memory.supersedes for r in results if r.memory.supersedes
+        }
+        if not superseded_ids:
+            return results
+        return [r for r in results if r.memory.id not in superseded_ids]
+
+
+def create_memory_service(
+    instance_id: Optional[str] = None,
+    db_path: Optional[str] = None,
+    openai_api_key: Optional[str] = None,
+    api_base: Optional[str] = None,
+    embedding_model: Optional[str] = None,
+    embedding_dimensions: Optional[int] = None,
+) -> TribalMemoryService:
+    """Factory function to create a memory service with sensible defaults.
+    
+    Args:
+        instance_id: Unique identifier for this agent instance.
+        db_path: Path for LanceDB persistent storage. If None, uses in-memory.
+        openai_api_key: API key. Falls back to OPENAI_API_KEY env var.
+            Not required for local models (when api_base is set).
+        api_base: Base URL for the embedding API.
+            For Ollama: "http://localhost:11434/v1"
+        embedding_model: Embedding model name. Default: "text-embedding-3-small".
+        embedding_dimensions: Embedding output dimensions. Default: 1536.
+    
+    Returns:
+        Configured TribalMemoryService ready for use.
+    
+    Warning:
+        If db_path is provided but LanceDB is not installed, falls back to
+        in-memory storage. This means data will NOT persist across restarts.
+    """
+    import logging
+    
+    from .embeddings import OpenAIEmbeddingService
+    from .vector_store import InMemoryVectorStore, LanceDBVectorStore
+    
+    logger = logging.getLogger(__name__)
+    
+    if not instance_id:
+        instance_id = os.environ.get("TRIBAL_MEMORY_INSTANCE_ID", "default")
+    
+    kwargs: dict = {"api_key": openai_api_key}
+    if api_base is not None:
+        kwargs["api_base"] = api_base
+    if embedding_model is not None:
+        kwargs["model"] = embedding_model
+    if embedding_dimensions is not None:
+        kwargs["dimensions"] = embedding_dimensions
+    
+    embedding_service = OpenAIEmbeddingService(**kwargs)
+    
+    if db_path:
+        try:
+            vector_store = LanceDBVectorStore(
+                embedding_service=embedding_service,
+                db_path=db_path
+            )
+        except ImportError:
+            logger.warning(
+                "LanceDB not installed. Falling back to in-memory storage. "
+                "Data will NOT persist across restarts. Install with: pip install lancedb"
+            )
+            vector_store = InMemoryVectorStore(embedding_service)
+    else:
+        vector_store = InMemoryVectorStore(embedding_service)
+    
+    return TribalMemoryService(
+        instance_id=instance_id,
+        embedding_service=embedding_service,
+        vector_store=vector_store
+    )

tribalmemory/services/vector_store.py

@@ -0,0 +1,360 @@
+"""Vector Store implementations.
+
+Provides both LanceDB (persistent) and in-memory storage options.
+"""
+
+import json
+import os
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, Union
+
+from ..interfaces import (
+    IVectorStore,
+    IEmbeddingService,
+    MemoryEntry,
+    MemorySource,
+    RecallResult,
+    StoreResult,
+)
+
+
+class LanceDBVectorStore(IVectorStore):
+    """LanceDB-backed vector store for persistent storage.
+    
+    Supports both local file storage and LanceDB Cloud.
+    """
+    
+    TABLE_NAME = "memories"
+    
+    def __init__(
+        self,
+        embedding_service: IEmbeddingService,
+        db_path: Optional[Union[str, Path]] = None,
+        db_uri: Optional[str] = None,
+        api_key: Optional[str] = None,
+    ):
+        self.embedding_service = embedding_service
+        self.db_path = Path(db_path) if db_path else None
+        self.db_uri = db_uri
+        self.api_key = api_key or os.environ.get("LANCEDB_API_KEY")
+        
+        self._db = None
+        self._table = None
+        self._initialized = False
+    
+    async def _ensure_initialized(self):
+        """Lazily initialize database connection."""
+        if self._initialized:
+            return
+        
+        try:
+            import lancedb
+        except ImportError:
+            raise ImportError("LanceDB not installed. Run: pip install lancedb")
+        
+        if self.db_uri:
+            self._db = lancedb.connect(self.db_uri, api_key=self.api_key)
+        elif self.db_path:
+            self.db_path.mkdir(parents=True, exist_ok=True)
+            self._db = lancedb.connect(str(self.db_path))
+        else:
+            raise ValueError("Either db_path or db_uri must be provided")
+        
+        if self.TABLE_NAME in self._db.table_names():
+            self._table = self._db.open_table(self.TABLE_NAME)
+        else:
+            self._table = self._create_table()
+        
+        self._initialized = True
+    
+    def _create_table(self) -> "lancedb.table.Table":
+        """Create the memories table with the defined schema."""
+        import pyarrow as pa
+        
+        schema = pa.schema([
+            pa.field("id", pa.string()),
+            pa.field("content", pa.string()),
+            pa.field("vector", pa.list_(pa.float32(), self._get_embedding_dim())),
+            pa.field("source_instance", pa.string()),
+            pa.field("source_type", pa.string()),
+            pa.field("created_at", pa.string()),
+            pa.field("updated_at", pa.string()),
+            pa.field("tags", pa.string()),
+            pa.field("context", pa.string()),
+            pa.field("confidence", pa.float32()),
+            pa.field("supersedes", pa.string()),
+            pa.field("related_to", pa.string()),
+            pa.field("deleted", pa.bool_()),
+        ])
+        
+        return self._db.create_table(self.TABLE_NAME, schema=schema)
+    
+    def _get_embedding_dim(self) -> int:
+        """Get the expected embedding dimension from the embedding service."""
+        if hasattr(self.embedding_service, 'dimensions'):
+            return self.embedding_service.dimensions
+        return 1536  # Default for text-embedding-3-small
+    
+    async def store(self, entry: MemoryEntry) -> StoreResult:
+        await self._ensure_initialized()
+        
+        try:
+            if entry.embedding is None:
+                entry.embedding = await self.embedding_service.embed(entry.content)
+            
+            # Validate embedding dimensions
+            expected_dim = self._get_embedding_dim()
+            if len(entry.embedding) != expected_dim:
+                return StoreResult(
+                    success=False,
+                    error=f"Invalid embedding dimension: got {len(entry.embedding)}, expected {expected_dim}"
+                )
+            
+            row = {
+                "id": entry.id,
+                "content": entry.content,
+                "vector": entry.embedding,
+                "source_instance": entry.source_instance,
+                "source_type": entry.source_type.value,
+                "created_at": entry.created_at.isoformat(),
+                "updated_at": entry.updated_at.isoformat(),
+                "tags": json.dumps(entry.tags),
+                "context": entry.context or "",
+                "confidence": entry.confidence,
+                "supersedes": entry.supersedes or "",
+                "related_to": json.dumps(entry.related_to),
+                "deleted": False,
+            }
+            
+            self._table.add([row])
+            return StoreResult(success=True, memory_id=entry.id)
+            
+        except Exception as e:
+            return StoreResult(success=False, error=str(e))
+    
+    async def recall(
+        self,
+        query_embedding: list[float],
+        limit: int = 10,
+        min_similarity: float = 0.7,
+        filters: Optional[dict] = None,
+    ) -> list[RecallResult]:
+        await self._ensure_initialized()
+        
+        start = time.perf_counter()
+        
+        results = (
+            self._table.search(query_embedding)
+            .where("deleted = false")
+            .limit(limit * 2)
+            .to_list()
+        )
+        
+        elapsed_ms = (time.perf_counter() - start) * 1000
+        
+        recall_results = []
+        for row in results:
+            # LanceDB returns L2 distance. Convert to cosine similarity.
+            # For normalized vectors (which OpenAI embeddings are):
+            # L2_distance² = 2 * (1 - cosine_similarity)
+            # Therefore: cosine_similarity = 1 - (L2_distance² / 2)
+            distance = row.get("_distance", 0)
+            similarity = max(0, 1 - (distance * distance / 2))
+            
+            if similarity < min_similarity:
+                continue
+            
+            entry = self._row_to_entry(row)
+            
+            # Apply filters
+            if filters:
+                if "tags" in filters and filters["tags"]:
+                    if not any(t in entry.tags for t in filters["tags"]):
+                        continue
+            
+            recall_results.append(RecallResult(
+                memory=entry,
+                similarity_score=similarity,
+                retrieval_time_ms=elapsed_ms
+            ))
+        
+        recall_results.sort(key=lambda x: x.similarity_score, reverse=True)
+        return recall_results[:limit]
+    
+    async def get(self, memory_id: str) -> Optional[MemoryEntry]:
+        await self._ensure_initialized()
+        
+        # Sanitize memory_id to prevent SQL injection
+        safe_id = self._sanitize_id(memory_id)
+        
+        results = (
+            self._table.search()
+            .where(f"id = '{safe_id}' AND deleted = false")
+            .limit(1)
+            .to_list()
+        )
+        
+        if not results:
+            return None
+        return self._row_to_entry(results[0])
+    
+    async def delete(self, memory_id: str) -> bool:
+        await self._ensure_initialized()
+        
+        # Sanitize memory_id to prevent SQL injection
+        safe_id = self._sanitize_id(memory_id)
+        
+        try:
+            self._table.update(
+                where=f"id = '{safe_id}'",
+                values={"deleted": True, "updated_at": datetime.utcnow().isoformat()}
+            )
+            return True
+        except Exception:
+            return False
+    
+    def _sanitize_id(self, memory_id: str) -> str:
+        """Sanitize memory_id to prevent SQL injection.
+        
+        UUIDs should only contain alphanumeric characters and hyphens.
+        Rejects any input containing quotes, semicolons, or other SQL metacharacters.
+        """
+        import re
+        # UUID pattern: only allow alphanumeric and hyphens
+        if not re.match(r'^[a-zA-Z0-9\-]+$', memory_id):
+            raise ValueError(f"Invalid memory_id format: {memory_id[:20]}...")
+        return memory_id
+    
+    async def list(
+        self,
+        limit: int = 1000,
+        offset: int = 0,
+        filters: Optional[dict] = None,
+    ) -> list[MemoryEntry]:
+        await self._ensure_initialized()
+        
+        results = (
+            self._table.search()
+            .where("deleted = false")
+            .limit(limit + offset)
+            .to_list()
+        )
+        
+        entries = [self._row_to_entry(row) for row in results[offset:offset + limit]]
+        
+        if filters and "tags" in filters and filters["tags"]:
+            entries = [e for e in entries if any(t in e.tags for t in filters["tags"])]
+        
+        return entries
+    
+    async def count(self, filters: Optional[dict] = None) -> int:
+        entries = await self.list(limit=100000, filters=filters)
+        return len(entries)
+    
+    def _row_to_entry(self, row: dict) -> MemoryEntry:
+        return MemoryEntry(
+            id=row["id"],
+            content=row["content"],
+            embedding=row.get("vector"),
+            source_instance=row.get("source_instance", "unknown"),
+            source_type=MemorySource(row.get("source_type", "unknown")),
+            created_at=datetime.fromisoformat(row["created_at"]) if row.get("created_at") else datetime.utcnow(),
+            updated_at=datetime.fromisoformat(row["updated_at"]) if row.get("updated_at") else datetime.utcnow(),
+            tags=json.loads(row.get("tags", "[]")),
+            context=row.get("context") or None,
+            confidence=row.get("confidence", 1.0),
+            supersedes=row.get("supersedes") or None,
+            related_to=json.loads(row.get("related_to", "[]")),
+        )
+
+
+class InMemoryVectorStore(IVectorStore):
+    """Simple in-memory vector store for testing."""
+    
+    def __init__(self, embedding_service: IEmbeddingService):
+        self.embedding_service = embedding_service
+        self._store: dict[str, MemoryEntry] = {}
+        self._deleted: set[str] = set()
+    
+    async def store(self, entry: MemoryEntry) -> StoreResult:
+        if entry.embedding is None:
+            entry.embedding = await self.embedding_service.embed(entry.content)
+        
+        self._store[entry.id] = entry
+        return StoreResult(success=True, memory_id=entry.id)
+    
+    async def recall(
+        self,
+        query_embedding: list[float],
+        limit: int = 10,
+        min_similarity: float = 0.7,
+        filters: Optional[dict] = None,
+    ) -> list[RecallResult]:
+        start = time.perf_counter()
+        
+        results = []
+        for entry in self._store.values():
+            if entry.id in self._deleted:
+                continue
+            if entry.embedding is None:
+                continue
+            
+            # Apply filters
+            if filters and "tags" in filters and filters["tags"]:
+                if not any(t in entry.tags for t in filters["tags"]):
+                    continue
+            
+            sim = self.embedding_service.similarity(query_embedding, entry.embedding)
+            if sim >= min_similarity:
+                results.append((entry, sim))
+        
+        results.sort(key=lambda x: x[1], reverse=True)
+        elapsed_ms = (time.perf_counter() - start) * 1000
+        
+        return [
+            RecallResult(memory=e, similarity_score=s, retrieval_time_ms=elapsed_ms)
+            for e, s in results[:limit]
+        ]
+    
+    async def get(self, memory_id: str) -> Optional[MemoryEntry]:
+        if memory_id in self._deleted:
+            return None
+        return self._store.get(memory_id)
+    
+    async def delete(self, memory_id: str) -> bool:
+        if memory_id in self._store:
+            self._deleted.add(memory_id)
+            return True
+        return False
+    
+    async def list(
+        self,
+        limit: int = 1000,
+        offset: int = 0,
+        filters: Optional[dict] = None,
+    ) -> list[MemoryEntry]:
+        entries = [
+            e for e in list(self._store.values())
+            if e.id not in self._deleted
+        ]
+        
+        if filters and "tags" in filters and filters["tags"]:
+            entries = [e for e in entries if any(t in e.tags for t in filters["tags"])]
+        
+        return entries[offset:offset + limit]
+
+    async def upsert(self, entry: MemoryEntry) -> StoreResult:
+        """Insert or replace, clearing any soft-delete tombstone."""
+        self._deleted.discard(entry.id)
+        if entry.embedding is None:
+            entry.embedding = (
+                await self.embedding_service.embed(entry.content)
+            )
+        self._store[entry.id] = entry
+        return StoreResult(success=True, memory_id=entry.id)
+
+    async def count(self, filters: Optional[dict] = None) -> int:
+        entries = await self.list(limit=100000, filters=filters)
+        return len(entries)

tribalmemory/testing/__init__.py

@@ -0,0 +1,22 @@
+"""Testing utilities for Tribal Memory."""
+
+from .mocks import (
+    MockEmbeddingService,
+    MockVectorStore,
+    MockMemoryService,
+    MockTimestampService,
+)
+from .metrics import LatencyTracker, SimilarityCalculator, TestResultLogger
+from .fixtures import load_test_data, TestDataSet
+
+__all__ = [
+    "MockEmbeddingService",
+    "MockVectorStore", 
+    "MockMemoryService",
+    "MockTimestampService",
+    "LatencyTracker",
+    "SimilarityCalculator",
+    "TestResultLogger",
+    "load_test_data",
+    "TestDataSet",
+]