semantic.py

"""
Semantic Memory Implementation

Semantic memory stores facts, preferences, and structured knowledge that persist
across conversation threads. This is ideal for:
- User preferences (language, style, interests)
- Factual knowledge (user's name, location, occupation)
- Structured information (product catalog, FAQ)

Examples:
    >>> from langgraph.store.memory import InMemoryStore
    >>> store = InMemoryStore()
    >>> memory = SemanticMemory(store, user_id="alice")
    >>> memory.store_fact("name", "Alice Smith")
    >>> memory.store_preference("cuisine", "Italian")
    >>> facts = memory.recall_all()
"""

from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass
from datetime import datetime
from langgraph.store.base import BaseStore


@dataclass
class Fact:
    """A single fact stored in semantic memory."""
    key: str
    value: Any
    category: str
    source: str = "conversation"
    confidence: float = 1.0
    created_at: str = None
    updated_at: str = None

    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now().isoformat()
        if self.updated_at is None:
            self.updated_at = self.created_at

    def to_dict(self) -> Dict:
        """Convert fact to dictionary for storage."""
        return {
            "value": self.value,
            "category": self.category,
            "source": self.source,
            "confidence": self.confidence,
            "created_at": self.created_at,
            "updated_at": self.updated_at
        }

    @classmethod
    def from_dict(cls, key: str, data: Dict) -> "Fact":
        """Create fact from stored dictionary."""
        return cls(
            key=key,
            value=data["value"],
            category=data["category"],
            source=data.get("source", "conversation"),
            confidence=data.get("confidence", 1.0),
            created_at=data.get("created_at"),
            updated_at=data.get("updated_at")
        )


