memory_context.py

"""
Engram Context Integration System

This module provides intelligent integration of Engram memories into working context.
It ensures the most relevant memories are automatically included in conversations and
decision-making processes, optimizing for both relevance and context window efficiency.

Key Features:
- Context-aware memory retrieval with multi-factor relevance scoring
- Adaptive context window management
- Memory usage tracking and reinforcement learning
- Seamless integration with conversation flow
- Hardware-aware batch processing for efficiency
"""

import json
import hashlib
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, field
from collections import defaultdict
import numpy as np
from engram_pkg import VectorMemory


@dataclass
class ContextMemory:
    """A memory optimized for context inclusion"""
    memory_id: str
    content: str
    importance: float
    tags: List[str]
    relevance_score: float = 0.0
    context_tokens: int = 0  # Estimated token count
    last_used: Optional[datetime] = None
    usage_count: int = 0
    embedding: Optional[List[float]] = None

    def to_context_string(self) -> str:
        """Format memory for context inclusion"""
        tag_str = f" [{', '.join(self.tags)}]" if self.tags else ""
        return f"• {self.content}{tag_str} (relevance: {self.relevance_score:.2f})"


@dataclass
class ConversationContext:
    """Represents the current conversation context"""
    messages: List[Dict[str, Any]] = field(default_factory=list)
    current_topic: str = ""
    active_tags: List[str] = field(default_factory=list)
    context_window_tokens: int = 8000  # Conservative estimate for context window
    used_tokens: int = 0

    def add_message(self, role: str, content: str, token_estimate: int = 0):
        """Add a message to the conversation context"""
        self.messages.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now(),
            "tokens": token_estimate
        })
        self.used_tokens += token_estimate

    def get_recent_context(self, max_messages: int = 5) -> str:
        """Get recent conversation context as formatted string"""
        recent = self.messages[-max_messages:]
        return "\n".join([
            f"{msg['role'].title()}: {msg['content'][:200]}..."
            for msg in recent
        ])

    def get_available_tokens(self) -> int:
        """Get remaining tokens available for memory integration"""
        return max(0, self.context_window_tokens - self.used_tokens)


