fleet-memory/hindsight-integrations/langgraph/hindsight_langgraph/nodes.py

"""Pre-built LangGraph nodes for Hindsight memory operations.

Provides node functions that can be added directly to a StateGraph to
inject memories at conversation start and store new memories after responses.
"""

import logging
from typing import Any, Optional

from hindsight_client import Hindsight
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph import MessagesState

from ._client import resolve_client

logger = logging.getLogger(__name__)


def _extract_text_content(content: Any) -> str:
    """Extract text from a message content field.

    Handles both plain string content and multimodal content lists
    (where each item may be a dict with "type" and "text" keys).
    Returns the concatenated text parts, or an empty string if no
    text content is found.
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for part in content:
            if isinstance(part, str):
                parts.append(part)
            elif isinstance(part, dict) and part.get("type") == "text":
                parts.append(part.get("text", ""))
        return " ".join(parts)
    return str(content) if content else ""


def create_recall_node(
    *,
    bank_id: Optional[str] = None,
    client: Optional[Hindsight] = None,
    hindsight_api_url: Optional[str] = None,
    api_key: Optional[str] = None,
    budget: str = "mid",
    max_tokens: int = 4096,
    max_results: int = 10,
    tags: Optional[list[str]] = None,
    tags_match: str = "any",
    bank_id_from_config: str = "user_id",
    output_key: Optional[str] = None,
):
    """Create a node that injects relevant memories into the conversation.

    This node extracts the latest user message, recalls relevant memories
    from Hindsight, and returns them either as a SystemMessage in the
    ``messages`` list (default) or as a plain string under a custom state
    key via ``output_key``.

    **Message ordering:** When using the default ``messages`` output,
    ``MessagesState`` uses ``add_messages`` as its reducer, which appends.
    The memory SystemMessage will appear after existing messages, not at
    position 0. If your LLM provider requires system messages first, use
    ``output_key`` to write the memory text to a separate state field and
    inject it into your prompt in the agent node.

    Example with ``output_key`` (recommended for correct ordering)::

        from typing import Optional
        from langgraph.graph import MessagesState

        class AgentState(MessagesState):
            memory_context: Optional[str] = None

        recall = create_recall_node(
            client=client, bank_id="user-123", output_key="memory_context"
        )
        # In your agent node, read state["memory_context"] and prepend
        # it to the system prompt.

    The bank_id can be provided directly or resolved dynamically from
    the graph's RunnableConfig via the ``bank_id_from_config`` key.

    Args:
        bank_id: Static Hindsight memory bank ID.
        client: Pre-configured Hindsight client.
        hindsight_api_url: API URL (used if no client provided).
        api_key: API key (used if no client provided).
        budget: Recall budget level (low/mid/high).
        max_tokens: Maximum tokens for recall results.
        max_results: Maximum number of memories to inject.
        tags: Tags to filter recall results.
        tags_match: Tag matching mode.
        bank_id_from_config: Config key to read bank_id from at runtime.
            Looked up in ``config["configurable"][bank_id_from_config]``.
            Only used when ``bank_id`` is not provided.
        output_key: If set, write the memory text to this state key as a
            plain string instead of appending a SystemMessage to ``messages``.
            Use this with a custom state type to control where memory context
            appears in your prompt.

    Returns:
        An async node function compatible with LangGraph StateGraph.
    """
    resolved_client = resolve_client(client, hindsight_api_url, api_key)

    async def recall_node(
        state: MessagesState, config: Optional[RunnableConfig] = None
    ) -> dict[str, Any]:
        resolved_bank_id = bank_id
        if resolved_bank_id is None and config:
            configurable = config.get("configurable", {})
            resolved_bank_id = configurable.get(bank_id_from_config)

        if not resolved_bank_id:
            logger.warning(
                "No bank_id available for recall node, skipping memory injection."
            )
            if output_key:
                return {output_key: None}
            return {"messages": []}

        # Extract query from the latest human message
        query = None
        for msg in reversed(state["messages"]):
            if isinstance(msg, HumanMessage):
                query = _extract_text_content(msg.content)
                break

        if not query:
            if output_key:
                return {output_key: None}
            return {"messages": []}

        try:
            recall_kwargs: dict[str, Any] = {
                "bank_id": resolved_bank_id,
                "query": query,
                "budget": budget,
                "max_tokens": max_tokens,
            }
            if tags:
                recall_kwargs["tags"] = tags
                recall_kwargs["tags_match"] = tags_match

            response = await resolved_client.arecall(**recall_kwargs)
            results = response.results[:max_results] if response.results else []

            if not results:
                if output_key:
                    return {output_key: None}
                return {"messages": []}

            lines = ["Relevant memories about this user:"]
            for i, result in enumerate(results, 1):
                lines.append(f"{i}. {result.text}")
            memory_text = "\n".join(lines)

            if output_key:
                return {output_key: memory_text}
            return {
                "messages": [
                    SystemMessage(content=memory_text, id="hindsight_memory_context")
                ]
            }
        except Exception as e:
            logger.error(f"Recall node failed: {e}")
            if output_key:
                return {output_key: None}
            return {"messages": []}

    return recall_node


def create_retain_node(
    *,
    bank_id: Optional[str] = None,
    client: Optional[Hindsight] = None,
    hindsight_api_url: Optional[str] = None,
    api_key: Optional[str] = None,
    tags: Optional[list[str]] = None,
    bank_id_from_config: str = "user_id",
    retain_human: bool = True,
    retain_ai: bool = False,
):
    """Create a node that stores conversation messages as memories.

    This node extracts messages from the conversation and stores them
    via Hindsight retain. It should be placed after the LLM response
    node in your graph.

    Args:
        bank_id: Static Hindsight memory bank ID.
        client: Pre-configured Hindsight client.
        hindsight_api_url: API URL (used if no client provided).
        api_key: API key (used if no client provided).
        tags: Tags to apply to stored memories.
        bank_id_from_config: Config key to read bank_id from at runtime.
        retain_human: Store human messages as memories.
        retain_ai: Store AI responses as memories.

    Returns:
        An async node function compatible with LangGraph StateGraph.
    """
    resolved_client = resolve_client(client, hindsight_api_url, api_key)

    async def retain_node(
        state: MessagesState, config: Optional[RunnableConfig] = None
    ) -> dict[str, Any]:
        resolved_bank_id = bank_id
        if resolved_bank_id is None and config:
            configurable = config.get("configurable", {})
            resolved_bank_id = configurable.get(bank_id_from_config)

        if not resolved_bank_id:
            logger.warning(
                "No bank_id available for retain node, skipping memory storage."
            )
            return {"messages": []}

        # Only retain the latest human and/or AI message to avoid
        # duplicating memories that were already stored in prior calls.
        messages_to_retain = []
        if retain_human:
            for msg in reversed(state["messages"]):
                if isinstance(msg, HumanMessage):
                    text = _extract_text_content(msg.content)
                    if text:
                        messages_to_retain.append(text)
                    break
        if retain_ai:
            for msg in reversed(state["messages"]):
                if isinstance(msg, AIMessage):
                    text = _extract_text_content(msg.content)
                    if text:
                        messages_to_retain.append(text)
                    break

        if not messages_to_retain:
            return {"messages": []}

        content = "\n\n".join(messages_to_retain)

        try:
            retain_kwargs: dict[str, Any] = {
                "bank_id": resolved_bank_id,
                "content": content,
            }
            if tags:
                retain_kwargs["tags"] = tags
            await resolved_client.aretain(**retain_kwargs)
        except Exception as e:
            logger.error(f"Retain node failed: {e}")

        return {"messages": []}

    return retain_node