tatlock/src/core/memory_service.py

"""
Memory service for direct key-based access.

Provides fast, LLM-free access to user memories for:
- Known-key lookups (location, timezone, preferences)
- Session context (current topic, recent entities)
- Structured storage (explicit user instructions)

This is the "direct access layer" - no LLM interpretation.
For semantic/fuzzy queries, use the Memory Agent instead.

Usage:
    from src.core.memory_service import memory_service

    # Get user's location (fast, no LLM)
    location = await memory_service.get_profile("location")

    # Set a preference
    await memory_service.set_preference("temperature_unit", "celsius")

    # Get session context
    ctx = await memory_service.get_session_context(conversation_id)
"""
from datetime import datetime, timezone
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field

from .config import config
from .context import get_user, get_conversation_id
from .embeddings import get_embedding_client
from .logging_config import get_logger
from .memory_cache import get_memory_cache
from .multi_tenancy import get_memory_collection_name
from .qdrant import get_qdrant_client

logger = get_logger(__name__)


class MemoryType(str, Enum):
    """Types of memories stored in Qdrant."""
    USER_PROFILE = "user_profile"      # Name, location, timezone
    PREFERENCE = "preference"          # Units, language, theme
    LEARNED_FACT = "learned_fact"      # "My car is a Tesla"


class MemoryRecord(BaseModel):
    """A memory record stored in Qdrant."""
    id: str
    type: MemoryType
    key: str                           # e.g., "location", "timezone", "car"
    value: str                         # The actual content
    keywords: list[str] = Field(default_factory=list)
    importance: float = 0.5            # 0.0 - 1.0
    source: str = "explicit"           # "explicit" | "inferred" | "conversation"
    created_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    updated_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat())