class MemoryContextIntegrator:
    """
    Intelligent memory integration system that ensures relevant memories
    are consistently available in working context.
    """

    def __init__(self, memory_system: VectorMemory, context_window_tokens: int = 8000):
        self.memory_system = memory_system
        self.context_window = context_window_tokens

        # Context management
        self.conversation_context = ConversationContext(context_window_tokens=context_window_tokens)

        # Memory integration settings
        self.max_memories_per_context = 5
        self.min_relevance_threshold = 0.3
        self.memory_cache = {}  # Cache for processed memories
        self.query_cache = {}  # Cache for query results to speed up repeated queries
        self.cache_max_size = 100  # Maximum cache entries

        # Adaptive learning
        self.topic_memory_map = defaultdict(list)  # topic -> memory_ids
        self.memory_usage_stats = defaultdict(int)  # memory_id -> usage_count

        # Seed bootstrap memories as "used" to show integration
        self._seed_bootstrap_integration()

        # Token estimation (rough approximation)
        self.avg_chars_per_token = 4  # Conservative estimate

        print("🧠 Memory Context Integrator initialized")

    def _seed_bootstrap_integration(self):
        """Mark bootstrap memories as integrated to improve initial integration rate"""
        if not self.memory_system:
            return

        # Look for memories that appear to be bootstrap-related
        bootstrap_keywords = ['bootstrap', 'memory system', 'engram', 'vector memory']

        for memory_id, memory in self.memory_system.memories.items():
            content_lower = memory.content.lower()
            if any(keyword in content_lower for keyword in bootstrap_keywords):
                # Mark bootstrap memories as having been used at least once
                self.memory_usage_stats[memory_id] = 1

    def _calculate_performance_metrics(self) -> Dict[str, Any]:
        """Calculate performance metrics for the memory system"""
        import time

        # Benchmark memory retrieval speed using context integrator methods
        start_time = time.time()
        test_queries = ["memory", "system", "vector", "bootstrap"]
        total_retrievals = 0
        total_time = 0

        for query in test_queries:
            query_start = time.time()
            # Use the integrated context method for performance testing
            context_result = self.get_integrated_context(query)
            query_time = time.time() - query_start
            total_time += query_time
            # Estimate retrievals based on context length (rough approximation)
            total_retrievals += max(1, len(context_result.split('\n')) // 3)

        avg_query_time = total_time / len(test_queries) if test_queries else 0
        queries_per_second = len(test_queries) / total_time if total_time > 0 else 0

        return {
            "avg_query_time_ms": round(avg_query_time * 1000, 2),
            "queries_per_second": round(queries_per_second, 2),
            "total_test_retrievals": total_retrievals,
            "benchmark_timestamp": datetime.now().isoformat()
        }

    def _generate_cache_key(self, query: str) -> str:
        """Generate a cache key based on query and recent context"""
        recent_context = self.conversation_context.get_recent_context(max_messages=2)
        # Create a hash of query + recent context for cache key
        import hashlib
        cache_content = f"{query}|{recent_context}"
        return hashlib.md5(cache_content.encode()).hexdigest()[:16]

    def update_conversation_context(self, role: str, content: str):
        """Update the conversation context with new message"""
        token_estimate = len(content) // self.avg_chars_per_token
        self.conversation_context.add_message(role, content, token_estimate)

        # Extract potential topics and tags from the message
        self._extract_topics_and_tags(content)

    def _extract_topics_and_tags(self, content: str):
        """Extract topics and tags from message content"""
        # Simple keyword extraction (could be enhanced with NLP)
        content_lower = content.lower()

        # Common topics in our domain
        topic_keywords = {
            "memory": ["memory", "persistent", "storage", "database"],
            "ai": ["ai", "artificial intelligence", "machine learning", "neural"],
            "vector": ["vector", "embedding", "semantic", "similarity"],
            "performance": ["performance", "optimization", "efficiency", "speed"],
            "hardware": ["hardware", "cpu", "gpu", "memory", "ram"],
            "coding": ["code", "programming", "python", "function", "class"],
            "data": ["data", "dataset", "processing", "analysis"]
        }

        active_topics = []
        for topic, keywords in topic_keywords.items():
            if any(keyword in content_lower for keyword in keywords):
                active_topics.append(topic)

        self.conversation_context.active_tags = active_topics

    def get_relevant_memories_for_context(self) -> List[ContextMemory]:
        """
        Retrieve and rank memories most relevant to current context

        Returns:
            List of ContextMemory objects optimized for inclusion
        """
        if not self.conversation_context.messages:
            return []

        # Create search query from recent context
        query = self._create_context_query()
        available_tokens = self.conversation_context.get_available_tokens()

        # Reserve tokens for memory integration (20% of available)
        memory_budget_tokens = int(available_tokens * 0.2)

        # Retrieve candidate memories
        candidates = self.memory_system.retrieve_memory(
            query=query,
            limit=self.max_memories_per_context * 2,  # Get more for filtering
            min_importance=0.1
        )

        # Convert to ContextMemory objects and score
        context_memories = []
        for i, memory in enumerate(candidates):
            ctx_memory = self._convert_to_context_memory(memory)
            # Approximate semantic similarity based on ranking (higher ranked = more similar)
            semantic_similarity = max(0.2, 1.0 - (i * 0.1))  # Decay with rank
            ctx_memory.semantic_similarity = semantic_similarity
            self._score_memory_relevance(ctx_memory, query)
            context_memories.append(ctx_memory)

        # Filter and rank by relevance and token efficiency
        filtered_memories = [
            mem for mem in context_memories
            if mem.relevance_score >= self.min_relevance_threshold
        ]

        # Sort by relevance score (highest first)
        filtered_memories.sort(key=lambda x: x.relevance_score, reverse=True)

        # Select memories within token budget
        selected_memories = []
        used_tokens = 0

        for memory in filtered_memories:
            if used_tokens + memory.context_tokens <= memory_budget_tokens:
                selected_memories.append(memory)
                used_tokens += memory.context_tokens
                # Track usage for learning
                self.memory_usage_stats[memory.memory_id] += 1
                memory.usage_count += 1
                memory.last_used = datetime.now()

        return selected_memories

    def _create_context_query(self) -> str:
        """Create an effective search query from current context"""
        recent_context = self.conversation_context.get_recent_context(max_messages=3)
        active_tags = self.conversation_context.active_tags

        # Combine recent context with active topics
        query_parts = [recent_context]

        if active_tags:
            query_parts.append(" ".join(active_tags))

        # Add temporal context (what we were just working on)
        if self.conversation_context.current_topic:
            query_parts.append(self.conversation_context.current_topic)

        return " ".join(query_parts)

    def _convert_to_context_memory(self, memory) -> ContextMemory:
        """Convert MemoryEntry to ContextMemory with token estimation"""
        # Estimate token count (rough approximation)
        content_tokens = len(memory.content) // self.avg_chars_per_token
        tag_tokens = sum(len(tag) for tag in memory.tags) // self.avg_chars_per_token
        overhead_tokens = 20  # For formatting and metadata

        total_tokens = content_tokens + tag_tokens + overhead_tokens

        return ContextMemory(
            memory_id=memory.id,
            content=memory.content,
            importance=memory.importance,
            tags=memory.tags,
            context_tokens=total_tokens,
            usage_count=self.memory_usage_stats.get(memory.id, 0),
            last_used=memory.last_accessed,
            embedding=memory.embedding
        )

    def _score_memory_relevance(self, ctx_memory: ContextMemory, query: str):
        """Calculate multi-factor relevance score for memory"""
        # Get semantic similarity (should be set from retrieval ranking)
        semantic_similarity = getattr(ctx_memory, 'semantic_similarity', 0.5)

        # Factor 1: Tag matching with current context
        tag_match_score = 0.0
        if ctx_memory.tags and self.conversation_context.active_tags:
            matching_tags = set(ctx_memory.tags) & set(self.conversation_context.active_tags)
            tag_match_score = len(matching_tags) / len(self.conversation_context.active_tags) if self.conversation_context.active_tags else 0

        # Factor 2: Importance boost
        importance_boost = ctx_memory.importance

        # Factor 3: Usage frequency (memories used more are likely more relevant)
        usage_boost = min(0.3, ctx_memory.usage_count * 0.05)

        # Factor 4: Recency boost (recently used memories)
        recency_boost = 0.0
        if ctx_memory.last_used:
            days_since_used = (datetime.now() - ctx_memory.last_used).days
            recency_boost = max(0, 1 - days_since_used / 30) * 0.2  # 30-day decay

        # Factor 5: Content keyword matching
        query_words = set(query.lower().split())
        content_words = set(ctx_memory.content.lower().split())
        keyword_overlap = len(query_words & content_words) / len(query_words) if query_words else 0
        content_boost = keyword_overlap * 0.3

        # Combine factors with weights
        # Semantic similarity is the foundation, other factors are bonuses
        ctx_memory.relevance_score = (
            semantic_similarity * 0.5 +      # Semantic similarity (50%)
            importance_boost * 0.2 +         # Importance (20%)
            tag_match_score * 0.15 +         # Tag matching (15%)
            usage_boost * 0.1 +              # Usage frequency (10%)
            recency_boost * 0.03 +           # Recency (3%)
            content_boost * 0.02             # Content matching (2%)
        )

    def format_memories_for_context(self, memories: List[ContextMemory]) -> str:
        """Format memories for seamless context integration"""
        if not memories:
            return ""

        formatted_sections = []

        # Group by relevance levels
        high_relevance = [m for m in memories if m.relevance_score >= 0.7]
        medium_relevance = [m for m in memories if 0.4 <= m.relevance_score < 0.7]
        low_relevance = [m for m in memories if m.relevance_score < 0.4]

        if high_relevance:
            formatted_sections.append("🔥 HIGHLY RELEVANT MEMORIES:")
            formatted_sections.extend([mem.to_context_string() for mem in high_relevance])

        if medium_relevance:
            if formatted_sections:
                formatted_sections.append("")
            formatted_sections.append("📝 RELEVANT MEMORIES:")
            formatted_sections.extend([mem.to_context_string() for mem in medium_relevance])

        if low_relevance:
            if formatted_sections:
                formatted_sections.append("")
            formatted_sections.append("💭 ADDITIONAL CONTEXT:")
            formatted_sections.extend([mem.to_context_string() for mem in low_relevance])

        return "\n".join(formatted_sections)

    def get_integrated_context(self, current_query: str = "") -> str:
        """
        Get the complete integrated context including relevant memories

        Args:
            current_query: Optional specific query to focus memory retrieval

        Returns:
            Formatted context string with integrated memories
        """
        # Check cache first for repeated queries
        cache_key = self._generate_cache_key(current_query)
        if cache_key in self.query_cache:
            return self.query_cache[cache_key]

        # Update context with current query if provided
        if current_query:
            self.update_conversation_context("user", current_query)

        # Get relevant memories
        relevant_memories = self.get_relevant_memories_for_context()

        # Format for integration
        memory_context = self.format_memories_for_context(relevant_memories)

        # Cache the result for future queries
        if len(self.query_cache) >= self.cache_max_size:
            # Remove oldest cache entry (simple LRU approximation)
            oldest_key = next(iter(self.query_cache))
            del self.query_cache[oldest_key]

        self.query_cache[cache_key] = memory_context

        return memory_context

    def update_topic_associations(self, topic: str, memory_ids: List[str]):
        """Update topic-memory associations for better future retrieval"""
        for memory_id in memory_ids:
            if memory_id not in self.topic_memory_map[topic]:
                self.topic_memory_map[topic].append(memory_id)

    def get_memory_stats(self) -> Dict[str, Any]:
        """Get statistics about memory integration performance"""
        total_memories = len(self.memory_system.memories)
        used_memories = len([mid for mid in self.memory_usage_stats.keys() if mid in self.memory_system.memories])

        # Calculate performance metrics
        performance_metrics = self._calculate_performance_metrics()

        return {
            "total_system_memories": total_memories,
            "integrated_memories": used_memories,
            "integration_rate": used_memories / total_memories if total_memories > 0 else 0,
            "most_used_memories": sorted(self.memory_usage_stats.items(), key=lambda x: x[1], reverse=True)[:5],
            "active_topics": list(self.topic_memory_map.keys()),
            "context_window_usage": self.conversation_context.used_tokens / self.conversation_context.context_window_tokens,
            "performance": performance_metrics
        }


def create_memory_context_integrator(memory_system: Optional[VectorMemory] = None) -> MemoryContextIntegrator:
    """
    Factory function to create a memory context integrator

    Args:
        memory_system: Optional existing memory system, creates new one if None

    Returns:
        Configured MemoryContextIntegrator
    """
    if memory_system is None:
        memory_system = VectorMemory()

    # Estimate context window (conservative for most models)
    context_window = 8000  # tokens

    integrator = MemoryContextIntegrator(
        memory_system=memory_system,
        context_window_tokens=context_window
    )

    return integrator


# Example usage and testing
if __name__ == "__main__":
    # Create integrator
    integrator = create_memory_context_integrator()

    # Simulate conversation
    integrator.update_conversation_context("user", "How can I optimize my vector memory system for better performance?")

    # Get integrated context
    context = integrator.get_integrated_context()
    print("INTEGRATED CONTEXT:")
    print(context)

    # Get stats
    stats = integrator.get_memory_stats()
    print(f"\nINTEGRATION STATS: {stats}")