From 395e41041a31c6cc239fbb9e9121c5054ad07744 Mon Sep 17 00:00:00 2001 From: JiehoonKwak Date: Mon, 10 Aug 2026 20:36:09 +0900 Subject: [PATCH] fix(search): keep mental-model derived state current Build embeddings and VectorChord lexical projections from the same canonical name-plus-content document for partial updates and clears. Keep Knowledge Page renames synchronized with their backing mental model so dense and lexical search state cannot silently retain the old page name.\n\nThis narrows the former PR after upstream #3318: backend query dispatch and schema reconciliation are intentionally excluded and will be handled separately.\n\nContext:\n- #3318 now owns backend-neutral Knowledge Page search dispatch.\n- Name-only and content-only writes previously embedded incomplete text.\n- Clear and page rename left dense embeddings stale.\n- Provider calls remain outside pooled database connections. --- .../hindsight_api/engine/memory_engine.py | 170 ++++++++++++------ .../tests/test_knowledge_base.py | 40 +++++ .../tests/test_mental_models.py | 21 ++- hindsight-api-slim/tests/test_reflections.py | 53 ++++++ 4 files changed, 231 insertions(+), 53 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index c4c9338b02..cee767621f 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -12792,34 +12792,43 @@ async def update_mental_model( await self._validate_operation(self._operation_validator.validate_bank_write(ctx)) backend = await self._get_backend() - # Compute the new embedding BEFORE acquiring a pooled connection: a slow - # embedder must never pin a DB connection. The embedding text depends only - # on the incoming name/content, never on DB state, so it can be done here. + # Resolve the complete searchable document before embedding it. A name-only + # or content-only update must include the stored half; otherwise the dense + # embedding and lexical projection describe different text. Keep the slow + # provider call outside pooled connections. + previous_content: str | None = None + previous_reflect_response: dict[str, Any] | None = None + effective_name = "" + effective_content = "" new_embedding_str: str | None = None - if content is not None: - embedding_text = f"{name or ''} {content}" - embedding = await embedding_utils.generate_embeddings_batch(self.embeddings, [embedding_text]) - if embedding: - new_embedding_str = str(embedding[0]) - - async with acquire_with_retry(backend) as conn: - # If content is changing, fetch current content + reflect_response to record history - previous_content: str | None = None - previous_reflect_response: dict[str, Any] | None = None - if content is not None: + document_changed = name is not None or content is not None + if document_changed: + async with acquire_with_retry(backend) as conn: current_row = await conn.fetchrow( - f"SELECT content, reflect_response FROM {fq_table('mental_models')} WHERE bank_id = $1 AND id = $2", + f"SELECT name, content, reflect_response FROM {fq_table('mental_models')} " + "WHERE bank_id = $1 AND id = $2", bank_id, mental_model_id, ) - if current_row: - previous_content = current_row["content"] - raw_rr = current_row["reflect_response"] - if isinstance(raw_rr, str): - previous_reflect_response = json.loads(raw_rr) if raw_rr else None - else: - previous_reflect_response = raw_rr + if current_row is None: + return None + + effective_name = name if name is not None else (current_row["name"] or "") + effective_content = content if content is not None else (current_row["content"] or "") + if content is not None: + previous_content = current_row["content"] + raw_rr = current_row["reflect_response"] + if isinstance(raw_rr, str): + previous_reflect_response = json.loads(raw_rr) if raw_rr else None + else: + previous_reflect_response = raw_rr + + embedding_text = f"{effective_name} {effective_content}" + embedding = await embedding_utils.generate_embeddings_batch(self.embeddings, [embedding_text]) + if embedding: + new_embedding_str = str(embedding[0]) + async with acquire_with_retry(backend) as conn: # Build dynamic update updates = [] params: list[Any] = [bank_id, mental_model_id] @@ -12829,22 +12838,14 @@ async def update_mental_model( record_mm_history = False slim_reflect_response: dict[str, Any] | None = None - # Track the SQL for the search_vector source columns: the new bind - # placeholder when the field is being updated, else the existing column - # (unchanged). Used to re-tokenize search_vector for vchord below. - name_sql = "name" - content_sql = "content" - if name is not None: updates.append(f"name = ${param_idx}") params.append(name) - name_sql = f"${param_idx}" param_idx += 1 if content is not None: updates.append(f"content = ${param_idx}") params.append(content) - content_sql = f"${param_idx}" param_idx += 1 if refresh_watermark is None: updates.append("last_refreshed_at = NOW()") @@ -12872,11 +12873,6 @@ async def update_mental_model( slim["trace"] = previous_trace slim_reflect_response = slim or None record_mm_history = True - # Apply the embedding computed above (off-connection). - if new_embedding_str is not None: - updates.append(f"embedding = ${param_idx}") - params.append(new_embedding_str) - param_idx += 1 elif refresh_watermark is not None: # A successful delta refresh can find no topic-relevant facts even though # the coarse staleness query found new rows. Advance the watermark to the @@ -12887,6 +12883,13 @@ async def update_mental_model( params.append(refresh_watermark) param_idx += 1 + if document_changed: + # If the provider returns no vector, NULL is safer than retaining + # an embedding for the previous document. + updates.append(f"embedding = ${param_idx}") + params.append(new_embedding_str) + param_idx += 1 + if reflect_response is not None: updates.append(f"reflect_response = ${param_idx}") params.append(json.dumps(reflect_response)) @@ -12922,16 +12925,20 @@ async def update_mental_model( params.append(json.dumps(structured_content)) param_idx += 1 - # Re-tokenize search_vector when the searchable text (name/content) - # changed, but only for vchord — its bm25vector column is written - # inline (native is a GENERATED column that updates itself; the other - # backends index base columns). Same helper as the insert/recall paths. - if name is not None or content is not None: + # Re-tokenize VectorChord from the same canonical document embedded + # above. Native is generated; other backends index base columns. + if document_changed: sv_expr = pg_search_vector_expr( - get_config(), text_col=name_sql, context_col=content_sql, signals_col=None, native_inline=False + get_config(), + text_col=f"${param_idx}", + context_col=f"${param_idx + 1}", + signals_col=None, + native_inline=False, ) if sv_expr: updates.append(f"search_vector = {sv_expr}") + params.extend([effective_name, effective_content]) + param_idx += 2 if not updates: return None @@ -13037,9 +13044,23 @@ async def clear_mental_model( await self._validate_operation(self._operation_validator.validate_bank_write(ctx)) backend = await self._get_backend() - # Content is cleared to '', so re-tokenize search_vector from the name - # alone — vchord only (see update_mental_model). Non-vchord backends leave - # the column untouched (generated / base-column indexed). + async with acquire_with_retry(backend) as conn: + current = await conn.fetchrow( + f"SELECT name FROM {fq_table('mental_models')} WHERE bank_id = $1 AND id = $2", + bank_id, + mental_model_id, + ) + if current is None: + return None + + embedding = await embedding_utils.generate_embeddings_batch( + self.embeddings, + [f"{current['name'] or ''} "], + ) + embedding_str = str(embedding[0]) if embedding else None + + # Content is cleared to '', so re-tokenize VectorChord from the name alone. + # Native is generated; other backends index base columns. sv_expr = pg_search_vector_expr( get_config(), text_col="name", context_col="''", signals_col=None, native_inline=False ) @@ -13050,7 +13071,8 @@ async def clear_mental_model( UPDATE {fq_table("mental_models")} SET content = '', structured_content = NULL, - last_refreshed_source_query = NULL{sv_clause} + last_refreshed_source_query = NULL, + embedding = $3{sv_clause} WHERE bank_id = $1 AND id = $2 RETURNING id, bank_id, name, source_query, content, tags, last_refreshed_at, created_at, reflect_response, @@ -13058,6 +13080,7 @@ async def clear_mental_model( """, bank_id, mental_model_id, + embedding_str, ) return self._row_to_mental_model(row) if row else None @@ -13491,21 +13514,64 @@ async def search_knowledge_pages( async def rename_knowledge_node( self, bank_id: str, node_id: str, name: str, *, request_context: "RequestContext" ) -> dict[str, Any] | None: - """Rename a folder or page node.""" + """Rename a folder or a page and its backing mental-model document.""" await self._authenticate_tenant(request_context) backend = await self._get_backend() + async with acquire_with_retry(backend) as conn: - row = await conn.fetchrow( + current = await conn.fetchrow( f""" - UPDATE {fq_table("knowledge_pages")} - SET name = $3, updated_at = now() - WHERE bank_id = $1 AND id = $2 - RETURNING {self._KP_COLUMNS} + SELECT kp.kind, kp.mental_model_id, mm.content AS mm_content + FROM {fq_table("knowledge_pages")} kp + LEFT JOIN {fq_table("mental_models")} mm + ON mm.id = kp.mental_model_id AND mm.bank_id = kp.bank_id + WHERE kp.bank_id = $1 AND kp.id = $2 """, bank_id, node_id, - name, ) + if current is None: + return None + + is_backed_page = current["kind"] == "page" and current["mental_model_id"] is not None + embedding_str: str | None = None + if is_backed_page: + embedding = await embedding_utils.generate_embeddings_batch( + self.embeddings, + [f"{name} {current['mm_content'] or ''}"], + ) + embedding_str = str(embedding[0]) if embedding else None + + sv_expr = pg_search_vector_expr( + get_config(), text_col="$3", context_col="content", signals_col=None, native_inline=False + ) + sv_clause = f", search_vector = {sv_expr}" if sv_expr else "" + + async with acquire_with_retry(backend) as conn: + async with conn.transaction(): + row = await conn.fetchrow( + f""" + UPDATE {fq_table("knowledge_pages")} + SET name = $3, updated_at = now() + WHERE bank_id = $1 AND id = $2 + RETURNING {self._KP_COLUMNS} + """, + bank_id, + node_id, + name, + ) + if row is not None and is_backed_page: + await conn.execute( + f""" + UPDATE {fq_table("mental_models")} + SET name = $3, embedding = $4{sv_clause} + WHERE bank_id = $1 AND id = $2 + """, + bank_id, + current["mental_model_id"], + name, + embedding_str, + ) return self._row_to_knowledge_node(row) if row else None async def update_knowledge_page( diff --git a/hindsight-api-slim/tests/test_knowledge_base.py b/hindsight-api-slim/tests/test_knowledge_base.py index b30fb7c471..a3294f8a55 100644 --- a/hindsight-api-slim/tests/test_knowledge_base.py +++ b/hindsight-api-slim/tests/test_knowledge_base.py @@ -12,6 +12,7 @@ import pytest_asyncio from hindsight_api.engine.memory_engine import MemoryEngine, _may_need_refresh +from hindsight_api.engine.retain import embedding_utils def _enc(bank_id: str) -> str: @@ -319,6 +320,45 @@ async def test_rename(self, api_client, kb_bank): assert resp.status_code == 200, resp.text assert resp.json()["name"] == "Compliance" + async def test_page_rename_updates_backing_search_document( + self, api_client, kb_bank, memory: MemoryEngine, request_context, monkeypatch + ): + bank_id, ids = kb_bank + async with memory._pool.acquire() as conn: + old_embedding = await conn.fetchval( + "SELECT embedding::text FROM mental_models WHERE bank_id = $1 AND id = $2", + bank_id, + ids.orders_mm, + ) + + original_generate = embedding_utils.generate_embeddings_batch + embedded_documents: list[list[str]] = [] + + async def recording_generate(*args, **kwargs): + embedded_documents.append(args[1]) + return await original_generate(*args, **kwargs) + + monkeypatch.setattr(embedding_utils, "generate_embeddings_batch", recording_generate) + resp = await api_client.patch( + f"/v1/default/banks/{_enc(bank_id)}/knowledge-base/nodes/{ids.orders}", + json={"name": "Order Operations"}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["name"] == "Order Operations" + assert embedded_documents == [["Order Operations # Orders\n\nOne row per order."]] + + async with memory._pool.acquire() as conn: + backing = await conn.fetchrow( + "SELECT name, embedding::text AS embedding, search_vector::text AS search_vector " + "FROM mental_models WHERE bank_id = $1 AND id = $2", + bank_id, + ids.orders_mm, + ) + assert backing["name"] == "Order Operations" + assert backing["embedding"] != old_embedding + assert "oper" in backing["search_vector"] + assert "order" in backing["search_vector"] + async def test_update_page_options(self, api_client, kb_bank): bank_id, ids = kb_bank resp = await api_client.patch( diff --git a/hindsight-api-slim/tests/test_mental_models.py b/hindsight-api-slim/tests/test_mental_models.py index 0a20d61459..1778f2a580 100644 --- a/hindsight-api-slim/tests/test_mental_models.py +++ b/hindsight-api-slim/tests/test_mental_models.py @@ -2255,8 +2255,15 @@ def test_trigger_tag_groups_rejects_invalid(self): class TestClearMentalModel: """Test clear_mental_model resets content so next refresh is full.""" - async def test_clear_resets_content(self, memory: MemoryEngine, request_context): + async def test_clear_resets_content(self, memory: MemoryEngine, request_context, monkeypatch): """Clear sets content to empty string and nulls structured/tracking fields.""" + original_generate = embedding_utils.generate_embeddings_batch + embedded_documents: list[list[str]] = [] + + async def recording_generate(*args, **kwargs): + embedded_documents.append(args[1]) + return await original_generate(*args, **kwargs) + bank_id = f"test-mm-clear-{uuid.uuid4().hex[:8]}" await memory.get_bank_profile(bank_id, request_context=request_context) @@ -2269,6 +2276,8 @@ async def test_clear_resets_content(self, memory: MemoryEngine, request_context) ) assert mm["content"] == "Some existing content" + monkeypatch.setattr(embedding_utils, "generate_embeddings_batch", recording_generate) + cleared = await memory.clear_mental_model( bank_id=bank_id, mental_model_id=mm["id"], @@ -2282,6 +2291,16 @@ async def test_clear_resets_content(self, memory: MemoryEngine, request_context) # Re-fetch to confirm persistence fetched = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context) assert fetched["content"] == "" + assert embedded_documents == [["Test Model "]] + async with memory._pool.acquire() as conn: + search_vector = await conn.fetchval( + "SELECT search_vector::text FROM mental_models WHERE bank_id = $1 AND id = $2", + bank_id, + mm["id"], + ) + assert "exist" not in search_vector + assert "content" not in search_vector + assert "test" in search_vector and "model" in search_vector await memory.delete_bank(bank_id, request_context=request_context) diff --git a/hindsight-api-slim/tests/test_reflections.py b/hindsight-api-slim/tests/test_reflections.py index 838b4acc61..35c1b2407f 100644 --- a/hindsight-api-slim/tests/test_reflections.py +++ b/hindsight-api-slim/tests/test_reflections.py @@ -194,6 +194,59 @@ async def test_update_mental_model(self, memory: MemoryEngine, request_context): # Cleanup await memory.delete_bank(bank_id, request_context=request_context) + @pytest.mark.asyncio + async def test_partial_updates_rebuild_canonical_search_document( + self, memory: MemoryEngine, request_context, monkeypatch + ): + """Each partial edit embeds the other stored half and refreshes FTS.""" + from hindsight_api.engine.retain import embedding_utils + + bank_id = f"test-mental-model-name-{uuid.uuid4().hex[:8]}" + mental_model = await memory.create_mental_model( + bank_id=bank_id, + name="Original Name", + source_query="Original Query", + content="Canonical Body", + request_context=request_context, + ) + + original_generate = embedding_utils.generate_embeddings_batch + embedded_documents: list[list[str]] = [] + + async def recording_generate(*args, **kwargs): + embedded_documents.append(args[1]) + return await original_generate(*args, **kwargs) + + monkeypatch.setattr(embedding_utils, "generate_embeddings_batch", recording_generate) + updated = await memory.update_mental_model( + bank_id=bank_id, + mental_model_id=mental_model["id"], + name="Updated Name", + request_context=request_context, + ) + + assert updated is not None and updated["name"] == "Updated Name" + assert embedded_documents == [["Updated Name Canonical Body"]] + + updated = await memory.update_mental_model( + bank_id=bank_id, + mental_model_id=mental_model["id"], + content="Revised Body", + request_context=request_context, + ) + assert updated is not None and updated["content"] == "Revised Body" + assert embedded_documents == [["Updated Name Canonical Body"], ["Updated Name Revised Body"]] + + async with memory._pool.acquire() as conn: + search_vector = await conn.fetchval( + "SELECT search_vector::text FROM mental_models WHERE bank_id = $1 AND id = $2", + bank_id, + mental_model["id"], + ) + assert "updat" in search_vector + assert "revis" in search_vector + await memory.delete_bank(bank_id, request_context=request_context) + @pytest.mark.asyncio async def test_delete_mental_model(self, memory: MemoryEngine, request_context): """Test deleting a mental model."""