class SemanticMemory:
    """
    Semantic memory manager for long-term fact storage.

    Stores and retrieves facts, preferences, and structured knowledge using
    LangGraph's Store interface.

    Args:
        store: LangGraph BaseStore instance (e.g., InMemoryStore, PostgresStore)
        user_id: Unique identifier for the user
        namespace_prefix: Optional prefix for memory namespace

    Examples:
        >>> memory = SemanticMemory(store, user_id="user_123")
        >>> memory.store_fact("email", "alice@example.com", category="contact")
        >>> email = memory.recall("email")
    """

    def __init__(
        self,
        store: BaseStore,
        user_id: str,
        namespace_prefix: str = "semantic"
    ):
        self.store = store
        self.user_id = user_id
        self.namespace = (namespace_prefix, user_id)

    def store_fact(
        self,
        key: str,
        value: Any,
        category: str = "general",
        source: str = "conversation",
        confidence: float = 1.0
    ) -> None:
        """
        Store a fact in semantic memory.

        Args:
            key: Unique identifier for the fact
            value: The fact's value (can be any JSON-serializable type)
            category: Category for organizing facts (e.g., "preference", "contact")
            source: Source of the fact (e.g., "conversation", "profile", "inference")
            confidence: Confidence level (0.0 to 1.0)

        Examples:
            >>> memory.store_fact("name", "Alice")
            >>> memory.store_fact("age", 30, category="demographics")
            >>> memory.store_fact("likes_coffee", True, category="preference", confidence=0.8)
        """
        fact = Fact(
            key=key,
            value=value,
            category=category,
            source=source,
            confidence=confidence
        )

        self.store.put(
            namespace=self.namespace,
            key=key,
            value=fact.to_dict()
        )

    def store_preference(
        self,
        preference_type: str,
        value: Any,
        confidence: float = 1.0
    ) -> None:
        """
        Store a user preference (convenience method).

        Args:
            preference_type: Type of preference (e.g., "cuisine", "language", "theme")
            value: Preference value
            confidence: Confidence level

        Examples:
            >>> memory.store_preference("language", "Python")
            >>> memory.store_preference("theme", "dark_mode")
        """
        self.store_fact(
            key=f"pref_{preference_type}",
            value=value,
            category="preference",
            confidence=confidence
        )

    def recall(self, key: str) -> Optional[Any]:
        """
        Recall a specific fact by key.

        Args:
            key: Fact identifier

        Returns:
            The fact's value, or None if not found

        Examples:
            >>> name = memory.recall("name")
            >>> if name:
            ...     print(f"User's name is {name}")
        """
        try:
            result = self.store.get(namespace=self.namespace, key=key)
            if result:
                return result.value.get("value")
        except Exception:
            pass
        return None

    def recall_with_metadata(self, key: str) -> Optional[Fact]:
        """
        Recall a fact with full metadata.

        Args:
            key: Fact identifier

        Returns:
            Fact object with metadata, or None if not found

        Examples:
            >>> fact = memory.recall_with_metadata("name")
            >>> print(f"Confidence: {fact.confidence}, Source: {fact.source}")
        """
        try:
            result = self.store.get(namespace=self.namespace, key=key)
            if result:
                return Fact.from_dict(key, result.value)
        except Exception:
            pass
        return None

    def recall_by_category(self, category: str) -> Dict[str, Any]:
        """
        Recall all facts in a category.

        Args:
            category: Category name

        Returns:
            Dictionary of {key: value} for facts in the category

        Examples:
            >>> preferences = memory.recall_by_category("preference")
            >>> for key, value in preferences.items():
            ...     print(f"{key}: {value}")
        """
        all_facts = self.recall_all()
        return {
            key: fact.value
            for key, fact in all_facts.items()
            if fact.category == category
        }

    def recall_all(self) -> Dict[str, Fact]:
        """
        Recall all facts from semantic memory.

        Returns:
            Dictionary of {key: Fact} for all stored facts

        Examples:
            >>> all_facts = memory.recall_all()
            >>> for key, fact in all_facts.items():
            ...     print(f"{key}: {fact.value} ({fact.category})")
        """
        try:
            results = self.store.search(namespace=self.namespace)
            return {
                item.key: Fact.from_dict(item.key, item.value)
                for item in results
            }
        except Exception:
            return {}

    def search(self, query: str, limit: int = 5) -> List[Tuple[str, Any, float]]:
        """
        Search semantic memory using similarity search.

        Note: Requires store to be configured with embeddings for semantic search.

        Args:
            query: Search query
            limit: Maximum number of results

        Returns:
            List of (key, value, similarity_score) tuples

        Examples:
            >>> results = memory.search("What food does the user like?")
            >>> for key, value, score in results:
            ...     print(f"{key}: {value} (similarity: {score:.2f})")
        """
        try:
            results = self.store.search(
                namespace=self.namespace,
                query=query,
                limit=limit
            )
            return [
                (item.key, item.value.get("value"), getattr(item, "score", 1.0))
                for item in results
            ]
        except Exception:
            return []

    def update_fact(
        self,
        key: str,
        value: Any = None,
        confidence: float = None
    ) -> bool:
        """
        Update an existing fact.

        Args:
            key: Fact identifier
            value: New value (if None, keeps existing)
            confidence: New confidence (if None, keeps existing)

        Returns:
            True if fact was updated, False if not found

        Examples:
            >>> memory.update_fact("name", "Alice Johnson")
            >>> memory.update_fact("age", confidence=0.9)
        """
        fact = self.recall_with_metadata(key)
        if not fact:
            return False

        if value is not None:
            fact.value = value
        if confidence is not None:
            fact.confidence = confidence

        fact.updated_at = datetime.now().isoformat()

        self.store.put(
            namespace=self.namespace,
            key=key,
            value=fact.to_dict()
        )
        return True

    def delete_fact(self, key: str) -> bool:
        """
        Delete a fact from semantic memory.

        Args:
            key: Fact identifier

        Returns:
            True if deleted, False if not found

        Examples:
            >>> memory.delete_fact("old_preference")
        """
        try:
            self.store.delete(namespace=self.namespace, key=key)
            return True
        except Exception:
            return False

    def clear_category(self, category: str) -> int:
        """
        Delete all facts in a category.

        Args:
            category: Category to clear

        Returns:
            Number of facts deleted

        Examples:
            >>> count = memory.clear_category("temporary")
            >>> print(f"Deleted {count} temporary facts")
        """
        facts = self.recall_all()
        count = 0
        for key, fact in facts.items():
            if fact.category == category:
                self.delete_fact(key)
                count += 1
        return count

    def get_context_string(self, categories: Optional[List[str]] = None) -> str:
        """
        Get a formatted string of facts for use in LLM context.

        Args:
            categories: Optional list of categories to include. If None, includes all.

        Returns:
            Formatted string describing stored facts

        Examples:
            >>> context = memory.get_context_string(categories=["preference"])
            >>> print(context)
            # Output:
            # User Preferences:
            # - cuisine: Italian
            # - theme: dark_mode
        """
        all_facts = self.recall_all()

        if categories:
            facts = {k: v for k, v in all_facts.items() if v.category in categories}
        else:
            facts = all_facts

        if not facts:
            return "No stored information available."

        # Group by category
        by_category: Dict[str, List[Tuple[str, Any]]] = {}
        for key, fact in facts.items():
            if fact.category not in by_category:
                by_category[fact.category] = []
            by_category[fact.category].append((key, fact.value))

        # Format output
        lines = []
        for category, items in sorted(by_category.items()):
            lines.append(f"{category.title()}:")
            for key, value in items:
                lines.append(f"  - {key}: {value}")

        return "\n".join(lines)


# Example usage
if __name__ == "__main__":
    from langgraph.store.memory import InMemoryStore

    # Create store and memory
    store = InMemoryStore()
    memory = SemanticMemory(store, user_id="alice")

    # Store facts
    print("Storing facts...")
    memory.store_fact("name", "Alice Smith")
    memory.store_fact("email", "alice@example.com", category="contact")
    memory.store_preference("cuisine", "Italian")
    memory.store_preference("language", "Python")

    # Recall
    print(f"\nName: {memory.recall('name')}")
    print(f"Email: {memory.recall('email')}")

    # Recall by category
    print("\nPreferences:")
    prefs = memory.recall_by_category("preference")
    for key, value in prefs.items():
        print(f"  {key}: {value}")

    # Get context
    print("\nContext string:")
    print(memory.get_context_string())