Refine paraphrase prompt; forward media modalities and support JSON mode

src/gabriel/api.py

@@ -1133,6 +1133,8 @@ async def paraphrase(
     reset_files: bool = False,
     reasoning_effort: Optional[str] = None,
     reasoning_summary: Optional[str] = None,
+    modality: str = "text",
+    search_context_size: str = "medium",
     n_rounds: int = 1,
     recursive_validation: Optional[bool] = None,
     n_initial_candidates: int = 1,
@@ -1171,6 +1173,10 @@ async def paraphrase(
         Whether to request JSON responses.
     web_search:
         Enable web search augmentation when supported by the model.
+    modality:
+        Modality of inputs (text, image, audio, pdf, entity, web).
+    search_context_size:
+        Web search context size when ``modality="web"``.
     n_parallels:
         Maximum concurrent paraphrase calls.
     use_dummy:
@@ -1220,6 +1226,8 @@ async def paraphrase(
         use_dummy=use_dummy,
         reasoning_effort=reasoning_effort,
         reasoning_summary=reasoning_summary,
+        modality=modality,
+        search_context_size=search_context_size,
         n_rounds=n_rounds,
         recursive_validation=recursive_validation,
         n_initial_candidates=n_initial_candidates,

src/gabriel/prompts/paraphrase_prompt.jinja2

@@ -1,17 +1,24 @@
-Consider the following text:
-BEGIN TEXT
-{{ text }}
-END TEXT
-Your task is to rewrite the above text in accordance with the following instructions:
-BEGIN INSTRUCTIONS
-{{ instructions }}
-END INSTRUCTIONS
-The text should be rewritten in a way that is consistent with the instructions.
-Be very diligent about following the instructions.
-Output only your revised text, nothing else, no preamble, no postamble, no explanation, no nothing other than the revised text.
-Consider the whole text, and follow the instructions carefully and completely.
-Be meticulous and exacting; your output of rewritten text must be as faithful as possible in following the instructions.
-Once again, the instructions are:
+{% import "snippets.jinja2" as snip %}
+{{ snip.single_entry(modality, text) }}
+
+Task: rewrite the provided content according to the instructions below. Be diligent, thorough, and faithful to the instructions in all cases.
+
 BEGIN INSTRUCTIONS
 {{ instructions }}
 END INSTRUCTIONS
+
+Modality guidance:
+- text: paraphrase the full passage, preserving meaning and key details while applying the instructions.
+- entity: rewrite the entity name in a consistent way per the instructions.
+- web: paraphrase/consolidate the web content found for the search term per the instructions.
+- image/audio/pdf: paraphrase the observed content into text in accordance with the instructions.
+
+Output format (template):
+{% if json_mode %}
+Always output JSON only in the indicated format.
+{"revised_text": "<revised_text>"}
+{% else %}
+<revised_text>
+{% endif %}
+
+No preamble, no explanation, only the desired revised text alone, faithful to the provided content and in accordance with the instructions.

src/gabriel/tasks/paraphrase.py

@@ -12,6 +12,12 @@
 
 from ..core.prompt_template import PromptTemplate, resolve_template
 from ..utils.openai_utils import get_all_responses
+from ..utils import (
+    load_audio_inputs,
+    load_image_inputs,
+    load_pdf_inputs,
+    warn_if_modality_mismatch,
+)
 from ..utils.logging import announce_prompt_rendering
 
 # Import classifier utilities for recursive validation.  Importing from
@@ -50,6 +56,10 @@ class ParaphraseConfig:
     # augmentation.  ``None`` defers to downstream defaults while ``True`` and
     # ``False`` explicitly enable or disable the feature.
     web_search: Optional[bool] = None
+    # What kind of input is being paraphrased (text, image, audio, pdf, entity, web).
+    modality: str = "text"
+    # Web search context size when modality is ``web``.
+    search_context_size: str = "medium"
     # Maximum number of parallel requests that will be sent to the
     # underlying API.  Note that classification and paraphrasing share
     # this value for simplicity.
@@ -155,7 +165,9 @@ async def run(
         # Convert the target column into a list of strings.  We coerce
         # values to strings so that non-string columns (e.g. numbers) are
         # handled gracefully.
-        texts: List[str] = df_proc[column_name].astype(str).tolist()
+        values = df_proc[column_name].tolist()
+        texts: List[str] = [str(v) for v in values]
+        warn_if_modality_mismatch(values, self.cfg.modality, column_name=column_name)
         # Determine the base name for the revised column(s).
         base_col = self.cfg.revised_column_name or f"{column_name}_revised"
         # Determine how many paraphrases to produce per passage.  A value
@@ -190,19 +202,67 @@ async def run(
             prompts: List[str] = []
             identifiers: List[str] = []
             for idx, text in enumerate(texts):
+                prompt_text = (
+                    text
+                    if self.cfg.modality in {"text", "entity", "web"}
+                    else ""
+                )
                 for j in range(1, n + 1):
                     prompts.append(
-                        self.template.render(text=text, instructions=self.cfg.instructions)
+                        self.template.render(
+                            text=prompt_text,
+                            instructions=self.cfg.instructions,
+                            modality=self.cfg.modality,
+                            json_mode=self.cfg.json_mode,
+                        )
                     )
                     identifiers.append(f"row_{idx}_rev{j}")
             save_path = os.path.join(self.cfg.save_dir, self.cfg.file_name)
+            prompt_images: Optional[Dict[str, List[str]]] = None
+            prompt_audio: Optional[Dict[str, List[Dict[str, str]]]] = None
+            prompt_pdfs: Optional[Dict[str, List[Dict[str, str]]]] = None
+
+            if self.cfg.modality == "image":
+                tmp: Dict[str, List[str]] = {}
+                for idx, val in enumerate(values):
+                    imgs = load_image_inputs(val)
+                    if imgs:
+                        for j in range(1, n + 1):
+                            tmp[f"row_{idx}_rev{j}"] = imgs
+                prompt_images = tmp or None
+            elif self.cfg.modality == "audio":
+                tmp_a: Dict[str, List[Dict[str, str]]] = {}
+                for idx, val in enumerate(values):
+                    auds = load_audio_inputs(val)
+                    if auds:
+                        for j in range(1, n + 1):
+                            tmp_a[f"row_{idx}_rev{j}"] = auds
+                prompt_audio = tmp_a or None
+            elif self.cfg.modality == "pdf":
+                tmp_p: Dict[str, List[Dict[str, str]]] = {}
+                for idx, val in enumerate(values):
+                    pdfs = load_pdf_inputs(val)
+                    if pdfs:
+                        for j in range(1, n + 1):
+                            tmp_p[f"row_{idx}_rev{j}"] = pdfs
+                prompt_pdfs = tmp_p or None
+
+            web_search = (
+                self.cfg.web_search
+                if self.cfg.web_search is not None
+                else self.cfg.modality == "web"
+            )
+            kwargs.setdefault("web_search", web_search)
+            kwargs.setdefault("search_context_size", self.cfg.search_context_size)
             resp_df = await get_all_responses(
                 prompts=prompts,
                 identifiers=identifiers,
+                prompt_images=prompt_images,
+                prompt_audio=prompt_audio,
+                prompt_pdfs=prompt_pdfs,
                 save_path=save_path,
                 model=self.cfg.model,
                 json_mode=self.cfg.json_mode,
-                web_search=self.cfg.web_search,
                 n_parallels=self.cfg.n_parallels,
                 use_dummy=self.cfg.use_dummy,
                 reset_files=reset_files,
@@ -227,6 +287,7 @@ async def run(
             resp_map: Dict[Tuple[int, int], str] = {}
             await self._recursive_validate(
                 texts,
+                values,
                 resp_map,
                 approval_map,
                 reset_files=reset_files,
@@ -261,6 +322,7 @@ async def run(
     async def _recursive_validate(
         self,
         original_texts: List[str],
+        original_values: List[Any],
         resp_map: Dict[Tuple[int, int], str],
         approval_map: Dict[Tuple[int, int], bool],
         *,
@@ -358,13 +420,27 @@ async def _recursive_validate(
                 # exists for this key, use that paraphrase as the base
                 # for regeneration.  Otherwise, continue to use the
                 # original.
-                if round_number > 0 and self.cfg.use_modified_source and key in resp_map:
+                if (
+                    round_number > 0
+                    and self.cfg.use_modified_source
+                    and self.cfg.modality in {"text", "entity", "web"}
+                    and key in resp_map
+                ):
                     base_text = resp_map[key]
                 else:
-                    base_text = original_texts[row_idx]
+                    base_text = (
+                        original_texts[row_idx]
+                        if self.cfg.modality in {"text", "entity", "web"}
+                        else ""
+                    )
                 for cand_idx in range(candidates_per_key):
                     prompts.append(
-                        self.template.render(text=base_text, instructions=self.cfg.instructions)
+                        self.template.render(
+                            text=base_text,
+                            instructions=self.cfg.instructions,
+                            modality=self.cfg.modality,
+                            json_mode=self.cfg.json_mode,
+                        )
                     )
                     # Encode row, revision, round and candidate index in
                     # the identifier.  Revision numbers are stored one-
@@ -389,13 +465,60 @@ async def _recursive_validate(
                 self.cfg.save_dir,
                 f"{os.path.splitext(self.cfg.file_name)[0]}_round{round_number}.csv",
             )
+            prompt_images: Optional[Dict[str, List[str]]] = None
+            prompt_audio: Optional[Dict[str, List[Dict[str, str]]]] = None
+            prompt_pdfs: Optional[Dict[str, List[Dict[str, str]]]] = None
+
+            if self.cfg.modality == "image":
+                tmp: Dict[str, List[str]] = {}
+                for key in to_check:
+                    row_idx, rev_idx = key
+                    imgs = load_image_inputs(original_values[row_idx])
+                    if imgs:
+                        for cand_idx in range(candidates_per_key):
+                            tmp[
+                                f"row_{row_idx}_rev{rev_idx + 1}_round{round_number}_cand{cand_idx}"
+                            ] = imgs
+                prompt_images = tmp or None
+            elif self.cfg.modality == "audio":
+                tmp_a: Dict[str, List[Dict[str, str]]] = {}
+                for key in to_check:
+                    row_idx, rev_idx = key
+                    auds = load_audio_inputs(original_values[row_idx])
+                    if auds:
+                        for cand_idx in range(candidates_per_key):
+                            tmp_a[
+                                f"row_{row_idx}_rev{rev_idx + 1}_round{round_number}_cand{cand_idx}"
+                            ] = auds
+                prompt_audio = tmp_a or None
+            elif self.cfg.modality == "pdf":
+                tmp_p: Dict[str, List[Dict[str, str]]] = {}
+                for key in to_check:
+                    row_idx, rev_idx = key
+                    pdfs = load_pdf_inputs(original_values[row_idx])
+                    if pdfs:
+                        for cand_idx in range(candidates_per_key):
+                            tmp_p[
+                                f"row_{row_idx}_rev{rev_idx + 1}_round{round_number}_cand{cand_idx}"
+                            ] = pdfs
+                prompt_pdfs = tmp_p or None
+
+            web_search = (
+                self.cfg.web_search
+                if self.cfg.web_search is not None
+                else self.cfg.modality == "web"
+            )
             new_resp_df = await get_all_responses(
                 prompts=prompts,
                 identifiers=identifiers,
+                prompt_images=prompt_images,
+                prompt_audio=prompt_audio,
+                prompt_pdfs=prompt_pdfs,
                 save_path=tmp_save_path,
                 model=self.cfg.model,
                 json_mode=self.cfg.json_mode,
-                web_search=self.cfg.web_search,
+                web_search=web_search,
+                search_context_size=self.cfg.search_context_size,
                 n_parallels=self.cfg.n_parallels,
                 use_dummy=self.cfg.use_dummy,
                 reset_files=reset_files,

tests/test_basic.py

@@ -1582,6 +1582,73 @@ def test_paraphrase_api(tmp_path):
     assert "txt_revised_1" in df_multi.columns and "txt_revised_2" in df_multi.columns
 
 
+def test_paraphrase_modalities_forward_media(monkeypatch, tmp_path):
+    captured = {}
+
+    async def fake_get_all_responses(*, prompts, identifiers, **kwargs):
+        captured["prompts"] = list(prompts)
+        captured["identifiers"] = list(identifiers)
+        captured["prompt_images"] = kwargs.get("prompt_images")
+        captured["prompt_audio"] = kwargs.get("prompt_audio")
+        captured["prompt_pdfs"] = kwargs.get("prompt_pdfs")
+        captured["web_search"] = kwargs.get("web_search")
+        captured["search_context_size"] = kwargs.get("search_context_size")
+        return pd.DataFrame({"Identifier": identifiers, "Response": prompts})
+
+    monkeypatch.setattr("gabriel.tasks.paraphrase.get_all_responses", fake_get_all_responses)
+
+    df_image = pd.DataFrame({"media": ["data:image/png;base64,xyz"]})
+    asyncio.run(
+        gabriel.paraphrase(
+            df_image,
+            "media",
+            instructions="caption it",
+            save_dir=str(tmp_path / "para_image"),
+            modality="image",
+        )
+    )
+    assert captured["prompt_images"]["row_0_rev1"] == ["data:image/png;base64,xyz"]
+
+    df_audio = pd.DataFrame({"media": [{"format": "mp3", "data": "abcd"}]})
+    asyncio.run(
+        gabriel.paraphrase(
+            df_audio,
+            "media",
+            instructions="summarize it",
+            save_dir=str(tmp_path / "para_audio"),
+            modality="audio",
+        )
+    )
+    assert captured["prompt_audio"]["row_0_rev1"][0]["format"] == "mp3"
+
+    df_pdf = pd.DataFrame({"media": ["data:application/pdf;base64,abcd"]})
+    asyncio.run(
+        gabriel.paraphrase(
+            df_pdf,
+            "media",
+            instructions="rewrite",
+            save_dir=str(tmp_path / "para_pdf"),
+            modality="pdf",
+        )
+    )
+    assert captured["prompt_pdfs"]["row_0_rev1"][0]["file_data"].startswith(
+        "data:application/pdf"
+    )
+
+    df_web = pd.DataFrame({"query": ["Mount Fuji"]})
+    asyncio.run(
+        gabriel.paraphrase(
+            df_web,
+            "query",
+            instructions="summarize",
+            save_dir=str(tmp_path / "para_web"),
+            modality="web",
+        )
+    )
+    assert captured["web_search"] is True
+    assert captured["search_context_size"] == "medium"
+
+
 def test_paraphrase_n_rounds_not_forwarded(monkeypatch, tmp_path):
     captured = {}
 
@@ -1607,7 +1674,16 @@ async def fake_get_all_responses(*, prompts, identifiers, **kwargs):
 
 
 def test_paraphrase_approval_column_recursive(monkeypatch, tmp_path):
-    async def fake_recursive_validate(self, original_texts, resp_map, approval_map, *, reset_files, max_rounds):
+    async def fake_recursive_validate(
+        self,
+        original_texts,
+        original_values,
+        resp_map,
+        approval_map,
+        *,
+        reset_files,
+        max_rounds,
+    ):
         for idx in range(len(original_texts)):
             resp_map[(idx, 0)] = f"approved {idx}"
             approval_map[(idx, 0)] = (idx % 2 == 0)