feat: add kwargs for long-term-memory viking (#340)

Author: zakahan

veadk/memory/long_term_memory.py

@@ -249,6 +249,7 @@ def _filter_and_convert_events(self, events: list[Event]) -> list[str]:
     async def add_session_to_memory(
         self,
         session: Session,
+        **kwargs,
     ):
         """Add a chat session's events to the long-term memory backend.
 
@@ -283,7 +284,12 @@ async def add_session_to_memory(
         logger.info(
             f"Adding {len(event_strings)} events to long term memory: index={self.index}"
         )
-        self._backend.save_memory(user_id=user_id, event_strings=event_strings)
+        if self.backend == "viking":
+            self._backend.save_memory(
+                user_id=user_id, event_strings=event_strings, **kwargs
+            )
+        else:
+            self._backend.save_memory(user_id=user_id, event_strings=event_strings)
         logger.info(
             f"Added {len(event_strings)} events to long term memory: index={self.index}, user_id={user_id}"
         )

veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py

@@ -120,6 +120,7 @@ def _get_client(self) -> VikingDBMemoryClient:
 
     @override
     def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
+        assistant_id = kwargs.get("assistant_id", "assistant")
         session_id = str(uuid.uuid1())
         messages = []
         for raw_events in event_strings:
@@ -131,7 +132,7 @@ def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
             messages.append({"role": role, "content": content})
         metadata = {
             "default_user_id": user_id,
-            "default_assistant_id": "assistant",
+            "default_assistant_id": assistant_id,
             "time": int(time.time() * 1000),
         }
 

veadk/runner.py

@@ -674,7 +674,7 @@ async def save_eval_set(self, session_id: str, eval_set_id: str = "default") ->
         return eval_set_path
 
     async def save_session_to_long_term_memory(
-        self, session_id: str, user_id: str = "", app_name: str = ""
+        self, session_id: str, user_id: str = "", app_name: str = "", **kwargs
     ) -> None:
         """Save the specified session to long-term memory.
 
@@ -730,5 +730,5 @@ async def save_session_to_long_term_memory(
             )
             return
 
-        await self.long_term_memory.add_session_to_memory(session)
+        await self.long_term_memory.add_session_to_memory(session, kwargs=kwargs)
         logger.info(f"Add session `{session.id}` to long term memory.")