tatlock/src/agents/biographer/tools.py

"""
Biographer tools for PydanticAI agent.

These tools enable The Biographer to record and recall the user's story:
- recall_semantic: Find memories by meaning/concept
- store_insight: Record new facts about the user
- list_memories: Browse recorded memories by type
- forget_memory: Remove specific memories

For direct key-based access (get/set profile, preferences),
use memory_service directly - these tools are for semantic queries.
"""

from src.core.context import get_user
from src.core.embeddings import get_embedding_client
from src.core.logging_config import get_logger
from src.core.memory_service import MemoryType, memory_service
from src.core.qdrant import get_qdrant_client

logger = get_logger(__name__)


# ============================================================================
# Semantic Recall
# ============================================================================


async def recall_semantic(
    query: str,
    memory_type: str = "",
    limit: int = 5,
) -> str:
    """
    Search memories by semantic similarity.

    Use this to find memories that are conceptually related to
    the query, even if exact words don't match. This is the main
    tool for answering questions like "What car do I drive?" or
    "What did I mention about my job?"

    Args:
        query: Natural language query to search for
        memory_type: Optional filter: "user_profile", "preference", "learned_fact"
        limit: Maximum memories to return (default: 5)

    Returns:
        Matching memories with their content and relevance scores

    Examples:
        recall_semantic("What is my car?")
        recall_semantic("work preferences", memory_type="preference")
        recall_semantic("family members")
    """
    try:
        user = get_user()
        embedding_client = get_embedding_client()
        qdrant = get_qdrant_client()

        # Generate embedding for query
        query_vector = await embedding_client.embed(query)
        if not query_vector:
            return "Unable to process query - embedding generation failed"

        # Search memories
        results = await qdrant.search_memories(
            user=user,
            query_vector=query_vector,
            limit=limit,
            memory_type=memory_type if memory_type else None,
        )

        if not results:
            return f"No memories found related to '{query}'"

        output_parts = [f"## Memories matching: {query}\n"]

        for i, memory in enumerate(results, 1):
            mem_type = memory.get("type", "unknown")
            key = memory.get("key", "")
            value = memory.get("value", "")
            score = memory.get("score", 0.0)
            source = memory.get("source", "unknown")

            type_icon = {
                "user_profile": "👤",
                "preference": "⚙️",
                "learned_fact": "💡",
            }.get(mem_type, "📝")

            output_parts.append(f"{i}. {type_icon} **{key}** (relevance: {score:.2f})")
            output_parts.append(f"   {value}")
            output_parts.append(f"   _Type: {mem_type}, Source: {source}_")
            output_parts.append("")

        logger.info(
            "memory_recall_semantic",
            query=query[:50],
            result_count=len(results),
            user=user,
        )

        return "\n".join(output_parts)

    except Exception as e:
        logger.error("memory_recall_semantic_error", error=str(e), query=query[:50])
        return f"Error searching memories: {str(e)}"


# ============================================================================
# Store Memory
# ============================================================================


async def store_insight(
    key: str,
    value: str,
    importance: float = 0.5,
) -> str:
    """
    Store a new insight or learned fact about the user.

    Use this when:
    - User explicitly asks to remember something
    - User shares personal information worth remembering
    - You learn something from conversation that should persist

    The memory will be stored with vector embedding for semantic search
    and can be recalled later using recall_semantic.

    Args:
        key: Short identifier for the memory (e.g., "car", "employer", "pet")
        value: The actual information to remember
        importance: How important is this? 0.0 (trivial) to 1.0 (critical)

    Returns:
        Confirmation of stored memory

    Examples:
        store_insight("car", "User drives a Tesla Model 3")
        store_insight("employer", "Works at Acme Corp as software engineer", importance=0.8)
    """
    try:
        # Auto-generate keywords from key and value
        keywords = [key]
        words = value.lower().split()
        keywords.extend([w for w in words if len(w) > 4][:5])

        success = await memory_service.store_fact(
            key=key,
            value=value,
            keywords=keywords,
            importance=importance,
            source="conversation",
        )

        if success:
            output_parts = [
                "## Memory Stored",
                f"**Key:** {key}",
                f"**Value:** {value}",
                f"**Keywords:** {', '.join(keywords)}",
                f"**Importance:** {importance:.1f}",
                "",
                "_Memory is now searchable via semantic recall._",
            ]

            logger.info(
                "memory_store_insight",
                key=key,
                importance=importance,
                user=get_user(),
            )

            return "\n".join(output_parts)
        else:
            return f"Failed to store memory for key '{key}'"

    except Exception as e:
        logger.error("memory_store_insight_error", error=str(e), key=key)
        return f"Error storing memory: {str(e)}"


async def update_profile(
    key: str,
    value: str,
) -> str:
    """
    Update user profile information.

    Use this for core identity information:
    - name, location, timezone
    - language preferences
    - occupation

    Profile data has high importance and is used for context
    by the Steward during request analysis.

    Args:
        key: Profile field (e.g., "name", "location", "timezone")
        value: The value to set

    Returns:
        Confirmation of profile update

    Examples:
        update_profile("location", "Amsterdam, Netherlands")
        update_profile("timezone", "Europe/Amsterdam")
        update_profile("name", "John")
    """
    try:
        success = await memory_service.set_profile(
            key=key,
            value=value,
            keywords=[key, "profile"],
        )

        if success:
            output_parts = [
                "## Profile Updated",
                f"**{key}:** {value}",
                "",
                "_Profile data is automatically included in context._",
            ]

            logger.info(
                "memory_update_profile",
                key=key,
                user=get_user(),
            )

            return "\n".join(output_parts)
        else:
            return f"Failed to update profile field '{key}'"

    except Exception as e:
        logger.error("memory_update_profile_error", error=str(e), key=key)
        return f"Error updating profile: {str(e)}"


