Skip to content

Performance optimization for searching memories#2142

Merged
bittergreen merged 5 commits into
MemTensor:dev-v2.0.25from
endxxxx:feat/performance-optimization-no-routing-v2.0.25
Jul 23, 2026
Merged

Performance optimization for searching memories#2142
bittergreen merged 5 commits into
MemTensor:dev-v2.0.25from
endxxxx:feat/performance-optimization-no-routing-v2.0.25

Conversation

@endxxxx

@endxxxx endxxxx commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

Description

Please include a summary of the change, the problem it solves, the implementation approach, and relevant context. List any dependencies required for this change.

Related Issue (Required): Fixes #issue_number

Type of change

Please delete options that are not relevant.

  • Bug fix (non-breaking change which fixes an issue)
  • Refactor (does not change functionality, e.g. code style improvements, linting)

How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

  • Unit Test
  • Test Script Or Test Steps (please provide)
  • Pipeline Automated API Test (please provide)

Checklist

  • I have performed a self-review of my own code | 我已自行检查了自己的代码
  • I have commented my code in hard-to-understand areas | 我已在难以理解的地方对代码进行了注释
  • I have added tests that prove my fix is effective or that my feature works | 我已添加测试以证明我的修复有效或功能正常
  • I have created related documentation issue/PR in MemOS-Docs (if applicable) | 我已在 MemOS-Docs 中创建了相关的文档 issue/PR(如果适用)
  • I have linked the issue to this PR (if applicable) | 我已将 issue 链接到此 PR(如果适用)
  • I have mentioned the person who will review this PR | 我已提及将审查此 PR 的人

Reviewer Checklist

  • closes #xxxx (Replace xxxx with the GitHub issue number)
  • Made sure Checks passed
  • Tests have been provided

@Memtensor-AI Memtensor-AI added area:api 云服务 / FastAPI / OpenAPI / MCP area:database graph_db + vector_db | 图数据库与向量数据库 area:memcube GeneralMemCube / cube 生命周期 / cube 配置 area:memory 记忆存储、检索、更新、召回逻辑 area:model llm + embedder + reranker status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Jul 22, 2026
@Memtensor-AI

Copy link
Copy Markdown
Collaborator

🤖 Open Code Review

Target: PR #2142
Task: manual-ocr:MemTensor/MemOS#2142
Base: dev-v2.0.25
Head: feat/performance-optimization-no-routing-v2.0.25

🔍 OpenCodeReview found 19 issue(s) in this PR.

⚠️ 1 warning(s) occurred during review.


1. src/memos/embedders/base.py (L55-L58)

Embedding failures are logged at logger.info instead of logger.error (or at minimum logger.warning). Since ark, ollama, and sentence_transformer embedders have no independent error-level logging for embed() failures, operators filtering on ERROR/WARNING will never see these failures from this decorator. Change the error branch to logger.error.

💡 Suggested Change

Before:

            if error_type is None:
                logger.info(log_message, *log_values)
            else:
                logger.info(log_message + " error_type=%s", *log_values, error_type)

After:

            if error_type is None:
                logger.info(log_message, *log_values)
            else:
                logger.error(log_message + " error_type=%s", *log_values, error_type)

2. src/memos/embedders/base.py (L22-L32)

list(texts or []) materializes the input into normalized_texts for logging purposes, but func is then called with the original texts. If a caller ever passes a generator or iterator (instead of the declared list[str]), the iterator will be exhausted by the list(...) call and func will receive an empty iterator, silently producing no embeddings. Pass the materialized normalized_texts to func as well to be safe.

💡 Suggested Change

Before:

        normalized_texts = [texts] if isinstance(texts, str) else list(texts or [])
        text_lengths = [len(str(text or "")) for text in normalized_texts]
        config = getattr(self, "config", None)
        model = getattr(config, "model_name_or_path", None) or "unknown"
        backup_model = getattr(config, "backup_model_name_or_path", None) or "none"
        backup_enabled = bool(getattr(self, "use_backup_client", False))
        started_at = time.perf_counter()
        status = "success"
        error_type = None
        try:
            return func(self, texts, *args, **kwargs)

