diff --git a/application/database/db.py b/application/database/db.py index 22e7358be..5db13c1c6 100644 --- a/application/database/db.py +++ b/application/database/db.py @@ -2382,6 +2382,7 @@ def delete_gapanalysis_results_for(self, node_name): return res def add_cre(self, cre: cre_defs.CRE) -> CRE: + """Add a CRE with metadata, or fill missing fields on an existing CRE.""" entry: CRE = None query = self.session.query(CRE).filter(func.lower(CRE.name) == cre.name.lower()) if cre.id: @@ -2412,6 +2413,7 @@ def add_cre(self, cre: cre_defs.CRE) -> CRE: name=cre.name, external_id=cre.id, tags=",".join([str(t) for t in cre.tags]), + metadata_json=cre.metadata or {}, ) self.session.add(entry) self.session.commit() diff --git a/application/prompt_client/prompt_client.py b/application/prompt_client/prompt_client.py index 81dddcd95..6e015fd2c 100644 --- a/application/prompt_client/prompt_client.py +++ b/application/prompt_client/prompt_client.py @@ -885,14 +885,14 @@ def query_llm(self, raw_question: str) -> str: return _extract_content_text(resp) def generate_embeddings_for(self, item_name: str): - # CRE embeddings are generated from the CRE's textual fields only - # (name/description/id). That path does not require fetching remote - # content via Playwright, so we can skip browser launch entirely and - # batch requests to avoid provider rate limits. + """Generate missing or changed embeddings using provider-ready CRE text.""" + # CRE fields and metadata do not need remote content, so skip + # Playwright and batch embedding requests to avoid provider rate limits. if item_name == cre_defs.Credoctypes.CRE.value: # list_cre_ids() returns plain CRE database ids (strings). cre_ids = self.database.list_cre_ids() pending: List[str] = [] + content_by_id: Dict[str, str] = {} cre_by_id: Dict[str, cre_defs.CRE] = {} for cid in cre_ids: @@ -903,7 +903,9 @@ def generate_embeddings_for(self, item_name: str): if getattr(cre, "metadata", None): metadata_json = stable_json(getattr(cre, "metadata", None)) embedding_text = f"{embedding_text}\nmetadata:{metadata_json}" - embedding_text = normalize_embeddings_content(embedding_text) + embedding_text = self._truncate_one( + normalize_embeddings_content(embedding_text) + ).rstrip() existing = self.database.get_embedding(cid) if ( existing @@ -914,6 +916,7 @@ def generate_embeddings_for(self, item_name: str): ): continue pending.append(cid) + content_by_id[cid] = embedding_text cre_by_id[cid] = cre if not pending: @@ -928,13 +931,7 @@ def generate_embeddings_for(self, item_name: str): for i in range(0, len(pending), batch_size): batch_ids = pending[i : i + batch_size] - contents: List[str] = [] - for cid in batch_ids: - cre = cre_by_id[cid] - contents.append( - f"{cre.doctype}\n name:{cre.name}\n description:{cre.description}\n id:{cre.id}\n " - ) - contents = [normalize_embeddings_content(c) for c in contents] + contents = [content_by_id[cid] for cid in batch_ids] embeddings = self.ai_client.get_text_embeddings(contents) # type: ignore[arg-type] if not embeddings: diff --git a/application/tests/db_test.py b/application/tests/db_test.py index 68362cfc2..11b440ea2 100644 --- a/application/tests/db_test.py +++ b/application/tests/db_test.py @@ -76,6 +76,40 @@ def setUp(self) -> None: self.collection = collection + def test_add_cre_preserves_metadata(self) -> None: + """Round-trip nested metadata when adding a new CRE.""" + metadata = {"domain": "authentication", "details": {"levels": [1, 2]}} + row = self.collection.add_cre( + defs.CRE(id="123-456", name="Metadata round trip", metadata=metadata) + ) + self.collection.session.expire_all() + + self.assertEqual(row.metadata_json, metadata) + self.assertEqual(self.collection.get_cre_by_db_id(row.id).metadata, metadata) + + def test_add_cre_without_metadata(self) -> None: + """Persist an empty metadata object when none is supplied.""" + row = self.collection.add_cre(defs.CRE(id="123-457", name="No metadata")) + self.collection.session.expire_all() + + self.assertEqual(self.collection.get_cre_by_db_id(row.id).metadata, {}) + + def test_add_cre_keeps_existing_metadata(self) -> None: + """Keep stored metadata unchanged when an existing CRE is added again.""" + row = self.collection.add_cre(defs.CRE(id="123-458", name="Existing metadata")) + row.metadata_json = {"source": "original"} + self.collection.session.commit() + + updated = self.collection.add_cre( + defs.CRE( + id="123-458", name="Existing metadata", metadata={"source": "incoming"} + ) + ) + self.collection.session.expire_all() + + self.assertEqual(updated.id, row.id) + self.assertEqual(updated.metadata_json, {"source": "original"}) + def test_get_by_tags(self) -> None: """ Given: A CRE with no links and a combination of possible tags: diff --git a/application/tests/prompt_client_incremental_embeddings_test.py b/application/tests/prompt_client_incremental_embeddings_test.py index 0ef70785f..109e3056a 100644 --- a/application/tests/prompt_client_incremental_embeddings_test.py +++ b/application/tests/prompt_client_incremental_embeddings_test.py @@ -4,7 +4,7 @@ import unittest from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch from application.defs import cre_defs from application.prompt_client import prompt_client @@ -157,6 +157,138 @@ def test_prompt_handler_cre_embeddings_skip_unchanged(self): args = fake_db.add_embedding.call_args[0] self.assertEqual(args[0].id, "db-4") + def test_prompt_handler_cre_metadata_cache(self): + """Reuse CRE embeddings until the metadata sent to the provider changes.""" + for metadata in ( + {}, + {"domain": "authentication", "details": {"levels": [1, 2]}}, + ): + with self.subTest(metadata=metadata): + fake_db = _FakeDB() + cre = cre_defs.CRE( + id="333-333", name="CRE-3", description="desc-3", metadata=metadata + ) + fake_db._cre_by_id = {"db-3": cre} + + def store(db_obj, _doc_type, _emb, text, fake_db=fake_db): + fake_db._emb_by_id[db_obj.id] = SimpleNamespace( + embeddings_content=text + ) + + fake_db.add_embedding.side_effect = store + ph = prompt_client.PromptHandler.__new__(prompt_client.PromptHandler) + ph.database = fake_db + ph.ai_client = Mock(spec=["get_max_batch_size", "get_text_embeddings"]) + ph.ai_client.get_max_batch_size.return_value = 16 + ph.ai_client.get_text_embeddings.return_value = [[0.1, 0.2]] + + ph.generate_embeddings_for(cre_defs.Credoctypes.CRE.value) + text = ph.ai_client.get_text_embeddings.call_args[0][0][0] + self.assertEqual(fake_db._emb_by_id["db-3"].embeddings_content, text) + if metadata: + self.assertIn( + "metadata:" + prompt_client.stable_json(metadata), text + ) + + ph.generate_embeddings_for(cre_defs.Credoctypes.CRE.value) + self.assertEqual(ph.ai_client.get_text_embeddings.call_count, 1) + + cre.metadata = {"domain": "cryptography"} + ph.generate_embeddings_for(cre_defs.Credoctypes.CRE.value) + self.assertEqual(ph.ai_client.get_text_embeddings.call_count, 2) + updated_text = ph.ai_client.get_text_embeddings.call_args[0][0][0] + self.assertIn('metadata:{"domain":"cryptography"}', updated_text) + self.assertEqual( + fake_db._emb_by_id["db-3"].embeddings_content, updated_text + ) + ph.generate_embeddings_for(cre_defs.Credoctypes.CRE.value) + self.assertEqual(ph.ai_client.get_text_embeddings.call_count, 2) + + def test_prompt_handler_cre_provider_limit(self): + """Cache and provider text stay identical at the 8,000-character limit.""" + for length in (7999, 8000, 8001, 8200): + with self.subTest(length=length): + cre = cre_defs.CRE(id="333-333", name="CRE-3", description="") + prefix = prompt_client.normalize_embeddings_content( + f"{cre.doctype} name:{cre.name} description: id:{cre.id}" + ) + cre.description = "x" * (length - len(prefix)) + if length == 8200: + # Put a space at the cutoff to check normalized cache hits. + index = 7999 - len(f"{cre.doctype} name:{cre.name} description:") + cre.description = ( + cre.description[:index] + " " + cre.description[index + 1 :] + ) + fake_db = _FakeDB() + fake_db._cre_by_id = {"db-3": cre} + + def store(db_obj, _doc_type, _emb, text, fake_db=fake_db): + fake_db._emb_by_id[db_obj.id] = SimpleNamespace( + embeddings_content=text + ) + + fake_db.add_embedding.side_effect = store + ph = prompt_client.PromptHandler.__new__(prompt_client.PromptHandler) + ph.database = fake_db + ph.ai_client = ph + ph._litellm = Mock() + ph.embed_model = "test-model" + ph._llm_max_retries = 0 + ph._llm_retry_sleep_seconds = 0 + ph._expected_embed_dim = None + with patch.object( + prompt_client.litellm_router, + "embedding", + return_value={"data": [{"embedding": [0.1, 0.2]}]}, + ) as provider: + ph.generate_embeddings_for(cre_defs.Credoctypes.CRE.value) + stored = fake_db._emb_by_id["db-3"].embeddings_content + expected_length = 7999 if length == 8200 else min(length, 8000) + self.assertEqual(len(stored), expected_length) + self.assertEqual(provider.call_args.kwargs["input"], [stored]) + ph.generate_embeddings_for(cre_defs.Credoctypes.CRE.value) + self.assertEqual(provider.call_count, 1) + + def test_prompt_handler_cre_metadata_beyond_provider_limit(self): + """Changes outside the provider text do not invalidate a cached vector.""" + fake_db = _FakeDB() + cre = cre_defs.CRE( + id="333-333", + name="CRE-3", + description="x" * 8100, + metadata={"value": "before"}, + ) + fake_db._cre_by_id = {"db-3": cre} + + def store(db_obj, _doc_type, _emb, text): + fake_db._emb_by_id[db_obj.id] = SimpleNamespace(embeddings_content=text) + + fake_db.add_embedding.side_effect = store + ph = prompt_client.PromptHandler.__new__(prompt_client.PromptHandler) + ph.database = fake_db + ph.ai_client = ph + ph._litellm = Mock() + ph.embed_model = "test-model" + ph._llm_max_retries = 0 + ph._llm_retry_sleep_seconds = 0 + ph._expected_embed_dim = None + with patch.object( + prompt_client.litellm_router, + "embedding", + return_value={"data": [{"embedding": [0.1, 0.2]}]}, + ) as provider: + ph.generate_embeddings_for(cre_defs.Credoctypes.CRE.value) + cre.metadata = {"value": "after"} + ph.generate_embeddings_for(cre_defs.Credoctypes.CRE.value) + self.assertEqual(provider.call_count, 1) + self.assertEqual(fake_db.add_embedding.call_count, 1) + + cre.description = "changed" + cre.description + ph.generate_embeddings_for(cre_defs.Credoctypes.CRE.value) + self.assertEqual(provider.call_count, 2) + stored = fake_db._emb_by_id["db-3"].embeddings_content + self.assertEqual(provider.call_args.kwargs["input"], [stored]) + def test_generate_embeddings_url_without_remote_text_uses_stored_fields(self): """If ``get_content`` returns nothing, embed from DB-backed node fields (no skip).""" fake_db = _FakeDB()