Skip to content

Commit fe0624e

Browse files
authored
feat: mos add load sdk for user (#263)
* feat: mos add load sdk for user * feat: remove think for reason model
1 parent e729218 commit fe0624e

File tree

2 files changed

+26
-0
lines changed

2 files changed

+26
-0
lines changed

src/memos/llms/vllm.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ def _generate_with_api_client(self, messages: list[MessageDict]) -> str:
105105
"temperature": float(getattr(self.config, "temperature", 0.8)),
106106
"max_tokens": int(getattr(self.config, "max_tokens", 1024)),
107107
"top_p": float(getattr(self.config, "top_p", 0.9)),
108+
"extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
108109
}
109110

110111
response = self.client.chat.completions.create(**completion_kwargs)
@@ -142,6 +143,7 @@ def generate_stream(self, messages: list[MessageDict]):
142143
"max_tokens": int(getattr(self.config, "max_tokens", 1024)),
143144
"top_p": float(getattr(self.config, "top_p", 0.9)),
144145
"stream": True, # Enable streaming
146+
"extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
145147
}
146148

147149
stream = self.client.chat.completions.create(**completion_kwargs)

src/memos/mem_os/core.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -960,6 +960,30 @@ def dump(
960960
self.mem_cubes[mem_cube_id].dump(dump_dir)
961961
logger.info(f"MemCube {mem_cube_id} dumped to {dump_dir}")
962962

963+
def load(
964+
self,
965+
load_dir: str,
966+
user_id: str | None = None,
967+
mem_cube_id: str | None = None,
968+
memory_types: list[Literal["text_mem", "act_mem", "para_mem"]] | None = None,
969+
) -> None:
970+
"""Dump the MemCube to a dictionary.
971+
Args:
972+
load_dir (str): The directory to load the MemCube from.
973+
user_id (str, optional): The identifier of the user to load the MemCube from.
974+
If None, the default user is used.
975+
mem_cube_id (str, optional): The identifier of the MemCube to load.
976+
If None, the default MemCube for the user is used.
977+
"""
978+
target_user_id = user_id if user_id is not None else self.user_id
979+
accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
980+
if not mem_cube_id:
981+
mem_cube_id = accessible_cubes[0].cube_id
982+
if mem_cube_id not in self.mem_cubes:
983+
raise ValueError(f"MemCube with ID {mem_cube_id} does not exist. please regiester")
984+
self.mem_cubes[mem_cube_id].load(load_dir, memory_types=memory_types)
985+
logger.info(f"MemCube {mem_cube_id} loaded from {load_dir}")
986+
963987
def get_user_info(self) -> dict[str, Any]:
964988
"""Get current user information including accessible cubes.
965989

0 commit comments

Comments
 (0)