After:

        normalized_texts = [texts] if isinstance(texts, str) else list(texts or [])
        text_lengths = [len(str(text or "")) for text in normalized_texts]
        config = getattr(self, "config", None)
        model = getattr(config, "model_name_or_path", None) or "unknown"
        backup_model = getattr(config, "backup_model_name_or_path", None) or "none"
        backup_enabled = bool(getattr(self, "use_backup_client", False))
        started_at = time.perf_counter()
        status = "success"
        error_type = None
        try:
            return func(self, normalized_texts, *args, **kwargs)

3. src/memos/embedders/base.py (L51-L53)

text_hash is called unconditionally on every embedding invocation — even when the INFO log level is disabled — computing a SHA-256 over the entire text batch (O(total characters)). This is wasted CPU on every hot-path embedding call when INFO is suppressed. Guard it with a level check, or defer the computation using a lazy wrapper.

💡 Suggested Change

Before:

                text_hash(normalized_texts),
                elapsed_ms,
                status,

After:

                text_hash(normalized_texts) if logger.isEnabledFor(logging.INFO) else "skipped",
                elapsed_ms,
                status,

4. src/memos/api/handlers/search_handler.py (L608-L620)

Silent data loss when top_k=0: When target_top_k is 0, candidate_limit becomes 0, so the if candidate_limit > 0 branch is skipped, selected stays [], and bucket['memories'] is overwritten with an empty list — silently dropping all candidates in that bucket before the MMR step.

Although the current call site passes search_req.top_k (which is typically positive), a defensive guard is needed to avoid catastrophic data loss when the value is zero. Consider treating candidate_limit == 0 as "no limit" (i.e., skip pruning for that bucket) rather than "keep nothing".

💡 Suggested Change

Before:

                selected: list[dict[str, Any]] = []
                if candidate_limit > 0:
                    for typed_memories in memories_by_type.values():
                        selected.extend(
                            sorted(
                                typed_memories,
                                key=self._mmr_candidate_score,
                                reverse=True,
                            )[:candidate_limit]
                        )
                    selected.sort(key=self._mmr_candidate_score, reverse=True)

                bucket["memories"] = selected

After:

                if candidate_limit == 0:
                    # top_k=0 means "no limit" — keep all candidates untouched
                    continue

                selected: list[dict[str, Any]] = []
                for typed_memories in memories_by_type.values():
                    selected.extend(
                        sorted(
                            typed_memories,
                            key=self._mmr_candidate_score,
                            reverse=True,
                        )[:candidate_limit]
                    )
                selected.sort(key=self._mmr_candidate_score, reverse=True)

                bucket["memories"] = selected

5. src/memos/api/handlers/search_handler.py (L606)

memory_type or result_key conflates empty-string type with untyped memories: The expression memory_type or result_key falls back to result_key not only when memory_type is None, but also when it is an empty string "" (falsy). If memory_type="" is a valid distinct type, it would be incorrectly grouped together with None-typed memories under the result_key bucket, and they would share the same candidate_limit, conflating two separate categories.

Prefer an explicit None check:

💡 Suggested Change

Before:

                    memories_by_type.setdefault(str(memory_type or result_key), []).append(memory)

After:

                    memories_by_type.setdefault(str(memory_type if memory_type is not None else result_key), []).append(memory)

6. src/memos/log.py (L76-L77)

The or [] fallback is dead code. The enclosing condition isinstance(bucket.get("memories"), list) already guarantees the value is a list before this branch executes. An empty list [] is falsy, so [] or [] still evaluates to [] — the fallback never produces a different result. Replace with len(bucket.get("memories")) to avoid misleading readers into thinking the value could be None at this point.

💡 Suggested Change

Before:

            len(bucket.get("memories") or [])
            if isinstance(bucket, Mapping) and isinstance(bucket.get("memories"), list)

After:

            len(bucket.get("memories"))
            if isinstance(bucket, Mapping) and isinstance(bucket.get("memories"), list)

7. src/memos/embedders/cache.py (L88)

Bug (medium): _request_cache_text_limit uses the wrong config variable.

cache_max_size is the maximum number of texts in the global TTL cache, but it's reused here to cap each per-request LRUCache. These are different dimensions: the outer TTLCache is sized by request_cache_max_requests, while each inner LRUCache should have its own independently configurable size. With the current code, setting MEMOS_EMBEDDING_CACHE_MAX_SIZE=64 (intended to bound global TTL memory) also silently limits every per-request LRU to 64 entries — causing unexpected within-request cache misses if a single request embeds more than 64 unique texts. The variable request_cache_max_requests is read from env just above but is only used for the outer TTLCache.

