"""Hindsight BaseMemory implementation for LlamaIndex. Provides automatic memory for LlamaIndex agents: - ``put()`` retains messages to Hindsight for long-term storage - ``get()`` recalls relevant memories and prepends them as context - Chat history is kept in-memory for the current session """ import logging import re import time import uuid from typing import Any, Optional from hindsight_client import Hindsight from llama_index.core.bridge.pydantic import Field, PrivateAttr from llama_index.core.llms import ChatMessage, MessageRole from llama_index.core.memory.types import BaseMemory logger = logging.getLogger(__name__) DEFAULT_SYSTEM_PROMPT = ( "Below are relevant memories from previous conversations:\n{memories}\n" "Use these memories to provide more personalized and contextual responses." ) # Patterns for detecting ReAct-style reasoning traces in assistant messages _REACT_PATTERN = re.compile( r"^(Thought|Action|Action Input|Observation)\s*:", re.MULTILINE ) _ANSWER_PATTERN = re.compile(r"^Answer\s*:\s*", re.MULTILINE) class HindsightMemory(BaseMemory): """Automatic long-term memory for LlamaIndex agents via Hindsight. On ``put()``, user and assistant messages are automatically retained to Hindsight. On ``get()``, relevant memories are recalled and prepended as a system message to enrich the agent's context. This follows the same pattern as Mem0's LlamaIndex integration: a local chat buffer for the current session, with Hindsight providing cross-session long-term memory. Args: bank_id: Hindsight memory bank to operate on. context: Source label for retain operations. budget: Recall budget level (low/mid/high). max_tokens: Maximum tokens for recall results. tags: Tags applied when storing memories. recall_tags: Tags to filter when recalling. recall_tags_match: Tag matching mode. system_prompt: Template for the memory system message. Must contain ``{memories}`` placeholder. chat_history_limit: Max messages to keep in local buffer. Oldest messages are dropped when exceeded. Example:: from hindsight_client import Hindsight from hindsight_llamaindex import HindsightMemory client = Hindsight(base_url="http://localhost:8888") memory = HindsightMemory.from_client( client=client, bank_id="user-123", mission="Track user preferences", ) # Use with any LlamaIndex agent — pass memory to run(), not the constructor agent = ReActAgent(tools=[], llm=llm) response = await agent.run("Hello!", memory=memory) """ bank_id: str = Field(description="Hindsight memory bank ID") context: str = Field(default="llamaindex", description="Source label for retain") budget: str = Field(default="mid", description="Recall budget level") max_tokens: int = Field(default=4096, description="Max tokens for recall") tags: Optional[list[str]] = Field(default=None, description="Tags for retain") recall_tags: Optional[list[str]] = Field( default=None, description="Tags to filter recall" ) recall_tags_match: str = Field(default="any", description="Tag matching mode") system_prompt: str = Field( default=DEFAULT_SYSTEM_PROMPT, description="Memory system message template" ) chat_history_limit: int = Field( default=100, description="Max messages in local buffer" ) _client: Hindsight = PrivateAttr() _chat_history: list[ChatMessage] = PrivateAttr(default_factory=list) _session_id: str = PrivateAttr() _bank_initialized: bool = PrivateAttr(default=False) _mission: Optional[str] = PrivateAttr(default=None) def __init__(self, client: Hindsight, mission: Optional[str] = None, **kwargs: Any): super().__init__(**kwargs) self._client = client self._session_id = str(uuid.uuid4())[:8] self._mission = mission self._chat_history = [] self._bank_initialized = False @classmethod def class_name(cls) -> str: return "HindsightMemory" @classmethod def from_defaults(cls, **kwargs: Any) -> "HindsightMemory": """Create from defaults. Prefer ``from_client()`` instead.""" raise NotImplementedError( "Use HindsightMemory.from_client() or HindsightMemory.from_url() instead." ) @classmethod def from_client( cls, client: Hindsight, bank_id: str, *, mission: Optional[str] = None, context: str = "llamaindex", budget: str = "mid", max_tokens: int = 4096, tags: Optional[list[str]] = None, recall_tags: Optional[list[str]] = None, recall_tags_match: str = "any", system_prompt: str = DEFAULT_SYSTEM_PROMPT, chat_history_limit: int = 100, ) -> "HindsightMemory": """Create a HindsightMemory with a pre-configured client. Args: client: Hindsight client instance. bank_id: Memory bank ID. mission: Bank mission (creates bank on first use if set). context: Source label for retain operations. budget: Recall budget level. max_tokens: Max recall tokens. tags: Tags for retain operations. recall_tags: Tags to filter recall. recall_tags_match: Tag matching mode. system_prompt: Memory system message template. chat_history_limit: Max local buffer size. """ return cls( client=client, bank_id=bank_id, mission=mission, context=context, budget=budget, max_tokens=max_tokens, tags=tags, recall_tags=recall_tags, recall_tags_match=recall_tags_match, system_prompt=system_prompt, chat_history_limit=chat_history_limit, ) @classmethod def from_url( cls, hindsight_api_url: str, bank_id: str, *, api_key: Optional[str] = None, **kwargs: Any, ) -> "HindsightMemory": """Create a HindsightMemory from an API URL. Args: hindsight_api_url: Hindsight API URL. bank_id: Memory bank ID. api_key: Optional API key. **kwargs: Additional arguments passed to ``from_client()``. """ client_kwargs: dict[str, Any] = {"base_url": hindsight_api_url, "timeout": 30.0} if api_key: client_kwargs["api_key"] = api_key client = Hindsight(**client_kwargs) return cls.from_client(client=client, bank_id=bank_id, **kwargs) def _ensure_bank(self) -> None: if self._bank_initialized or not self._mission: return try: self._client.create_bank( bank_id=self.bank_id, name=self.bank_id, mission=self._mission, ) except Exception as e: logger.debug(f"Bank creation for {self.bank_id}: {e}") self._bank_initialized = True async def _aensure_bank(self) -> None: if self._bank_initialized or not self._mission: return try: await self._client.acreate_bank( bank_id=self.bank_id, name=self.bank_id, mission=self._mission, ) except Exception as e: logger.debug(f"Bank creation for {self.bank_id}: {e}") self._bank_initialized = True def _generate_document_id(self) -> str: return f"{self._session_id}-{uuid.uuid4().hex[:12]}" @staticmethod def _extract_clean_content(content: str, role: MessageRole) -> str: """Extract clean content from a message, stripping ReAct traces. For assistant messages containing ReAct reasoning (Thought:/Action:/ Observation: prefixes), extracts only the final Answer: text. Returns empty string if the message is purely reasoning with no answer. User messages are returned as-is. """ if role != MessageRole.ASSISTANT: return content # Check if this looks like ReAct reasoning if not _REACT_PATTERN.search(content): return content # Extract the final Answer: block answer_match = list(_ANSWER_PATTERN.finditer(content)) if answer_match: # Use the last Answer: block (final answer after reasoning) last_answer = answer_match[-1] return content[last_answer.end():].strip() # ReAct traces with no Answer: — skip retention return "" def _retain_message(self, message: ChatMessage) -> None: """Retain a message to Hindsight (sync).""" if message.role not in (MessageRole.USER, MessageRole.ASSISTANT): return content = str(message.content) if message.content else "" content = self._extract_clean_content(content, message.role) if not content.strip(): return try: self._ensure_bank() kwargs: dict[str, Any] = { "bank_id": self.bank_id, "content": content, "context": self.context, "document_id": self._generate_document_id(), "metadata": {"role": message.role.value, "source": "llamaindex"}, } if self.tags: kwargs["tags"] = self.tags self._client.retain(**kwargs) except Exception as e: logger.error(f"Failed to retain message: {e}") async def _aretain_message(self, message: ChatMessage) -> None: """Retain a message to Hindsight (async).""" if message.role not in (MessageRole.USER, MessageRole.ASSISTANT): return content = str(message.content) if message.content else "" content = self._extract_clean_content(content, message.role) if not content.strip(): return try: await self._aensure_bank() kwargs: dict[str, Any] = { "bank_id": self.bank_id, "content": content, "context": self.context, "document_id": self._generate_document_id(), "metadata": {"role": message.role.value, "source": "llamaindex"}, } if self.tags: kwargs["tags"] = self.tags await self._client.aretain(**kwargs) except Exception as e: logger.error(f"Failed to retain message: {e}") def _recall_memories(self, query: str) -> str: """Recall relevant memories (sync).""" try: self._ensure_bank() kwargs: dict[str, Any] = { "bank_id": self.bank_id, "query": query, "budget": self.budget, "max_tokens": self.max_tokens, } if self.recall_tags: kwargs["tags"] = self.recall_tags kwargs["tags_match"] = self.recall_tags_match response = self._client.recall(**kwargs) if not response.results: return "" lines = [r.text for r in response.results] return "\n".join(f"- {line}" for line in lines) except Exception as e: logger.error(f"Failed to recall memories: {e}") return "" async def _arecall_memories(self, query: str) -> str: """Recall relevant memories (async).""" try: await self._aensure_bank() kwargs: dict[str, Any] = { "bank_id": self.bank_id, "query": query, "budget": self.budget, "max_tokens": self.max_tokens, } if self.recall_tags: kwargs["tags"] = self.recall_tags kwargs["tags_match"] = self.recall_tags_match response = await self._client.arecall(**kwargs) if not response.results: return "" lines = [r.text for r in response.results] return "\n".join(f"- {line}" for line in lines) except Exception as e: logger.error(f"Failed to recall memories: {e}") return "" # -- BaseMemory interface -- def get(self, input: Optional[str] = None, **kwargs: Any) -> list[ChatMessage]: """Get chat history, enriched with recalled Hindsight memories. If ``input`` is provided, relevant memories are recalled and prepended as a system message. """ messages: list[ChatMessage] = [] if input: memories_text = self._recall_memories(input) if memories_text: system_content = self.system_prompt.format(memories=memories_text) messages.append( ChatMessage(role=MessageRole.SYSTEM, content=system_content) ) messages.extend(self._chat_history) return messages async def aget( self, input: Optional[str] = None, **kwargs: Any ) -> list[ChatMessage]: """Async version of get().""" messages: list[ChatMessage] = [] if input: memories_text = await self._arecall_memories(input) if memories_text: system_content = self.system_prompt.format(memories=memories_text) messages.append( ChatMessage(role=MessageRole.SYSTEM, content=system_content) ) messages.extend(self._chat_history) return messages def get_all(self) -> list[ChatMessage]: """Get all messages in the local chat buffer.""" return list(self._chat_history) def put(self, message: ChatMessage) -> None: """Store a message in local buffer and retain to Hindsight.""" self._chat_history.append(message) # Trim to limit if len(self._chat_history) > self.chat_history_limit: self._chat_history = self._chat_history[-self.chat_history_limit :] self._retain_message(message) async def aput(self, message: ChatMessage) -> None: """Async version of put().""" self._chat_history.append(message) if len(self._chat_history) > self.chat_history_limit: self._chat_history = self._chat_history[-self.chat_history_limit :] await self._aretain_message(message) def set(self, messages: list[ChatMessage]) -> None: """Set the chat history, retaining new messages to Hindsight.""" existing_len = len(self._chat_history) self._chat_history = list(messages) # Retain only new messages (beyond previous length) for msg in messages[existing_len:]: self._retain_message(msg) async def aset(self, messages: list[ChatMessage]) -> None: """Async version of set().""" existing_len = len(self._chat_history) self._chat_history = list(messages) for msg in messages[existing_len:]: await self._aretain_message(msg) def reset(self) -> None: """Reset the local chat buffer. Does not clear Hindsight memories.""" self._chat_history = []