class MemoryService:
    """
    Direct access to user memories without LLM overhead.

    Use this for:
    - Known-key lookups: get_profile("location"), get_preference("units")
    - Explicit storage: set_preference("theme", "dark")
    - Session context: get_session_context(), update_session_context()

    Do NOT use for:
    - Fuzzy queries: "What car do I drive?" → Use Memory Agent
    - Semantic recall: "What did I mention about X?" → Use Memory Agent
    """

    def __init__(self):
        """Initialize memory service with lazy client loading."""
        self._qdrant = None
        self._embedding = None
        self._cache = None

    @property
    def qdrant(self):
        """Lazy-load Qdrant client."""
        if self._qdrant is None:
            self._qdrant = get_qdrant_client()
        return self._qdrant

    @property
    def embedding(self):
        """Lazy-load embedding client."""
        if self._embedding is None:
            self._embedding = get_embedding_client()
        return self._embedding

    @property
    def cache(self):
        """Lazy-load Redis cache."""
        if self._cache is None:
            self._cache = get_memory_cache()
        return self._cache

    # =========================================================================
    # Profile Methods (user_profile type)
    # =========================================================================

    async def get_profile(self, key: str, user: str | None = None) -> str | None:
        """
        Get a user profile value by key.

        Args:
            key: Profile key (e.g., "location", "timezone", "name")
            user: User ID (defaults to current request context)

        Returns:
            Profile value or None if not found

        Example:
            >>> location = await memory_service.get_profile("location")
            >>> location
            "Amsterdam, Netherlands"
        """
        user = user or get_user()
        return await self._get_memory(user, MemoryType.USER_PROFILE, key)

    async def set_profile(
        self,
        key: str,
        value: str,
        user: str | None = None,
        keywords: list[str] | None = None,
    ) -> bool:
        """
        Set a user profile value.

        Args:
            key: Profile key (e.g., "location", "timezone")
            value: Profile value
            user: User ID (defaults to current request context)
            keywords: Optional keywords for semantic search

        Returns:
            True if successful

        Example:
            >>> await memory_service.set_profile("location", "Amsterdam, Netherlands")
            True
        """
        user = user or get_user()
        return await self._set_memory(
            user=user,
            memory_type=MemoryType.USER_PROFILE,
            key=key,
            value=value,
            keywords=keywords or [key],
            importance=0.9,  # Profile data is important
        )

    # =========================================================================
    # Preference Methods (preference type)
    # =========================================================================

    async def get_preference(self, key: str, user: str | None = None) -> str | None:
        """
        Get a user preference by key.

        Args:
            key: Preference key (e.g., "temperature_unit", "language", "theme")
            user: User ID (defaults to current request context)

        Returns:
            Preference value or None if not found

        Example:
            >>> units = await memory_service.get_preference("temperature_unit")
            >>> units
            "celsius"
        """
        user = user or get_user()
        return await self._get_memory(user, MemoryType.PREFERENCE, key)

    async def set_preference(
        self,
        key: str,
        value: str,
        user: str | None = None,
    ) -> bool:
        """
        Set a user preference.

        Args:
            key: Preference key
            value: Preference value
            user: User ID (defaults to current request context)

        Returns:
            True if successful

        Example:
            >>> await memory_service.set_preference("theme", "dark")
            True
        """
        user = user or get_user()
        return await self._set_memory(
            user=user,
            memory_type=MemoryType.PREFERENCE,
            key=key,
            value=value,
            keywords=[key, "preference"],
            importance=0.7,
        )

    async def get_all_preferences(self, user: str | None = None) -> dict[str, str]:
        """
        Get all preferences for a user.

        Returns:
            Dict of key -> value for all preferences
        """
        user = user or get_user()
        memories = await self._get_all_by_type(user, MemoryType.PREFERENCE)
        return {m["key"]: m["value"] for m in memories}

    # =========================================================================
    # Learned Facts (learned_fact type) - for direct storage only
    # =========================================================================

    async def store_fact(
        self,
        key: str,
        value: str,
        user: str | None = None,
        keywords: list[str] | None = None,
        importance: float = 0.5,
        source: str = "explicit",
    ) -> bool:
        """
        Store a learned fact about the user.

        Use this for explicit user statements like:
        - "Remember that my car is a Tesla"
        - "I work at Acme Corp"

        For semantic extraction from conversation, use the Memory Agent.

        Args:
            key: Fact identifier (e.g., "car", "employer")
            value: The fact content
            user: User ID
            keywords: Keywords for semantic search
            importance: 0.0-1.0 importance score
            source: "explicit" | "inferred" | "conversation"

        Returns:
            True if successful
        """
        user = user or get_user()
        return await self._set_memory(
            user=user,
            memory_type=MemoryType.LEARNED_FACT,
            key=key,
            value=value,
            keywords=keywords or [key],
            importance=importance,
            source=source,
        )

    async def get_fact(self, key: str, user: str | None = None) -> str | None:
        """
        Get a specific fact by key.

        For semantic/fuzzy queries, use the Memory Agent.
        """
        user = user or get_user()
        return await self._get_memory(user, MemoryType.LEARNED_FACT, key)

    # =========================================================================
    # Session Context (Redis-backed, 24h TTL)
    # =========================================================================

    async def get_session_context(
        self,
        conversation_id: str | None = None,
        user: str | None = None,
    ) -> dict[str, Any] | None:
        """
        Get session context for current conversation.

        Args:
            conversation_id: Conversation ID (defaults to current context)
            user: User ID (defaults to current context)

        Returns:
            Session context dict or None
        """
        user = user or get_user()
        conversation_id = conversation_id or get_conversation_id()

        if not conversation_id:
            return None

        return await self.cache.get_session_context(user, conversation_id)

    async def set_session_context(
        self,
        context: dict[str, Any],
        conversation_id: str | None = None,
        user: str | None = None,
    ) -> bool:
        """
        Set session context for current conversation.

        Args:
            context: Context data to store
            conversation_id: Conversation ID
            user: User ID

        Returns:
            True if successful
        """
        user = user or get_user()
        conversation_id = conversation_id or get_conversation_id()

        if not conversation_id:
            logger.warning("memory_service_no_conversation_id")
            return False

        return await self.cache.set_session_context(user, conversation_id, context)

    async def update_session_context(
        self,
        updates: dict[str, Any],
        conversation_id: str | None = None,
        user: str | None = None,
    ) -> bool:
        """
        Update session context (merge with existing).

        Args:
            updates: Fields to update
            conversation_id: Conversation ID
            user: User ID

        Returns:
            True if successful
        """
        user = user or get_user()
        conversation_id = conversation_id or get_conversation_id()

        if not conversation_id:
            return False

        return await self.cache.update_session_context(user, conversation_id, updates)

    async def get_recent_entities(
        self,
        conversation_id: str | None = None,
        user: str | None = None,
    ) -> list[str]:
        """
        Get recently mentioned entities in conversation.

        Returns:
            List of entity names
        """
        user = user or get_user()
        conversation_id = conversation_id or get_conversation_id()

        if not conversation_id:
            return []

        return await self.cache.get_recent_entities(user, conversation_id)

    async def add_recent_entities(
        self,
        entities: list[str],
        conversation_id: str | None = None,
        user: str | None = None,
    ) -> bool:
        """
        Add entities to recent entities set.

        Args:
            entities: Entity names to add
            conversation_id: Conversation ID
            user: User ID

        Returns:
            True if successful
        """
        user = user or get_user()
        conversation_id = conversation_id or get_conversation_id()

        if not conversation_id:
            return False

        return await self.cache.add_recent_entities(user, conversation_id, entities)

    # =========================================================================
    # Bulk / Pre-fetch Methods (for Steward)
    # =========================================================================

    async def prefetch_context(
        self,
        user: str | None = None,
        include_profile: bool = True,
        include_preferences: bool = True,
        profile_keys: list[str] | None = None,
    ) -> dict[str, Any]:
        """
        Pre-fetch commonly needed context for Steward.

        This is the main entry point for Steward to get user context
        before analyzing a request.

        Args:
            user: User ID
            include_profile: Include profile data
            include_preferences: Include preferences
            profile_keys: Specific profile keys to fetch (None = common ones)

        Returns:
            Dict with profile and preferences data

        Example:
            >>> ctx = await memory_service.prefetch_context()
            >>> ctx
            {
                "profile": {"location": "Amsterdam", "timezone": "Europe/Amsterdam"},
                "preferences": {"temperature_unit": "celsius"}
            }
        """
        user = user or get_user()
        result: dict[str, Any] = {}

        if include_profile:
            profile_keys = profile_keys or ["location", "timezone", "name"]
            profile = {}
            for key in profile_keys:
                value = await self.get_profile(key, user)
                if value:
                    profile[key] = value
            if profile:
                result["profile"] = profile

        if include_preferences:
            preferences = await self.get_all_preferences(user)
            if preferences:
                result["preferences"] = preferences

        logger.debug(
            "memory_service_prefetch",
            user=user,
            profile_keys=list(result.get("profile", {}).keys()),
            preference_keys=list(result.get("preferences", {}).keys()),
        )

        return result

    # =========================================================================
    # Internal Methods
    # =========================================================================

    async def _get_memory(
        self,
        user: str,
        memory_type: MemoryType,
        key: str,
    ) -> str | None:
        """Get a memory by type and key (exact match)."""
        collection = get_memory_collection_name(user)

        try:
            # Search with filter for exact type + key match
            # We use a dummy vector since we're filtering by payload
            results = self.qdrant._client.scroll(
                collection_name=collection,
                scroll_filter={
                    "must": [
                        {"key": "type", "match": {"value": memory_type.value}},
                        {"key": "key", "match": {"value": key}},
                    ]
                },
                limit=1,
                with_payload=True,
                with_vectors=False,
            )

            points, _ = results
            if points:
                return points[0].payload.get("value")
            return None

        except Exception as e:
            logger.warning(
                "memory_service_get_failed",
                user=user,
                type=memory_type.value,
                key=key,
                error=str(e),
            )
            return None

    async def _set_memory(
        self,
        user: str,
        memory_type: MemoryType,
        key: str,
        value: str,
        keywords: list[str],
        importance: float = 0.5,
        source: str = "explicit",
    ) -> bool:
        """Set a memory (upsert by type + key)."""
        try:
            # Generate embedding for semantic search
            embedding = await self.embedding.embed(f"{key}: {value}")
            if not embedding:
                logger.error("memory_service_embedding_failed", key=key)
                return False

            # Create memory ID from type + key for idempotent upserts
            memory_id = f"{memory_type.value}:{key}"

            payload = {
                "type": memory_type.value,
                "key": key,
                "value": value,
                "keywords": keywords,
                "importance": importance,
                "source": source,
                "updated_at": datetime.now(timezone.utc).isoformat(),
            }

            result = await self.qdrant.upsert_memory(
                user=user,
                memory_id=memory_id,
                vector=embedding,
                payload=payload,
            )

            if result:
                logger.debug(
                    "memory_service_set",
                    user=user,
                    type=memory_type.value,
                    key=key,
                )
                return True
            return False

        except Exception as e:
            logger.error(
                "memory_service_set_failed",
                user=user,
                type=memory_type.value,
                key=key,
                error=str(e),
            )
            return False

    async def _get_all_by_type(
        self,
        user: str,
        memory_type: MemoryType,
        limit: int = 100,
    ) -> list[dict[str, Any]]:
        """Get all memories of a specific type."""
        collection = get_memory_collection_name(user)

        try:
            results = self.qdrant._client.scroll(
                collection_name=collection,
                scroll_filter={
                    "must": [
                        {"key": "type", "match": {"value": memory_type.value}},
                    ]
                },
                limit=limit,
                with_payload=True,
                with_vectors=False,
            )

            points, _ = results
            return [p.payload for p in points]

        except Exception as e:
            logger.warning(
                "memory_service_get_all_failed",
                user=user,
                type=memory_type.value,
                error=str(e),
            )
            return []

    async def delete_memory(
        self,
        key: str,
        memory_type: MemoryType,
        user: str | None = None,
    ) -> bool:
        """
        Delete a specific memory.

        Args:
            key: Memory key
            memory_type: Type of memory
            user: User ID

        Returns:
            True if deleted
        """
        user = user or get_user()
        memory_id = f"{memory_type.value}:{key}"

        return await self.qdrant.delete_memory(user, memory_id)


# Global service instance
memory_service = MemoryService()