Suggested fix: introduce a dedicated env variable (e.g. MEMOS_EMBEDDING_REQUEST_CACHE_TEXT_LIMIT) with its own default, or at minimum use the existing _DEFAULT_CACHE_MAX_SIZE explicitly as a documented fallback rather than re-referencing cache_max_size.


8. src/memos/embedders/cache.py (L30)

Maintainability (low): "trace-id" is a hardcoded copy of a default defined in another module.

In src/memos/context/context.py line 41, RequestContext.__init__ defaults trace_id to "trace-id" when none is supplied: self.trace_id = trace_id or "trace-id". This cache module replicates that sentinel as a magic string. If the default placeholder in context.py ever changes, this set will silently fail to exclude uninstrumented requests, causing all of them to share the same per-request LRU cache entry — every no-trace-id request would cross-contaminate each other's cache.

Suggested fix: export the sentinel constant from context.py (e.g., DEFAULT_TRACE_ID = "trace-id") and import it here, or expose a predicate like is_placeholder_trace_id(tid) to keep the two modules in sync.


9. src/memos/graph_dbs/polardb.py (L42-L43)

TypeError is dead code here: json.loads raises TypeError only when its argument is not str, bytes, or bytearray, but the surrounding isinstance(value, str) guard already ensures the argument is always a str. More importantly, both exception types are swallowed silently with no log entry. If a field is stored as malformed agtype JSON, the caller silently receives a raw string where it may expect a list or dict, producing subtle downstream type errors that are very hard to diagnose. Add at minimum a logger.debug (or logger.warning) when falling back, and drop the dead TypeError clause.

💡 Suggested Change

Before:

        except (json.JSONDecodeError, TypeError):
            return value

After:

        except json.JSONDecodeError:
            logger.debug("Failed to JSON-decode agtype value %r; returning raw string", value)
            return value

10. src/memos/graph_dbs/polardb.py (L1957-L1959)

Minor log-formatting inconsistency: the conditional branch that logs the full query still has a leading space (" search_by_embedding query: %s"), while this new branch omits it. Either add the leading space here or remove it from the other branch for consistency.


11. tests/embedders/test_base.py (L27-L29)

mock_logger.info.call_args is None when the mock was never invoked (documented unittest.mock behaviour). If a regression in log_embedding_call causes it to skip the logger.info call, .call_args.args raises AttributeError: 'NoneType' object has no attribute 'args' instead of a clean assertion failure, making the root cause hard to diagnose. Add mock_logger.info.assert_called_once() before dereferencing call_args in all three tests.

💡 Suggested Change

Before:

    log_args = mock_logger.info.call_args.args
    rendered = log_args[0] % log_args[1:]
    assert "model=embedding-model" in rendered

After:

    mock_logger.info.assert_called_once()
    log_args = mock_logger.info.call_args.args
    rendered = log_args[0] % log_args[1:]
    assert "model=embedding-model" in rendered

12. tests/embedders/test_base.py (L55-L57)

Same call_args dereference without a prior assert_called_once() guard. If the error-path branch of log_embedding_call skips logger.info, the failure surfaces as AttributeError rather than a clear test failure.

💡 Suggested Change

Before:

    log_args = mock_logger.info.call_args.args
    rendered = log_args[0] % log_args[1:]
    assert "status=failed" in rendered

After:

    mock_logger.info.assert_called_once()
    log_args = mock_logger.info.call_args.args
    rendered = log_args[0] % log_args[1:]
    assert "status=failed" in rendered

13. tests/embedders/test_base.py (L77-L79)

Same call_args dereference without a prior assert_called_once() guard. A decorator regression that suppresses the logger.info call would yield a confusing AttributeError here instead of a meaningful assertion failure.

💡 Suggested Change

Before:

    log_args = mock_logger.info.call_args.args
    rendered = log_args[0] % log_args[1:]
    assert "model=primary-model" in rendered

After:

    mock_logger.info.assert_called_once()
    log_args = mock_logger.info.call_args.args
    rendered = log_args[0] % log_args[1:]
    assert "model=primary-model" in rendered

14. tests/api/test_mmr_candidate_pruning.py (L99-L100)

