headroom/proxy/memory_handler.py

"""Memory integration handler for the proxy server.

This module provides memory capabilities for the Headroom proxy:
1. MemoryHandler - Unified handler for memory operations
   - inject_tools() - Add memory tools to requests
   - search_and_format_context() - Search memories, format for injection
   - has_memory_tool_calls() - Detect memory tool usage in response
   - handle_memory_tool_calls() - Execute tools, return results

Usage:
    config = MemoryConfig(enabled=True, backend="local")
    handler = MemoryHandler(config)

    # Inject tools into request
    tools, was_injected = handler.inject_tools(existing_tools, "anthropic")

    # Search and inject context
    context = await handler.search_and_format_context(user_id, messages)

    # Handle tool calls in response
    if handler.has_memory_tool_calls(response, "anthropic"):
        results = await handler.handle_memory_tool_calls(response, user_id, "anthropic")
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal

if TYPE_CHECKING:
    from headroom.memory.backends.local import LocalBackend

logger = logging.getLogger(__name__)

# Memory tool names for detection
MEMORY_TOOL_NAMES = {"memory_save", "memory_search", "memory_update", "memory_delete"}


@dataclass
class MemoryConfig:
    """Configuration for memory handler."""

    enabled: bool = False
    backend: Literal["local", "qdrant-neo4j"] = "local"
    db_path: str = "headroom_memory.db"
    inject_tools: bool = True
    inject_context: bool = True
    top_k: int = 10
    min_similarity: float = 0.3
    # Qdrant+Neo4j config
    qdrant_host: str = "localhost"
    qdrant_port: int = 6333
    neo4j_uri: str = "neo4j://localhost:7687"
    neo4j_user: str = "neo4j"
    neo4j_password: str = "password"


class MemoryHandler:
    """Unified handler for memory operations in the proxy.

    Responsibilities:
    1. Initialize and manage memory backend
    2. Inject memory tools into requests
    3. Search and inject relevant memories as context
    4. Handle memory tool calls in responses
    """

    def __init__(self, config: MemoryConfig) -> None:
        self.config = config
        self._backend: LocalBackend | Any = None
        self._initialized = False
        self._memory_tools: list[dict[str, Any]] | None = None

    async def _ensure_initialized(self) -> None:
        """Lazy initialization of memory backend."""
        if self._initialized:
            return

        if not self.config.enabled:
            return

        if self.config.backend == "local":
            from headroom.memory.backends.local import LocalBackend, LocalBackendConfig

            backend_config = LocalBackendConfig(db_path=self.config.db_path)
            self._backend = LocalBackend(backend_config)
            await self._backend._ensure_initialized()
            logger.info(f"Memory: Initialized LocalBackend at {self.config.db_path}")

        elif self.config.backend == "qdrant-neo4j":
            try:
                from headroom.memory.backends.direct_mem0 import (
                    DirectMem0Adapter,
                    Mem0Config,
                )

                mem0_config = Mem0Config(
                    qdrant_host=self.config.qdrant_host,
                    qdrant_port=self.config.qdrant_port,
                    neo4j_uri=self.config.neo4j_uri,
                    neo4j_user=self.config.neo4j_user,
                    neo4j_password=self.config.neo4j_password,
                    enable_graph=True,
                )
                self._backend = DirectMem0Adapter(mem0_config)
                logger.info(
                    f"Memory: Initialized Qdrant+Neo4j backend "
                    f"({self.config.qdrant_host}:{self.config.qdrant_port})"
                )
            except ImportError as e:
                logger.error(
                    f"Memory: Failed to import qdrant-neo4j dependencies: {e}. "
                    "Install with: pip install mem0ai qdrant-client neo4j"
                )
                raise
        else:
            raise ValueError(f"Unknown memory backend: {self.config.backend}")

        self._initialized = True

    def _get_memory_tools(self) -> list[dict[str, Any]]:
        """Get memory tool definitions (cached)."""
        if self._memory_tools is None:
            from headroom.memory.tools import get_memory_tools_optimized

            self._memory_tools = get_memory_tools_optimized()
        return self._memory_tools

    def inject_tools(
        self,
        tools: list[dict[str, Any]] | None,
        provider: str = "anthropic",
    ) -> tuple[list[dict[str, Any]], bool]:
        """Inject memory tools into tools list.

        Args:
            tools: Existing tools list (may be None).
            provider: Provider for tool format ("anthropic" or "openai").

        Returns:
            Tuple of (updated_tools, was_injected).
        """
        if not self.config.inject_tools:
            return tools or [], False

        tools = list(tools) if tools else []

        # Check which tools are already present
        existing_names: set[str] = set()
        for tool in tools:
            name = tool.get("name") or tool.get("function", {}).get("name")
            if name:
                existing_names.add(name)

        # Add missing memory tools
        was_injected = False
        for memory_tool in self._get_memory_tools():
            tool_name = memory_tool["function"]["name"]
            if tool_name in existing_names:
                continue

            # Convert to provider format
            if provider == "anthropic":
                tools.append(
                    {
                        "name": tool_name,
                        "description": memory_tool["function"]["description"],
                        "input_schema": memory_tool["function"]["parameters"],
                    }
                )
            else:
                # OpenAI format
                tools.append(memory_tool)

            was_injected = True

        return tools, was_injected

    async def search_and_format_context(
        self,
        user_id: str,
        messages: list[dict[str, Any]],
    ) -> str | None:
        """Search memories and format as context injection.

        Args:
            user_id: User identifier for memory scoping.
            messages: Conversation messages (used to extract query).

        Returns:
            Formatted context string, or None if no relevant memories.
        """
        if not self.config.inject_context:
            return None

        await self._ensure_initialized()
        if not self._backend:
            return None

        # Extract query from last user message
        query = self._extract_user_query(messages)
        if not query:
            logger.debug("Memory: No user query found for context search")
            return None

        try:
            # Search memories
            results = await self._backend.search_memories(
                query=query,
                user_id=user_id,
                top_k=self.config.top_k,
                include_related=True,
            )

            if not results:
                logger.debug(f"Memory: No memories found for user {user_id}")
                return None

            # Filter by minimum similarity
            filtered_results = [r for r in results if r.score >= self.config.min_similarity]

            if not filtered_results:
                logger.debug(
                    f"Memory: {len(results)} memories found but none above threshold "
                    f"{self.config.min_similarity}"
                )
                return None

            # Format as context
            memory_lines = []
            for i, result in enumerate(filtered_results, 1):
                memory_lines.append(f"{i}. {result.memory.content}")
                if hasattr(result, "related_entities") and result.related_entities:
                    entities_str = ", ".join(result.related_entities[:3])
                    memory_lines.append(f"   (Related: {entities_str})")

            context = f"""## Relevant Memories for This User

