Merge pull request #14774 from rragundez/images-from-db

fix: Store and load code interpreter generated images from a central location (DB and/or cloud storage)
This commit is contained in:
Tim Jaeryang Baek 2025-06-08 19:09:28 +04:00 committed by GitHub
commit 6cb519ca0e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 31 additions and 43 deletions

View File

@ -420,7 +420,7 @@ def load_b64_image_data(b64_str):
try: try:
if "," in b64_str: if "," in b64_str:
header, encoded = b64_str.split(",", 1) header, encoded = b64_str.split(",", 1)
mime_type = header.split(";")[0] mime_type = header.split(";")[0].lstrip("data:")
img_data = base64.b64decode(encoded) img_data = base64.b64decode(encoded)
else: else:
mime_type = "image/png" mime_type = "image/png"
@ -428,7 +428,7 @@ def load_b64_image_data(b64_str):
return img_data, mime_type return img_data, mime_type
except Exception as e: except Exception as e:
log.exception(f"Error loading image data: {e}") log.exception(f"Error loading image data: {e}")
return None return None, None
def load_url_image_data(url, headers=None): def load_url_image_data(url, headers=None):

View File

@ -37,7 +37,12 @@ from open_webui.routers.tasks import (
generate_chat_tags, generate_chat_tags,
) )
from open_webui.routers.retrieval import process_web_search, SearchForm from open_webui.routers.retrieval import process_web_search, SearchForm
from open_webui.routers.images import image_generations, GenerateImageForm from open_webui.routers.images import (
load_b64_image_data,
image_generations,
GenerateImageForm,
upload_image,
)
from open_webui.routers.pipelines import ( from open_webui.routers.pipelines import (
process_pipeline_inlet_filter, process_pipeline_inlet_filter,
process_pipeline_outlet_filter, process_pipeline_outlet_filter,
@ -2259,28 +2264,19 @@ async def process_chat_response(
stdoutLines = stdout.split("\n") stdoutLines = stdout.split("\n")
for idx, line in enumerate(stdoutLines): for idx, line in enumerate(stdoutLines):
if "data:image/png;base64" in line: if "data:image/png;base64" in line:
id = str(uuid4()) image_url = ""
# Extract base64 image data from the line
# ensure the path exists image_data, content_type = load_b64_image_data(line)
os.makedirs( if image_data is not None:
os.path.join(CACHE_DIR, "images"), image_url = upload_image(
exist_ok=True, request,
image_data,
content_type,
metadata,
user,
) )
image_path = os.path.join(
CACHE_DIR,
f"images/{id}.png",
)
with open(image_path, "wb") as f:
f.write(
base64.b64decode(
line.split(",")[1]
)
)
stdoutLines[idx] = ( stdoutLines[idx] = (
f"![Output Image {idx}](/cache/images/{id}.png)" f"![Output Image]({image_url})"
) )
output["stdout"] = "\n".join(stdoutLines) output["stdout"] = "\n".join(stdoutLines)
@ -2291,30 +2287,22 @@ async def process_chat_response(
resultLines = result.split("\n") resultLines = result.split("\n")
for idx, line in enumerate(resultLines): for idx, line in enumerate(resultLines):
if "data:image/png;base64" in line: if "data:image/png;base64" in line:
id = str(uuid4()) image_url = ""
# Extract base64 image data from the line
# ensure the path exists image_data, content_type = (
os.makedirs( load_b64_image_data(line)
os.path.join(CACHE_DIR, "images"),
exist_ok=True,
) )
if image_data is not None:
image_path = os.path.join( image_url = upload_image(
CACHE_DIR, request,
f"images/{id}.png", image_data,
content_type,
metadata,
user,
) )
with open(image_path, "wb") as f:
f.write(
base64.b64decode(
line.split(",")[1]
)
)
resultLines[idx] = ( resultLines[idx] = (
f"![Output Image {idx}](/cache/images/{id}.png)" f"![Output Image]({image_url})"
) )
output["result"] = "\n".join(resultLines) output["result"] = "\n".join(resultLines)
except Exception as e: except Exception as e:
output = str(e) output = str(e)