fleet-memory/hindsight-integrations/litellm/hindsight_litellm/callbacks.py

"""LiteLLM callback handlers for Hindsight memory integration.

This module implements LiteLLM's CustomLogger interface to intercept
LLM calls and integrate with Hindsight for memory injection and storage.

Uses direct HTTP calls via requests/httpx to avoid async event loop conflicts
when the hindsight_client's async methods are called from LiteLLM callbacks.
"""

import logging
import fnmatch
import hashlib
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Set
import asyncio
import threading
import concurrent.futures

from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import ModelResponse

from .config import get_config, is_configured, HindsightConfig, MemoryInjectionMode

# Use requests for sync HTTP calls to avoid async event loop issues
try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

try:
    import httpx
    HAS_HTTPX = True
except ImportError:
    HAS_HTTPX = False


logger = logging.getLogger(__name__)

# Thread pool for running async operations in background
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=4, thread_name_prefix="hindsight-")


class HindsightCallback(CustomLogger):
    """LiteLLM custom logger that integrates with Hindsight memory system.

    This callback handler:
    1. Injects relevant memories into prompts before LLM calls
    2. Stores conversations to Hindsight after successful LLM calls

    Features:
    - Works with 100+ LLM providers via LiteLLM
    - Deduplication to avoid storing duplicate conversations
    - Configurable memory injection modes
    - Support for entity observations in recall

    Usage:
        >>> from hindsight_litellm import configure, enable
        >>> configure(bank_id="my-agent", hindsight_api_url="http://localhost:8888")
        >>> enable()
        >>>
        >>> # Now all LiteLLM calls will have memory integration
        >>> import litellm
        >>> response = litellm.completion(
        ...     model="gpt-4",
        ...     messages=[{"role": "user", "content": "What did we discuss?"}]
        ... )
    """

    def __init__(self):
        """Initialize the Hindsight callback handler."""
        super().__init__()
        self._http_session = None
        self._http_lock = threading.Lock()
        # Track recently stored conversation hashes for deduplication
        self._recent_hashes: Set[str] = set()
        self._max_hash_cache = 1000

    def _get_http_session(self):
        """Get or create a requests Session (thread-safe)."""
        if self._http_session is None:
            with self._http_lock:
                if self._http_session is None:
                    if HAS_REQUESTS:
                        self._http_session = requests.Session()
                    elif HAS_HTTPX:
                        self._http_session = httpx.Client(timeout=30.0)
                    else:
                        raise RuntimeError(
                            "Neither 'requests' nor 'httpx' is installed. "
                            "Please install one: pip install requests"
                        )
        return self._http_session

    def _http_post(self, url: str, json_data: dict, config: HindsightConfig) -> Optional[dict]:
        """Make a synchronous HTTP POST request."""
        try:
            session = self._get_http_session()
            headers = {"Content-Type": "application/json"}
            if config.api_key:
                headers["Authorization"] = f"Bearer {config.api_key}"

            if HAS_REQUESTS:
                response = session.post(url, json=json_data, headers=headers, timeout=30)
                response.raise_for_status()
                return response.json()
            elif HAS_HTTPX:
                response = session.post(url, json=json_data, headers=headers)
                response.raise_for_status()
                return response.json()
        except Exception as e:
            if config.verbose:
                logger.warning(f"HTTP POST failed: {e}")
            return None

    def _should_skip_model(self, model: str, config: HindsightConfig) -> bool:
        """Check if this model should be excluded from interception."""
        for pattern in config.excluded_models:
            if fnmatch.fnmatch(model.lower(), pattern.lower()):
                return True
        return False

    def _extract_user_query(self, messages: List[Dict[str, Any]]) -> Optional[str]:
        """Extract the user's query from the last user message."""
        for msg in reversed(messages):
            role = msg.get("role", "")
            if role == "user":
                content = msg.get("content")
                if isinstance(content, str):
                    return content
                elif isinstance(content, list):
                    # Handle structured content (e.g., vision messages)
                    text_parts = []
                    for item in content:
                        if isinstance(item, dict) and item.get("type") == "text":
                            text_parts.append(item.get("text", ""))
                    if text_parts:
                        return " ".join(text_parts)
        return None

    def _compute_conversation_hash(
        self,
        user_input: str,
        assistant_output: str,
    ) -> str:
        """Compute a hash for deduplication."""
        content = f"{user_input.strip().lower()}|{assistant_output.strip().lower()}"
        return hashlib.md5(content.encode()).hexdigest()[:16]

    def _is_duplicate(self, conv_hash: str) -> bool:
        """Check if this conversation was recently stored."""
        if conv_hash in self._recent_hashes:
            return True

        # Add to cache, evict oldest if full
        self._recent_hashes.add(conv_hash)
        if len(self._recent_hashes) > self._max_hash_cache:
            # Remove oldest (arbitrary since set, but good enough)
            self._recent_hashes.pop()

        return False

    def _format_memories(
        self,
        results: List[Any],
        config: HindsightConfig
    ) -> str:
        """Format memory recall results into a context string.

        Results can be RecallResult objects (with .text, .type attributes)
        or dicts (with get() method).
        """
        if not results:
            return ""

        # Apply limit if set, otherwise use all results
        results_to_use = results[:config.max_memories] if config.max_memories else results
        memory_lines = []
        for i, result in enumerate(results_to_use, 1):
            # Handle both RecallResult objects and dicts
            if hasattr(result, 'text'):
                text = result.text or ""
                fact_type = getattr(result, 'type', 'world') or "world"
                weight = getattr(result, 'weight', 0.0) or 0.0
            else:
                text = result.get("text", "")
                fact_type = result.get("type", result.get("fact_type", "world"))
                weight = result.get("weight", 0.0)

            if text:
                # Include metadata for context
                type_label = fact_type.upper() if fact_type else "MEMORY"
                line = f"{i}. [{type_label}] {text}"
                if weight > 0 and config.verbose:
                    line += f" (relevance: {weight:.2f})"
                memory_lines.append(line)

        if not memory_lines:
            return ""

        return (
            "# Relevant Memories\n"
            "The following information from memory may be relevant:\n\n"
            + "\n".join(memory_lines)
        )

    def _inject_memories_into_messages(
        self,
        messages: List[Dict[str, Any]],
        memory_context: str,
        config: HindsightConfig,
    ) -> List[Dict[str, Any]]:
        """Inject memory context into the messages list."""
        if not memory_context:
            return messages

        updated_messages = list(messages)  # Make a copy

        if config.injection_mode == MemoryInjectionMode.SYSTEM_MESSAGE:
            # Find existing system message or create new one
            for i, msg in enumerate(updated_messages):
                if msg.get("role") == "system":
                    # Append to existing system message
                    existing_content = msg.get("content", "")
                    updated_messages[i] = {
                        **msg,
                        "content": f"{existing_content}\n\n{memory_context}"
                    }
                    return updated_messages

            # No system message found, prepend one
            updated_messages.insert(0, {
                "role": "system",
                "content": memory_context
            })

        elif config.injection_mode == MemoryInjectionMode.PREPEND_USER:
            # Find the last user message and prepend context
            for i in range(len(updated_messages) - 1, -1, -1):
                if updated_messages[i].get("role") == "user":
                    original_content = updated_messages[i].get("content", "")
                    if isinstance(original_content, str):
                        updated_messages[i] = {
                            **updated_messages[i],
                            "content": f"{memory_context}\n\n---\n\n{original_content}"
                        }
                    break

        return updated_messages

    def _get_bank_id(self, config: HindsightConfig) -> str:
        """Get the bank_id for API calls."""
        return config.bank_id

    def _recall_memories_sync(
        self,
        query: str,
        config: HindsightConfig
    ) -> List[Dict[str, Any]]:
        """Recall relevant memories from Hindsight (sync) using direct HTTP."""
        try:
            bank_id = self._get_bank_id(config)
            url = f"{config.hindsight_api_url}/v1/default/banks/{bank_id}/memories/recall"

            request_data = {
                "query": query,
                "budget": config.recall_budget or "mid",
                "max_tokens": config.max_memory_tokens or 4096,
            }
            if config.fact_types:
                request_data["types"] = config.fact_types

            response = self._http_post(url, request_data, config)
            if response and "results" in response:
                return response["results"]
            return []

        except Exception as e:
            if config.verbose:
                logger.warning(f"Failed to recall memories: {e}")
            return []

    async def _recall_memories_async(
        self,
        query: str,
        config: HindsightConfig
    ) -> List[Any]:
        """Recall relevant memories from Hindsight (async).

        Uses thread pool executor with sync HTTP to avoid event loop conflicts.
        """
        try:
            loop = asyncio.get_running_loop()
            results = await loop.run_in_executor(
                _executor,
                self._recall_memories_sync,
                query,
                config
            )

            return results if isinstance(results, list) else []

        except Exception as e:
            if config.verbose:
                logger.warning(f"Failed to recall memories: {e}")
            return []

    def _store_conversation_sync(
        self,
        messages: List[Dict[str, Any]],
        response: ModelResponse,
        model: str,
        config: HindsightConfig,
    ) -> None:
        """Store the conversation to Hindsight (sync) using direct HTTP.

        By default, stores the full conversation history passed to the LLM.
        Each message is stored as a separate item, all linked by document_id.

        Hindsight will process the document as a whole for memory extraction.
        """
        try:
            # Extract assistant response from the LLM response
            assistant_output = ""
            if response.choices and len(response.choices) > 0:
                choice = response.choices[0]
                if hasattr(choice, "message") and choice.message:
                    assistant_output = choice.message.content or ""

            if not assistant_output:
                return

            # Build conversation items - each message becomes a separate item
            # All linked by document_id for Hindsight to process together
            items = []
            for msg in messages:
                role = msg.get("role", "").upper()
                content = msg.get("content", "")

                # Skip system messages - they're instructions, not conversation
                if role == "SYSTEM":
                    continue

                # Skip if this looks like our injected memory context
                if isinstance(content, str) and content.startswith("# Relevant Memories"):
                    continue

                # Handle structured content (e.g., vision messages)
                if isinstance(content, list):
                    text_parts = []
                    for item in content:
                        if isinstance(item, dict) and item.get("type") == "text":
                            text_parts.append(item.get("text", ""))
                    content = " ".join(text_parts)

                if content:
                    # Map roles to clearer labels
                    label = "USER" if role == "USER" else "ASSISTANT"
                    items.append(f"{label}: {content}")

            # Add the new assistant response
            items.append(f"ASSISTANT: {assistant_output}")

            if not items:
                return

            # Use last user message for deduplication hash
            user_input = self._extract_user_query(messages) or ""

            # Deduplication check
            conv_hash = self._compute_conversation_hash(user_input, assistant_output)
            if self._is_duplicate(conv_hash):
                if config.verbose:
                    logger.debug(f"Skipping duplicate conversation: {conv_hash}")
                return

            # Build the full conversation as a single item for now
            # (Future: could store each message as separate item in same document)
            conversation_text = "\n\n".join(items)

            # Build metadata
            metadata = {
                "source": "litellm",
                "model": model,
            }

            # Add token usage if available
            if hasattr(response, "usage") and response.usage:
                if hasattr(response.usage, "total_tokens"):
                    metadata["tokens"] = str(response.usage.total_tokens)

            bank_id = self._get_bank_id(config)
            url = f"{config.hindsight_api_url}/v1/default/banks/{bank_id}/memories"

            request_data = {
                "items": [
                    {
                        "content": conversation_text,
                        "context": f"conversation:litellm:{model}",
                        "metadata": metadata,
                        "document_id": config.document_id,  # Group by document
                    }
                ],
            }

            self._http_post(url, request_data, config)

            if config.verbose:
                logger.info(f"Stored conversation to Hindsight bank: {config.bank_id}")

        except Exception as e:
            if config.verbose:
                logger.warning(f"Failed to store conversation: {e}")

    async def _store_conversation_async(
        self,
        messages: List[Dict[str, Any]],
        response: ModelResponse,
        model: str,
        config: HindsightConfig,
    ) -> None:
        """Store the conversation to Hindsight (async).

        Uses thread pool executor with sync HTTP to avoid event loop conflicts.
        """
        try:
            loop = asyncio.get_running_loop()
            await loop.run_in_executor(
                _executor,
                self._store_conversation_sync,
                messages,
                response,
                model,
                config
            )
        except Exception as e:
            if config.verbose:
                logger.warning(f"Failed to store conversation: {e}")

    # ========== LiteLLM CustomLogger Interface ==========

    def log_pre_api_call(
        self,
        model: str,
        messages: List[Dict[str, Any]],
        kwargs: Dict[str, Any],
    ) -> None:
        """Called before making the API call (sync).

        This is where we inject memories into the messages.
        """
        if not is_configured():
            return

        config = get_config()
        if not config or not config.enabled or not config.inject_memories:
            return

        if self._should_skip_model(model, config):
            return

        # Extract user query
        user_query = self._extract_user_query(messages)
        if not user_query:
            return

        # Recall relevant memories
        memories = self._recall_memories_sync(user_query, config)
        if not memories:
            return

        # Format and inject memories
        memory_context = self._format_memories(memories, config)
        updated_messages = self._inject_memories_into_messages(
            messages, memory_context, config
        )

        # Modify messages list IN-PLACE (don't just reassign kwargs)
        messages.clear()
        messages.extend(updated_messages)

        if config.verbose:
            logger.info(f"Injected {len(memories)} memories into prompt")

    async def async_log_pre_api_call(
        self,
        model: str,
        messages: List[Dict[str, Any]],
        kwargs: Dict[str, Any],
    ) -> None:
        """Called before making the API call (async).

        This is where we inject memories into the messages.
        """
        if not is_configured():
            return

        config = get_config()
        if not config or not config.enabled or not config.inject_memories:
            return

        if self._should_skip_model(model, config):
            return

        # Extract user query
        user_query = self._extract_user_query(messages)
        if not user_query:
            return

        # Recall relevant memories
        memories = await self._recall_memories_async(user_query, config)
        if not memories:
            return

        # Format and inject memories
        memory_context = self._format_memories(memories, config)
        updated_messages = self._inject_memories_into_messages(
            messages, memory_context, config
        )

        # Modify messages list IN-PLACE (don't just reassign kwargs)
        messages.clear()
        messages.extend(updated_messages)

        if config.verbose:
            logger.info(f"Injected {len(memories)} memories into prompt")

    def log_success_event(
        self,
        kwargs: Dict[str, Any],
        response_obj: Any,
        start_time: float,
        end_time: float,
    ) -> None:
        """Called after successful API call (sync).

        This is where we store the conversation.
        """
        if not is_configured():
            return

        config = get_config()
        if not config or not config.enabled or not config.store_conversations:
            return

        model = kwargs.get("model", "unknown")
        if self._should_skip_model(model, config):
            return

        messages = kwargs.get("messages", [])
        if not messages:
            return

        # Store the conversation
        self._store_conversation_sync(messages, response_obj, model, config)

    async def async_log_success_event(
        self,
        kwargs: Dict[str, Any],
        response_obj: Any,
        start_time: float,
        end_time: float,
    ) -> None:
        """Called after successful API call (async).

        This is where we store the conversation.
        """
        if not is_configured():
            return

        config = get_config()
        if not config or not config.enabled or not config.store_conversations:
            return

        model = kwargs.get("model", "unknown")
        if self._should_skip_model(model, config):
            return

        messages = kwargs.get("messages", [])
        if not messages:
            return

        # Store the conversation
        await self._store_conversation_async(messages, response_obj, model, config)

    def log_failure_event(
        self,
        kwargs: Dict[str, Any],
        response_obj: Any,
        start_time: float,
        end_time: float,
    ) -> None:
        """Called after failed API call (sync)."""
        # We don't store failed conversations
        pass

    async def async_log_failure_event(
        self,
        kwargs: Dict[str, Any],
        response_obj: Any,
        start_time: float,
        end_time: float,
    ) -> None:
        """Called after failed API call (async)."""
        # We don't store failed conversations
        pass

    def close(self) -> None:
        """Clean up resources."""
        with self._http_lock:
            if self._http_session is not None:
                try:
                    if HAS_REQUESTS:
                        self._http_session.close()
                    elif HAS_HTTPX:
                        self._http_session.close()
                except Exception:
                    pass
                self._http_session = None
        self._recent_hashes.clear()


# Global callback instance
_callback: Optional[HindsightCallback] = None


def get_callback() -> HindsightCallback:
    """Get the global callback instance, creating it if necessary."""
    global _callback
    if _callback is None:
        _callback = HindsightCallback()
    return _callback


def cleanup_callback() -> None:
    """Clean up the global callback instance."""
    global _callback
    if _callback is not None:
        _callback.close()
        _callback = None