Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/memos/api/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -427,7 +427,7 @@ def get_reader_config() -> dict[str, Any]:
"config": {
"chunk_type": os.getenv("MEM_READER_CHAT_CHUNK_TYPE", "default"),
"chunk_length": int(os.getenv("MEM_READER_CHAT_CHUNK_TOKEN_SIZE", 1600)),
"chunk_session": int(os.getenv("MEM_READER_CHAT_CHUNK_SESS_SIZE", 20)),
"chunk_session": int(os.getenv("MEM_READER_CHAT_CHUNK_SESS_SIZE", 10)),
"chunk_overlap": int(os.getenv("MEM_READER_CHAT_CHUNK_OVERLAP", 2)),
},
}
Expand Down
2 changes: 1 addition & 1 deletion src/memos/configs/memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ class TreeTextMemoryConfig(BaseTextMemoryConfig):
),
)

search_strategy: dict[str, bool] | None = Field(
search_strategy: dict[str, Any] | None = Field(
default=None,
description=(
'Set search strategy for this memory configuration.{"bm25": true, "cot": false}'
Expand Down
15 changes: 14 additions & 1 deletion src/memos/mem_reader/strategy_struct.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def _get_llm_response(self, mem_str: str) -> dict:
template = STRATEGY_PROMPT_DICT["chat"][lang]
examples = STRATEGY_PROMPT_DICT["chat"][f"{lang}_example"]
prompt = template.replace("${conversation}", mem_str)
if self.config.remove_prompt_example:
if self.config.remove_prompt_example: # TODO unused
prompt = prompt.replace(examples, "")
messages = [{"role": "user", "content": prompt}]
try:
Expand Down Expand Up @@ -112,6 +112,19 @@ def get_scene_data_info(self, scene_data: list, type: str) -> list[str]:

results.append([overlap_item, item])
current_length = overlap_length + content_length
else:
cut_size, cut_overlap = (
self.chat_chunker["chunk_session"],
self.chat_chunker["chunk_overlap"],
)
for items in scene_data:
step = cut_size - cut_overlap
end = len(items) - cut_overlap
if end <= 0:
results.extend([items[:]])
else:
results.extend([items[i : i + cut_size] for i in range(0, end, step)])

elif type == "doc":
parser_config = ParserConfigFactory.model_validate(
{
Expand Down
4 changes: 3 additions & 1 deletion src/memos/memories/textual/simple_tree.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,9 @@ def __init__(
time_start_bm = time.time()
self.search_strategy = config.search_strategy
self.bm25_retriever = (
EnhancedBM25() if self.search_strategy and self.search_strategy["bm25"] else None
EnhancedBM25()
if self.search_strategy and self.search_strategy.get("bm25", False)
else None
)
logger.info(f"time init: bm25_retriever time is: {time.time() - time_start_bm}")

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,3 +13,4 @@ class ParsedTaskGoal:
rephrased_query: str | None = None
internet_search: bool = False
goal_type: str | None = None # e.g., 'default', 'explanation', etc.
context: str = ""
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ def find_project_root(marker=".git"):
if (current / marker).exists():
return current
current = current.parent
logger.warn(f"The project root directory tag file was not found: {marker}")
return Path(".")


PROJECT_ROOT = find_project_root()
Expand Down
17 changes: 7 additions & 10 deletions src/memos/memories/textual/tree_text_memory/retrieve/searcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@

logger = get_logger(__name__)
COT_DICT = {
"fast": {"en": COT_PROMPT, "zh": COT_PROMPT_ZH},
"fine": {"en": SIMPLE_COT_PROMPT, "zh": SIMPLE_COT_PROMPT_ZH},
"fine": {"en": COT_PROMPT, "zh": COT_PROMPT_ZH},
"fast": {"en": SIMPLE_COT_PROMPT, "zh": SIMPLE_COT_PROMPT_ZH},
}


Expand Down Expand Up @@ -59,12 +59,8 @@ def __init__(
# Create internet retriever from config if provided
self.internet_retriever = internet_retriever
self.moscube = moscube
self.vec_cot = (
search_strategy.get("vec_cot", "false") == "true" if search_strategy else False
)
self.use_fast_graph = (
search_strategy.get("fast_graph", "false") == "true" if search_strategy else False
)
self.vec_cot = search_strategy.get("cot", False) if search_strategy else False
self.use_fast_graph = search_strategy.get("fast_graph", False) if search_strategy else False

self._usage_executor = ContextThreadPoolExecutor(max_workers=4, thread_name_prefix="usage")

Expand Down Expand Up @@ -287,6 +283,7 @@ def _retrieve_paths(
search_filter,
user_name,
id_filter,
mode=mode,
)
)
tasks.append(
Expand Down Expand Up @@ -369,6 +366,7 @@ def _retrieve_from_long_term_and_user(
search_filter: dict | None = None,
user_name: str | None = None,
id_filter: dict | None = None,
mode: str = "fast",
):
"""Retrieve and rerank from LongTermMemory and UserMemory"""
results = []
Expand All @@ -377,7 +375,7 @@ def _retrieve_from_long_term_and_user(
# chain of thinking
cot_embeddings = []
if self.vec_cot:
queries = self._cot_query(query)
queries = self._cot_query(query, mode=mode, context=parsed_goal.context)
if len(queries) > 1:
cot_embeddings = self.embedder.embed(queries)
cot_embeddings.extend(query_embedding)
Expand Down Expand Up @@ -566,7 +564,6 @@ def _cot_query(
prompt = template.replace("${original_query}", query).replace(
"${split_num_threshold}", str(split_num)
)
logger.info("COT process")

messages = [{"role": "user", "content": prompt}]
try:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ def parse(
- mode == 'fine': use LLM to parse structured topic/keys/tags
"""
if mode == "fast":
return self._parse_fast(task_description, **kwargs)
return self._parse_fast(task_description, context=context, **kwargs)
elif mode == "fine":
if not self.llm:
raise ValueError("LLM not provided for slow mode.")
Expand All @@ -51,6 +51,7 @@ def _parse_fast(self, task_description: str, **kwargs) -> ParsedTaskGoal:
"""
Fast mode: simple jieba word split.
"""
context = kwargs.get("context", "")
use_fast_graph = kwargs.get("use_fast_graph", False)
if use_fast_graph:
desc_tokenized = self.tokenizer.tokenize_mixed(task_description)
Expand All @@ -61,6 +62,7 @@ def _parse_fast(self, task_description: str, **kwargs) -> ParsedTaskGoal:
goal_type="default",
rephrased_query=task_description,
internet_search=False,
context=context,
)
else:
return ParsedTaskGoal(
Expand All @@ -70,6 +72,7 @@ def _parse_fast(self, task_description: str, **kwargs) -> ParsedTaskGoal:
goal_type="default",
rephrased_query=task_description,
internet_search=False,
context=context,
)

def _parse_fine(
Expand All @@ -91,16 +94,17 @@ def _parse_fine(
logger.info(f"Parsing Goal... LLM input is {prompt}")
response = self.llm.generate(messages=[{"role": "user", "content": prompt}])
logger.info(f"Parsing Goal... LLM Response is {response}")
return self._parse_response(response)
return self._parse_response(response, context=context)
except Exception:
logger.warning(f"Fail to fine-parse query {query}: {traceback.format_exc()}")
return self._parse_fast(query)
return self._parse_fast(query, context=context)

def _parse_response(self, response: str) -> ParsedTaskGoal:
def _parse_response(self, response: str, **kwargs) -> ParsedTaskGoal:
"""
Parse LLM JSON output safely.
"""
try:
context = kwargs.get("context", "")
response = response.replace("```", "").replace("json", "").strip()
response_json = eval(response)
return ParsedTaskGoal(
Expand All @@ -110,6 +114,7 @@ def _parse_response(self, response: str) -> ParsedTaskGoal:
rephrased_query=response_json.get("rephrased_instruction", None),
internet_search=response_json.get("internet_search", False),
goal_type=response_json.get("goal_type", "default"),
context=context,
)
except Exception as e:
raise ValueError(f"Failed to parse LLM output: {e}\nRaw response:\n{response}") from e
Loading
Loading