Added media enrichment to embedding process, with a test

Author: jamesbraza

src/paperqa/docs.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import json
 import logging
 import os
@@ -481,7 +482,16 @@ async def aadd_texts(
         if embedding_model and texts[0].embedding is None:
             for t, t_embedding in zip(
                 texts,
-                await embedding_model.embed_documents(texts=[t.text for t in texts]),
+                await embedding_model.embed_documents(
+                    texts=await asyncio.gather(
+                        *(
+                            t.get_embeddable_text(
+                                all_settings.parsing.should_parse_and_enrich_media[1]
+                            )
+                            for t in texts
+                        )
+                    )
+                ),
                 strict=True,
             ):
                 t.embedding = t_embedding
@@ -535,14 +545,20 @@ def delete(
         self.deleted_dockeys.add(dockey)
         self.texts = list(filter(lambda x: x.doc.dockey != dockey, self.texts))
 
-    async def _build_texts_index(self, embedding_model: EmbeddingModel) -> None:
+    async def _build_texts_index(
+        self, embedding_model: EmbeddingModel, with_enrichment: bool = False
+    ) -> None:
         texts = [t for t in self.texts if t not in self.texts_index]
         # For any embeddings we are supposed to lazily embed, embed them now
         to_embed = [t for t in texts if t.embedding is None]
         if to_embed:
             for t, t_embedding in zip(
                 to_embed,
-                await embedding_model.embed_documents(texts=[t.text for t in to_embed]),
+                await embedding_model.embed_documents(
+                    texts=await asyncio.gather(
+                        *(t.get_embeddable_text(with_enrichment) for t in to_embed)
+                    )
+                ),
                 strict=True,
             ):
                 t.embedding = t_embedding
@@ -564,7 +580,10 @@ async def retrieve_texts(
         # TODO: should probably happen elsewhere
         self.texts_index.mmr_lambda = settings.texts_index_mmr_lambda
 
-        await self._build_texts_index(embedding_model)
+        await self._build_texts_index(
+            embedding_model,
+            with_enrichment=settings.parsing.should_parse_and_enrich_media[1],
+        )
         _k = k + len(self.deleted_dockeys)
         matches: list[Text] = cast(
             "list[Text]",

src/paperqa/types.py

@@ -173,6 +173,33 @@ def __eq__(self, other) -> bool:
     def __hash__(self) -> int:
         return hash((self.name, self.text))
 
+    async def get_embeddable_text(self, with_enrichment: bool = False) -> str:
+        """Get the text to embed, which may be different from the actual text content.
+
+        This method is async so subclassers could use custom enrichment logic here.
+
+        Args:
+            with_enrichment: Opt-in flag to include media enrichment in the return.
+                Media enrichment can improve placement in embedding space,
+                without affecting the text used for quotation.
+
+        Returns:
+            Content to embed.
+        """
+        if not with_enrichment:
+            return self.text
+        # Media enrichment can improve placement in embedding space,
+        # without affecting the text used for quotation
+        enriched_media = (
+            (
+                f"Media {m.index} from page {m.info.get('page_num', 'unknown')!s}'s"
+                f" enriched description:\n\n{m.info['enriched_description']!s}"
+            )
+            for m in self.media
+            if m.info.get("enriched_description")
+        )
+        return "\n\n".join((self.text, *enriched_media))
+
 
 # Sentinel to autopopulate a field within model_validator
 AUTOPOPULATE_VALUE = ""  # NOTE: this is falsy by design