The following information was previously saved about this user:

{chr(10).join(memory_lines)}

Use this context to provide personalized and contextually relevant responses."""

            logger.info(
                f"Memory: Injecting {len(filtered_results)} memories "
                f"({len(context)} chars) for user {user_id}"
            )
            return context

        except Exception as e:
            logger.warning(f"Memory: Search failed for user {user_id}: {e}")
            return None

    def _extract_user_query(self, messages: list[dict[str, Any]]) -> str:
        """Extract the user query from the last user message."""
        for msg in reversed(messages):
            if msg.get("role") != "user":
                continue

            content = msg.get("content", "")

            if isinstance(content, str):
                return content[:500]  # Limit query length

            if isinstance(content, list):
                for block in content:
                    if isinstance(block, dict) and block.get("type") == "text":
                        text = str(block.get("text", ""))
                        if text:
                            return text[:500]

        return ""

    def has_memory_tool_calls(
        self,
        response: dict[str, Any],
        provider: str = "anthropic",
    ) -> bool:
        """Check if response contains memory tool calls."""
        tool_calls = self._extract_tool_calls(response, provider)
        for tc in tool_calls:
            name = tc.get("name") or tc.get("function", {}).get("name")
            if name in MEMORY_TOOL_NAMES:
                return True
        return False

    def _extract_tool_calls(
        self,
        response: dict[str, Any],
        provider: str,
    ) -> list[dict[str, Any]]:
        """Extract tool calls from response based on provider format."""
        if provider == "anthropic":
            content = response.get("content", [])
            if isinstance(content, list):
                return [block for block in content if block.get("type") == "tool_use"]
            return []

        elif provider == "openai":
            choices = response.get("choices", [])
            if choices:
                message = choices[0].get("message", {})
                return list(message.get("tool_calls", []) or [])
            return []

        return []

    async def handle_memory_tool_calls(
        self,
        response: dict[str, Any],
        user_id: str,
        provider: str = "anthropic",
    ) -> list[dict[str, Any]]:
        """Execute memory tool calls and return results.

        Args:
            response: The API response containing tool calls.
            user_id: User identifier for memory operations.
            provider: Provider format ("anthropic" or "openai").

        Returns:
            List of tool results in provider format.
        """
        await self._ensure_initialized()
        if not self._backend:
            return []

        tool_calls = self._extract_tool_calls(response, provider)
        results: list[dict[str, Any]] = []

        for tc in tool_calls:
            tool_name = tc.get("name") or tc.get("function", {}).get("name")
            if tool_name not in MEMORY_TOOL_NAMES:
                continue

            tool_id = tc.get("id", "")

            # Parse input data
            if provider == "anthropic":
                input_data = tc.get("input", {})
            else:
                args_str = tc.get("function", {}).get("arguments", "{}")
                try:
                    input_data = json.loads(args_str)
                except json.JSONDecodeError:
                    input_data = {}

            # Execute the tool
            result_content = await self._execute_memory_tool(tool_name, input_data, user_id)

            # Format result based on provider
            if provider == "anthropic":
                results.append(
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_id,
                        "content": result_content,
                    }
                )
            else:
                results.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_id,
                        "content": result_content,
                    }
                )

            logger.info(f"Memory: Executed {tool_name} for user {user_id}")

        return results

    async def _execute_memory_tool(
        self,
        tool_name: str,
        input_data: dict[str, Any],
        user_id: str,
    ) -> str:
        """Execute a memory tool and return result string."""
        try:
            if tool_name == "memory_save":
                return await self._execute_save(input_data, user_id)
            elif tool_name == "memory_search":
                return await self._execute_search(input_data, user_id)
            elif tool_name == "memory_update":
                return await self._execute_update(input_data, user_id)
            elif tool_name == "memory_delete":
                return await self._execute_delete(input_data, user_id)
            else:
                return json.dumps({"error": f"Unknown tool: {tool_name}"})

        except Exception as e:
            logger.error(f"Memory: Tool {tool_name} failed: {e}")
            return json.dumps({"status": "error", "error": str(e)})

    async def _execute_save(self, input_data: dict[str, Any], user_id: str) -> str:
        """Execute memory_save tool."""
        content = input_data.get("content", "")
        if not content:
            return json.dumps({"status": "error", "error": "content is required"})

        # Extract parameters
        importance = input_data.get("importance", 0.5)
        facts = input_data.get("facts")
        entities = input_data.get("entities")
        extracted_entities = input_data.get("extracted_entities")
        relationships = input_data.get("relationships")
        extracted_relationships = input_data.get("extracted_relationships")

        # Call backend
        memory = await self._backend.save_memory(
            content=content,
            user_id=user_id,
            importance=importance,
            facts=facts,
            entities=entities,
            extracted_entities=extracted_entities,
            relationships=relationships,
            extracted_relationships=extracted_relationships,
        )

        return json.dumps(
            {
                "status": "saved",
                "memory_id": memory.id,
                "content": (
                    memory.content[:100] + "..." if len(memory.content) > 100 else memory.content
                ),
            }
        )

    async def _execute_search(self, input_data: dict[str, Any], user_id: str) -> str:
        """Execute memory_search tool."""
        query = input_data.get("query", "")
        if not query:
            return json.dumps({"status": "error", "error": "query is required"})

        top_k = input_data.get("top_k", 10)
        include_related = input_data.get("include_related", True)
        entities_filter = input_data.get("entities")

        results = await self._backend.search_memories(
            query=query,
            user_id=user_id,
            top_k=top_k,
            include_related=include_related,
            entities=entities_filter,
        )

        return json.dumps(
            {
                "status": "found",
                "count": len(results),
                "memories": [
                    {
                        "id": r.memory.id,
                        "content": r.memory.content,
                        "score": round(r.score, 3),
                        "entities": (
                            r.related_entities[:5]
                            if hasattr(r, "related_entities") and r.related_entities
                            else []
                        ),
                    }
                    for r in results
                ],
            }
        )

    async def _execute_update(self, input_data: dict[str, Any], user_id: str) -> str:
        """Execute memory_update tool."""
        memory_id = input_data.get("memory_id", "")
        new_content = input_data.get("new_content", "")

        if not memory_id:
            return json.dumps({"status": "error", "error": "memory_id is required"})
        if not new_content:
            return json.dumps({"status": "error", "error": "new_content is required"})

        reason = input_data.get("reason")

        # Check if backend has update_memory method
        if hasattr(self._backend, "update_memory"):
            memory = await self._backend.update_memory(
                memory_id=memory_id,
                new_content=new_content,
                reason=reason,
                user_id=user_id,
            )
            return json.dumps({"status": "updated", "memory_id": memory.id})
        else:
            # Fallback: delete old, save new
            await self._backend.delete_memory(memory_id)
            memory = await self._backend.save_memory(
                content=new_content,
                user_id=user_id,
                importance=0.5,
            )
            return json.dumps(
                {
                    "status": "updated",
                    "memory_id": memory.id,
                    "note": "Replaced via delete+save",
                }
            )

    async def _execute_delete(self, input_data: dict[str, Any], user_id: str) -> str:
        """Execute memory_delete tool."""
        memory_id = input_data.get("memory_id", "")
        if not memory_id:
            return json.dumps({"status": "error", "error": "memory_id is required"})

        deleted = await self._backend.delete_memory(memory_id)

        return json.dumps(
            {
                "status": "deleted" if deleted else "not_found",
                "memory_id": memory_id,
            }
        )

    async def close(self) -> None:
        """Close the memory backend."""
        if self._backend and hasattr(self._backend, "close"):
            await self._backend.close()
        self._backend = None
        self._initialized = False
        logger.info("Memory: Handler closed")