The assertion relies on the pruned list being returned in score-descending insertion order, but _prune_mmr_candidates_by_bucket accumulates results via selected.extend(sorted(...)) and then re-sorts by score — the ordering contract is sound here. However, the real fragility is that the test verifies exactly two elements without asserting which two were dropped. With text_top_k=1 and _MMR_CANDIDATE_MULTIPLIER=2, the pruner should keep 1 * 2 = 2 candidates. If _MMR_CANDIDATE_MULTIPLIER is changed (it is a module-level constant with no indirection), this assertion silently breaks (either wrong count or wrong items). Consider either importing and using _MMR_CANDIDATE_MULTIPLIER explicitly in the assertion, or parameterising the multiplier so the expected slice is derived rather than hardcoded.


15. tests/api/test_mmr_candidate_pruning.py (L67-L72)

Using a set literal hides ordering, which is intentional for multi-type mixing — good. But note that _prune_mmr_candidates_by_bucket always performs a final selected.sort(key=self._mmr_candidate_score, reverse=True) on the combined list. The pref_mem assertion below does use a list and relies on that sort. These two assertions are inconsistent: either both can assert an ordered list (since the sort is deterministic for distinct scores), or both should use sets. Using a list assertion throughout would catch if the sort contract is accidentally dropped.


16. tests/embedders/test_cache.py (L140)

Using assert to wait for backend_started is a synchronisation barrier that is stripped under python -O, turning the barrier into a no-op. Thread 2 is then submitted before Thread 1 has entered blocking_embed and registered an in-flight request, so Thread 2 may start its own independent backend call instead of joining Thread 1's. singleflight_joins never reaches 1 and the polling loop below hangs (or, with the deadline assertion also stripped, spins forever). Replace with an explicit check.

💡 Suggested Change

Before:

        assert backend_started.wait(timeout=2)

After:

        if not backend_started.wait(timeout=2):
            pytest.fail("backend thread did not start within deadline")

17. tests/embedders/test_cache.py (L145)

Using assert for deadline enforcement is stripped under python -O, turning the bounded wait into an infinite spin. If singleflight_joins never reaches 1 (e.g., because Thread 2 raced ahead of Thread 1 due to the stripped barrier above), the test process hangs indefinitely instead of reporting a clear failure.

💡 Suggested Change

Before:

            assert time.monotonic() < deadline

After:

            if time.monotonic() >= deadline:
                pytest.fail("singleflight join not observed within deadline")

18. tests/embedders/test_cache.py (L132)

Using assert to validate threading.Event.wait inside a worker thread is stripped under python -O. If release_backend is never set (because the main thread's polling loop hanged), the 2-second wait returns False and the worker silently continues, returning the correct value. The test then passes while the singleflight mechanism was never exercised.

💡 Suggested Change

Before:

        assert release_backend.wait(timeout=2)

After:

        if not release_backend.wait(timeout=2):
            raise RuntimeError("release_backend was not set within deadline")

19. tests/reranker/test_local_rerankers.py (L101-L105)

next() on the bare generator raises StopIteration (surfaced by pytest as a confusing error, not a clean assertion failure) when the expected log record is absent — e.g., if caplog does not capture the logger used inside CosineLocalReranker because its logger is not propagating to the root logger. Provide a sentinel default and assert on it:

message = next(
    (
        record.getMessage()
        for record in caplog.records
        if "CosineLocalReranker rerank result" in record.getMessage()
    ),
    None,
)
assert message is not None, "Expected 'CosineLocalReranker rerank result' log record not found"

Generated by cloud-assistant via Open Code Review.

@Memtensor-AI

Copy link
Copy Markdown
Collaborator

✅ Automated Test Results: PASSED

All tests passed (57/57 executed). memos_python_core/changed-repo-python: 57/57. Duration: 6s [advisory, non-gating] AI-generated tests on branch test/auto-gen-3d2915ca2b1b0243-20260723102845: 94/99 passed, 5 failed — these do NOT affect the PR verdict; review the branch manually.

Branch: feat/performance-optimization-no-routing-v2.0.25

@Memtensor-AI Memtensor-AI added status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 and removed status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Jul 23, 2026
@bittergreen
bittergreen merged commit 7945f28 into MemTensor:dev-v2.0.25 Jul 23, 2026
18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:api 云服务 / FastAPI / OpenAPI / MCP area:database graph_db + vector_db | 图数据库与向量数据库 area:memcube GeneralMemCube / cube 生命周期 / cube 配置 area:memory 记忆存储、检索、更新、召回逻辑 area:model llm + embedder + reranker status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants