Validate chunk image_base64 before doc-store write (#15364)

## Summary

Fixes [#15363](https://github.com/infiniflow/ragflow/issues/15363) —
`add_chunk` / `update_chunk` indexed chunks with `image_id` before
validating or storing `image_base64`, leaving orphan chunks on invalid
input.

## Related Issue

Fixes #15363

## Change Type

- [x] Bug fix
- [x] Regression tests

## What Changed

- Added `_decode_chunk_image_base64()` — strict base64 decode with
structured 4xx errors
- Added `_store_chunk_image_or_error()` — catches `store_chunk_image`
failures
- **`add_chunk` / `update_chunk`**: decode + store image **before**
`docStoreConn.insert` / `update`; only set `img_id` after successful
storage

## Files Changed

| File | Change |
|------|--------|
| `api/apps/restful_apis/chunk_api.py` | Helpers + reorder image
handling |
| `test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py`
| 3 regression tests |

## Validation

```bash
cd /root/gittensor/ragflow
pytest test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py::test_restful_add_chunk_invalid_image_base64_does_not_index_chunk -v
pytest test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py::test_restful_update_chunk_invalid_image_base64_does_not_update_chunk -v
pytest test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py::test_restful_add_chunk_valid_image_base64_stores_before_insert -v
pytest test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py -v
```

## Test Plan

- [x] Invalid `image_base64` on add → 4xx, no doc-store insert
- [x] Invalid `image_base64` on update → 4xx, no doc-store update
- [x] Valid PNG base64 on add → image stored, chunk indexed with
`img_id`
- [ ] CI green
This commit is contained in:
kpdev
2026-05-29 04:36:46 -07:00
committed by GitHub
parent 04aa8d04e8
commit cb1ea5a47f
2 changed files with 116 additions and 9 deletions

View File

@@ -14,6 +14,7 @@
# limitations under the License.
#
import base64
import binascii
import datetime
import logging
import re
@@ -66,6 +67,31 @@ DOC_STOP_PARSING_INVALID_STATE_MESSAGE = "Can't stop parsing document that has n
DOC_STOP_PARSING_INVALID_STATE_ERROR_CODE = "DOC_STOP_PARSING_INVALID_STATE"
def _decode_chunk_image_base64(image_base64):
if not isinstance(image_base64, str) or not image_base64.strip():
return None, "`image_base64` must be a non-empty string"
try:
image_binary = base64.b64decode(image_base64, validate=True)
except (binascii.Error, ValueError):
return None, "Invalid `image_base64`"
if not image_binary:
return None, "`image_base64` is empty"
return image_binary, None
def _store_chunk_image_or_error(dataset_id, chunk_id, image_binary):
try:
store_chunk_image(dataset_id, chunk_id, image_binary)
except Exception:
logging.exception(
"Failed to store chunk image. dataset_id=%s chunk_id=%s",
dataset_id,
chunk_id,
)
return "Failed to store chunk image"
return None
class Chunk(BaseModel):
id: str = ""
content: str = ""
@@ -507,8 +533,13 @@ async def add_chunk(tenant_id, dataset_id, document_id):
except ValueError as exc:
return get_error_data_result(f"`tag_feas` {exc}")
image_base64 = req.get("image_base64")
if image_base64:
if "image_base64" in req:
image_binary, image_err = _decode_chunk_image_base64(req.get("image_base64"))
if image_err:
return get_error_data_result(message=image_err)
store_err = _store_chunk_image_or_error(dataset_id, chunk_id, image_binary)
if store_err:
return get_error_data_result(message=store_err)
d["img_id"] = f"{dataset_id}-{chunk_id}"
d["doc_type_kwd"] = "image"
@@ -520,9 +551,6 @@ async def add_chunk(tenant_id, dataset_id, document_id):
d[f"q_{len(v)}_vec"] = v.tolist()
settings.docStoreConn.insert([d], search.index_name(dataset_tenant_id), dataset_id)
if image_base64:
store_chunk_image(dataset_id, chunk_id, base64.b64decode(image_base64))
DocumentService.increment_chunk_num(doc.id, doc.kb_id, c, 1, 0)
key_mapping = {
"id": "id",
@@ -647,8 +675,13 @@ async def update_chunk(tenant_id, dataset_id, document_id, chunk_id):
d["tag_feas"] = validate_tag_features(req["tag_feas"])
except ValueError as exc:
return get_error_data_result(f"`tag_feas` {exc}")
image_base64 = req.get("image_base64")
if image_base64:
if "image_base64" in req:
image_binary, image_err = _decode_chunk_image_base64(req.get("image_base64"))
if image_err:
return get_error_data_result(message=image_err)
store_err = _store_chunk_image_or_error(dataset_id, chunk_id, image_binary)
if store_err:
return get_error_data_result(message=store_err)
d["img_id"] = f"{dataset_id}-{chunk_id}"
d["doc_type_kwd"] = "image"
@@ -671,8 +704,6 @@ async def update_chunk(tenant_id, dataset_id, document_id, chunk_id):
v = 0.1 * v[0] + 0.9 * v[1] if doc.parser_id != ParserType.QA else v[1]
d[f"q_{len(v)}_vec"] = v.tolist()
settings.docStoreConn.update({"id": chunk_id}, d, search.index_name(dataset_tenant_id), dataset_id)
if image_base64:
store_chunk_image(dataset_id, chunk_id, base64.b64decode(image_base64))
return get_result()