fix(memory): honor prompt param in vector store extraction (#4914)
This commit is contained in:
+6
-5
@@ -656,10 +656,10 @@ class Memory(MemoryBase):
|
||||
else:
|
||||
messages = parse_vision_messages(messages)
|
||||
|
||||
vector_store_result = self._add_to_vector_store(messages, processed_metadata, effective_filters, infer)
|
||||
vector_store_result = self._add_to_vector_store(messages, processed_metadata, effective_filters, infer, prompt=prompt)
|
||||
return {"results": vector_store_result}
|
||||
|
||||
def _add_to_vector_store(self, messages, metadata, filters, infer):
|
||||
def _add_to_vector_store(self, messages, metadata, filters, infer, prompt=None):
|
||||
if not infer:
|
||||
returned_memories = []
|
||||
for message_dict in messages:
|
||||
@@ -726,7 +726,7 @@ class Memory(MemoryBase):
|
||||
if is_agent_scoped:
|
||||
system_prompt += AGENT_CONTEXT_SUFFIX
|
||||
|
||||
custom_instr = self.custom_instructions
|
||||
custom_instr = prompt or self.custom_instructions
|
||||
|
||||
user_prompt = generate_additive_extraction_prompt(
|
||||
existing_memories=existing_memories,
|
||||
@@ -2064,7 +2064,7 @@ class AsyncMemory(MemoryBase):
|
||||
else:
|
||||
messages = parse_vision_messages(messages)
|
||||
|
||||
vector_store_result = await self._add_to_vector_store(messages, processed_metadata, effective_filters, infer)
|
||||
vector_store_result = await self._add_to_vector_store(messages, processed_metadata, effective_filters, infer, prompt=prompt)
|
||||
return {"results": vector_store_result}
|
||||
|
||||
async def _add_to_vector_store(
|
||||
@@ -2073,6 +2073,7 @@ class AsyncMemory(MemoryBase):
|
||||
metadata: dict,
|
||||
effective_filters: dict,
|
||||
infer: bool,
|
||||
prompt: Optional[str] = None,
|
||||
):
|
||||
if not infer:
|
||||
returned_memories = []
|
||||
@@ -2141,7 +2142,7 @@ class AsyncMemory(MemoryBase):
|
||||
if is_agent_scoped:
|
||||
system_prompt += AGENT_CONTEXT_SUFFIX
|
||||
|
||||
custom_instr = self.custom_instructions
|
||||
custom_instr = prompt or self.custom_instructions
|
||||
|
||||
user_prompt = generate_additive_extraction_prompt(
|
||||
existing_memories=existing_memories,
|
||||
|
||||
@@ -78,6 +78,43 @@ class TestAddToVectorStoreErrors:
|
||||
assert result == [] # Should return empty list when no memories processed
|
||||
|
||||
|
||||
class TestPromptOverridesCustomInstructions:
|
||||
@pytest.fixture
|
||||
def mock_memory(self, mocker):
|
||||
mock_llm, _ = _setup_mocks(mocker)
|
||||
mock_llm.return_value.generate_response.return_value = '{"memory": []}'
|
||||
|
||||
memory = Memory()
|
||||
memory.custom_instructions = "config-level instructions"
|
||||
memory.db.get_last_messages = MagicMock(return_value=[])
|
||||
memory.db.save_messages = MagicMock()
|
||||
return memory
|
||||
|
||||
def test_prompt_overrides_custom_instructions(self, mock_memory):
|
||||
mock_memory._add_to_vector_store(
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
metadata={},
|
||||
filters={},
|
||||
infer=True,
|
||||
prompt="per-call override",
|
||||
)
|
||||
|
||||
user_prompt = mock_memory.llm.generate_response.call_args[1]["messages"][1]["content"]
|
||||
assert "per-call override" in user_prompt
|
||||
assert "config-level instructions" not in user_prompt
|
||||
|
||||
def test_falls_back_to_custom_instructions_when_no_prompt(self, mock_memory):
|
||||
mock_memory._add_to_vector_store(
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
metadata={},
|
||||
filters={},
|
||||
infer=True,
|
||||
)
|
||||
|
||||
user_prompt = mock_memory.llm.generate_response.call_args[1]["messages"][1]["content"]
|
||||
assert "config-level instructions" in user_prompt
|
||||
|
||||
|
||||
class TestAsyncUpdate:
|
||||
@pytest.fixture
|
||||
def mock_async_memory(self, mocker):
|
||||
|
||||
+1
-1
@@ -61,7 +61,7 @@ def test_add(memory_instance):
|
||||
assert result["results"] == [{"memory": "Test memory", "event": "ADD"}]
|
||||
|
||||
memory_instance._add_to_vector_store.assert_called_once_with(
|
||||
[{"role": "user", "content": "Test message"}], {"user_id": "test_user"}, {"user_id": "test_user"}, True
|
||||
[{"role": "user", "content": "Test message"}], {"user_id": "test_user"}, {"user_id": "test_user"}, True, prompt=None
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user