async def update_preference(
    key: str,
    value: str,
) -> str:
    """
    Update user preferences.

    Use this for settings and preferences:
    - temperature_unit (celsius/fahrenheit)
    - distance_unit (metric/imperial)
    - theme, language, etc.

    Preferences are used by agents to customize responses.

    Args:
        key: Preference name (e.g., "temperature_unit", "theme")
        value: Preference value

    Returns:
        Confirmation of preference update

    Examples:
        update_preference("temperature_unit", "celsius")
        update_preference("distance_unit", "metric")
        update_preference("theme", "dark")
    """
    try:
        success = await memory_service.set_preference(
            key=key,
            value=value,
        )

        if success:
            output_parts = [
                "## Preference Updated",
                f"**{key}:** {value}",
                "",
                "_Preference will be applied to future responses._",
            ]

            logger.info(
                "memory_update_preference",
                key=key,
                user=get_user(),
            )

            return "\n".join(output_parts)
        else:
            return f"Failed to update preference '{key}'"

    except Exception as e:
        logger.error("memory_update_preference_error", error=str(e), key=key)
        return f"Error updating preference: {str(e)}"


# ============================================================================
# List Memories
# ============================================================================


async def list_memories(
    memory_type: str = "learned_fact",
    limit: int = 20,
) -> str:
    """
    List stored memories of a specific type.

    Use this to browse what's stored in memory without
    a specific search query.

    Args:
        memory_type: Type to list: "user_profile", "preference", "learned_fact"
        limit: Maximum memories to return (default: 20)

    Returns:
        List of memories with their keys and values

    Examples:
        list_memories("user_profile")
        list_memories("preference")
        list_memories("learned_fact", limit=10)
    """
    try:
        user = get_user()
        qdrant = get_qdrant_client()

        # Convert string to MemoryType
        try:
            MemoryType(memory_type)  # validated for its ValueError; the value is unused
        except ValueError:
            return f"Invalid memory type '{memory_type}'. Use: user_profile, preference, or learned_fact"

        # Get all memories of type
        results = qdrant._client.scroll(
            collection_name=f"memories_{user}",
            scroll_filter={
                "must": [
                    {"key": "type", "match": {"value": memory_type}},
                ]
            },
            limit=limit,
            with_payload=True,
            with_vectors=False,
        )

        points, _ = results
        if not points:
            return f"No {memory_type} memories found"

        type_icon = {
            "user_profile": "👤",
            "preference": "⚙️",
            "learned_fact": "💡",
        }.get(memory_type, "📝")

        output_parts = [f"## {type_icon} {memory_type.replace('_', ' ').title()} Memories\n"]

        for point in points:
            payload = point.payload
            key = payload.get("key", "unknown")
            value = payload.get("value", "")
            importance = payload.get("importance", 0.5)

            output_parts.append(f"- **{key}**: {value}")
            if importance > 0.7:
                output_parts.append(f"  _(importance: {importance:.1f})_")

        logger.info(
            "memory_list",
            memory_type=memory_type,
            count=len(points),
            user=user,
        )

        return "\n".join(output_parts)

    except Exception as e:
        logger.error("memory_list_error", error=str(e), memory_type=memory_type)
        return f"Error listing memories: {str(e)}"


# ============================================================================
# Forget Memory
# ============================================================================


async def forget_memory(
    key: str,
    memory_type: str = "learned_fact",
) -> str:
    """
    Remove a specific memory.

    Use this when:
    - User asks to forget something
    - Information is outdated or incorrect
    - Privacy concerns

    Args:
        key: Key of the memory to forget
        memory_type: Type of memory: "user_profile", "preference", "learned_fact"

    Returns:
        Confirmation of deletion

    Examples:
        forget_memory("old_car")
        forget_memory("location", memory_type="user_profile")
        forget_memory("theme", memory_type="preference")
    """
    try:
        # Convert string to MemoryType
        try:
            mem_type = MemoryType(memory_type)
        except ValueError:
            return f"Invalid memory type '{memory_type}'. Use: user_profile, preference, or learned_fact"

        success = await memory_service.delete_memory(
            key=key,
            memory_type=mem_type,
        )

        if success:
            output_parts = [
                "## Memory Forgotten",
                f"**Key:** {key}",
                f"**Type:** {memory_type}",
                "",
                "_Memory has been removed._",
            ]

            logger.info(
                "memory_forget",
                key=key,
                memory_type=memory_type,
                user=get_user(),
            )

            return "\n".join(output_parts)
        else:
            return f"Memory '{key}' not found or already deleted"

    except Exception as e:
        logger.error("memory_forget_error", error=str(e), key=key)
        return f"Error forgetting memory: {str(e)}"


# ============================================================================
# Tool Collection for Registration
# ============================================================================

# All tools available to The Biographer
BIOGRAPHER_TOOLS = [
    # Recall
    recall_semantic,
    list_memories,
    # Record
    store_insight,
    update_profile,
    update_preference,
    # Manage
    forget_memory,
]