agent-restore-harness/scripts/ingest/memory_api_fix.py

"""Main Agent Memory class - the primary interface."""

import logging
from datetime import datetime, timedelta
from typing import Any
from uuid import UUID

import redis.asyncio as redis
from qdrant_client import AsyncQdrantClient, models
from qdrant_client.models import (
    NamedSparseVector,
    NamedVector,
    PointStruct,
    SparseVector,
)

from config import get_settings
from embeddings import BGEm3Client
from hybrid_search import HybridSearcher
from models import (
    Memory,
    MemoryCreate,
    MemorySource,
    MemoryType,
    SearchQuery,
    SearchResult,
)
from routing.router import QueryRouter
from routing.stats_store import RoutingStatsStore

logger = logging.getLogger(__name__)


class AgentMemory:
    """
    Main interface for agent memory operations.
    
    Provides:
    - Store memories with automatic embedding
    - Hybrid search (dense + sparse)
    - Time-based retrieval
    - Importance decay
    - Memory consolidation
    """
    
    def __init__(
        self,
        agent_id: str,
        qdrant: AsyncQdrantClient | None = None,
        embedding_client: BGEm3Client | None = None,
        redis_client: redis.Redis | None = None,
    ):
        self.agent_id = agent_id
        self.settings = get_settings()
        
        # Initialize clients
        self._qdrant = qdrant
        self._embeddings = embedding_client
        self._redis = redis_client
        self._searcher: HybridSearcher | None = None
        self._router: QueryRouter | None = None
        # Only close clients this instance created — not externally-shared ones
        self._owns_qdrant = qdrant is None
        self._owns_embeddings = embedding_client is None
        self._owns_redis = redis_client is None

        self._initialized = False
    
    async def _ensure_initialized(self):
        """Lazy initialization of clients."""
        if self._initialized:
            return
        
        # Qdrant client
        if self._qdrant is None:
            self._qdrant = AsyncQdrantClient(
                url=self.settings.qdrant_url,
                api_key=self.settings.qdrant_api_key,
            )
        
        # Redis client (optional)
        if self._redis is None and self.settings.cache_embeddings:
            try:
                self._redis = redis.from_url(self.settings.redis_url)
                await self._redis.ping()
            except Exception as e:
                logger.warning(f"Redis connection failed, caching disabled: {e}")
                self._redis = None
        
        # Embedding client
        if self._embeddings is None:
            self._embeddings = BGEm3Client(redis_client=self._redis)
        
        # Hybrid searcher
        self._searcher = HybridSearcher(
            qdrant=self._qdrant,
            embedding_client=self._embeddings,
        )

        # Query router (M4 smart routing)
        stats_store = None
        if self._redis:
            stats_store = RoutingStatsStore(self._redis)
        self._router = QueryRouter(
            qdrant=self._qdrant,
            embedding_client=self._embeddings,
            stats_store=stats_store,
        )

        self._initialized = True
    
    async def close(self):
        """Close connections owned by this instance (not externally-shared clients)."""
        if self._embeddings and self._owns_embeddings:
            await self._embeddings.close()
        if self._redis and self._owns_redis:
            await self._redis.close()
        if self._qdrant and self._owns_qdrant:
            await self._qdrant.close()
    
    async def __aenter__(self):
        await self._ensure_initialized()
        return self
    
    async def __aexit__(self, *args):
        await self.close()
    
    # =========================================================================
    # STORE OPERATIONS
    # =========================================================================
    
    async def store(
        self,
        content: str,
        memory_type: MemoryType = MemoryType.EPISODIC,
        importance: float | None = None,
        source: MemorySource = MemorySource.CONVERSATION,
        entities: list[str] | None = None,
        session_id: str | None = None,
        user_id: str | None = None,
        language: str = "en",
        metadata: dict[str, Any] | None = None,
    ) -> Memory:
        """
        Store a new memory with automatic embedding.
        
        Args:
            content: The memory content
            memory_type: Type of memory (episodic, semantic, working)
            importance: Importance score (0-1), defaults to settings default
            source: Source of the memory
            entities: Extracted entities
            session_id: Session ID for episodic memories
            user_id: Associated user ID
            language: ISO language code
            metadata: Additional metadata
            
        Returns:
            The created Memory object
        """
        await self._ensure_initialized()
        
        # Create memory object
        memory = Memory(
            content=content,
            memory_type=memory_type,
            agent_id=self.agent_id,
            session_id=session_id,
            user_id=user_id,
            importance=importance or self.settings.default_importance,
            initial_importance=importance or self.settings.default_importance,
            source=source,
            entities=entities or [],
            language=language,
            metadata=metadata or {},
        )

        # Episodic memories start unconsolidated
        if memory_type == MemoryType.EPISODIC:
            memory.consolidated = False
        
        # Generate embeddings
        embedding_response = await self._embeddings.embed_full(content)
        
        memory.has_dense = embedding_response.dense is not None
        memory.has_sparse = embedding_response.sparse is not None
        memory.has_colbert = embedding_response.colbert is not None
        
        # Build vectors dict for Qdrant
        vectors = {}
        
        if embedding_response.dense:
            vectors["dense"] = embedding_response.dense
        
        # Prepare point
        point = PointStruct(
            id=str(memory.id),
            vector=vectors,
            payload=memory.model_dump(mode="json"),
        )
        
        # If sparse embeddings, add as named sparse vector
        if embedding_response.sparse:
            sparse_indices = list(embedding_response.sparse.keys())
            sparse_values = list(embedding_response.sparse.values())
            
            # Qdrant requires sparse vectors in a specific format
            point.vector["sparse"] = SparseVector(
                indices=sparse_indices,
                values=sparse_values,
            )
        
        # Upsert to Qdrant
        await self._qdrant.upsert(
            collection_name=self.settings.collection_name,
            points=[point],
        )
        
        logger.info(f"Stored memory {memory.id} for agent {self.agent_id}")
        
        return memory
    
    async def store_batch(
        self,
        memories: list[MemoryCreate],
    ) -> list[Memory]:
        """Store multiple memories efficiently."""
        await self._ensure_initialized()
        
        if not memories:
            return []
        
        # Create memory objects
        memory_objects = []
        contents = []
        
        for mem_create in memories:
            memory = Memory(
                content=mem_create.content,
                memory_type=mem_create.memory_type,
                agent_id=self.agent_id,
                session_id=mem_create.session_id,
                user_id=mem_create.user_id,
                importance=mem_create.importance,
                initial_importance=mem_create.importance,
                source=mem_create.source,
                entities=mem_create.entities,
                language=mem_create.language,
                metadata=mem_create.metadata,
            )
            # Episodic memories start unconsolidated
            if mem_create.memory_type == MemoryType.EPISODIC:
                memory.consolidated = False
            memory_objects.append(memory)
            contents.append(mem_create.content)
        
        # Batch embed
        embedding_responses = await self._embeddings.embed_batch(
            contents,
            include_sparse=True,
            include_colbert=False,
        )
        
        # Build points
        points = []
        for memory, emb_response in zip(memory_objects, embedding_responses):
            memory.has_dense = emb_response.dense is not None
            memory.has_sparse = emb_response.sparse is not None
            
            vectors = {}
            if emb_response.dense:
                vectors["dense"] = emb_response.dense
            
            point = PointStruct(
                id=str(memory.id),
                vector=vectors,
                payload=memory.model_dump(mode="json"),
            )
            
            if emb_response.sparse:
                point.vector["sparse"] = SparseVector(
                    indices=list(emb_response.sparse.keys()),
                    values=list(emb_response.sparse.values()),
                )
            
            points.append(point)
        
        # Batch upsert
        await self._qdrant.upsert(
            collection_name=self.settings.collection_name,
            points=points,
        )
        
        logger.info(f"Stored {len(points)} memories for agent {self.agent_id}")
        
        return memory_objects
    
    # =========================================================================
    # SEARCH OPERATIONS
    # =========================================================================
    
    async def search(
        self,
        query: str,
        memory_types: list[MemoryType] | None = None,
        limit: int = 10,
        min_importance: float = 0.0,
        session_id: str | None = None,
        user_id: str | None = None,
        after: datetime | None = None,
        before: datetime | None = None,
        routing_strategy: str | None = None,
    ) -> list[SearchResult]:
        """
        Search memories using hybrid search.

        Args:
            query: Search query text
            memory_types: Filter by memory types
            limit: Maximum results
            min_importance: Minimum importance threshold
            session_id: Filter by session
            user_id: Filter by user
            after: Only memories after this time
            before: Only memories before this time
            routing_strategy: Smart routing strategy ('auto', 'lookup',
                'standard', 'deep', 'synthesis'). None = legacy hybrid search.

        Returns:
            List of SearchResult with relevance scores
        """
        await self._ensure_initialized()

        # --- Smart routing path (M4) ---
        if routing_strategy is not None:
            routing_result = await self._router.route(
                query=query,
                agent_id=self.agent_id,
                strategy=routing_strategy,
                memory_types=memory_types,
                limit=limit,
                min_importance=min_importance,
                session_id=session_id,
                user_id=user_id,
                after=after,
                before=before,
            )
            results = routing_result.strategy_result.results
            for result in results:
                await self._update_access(result.memory.id)
            return results

        # --- Legacy hybrid search path (default, backward compat) ---
        search_query = SearchQuery(
            query=query,
            memory_types=memory_types,
            limit=limit,
            min_importance=min_importance,
            session_id=session_id,
            user_id=user_id,
            after=after,
            before=before,
        )

        results = await self._searcher.search(search_query, self.agent_id)

        # Update access count and last_accessed for retrieved memories
        for result in results:
            await self._update_access(result.memory.id)

        return results
    
    async def _update_access(self, memory_id: UUID):
        """Update access tracking for a memory."""
        try:
            await self._qdrant.set_payload(
                collection_name=self.settings.collection_name,
                payload={
                    "last_accessed": datetime.utcnow().isoformat(),
                },
                points=[str(memory_id)],
            )
        except Exception as e:
            logger.warning(f"Failed to update access for {memory_id}: {e}")
    
    # =========================================================================
    # RETRIEVAL OPERATIONS
    # =========================================================================
    
    async def get(self, memory_id: UUID) -> Memory | None:
        """Get a specific memory by ID."""
        await self._ensure_initialized()
        
        results = await self._qdrant.retrieve(
            collection_name=self.settings.collection_name,
            ids=[str(memory_id)],
            with_payload=True,
        )
        
        if results:
            return Memory(**results[0].payload)
        return None
    
    async def get_recent(
        self,
        memory_type: MemoryType | None = None,
        hours: int = 24,
        limit: int = 50,
        session_id: str | None = None,
    ) -> list[Memory]:
        """Get recent memories by time."""
        await self._ensure_initialized()
        
        # Build filter
        must_conditions = [
            models.FieldCondition(
                key="agent_id",
                match=models.MatchValue(value=self.agent_id),
            ),
            models.FieldCondition(
                key="timestamp",
                range=models.Range(
                    gte=(datetime.utcnow() - timedelta(hours=hours)).timestamp(),
                ),
            ),
        ]
        
        if memory_type:
            must_conditions.append(
                models.FieldCondition(
                    key="memory_type",
                    match=models.MatchValue(value=memory_type.value),
                )
            )
        
        if session_id:
            must_conditions.append(
                models.FieldCondition(
                    key="session_id",
                    match=models.MatchValue(value=session_id),
                )
            )
        
        # Scroll through results
        results, _ = await self._qdrant.scroll(
            collection_name=self.settings.collection_name,
            scroll_filter=models.Filter(must=must_conditions),
            limit=limit,
            with_payload=True,
            order_by=models.OrderBy(
                key="timestamp",
                direction=models.Direction.DESC,
            ),
        )
        
        return [Memory(**r.payload) for r in results]
    
    async def get_by_entities(
        self,
        entities: list[str],
        memory_type: MemoryType | None = None,
        limit: int = 20,
    ) -> list[Memory]:
        """Get memories containing specific entities."""
        await self._ensure_initialized()
        
        must_conditions = [
            models.FieldCondition(
                key="agent_id",
                match=models.MatchValue(value=self.agent_id),
            ),
        ]
        
        # Match any of the entities
        for entity in entities:
            must_conditions.append(
                models.FieldCondition(
                    key="entities",
                    match=models.MatchValue(value=entity),
                )
            )
        
        if memory_type:
            must_conditions.append(
                models.FieldCondition(
                    key="memory_type",
                    match=models.MatchValue(value=memory_type.value),
                )
            )
        
        results, _ = await self._qdrant.scroll(
            collection_name=self.settings.collection_name,
            scroll_filter=models.Filter(must=must_conditions),
            limit=limit,
            with_payload=True,
        )
        
        return [Memory(**r.payload) for r in results]
    
    # =========================================================================
    # UPDATE OPERATIONS
    # =========================================================================
    
    async def update_importance(
        self,
        memory_id: UUID,
        importance: float,
    ):
        """Update importance score for a memory."""
        await self._ensure_initialized()
        
        await self._qdrant.set_payload(
            collection_name=self.settings.collection_name,
            payload={"importance": max(0.0, min(1.0, importance))},
            points=[str(memory_id)],
        )
    
    async def reinforce(self, memory_id: UUID, boost: float = 0.1):
        """Reinforce a memory (increase importance)."""
        memory = await self.get(memory_id)
        if memory:
            new_importance = min(1.0, memory.importance + boost)
            await self.update_importance(memory_id, new_importance)
    
    async def decay_importance(self, days_old: int = 7):
        """Apply importance decay to old memories."""
        await self._ensure_initialized()
        
        cutoff = datetime.utcnow() - timedelta(days=days_old)
        
        # Find old memories
        results, _ = await self._qdrant.scroll(
            collection_name=self.settings.collection_name,
            scroll_filter=models.Filter(
                must=[
                    models.FieldCondition(
                        key="agent_id",
                        match=models.MatchValue(value=self.agent_id),
                    ),
                    models.FieldCondition(
                        key="timestamp",
                        range=models.Range(lte=cutoff.timestamp()),
                    ),
                ]
            ),
            limit=1000,
            with_payload=True,
        )
        
        decay_rate = self.settings.importance_decay_rate
        
        for point in results:
            memory = Memory(**point.payload)
            days_since = (datetime.utcnow() - memory.timestamp).days
            decayed = memory.initial_importance * (1 - decay_rate * days_since)
            decayed = max(0.05, decayed)  # Minimum importance
            
            if abs(decayed - memory.importance) > 0.01:
                await self.update_importance(memory.id, decayed)
    
    # =========================================================================
    # DELETE OPERATIONS
    # =========================================================================
    
    async def delete(self, memory_id: UUID) -> bool:
        """Delete a specific memory."""
        await self._ensure_initialized()
        
        await self._qdrant.delete(
            collection_name=self.settings.collection_name,
            points_selector=models.PointIdsList(
                points=[str(memory_id)],
            ),
        )
        return True
    
    async def delete_session(self, session_id: str) -> int:
        """Delete all memories from a session."""
        await self._ensure_initialized()
        
        result = await self._qdrant.delete(
            collection_name=self.settings.collection_name,
            points_selector=models.FilterSelector(
                filter=models.Filter(
                    must=[
                        models.FieldCondition(
                            key="agent_id",
                            match=models.MatchValue(value=self.agent_id),
                        ),
                        models.FieldCondition(
                            key="session_id",
                            match=models.MatchValue(value=session_id),
                        ),
                    ]
                )
            ),
        )
        return result.status
    
    async def clear_all(self):
        """Delete all memories for this agent. Use with caution!"""
        await self._ensure_initialized()
        
        await self._qdrant.delete(
            collection_name=self.settings.collection_name,
            points_selector=models.FilterSelector(
                filter=models.Filter(
                    must=[
                        models.FieldCondition(
                            key="agent_id",
                            match=models.MatchValue(value=self.agent_id),
                        ),
                    ]
                )
            ),
        )
        logger.warning(f"Cleared all memories for agent {self.agent_id}")
    
    # =========================================================================
    # STATS
    # =========================================================================
    
    async def count(
        self,
        memory_type: MemoryType | None = None,
    ) -> int:
        """Count memories for this agent."""
        await self._ensure_initialized()
        
        must_conditions = [
            models.FieldCondition(
                key="agent_id",
                match=models.MatchValue(value=self.agent_id),
            ),
        ]
        
        if memory_type:
            must_conditions.append(
                models.FieldCondition(
                    key="memory_type",
                    match=models.MatchValue(value=memory_type.value),
                )
            )
        
        result = await self._qdrant.count(
            collection_name=self.settings.collection_name,
            count_filter=models.Filter(must=must_conditions),
        )
        
        return result.count