Add config option for vertex embedding tasks (#2266)

2025-02-28 18:50:05 +09:00
parent 8143f86be6
commit 6acb00731d
14 changed files with 141 additions and 48 deletions
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -119,7 +119,7 @@ def test_search(memory_instance, version, enable_graph):
    memory_instance.vector_store.search.assert_called_once_with(
        query=[0.1, 0.2, 0.3], limit=100, filters={"user_id": "test_user"}
    )
-    memory_instance.embedding_model.embed.assert_called_once_with("test query")
+    memory_instance.embedding_model.embed.assert_called_once_with("test query", "search")

    if enable_graph:
        memory_instance.graph.search.assert_called_once_with("test query", {"user_id": "test_user"}, 100)