56489a58d1
* refactor(llamaindex): merge two packages into single hindsight-llamaindex Merge `llama-index-tools-hindsight` and `llama-index-memory-hindsight` into a single `hindsight-llamaindex` package following our naming convention. - Rename package to `hindsight-llamaindex` (Python module: `hindsight_llamaindex`) - Move HindsightToolSpec and HindsightMemory into the same package - Delete `llamaindex-memory/` directory - Add CI test job for llamaindex integration - Update docs, blog post, and integrations.json * fix(blog): update llamaindex blog post for merged package - Move date to 2026-03-30 - Add HindsightMemory (automatic BaseMemory) pattern - Fix "bank must exist first" pitfall — mission auto-creates - Align all code examples with docs page - Update architecture diagram to show both patterns * fix(docs): add llamaindex/openai icons, rename Codex - Add llamaindex.png and openai.png icons - Rename "OpenAI Codex CLI" to "Codex" in integrations.json and docs - Use openai.png icon for Codex integration
348 lines
12 KiB
Python
348 lines
12 KiB
Python
"""Unit tests for Hindsight LlamaIndex memory."""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from llama_index.core.llms import ChatMessage, MessageRole
|
|
from hindsight_llamaindex import HindsightMemory
|
|
|
|
|
|
def _mock_client():
|
|
"""Create a mock Hindsight client."""
|
|
client = MagicMock()
|
|
client.retain = MagicMock()
|
|
client.recall = MagicMock()
|
|
client.create_bank = MagicMock()
|
|
client.aretain = AsyncMock()
|
|
client.arecall = AsyncMock()
|
|
client.acreate_bank = AsyncMock()
|
|
return client
|
|
|
|
|
|
def _mock_recall_response(texts: list[str]):
|
|
response = MagicMock()
|
|
results = []
|
|
for t in texts:
|
|
r = MagicMock()
|
|
r.text = t
|
|
results.append(r)
|
|
response.results = results
|
|
return response
|
|
|
|
|
|
class TestHindsightMemoryCreation:
|
|
def test_from_client(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(
|
|
client=client,
|
|
bank_id="test-bank",
|
|
)
|
|
assert memory.bank_id == "test-bank"
|
|
assert memory.context == "llamaindex"
|
|
assert memory.budget == "mid"
|
|
|
|
def test_from_client_with_options(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(
|
|
client=client,
|
|
bank_id="test-bank",
|
|
mission="Track preferences",
|
|
context="my-app",
|
|
budget="high",
|
|
tags=["source:chat"],
|
|
)
|
|
assert memory.bank_id == "test-bank"
|
|
assert memory.context == "my-app"
|
|
assert memory.budget == "high"
|
|
assert memory.tags == ["source:chat"]
|
|
|
|
def test_from_url(self):
|
|
with patch("hindsight_llamaindex.memory.Hindsight") as mock_cls:
|
|
mock_cls.return_value = _mock_client()
|
|
memory = HindsightMemory.from_url(
|
|
hindsight_api_url="http://localhost:8888",
|
|
bank_id="test-bank",
|
|
)
|
|
assert memory.bank_id == "test-bank"
|
|
mock_cls.assert_called_once_with(
|
|
base_url="http://localhost:8888", timeout=30.0
|
|
)
|
|
|
|
def test_from_defaults_raises(self):
|
|
with pytest.raises(NotImplementedError):
|
|
HindsightMemory.from_defaults()
|
|
|
|
def test_class_name(self):
|
|
assert HindsightMemory.class_name() == "HindsightMemory"
|
|
|
|
|
|
class TestPut:
|
|
def test_put_user_message_retains(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
msg = ChatMessage(role=MessageRole.USER, content="I like Python")
|
|
memory.put(msg)
|
|
|
|
client.retain.assert_called_once()
|
|
kwargs = client.retain.call_args[1]
|
|
assert kwargs["bank_id"] == "test"
|
|
assert kwargs["content"] == "I like Python"
|
|
assert kwargs["context"] == "llamaindex"
|
|
assert kwargs["metadata"]["role"] == "user"
|
|
|
|
def test_put_assistant_message_retains(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
msg = ChatMessage(role=MessageRole.ASSISTANT, content="Noted!")
|
|
memory.put(msg)
|
|
|
|
client.retain.assert_called_once()
|
|
kwargs = client.retain.call_args[1]
|
|
assert kwargs["metadata"]["role"] == "assistant"
|
|
|
|
def test_put_system_message_skipped(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
msg = ChatMessage(role=MessageRole.SYSTEM, content="You are helpful")
|
|
memory.put(msg)
|
|
|
|
client.retain.assert_not_called()
|
|
|
|
def test_put_empty_content_skipped(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
msg = ChatMessage(role=MessageRole.USER, content=" ")
|
|
memory.put(msg)
|
|
|
|
client.retain.assert_not_called()
|
|
|
|
def test_put_adds_to_local_history(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
msg = ChatMessage(role=MessageRole.USER, content="hello")
|
|
memory.put(msg)
|
|
|
|
assert len(memory.get_all()) == 1
|
|
assert memory.get_all()[0].content == "hello"
|
|
|
|
def test_put_trims_to_limit(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(
|
|
client=client, bank_id="test", chat_history_limit=3
|
|
)
|
|
for i in range(5):
|
|
memory.put(ChatMessage(role=MessageRole.USER, content=f"msg-{i}"))
|
|
|
|
history = memory.get_all()
|
|
assert len(history) == 3
|
|
assert history[0].content == "msg-2"
|
|
assert history[2].content == "msg-4"
|
|
|
|
def test_put_tags_passed(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(
|
|
client=client, bank_id="test", tags=["source:chat"]
|
|
)
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
|
|
kwargs = client.retain.call_args[1]
|
|
assert kwargs["tags"] == ["source:chat"]
|
|
|
|
def test_put_retain_failure_is_graceful(self):
|
|
client = _mock_client()
|
|
client.retain.side_effect = RuntimeError("connection refused")
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
# Should not raise
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
# Message still in local history
|
|
assert len(memory.get_all()) == 1
|
|
|
|
def test_put_generates_document_id(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
|
|
kwargs = client.retain.call_args[1]
|
|
doc_id = kwargs["document_id"]
|
|
parts = doc_id.rsplit("-", 1)
|
|
assert len(parts) == 2
|
|
assert parts[1].isdigit()
|
|
|
|
|
|
class TestGet:
|
|
def test_get_without_input_returns_history(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
memory.put(ChatMessage(role=MessageRole.ASSISTANT, content="hi there"))
|
|
|
|
messages = memory.get()
|
|
assert len(messages) == 2
|
|
client.recall.assert_not_called()
|
|
|
|
def test_get_with_input_recalls_memories(self):
|
|
client = _mock_client()
|
|
client.recall.return_value = _mock_recall_response(
|
|
["User likes Python", "User prefers dark mode"]
|
|
)
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
|
|
messages = memory.get(input="What are my preferences?")
|
|
|
|
# Should have system message + chat history
|
|
assert len(messages) == 2
|
|
assert messages[0].role == MessageRole.SYSTEM
|
|
assert "User likes Python" in str(messages[0].content)
|
|
assert "User prefers dark mode" in str(messages[0].content)
|
|
assert messages[1].content == "hello"
|
|
|
|
def test_get_with_input_no_memories(self):
|
|
client = _mock_client()
|
|
client.recall.return_value = _mock_recall_response([])
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
|
|
messages = memory.get(input="anything")
|
|
# No system message when no memories found
|
|
assert len(messages) == 1
|
|
assert messages[0].content == "hello"
|
|
|
|
def test_get_recall_passes_params(self):
|
|
client = _mock_client()
|
|
client.recall.return_value = _mock_recall_response(["fact"])
|
|
memory = HindsightMemory.from_client(
|
|
client=client,
|
|
bank_id="test",
|
|
budget="high",
|
|
max_tokens=2048,
|
|
recall_tags=["scope:user"],
|
|
recall_tags_match="all",
|
|
)
|
|
|
|
memory.get(input="query")
|
|
kwargs = client.recall.call_args[1]
|
|
assert kwargs["budget"] == "high"
|
|
assert kwargs["max_tokens"] == 2048
|
|
assert kwargs["tags"] == ["scope:user"]
|
|
assert kwargs["tags_match"] == "all"
|
|
|
|
def test_get_recall_failure_is_graceful(self):
|
|
client = _mock_client()
|
|
client.recall.side_effect = RuntimeError("timeout")
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
|
|
# Should not raise, returns history without memories
|
|
messages = memory.get(input="query")
|
|
assert len(messages) == 1
|
|
|
|
def test_get_custom_system_prompt(self):
|
|
client = _mock_client()
|
|
client.recall.return_value = _mock_recall_response(["fact1"])
|
|
memory = HindsightMemory.from_client(
|
|
client=client,
|
|
bank_id="test",
|
|
system_prompt="MEMORIES: {memories}",
|
|
)
|
|
|
|
messages = memory.get(input="query")
|
|
assert str(messages[0].content) == "MEMORIES: - fact1"
|
|
|
|
|
|
class TestSet:
|
|
def test_set_replaces_history(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="old"))
|
|
|
|
new_messages = [
|
|
ChatMessage(role=MessageRole.USER, content="new1"),
|
|
ChatMessage(role=MessageRole.ASSISTANT, content="new2"),
|
|
]
|
|
memory.set(new_messages)
|
|
|
|
history = memory.get_all()
|
|
assert len(history) == 2
|
|
assert history[0].content == "new1"
|
|
|
|
def test_set_retains_only_new_messages(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="existing"))
|
|
client.retain.reset_mock()
|
|
|
|
# Set with 3 messages (1 existing + 2 new)
|
|
messages = [
|
|
ChatMessage(role=MessageRole.USER, content="existing"),
|
|
ChatMessage(role=MessageRole.USER, content="new1"),
|
|
ChatMessage(role=MessageRole.ASSISTANT, content="new2"),
|
|
]
|
|
memory.set(messages)
|
|
|
|
# Should retain only the 2 new messages
|
|
assert client.retain.call_count == 2
|
|
|
|
|
|
class TestReset:
|
|
def test_reset_clears_local_history(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
assert len(memory.get_all()) == 1
|
|
|
|
memory.reset()
|
|
assert len(memory.get_all()) == 0
|
|
|
|
|
|
class TestBankMission:
|
|
def test_creates_bank_with_mission_on_put(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(
|
|
client=client, bank_id="test", mission="Track preferences"
|
|
)
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
|
|
client.create_bank.assert_called_once_with(
|
|
bank_id="test",
|
|
name="test",
|
|
mission="Track preferences",
|
|
)
|
|
|
|
def test_creates_bank_with_mission_on_get(self):
|
|
client = _mock_client()
|
|
client.recall.return_value = _mock_recall_response([])
|
|
memory = HindsightMemory.from_client(
|
|
client=client, bank_id="test", mission="Track preferences"
|
|
)
|
|
memory.get(input="query")
|
|
|
|
client.create_bank.assert_called_once()
|
|
|
|
def test_bank_creation_idempotent(self):
|
|
client = _mock_client()
|
|
client.recall.return_value = _mock_recall_response([])
|
|
memory = HindsightMemory.from_client(
|
|
client=client, bank_id="test", mission="mission"
|
|
)
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
memory.get(input="query")
|
|
|
|
assert client.create_bank.call_count == 1
|
|
|
|
def test_no_bank_creation_without_mission(self):
|
|
client = _mock_client()
|
|
memory = HindsightMemory.from_client(client=client, bank_id="test")
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
|
|
client.create_bank.assert_not_called()
|
|
|
|
def test_bank_creation_failure_is_graceful(self):
|
|
client = _mock_client()
|
|
client.create_bank.side_effect = RuntimeError("already exists")
|
|
memory = HindsightMemory.from_client(
|
|
client=client, bank_id="test", mission="mission"
|
|
)
|
|
# Should not raise
|
|
memory.put(ChatMessage(role=MessageRole.USER, content="hello"))
|
|
client.retain.assert_called_once()
|