fix(parsers): delete uploaded Mistral OCR files; fix patch.stopall test hygiene (#77)

Author: KylinMountain

openkb/parsers/mistral.py

@@ -42,26 +42,37 @@ def parse(self, src: Path) -> ParseResult:
             ) from exc
 
         client = Mistral(api_key=api_key)
-        uploaded = client.files.upload(
-            file={"file_name": src.name, "content": src.read_bytes()}, purpose="ocr"
-        )
-        signed = client.files.get_signed_url(file_id=uploaded.id)
-        resp = client.ocr.process(
-            model=self.model,
-            document={"type": "document_url", "document_url": signed.url},
-            include_image_base64=True,
-        )
+        uploaded = None
+        try:
+            uploaded = client.files.upload(
+                file={"file_name": src.name, "content": src.read_bytes()}, purpose="ocr"
+            )
+            signed = client.files.get_signed_url(file_id=uploaded.id)
+            resp = client.ocr.process(
+                model=self.model,
+                document={"type": "document_url", "document_url": signed.url},
+                include_image_base64=True,
+            )
 
-        parts: list[str] = []
-        images: dict[str, bytes] = {}
-        for page in resp.pages:
-            parts.append(page.markdown or "")
-            for img in getattr(page, "images", None) or []:
-                raw = img.image_base64 or ""
-                raw = _DATA_URI_RE.sub("", raw)
+            parts: list[str] = []
+            images: dict[str, bytes] = {}
+            for page in resp.pages:
+                parts.append(page.markdown or "")
+                for img in getattr(page, "images", None) or []:
+                    raw = img.image_base64 or ""
+                    raw = _DATA_URI_RE.sub("", raw)
+                    try:
+                        images[img.id] = base64.b64decode(raw, validate=True)
+                    except Exception:
+                        logger.warning("Skipping undecodable Mistral image: %s", getattr(img, "id", "?"))
+                        continue
+            return ParseResult(markdown="\n\n".join(parts), images=images)
+        finally:
+            if uploaded is not None:
                 try:
-                    images[img.id] = base64.b64decode(raw, validate=True)
+                    client.files.delete(file_id=uploaded.id)
                 except Exception:
-                    logger.warning("Skipping undecodable Mistral image: %s", getattr(img, "id", "?"))
-                    continue
-        return ParseResult(markdown="\n\n".join(parts), images=images)
+                    logger.warning(
+                        "Failed to delete uploaded Mistral OCR file %s",
+                        getattr(uploaded, "id", "?"),
+                    )

tests/test_parsers_local.py

@@ -39,13 +39,10 @@ def test_parse_other_uses_markitdown_and_extracts_base64(tmp_path):
     src = tmp_path / "deck.pptx"
     src.write_bytes(b"PK fake")
     images_dir = tmp_path / "img" / "deck"
-    fake_mid = patch("openkb.parsers.local.MarkItDown").start()
-    fake_mid.return_value.convert.return_value.text_content = "MARKITDOWN MD"
-    try:
-        with patch("openkb.parsers.local.extract_base64_images", return_value="CLEANED") as ex:
-            p = LocalParser(doc_name="deck", images_dir=images_dir, source_dir=tmp_path)
-            result = p.parse(src)
-        ex.assert_called_once_with("MARKITDOWN MD", "deck", images_dir)
-        assert result.markdown == "CLEANED"
-    finally:
-        patch.stopall()
+    with patch("openkb.parsers.local.MarkItDown") as fake_mid, \
+         patch("openkb.parsers.local.extract_base64_images", return_value="CLEANED") as ex:
+        fake_mid.return_value.convert.return_value.text_content = "MARKITDOWN MD"
+        p = LocalParser(doc_name="deck", images_dir=images_dir, source_dir=tmp_path)
+        result = p.parse(src)
+    ex.assert_called_once_with("MARKITDOWN MD", "deck", images_dir)
+    assert result.markdown == "CLEANED"

tests/test_parsers_mistral.py

@@ -86,3 +86,37 @@ def test_undecodable_image_logged_and_skipped(monkeypatch, tmp_path, caplog):
         result = MistralParser({}).parse(src)
     assert "bad.png" not in result.images
     assert any("bad.png" in r.message for r in caplog.records)
+
+
+def test_uploaded_file_is_deleted(monkeypatch, tmp_path):
+    import sys, types
+    from unittest.mock import MagicMock
+    monkeypatch.setenv("MISTRAL_API_KEY", "k")
+    client = MagicMock()
+    client.files.upload.return_value = MagicMock(id="file-1")
+    client.files.get_signed_url.return_value = MagicMock(url="https://signed")
+    client.ocr.process.return_value = MagicMock(pages=[])
+    mod = types.ModuleType("mistralai"); mod.Mistral = MagicMock(return_value=client)
+    monkeypatch.setitem(sys.modules, "mistralai", mod)
+    from openkb.parsers.mistral import MistralParser
+    src = tmp_path / "d.pdf"; src.write_bytes(b"%PDF")
+    MistralParser({}).parse(src)
+    client.files.delete.assert_called_once_with(file_id="file-1")
+
+
+def test_uploaded_file_deleted_even_on_ocr_error(monkeypatch, tmp_path):
+    import sys, types
+    from unittest.mock import MagicMock
+    import pytest
+    monkeypatch.setenv("MISTRAL_API_KEY", "k")
+    client = MagicMock()
+    client.files.upload.return_value = MagicMock(id="file-2")
+    client.files.get_signed_url.return_value = MagicMock(url="https://signed")
+    client.ocr.process.side_effect = RuntimeError("ocr boom")
+    mod = types.ModuleType("mistralai"); mod.Mistral = MagicMock(return_value=client)
+    monkeypatch.setitem(sys.modules, "mistralai", mod)
+    from openkb.parsers.mistral import MistralParser
+    src = tmp_path / "d.pdf"; src.write_bytes(b"%PDF")
+    with pytest.raises(RuntimeError):
+        MistralParser({}).parse(src)
+    client.files.delete.assert_called_once_with(file_id="file-2")