enh: image prompt enhancer

backend/open_webui/utils/middleware.py

@@ -28,6 +28,7 @@ from open_webui.socket.main import (
 from open_webui.routers.tasks import (
     generate_queries,
     generate_title,
+    generate_image_prompt,
     generate_chat_tags,
 )
 from open_webui.routers.retrieval import process_web_search, SearchForm
@@ -503,12 +504,44 @@ async def chat_image_generation_handler(
     messages = form_data["messages"]
     user_message = get_last_user_message(messages)
 
+    prompt = ""
+    negative_prompt = ""
+
+    try:
+        res = await generate_image_prompt(
+            request,
+            {
+                "model": form_data["model"],
+                "messages": messages,
+            },
+            user,
+        )
+
+        response = res["choices"][0]["message"]["content"]
+
+        try:
+            bracket_start = response.find("{")
+            bracket_end = response.rfind("}") + 1
+
+            if bracket_start == -1 or bracket_end == -1:
+                raise Exception("No JSON object found in the response")
+
+            response = response[bracket_start:bracket_end]
+            response = json.loads(response)
+            prompt = response.get("prompt", [])
+        except Exception as e:
+            prompt = user_message
+
+    except Exception as e:
+        log.exception(e)
+        prompt = user_message
+
     system_message_content = ""
 
     try:
         images = await image_generations(
             request=request,
-            form_data=GenerateImageForm(**{"prompt": user_message}),
+            form_data=GenerateImageForm(**{"prompt": prompt}),
             user=user,
         )