From 127af4e45c3a93eced8d3504f0620f3b5de98c18 Mon Sep 17 00:00:00 2001
From: Stephen Hu <stephenhu@seismic.com>
Date: Wed, 10 Sep 2025 15:55:33 +0800
Subject: [PATCH] Refactor:Improve BytesIO usage for image2base64 (#9997)

### What problem does this PR solve?

Improve BytesIO usage for image2base64

### Type of change

- [x] Refactoring
---
 rag/llm/cv_model.py | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py
index 4ca4ee190..83c79bbfd 100644
--- a/rag/llm/cv_model.py
+++ b/rag/llm/cv_model.py
@@ -124,17 +124,19 @@ class Base(ABC):
                 mime = "image/jpeg"
             b64 = base64.b64encode(data).decode("utf-8")
             return f"data:{mime};base64,{b64}"
-        buffered = BytesIO()
-        fmt = "JPEG"
-        try:
-            image.save(buffered, format="JPEG")
-        except Exception:
-            buffered = BytesIO()  # reset buffer before saving PNG
-            image.save(buffered, format="PNG")
-            fmt = "PNG"
-        data = buffered.getvalue()
-        b64 = base64.b64encode(data).decode("utf-8")
-        mime = f"image/{fmt.lower()}"
+        with BytesIO() as buffered:
+            fmt = "JPEG"
+            try:
+                image.save(buffered, format="JPEG")
+            except Exception:
+                 # reset buffer before saving PNG
+                buffered.seek(0)
+                buffered.truncate()
+                image.save(buffered, format="PNG")
+                fmt = "PNG"
+            data = buffered.getvalue()
+            b64 = base64.b64encode(data).decode("utf-8")
+            mime = f"image/{fmt.lower()}"
         return f"data:{mime};base64,{b64}"
 
     def prompt(self, b64):