diff --git a/src/atocore/memory/service.py b/src/atocore/memory/service.py index 3c9f607..2979f50 100644 --- a/src/atocore/memory/service.py +++ b/src/atocore/memory/service.py @@ -389,29 +389,37 @@ def get_memories_for_context( if not query_tokens: query_tokens = None - # Flat budget across types so paragraph-length project memories - # aren't starved by an even slice. Types are still walked in order - # (identity/preference first when they're the input), so earlier - # types still get first pick when the budget is tight. + # Collect ALL candidates across the requested types into one + # pool, then rank globally before the budget walk. Ranking per + # type and walking types in order would starve later types when + # the first type's candidates filled the budget — even if a + # later-type candidate matched the query perfectly. Type order + # is preserved as a stable tiebreaker inside + # ``_rank_memories_for_query`` via Python's stable sort. + pool: list[Memory] = [] + seen_ids: set[str] = set() for mtype in memory_types: - # Raise the fetch limit above the budget slice so query-relevance - # ordering has a real pool to rerank. Without a query, the extras - # just fall off the end harmlessly. - candidates = get_memories( + for mem in get_memories( memory_type=mtype, project=project, min_confidence=0.5, limit=30, - ) - if query_tokens is not None: - candidates = _rank_memories_for_query(candidates, query_tokens) - for mem in candidates: - entry = f"[{mem.memory_type}] {mem.content}" - entry_len = len(entry) + 1 - if entry_len > available - used: + ): + if mem.id in seen_ids: continue - selected_entries.append(entry) - used += entry_len + seen_ids.add(mem.id) + pool.append(mem) + + if query_tokens is not None: + pool = _rank_memories_for_query(pool, query_tokens) + + for mem in pool: + entry = f"[{mem.memory_type}] {mem.content}" + entry_len = len(entry) + 1 + if entry_len > available - used: + continue + selected_entries.append(entry) + used += entry_len if not selected_entries: return "", 0