From 5420c165c60edb23a841f1b469d2fec381775cbf Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Wed, 29 Jan 2025 14:40:36 -0800
Subject: [PATCH] refac: title generation

---
 backend/open_webui/config.py           | 26 ++++++++++++++++++--------
 backend/open_webui/routers/tasks.py    |  4 ++--
 backend/open_webui/utils/middleware.py | 22 +++++++++++++++-------
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py
index c37b831de..f65726c3c 100644
--- a/backend/open_webui/config.py
+++ b/backend/open_webui/config.py
@@ -1094,16 +1094,26 @@ TITLE_GENERATION_PROMPT_TEMPLATE = PersistentConfig(
     os.environ.get("TITLE_GENERATION_PROMPT_TEMPLATE", ""),
 )
 
-DEFAULT_TITLE_GENERATION_PROMPT_TEMPLATE = """Create a concise, 3-5 word title with an emoji as a title for the chat history, in the given language. Suitable Emojis for the summary can be used to enhance understanding but avoid quotation marks or special formatting. RESPOND ONLY WITH THE TITLE TEXT.
+DEFAULT_TITLE_GENERATION_PROMPT_TEMPLATE = """### Task:
+Generate a concise, 3-5 word title with an emoji summarizing the chat history.
+### Guidelines:
+- The title should clearly represent the main theme or subject of the conversation.
+- Use emojis that enhance understanding of the topic, but avoid quotation marks or special formatting.
+- Write the title in the chat's primary language; default to English if multilingual.
+- Prioritize accuracy over excessive creativity; keep it clear and simple.
 
-Examples of titles:
-📉 Stock Market Trends
-🍪 Perfect Chocolate Chip Recipe
-Evolution of Music Streaming
-Remote Work Productivity Tips
-Artificial Intelligence in Healthcare
-🎮 Video Game Development Insights
+### Examples:
+- 📉 Stock Market Trends
+- 🍪 Perfect Chocolate Chip Recipe
+- Evolution of Music Streaming
+- Remote Work Productivity Tips
+- Artificial Intelligence in Healthcare
+- 🎮 Video Game Development Insights
 
+### Output:
+JSON format: { "title": "your concise title here" }
+
+### Chat History:
 <chat_history>
 {{MESSAGES:END:2}}
 </chat_history>"""
diff --git a/backend/open_webui/routers/tasks.py b/backend/open_webui/routers/tasks.py
index 6d7343c8a..2dc58b4a5 100644
--- a/backend/open_webui/routers/tasks.py
+++ b/backend/open_webui/routers/tasks.py
@@ -175,10 +175,10 @@ async def generate_title(
         "messages": [{"role": "user", "content": content}],
         "stream": False,
         **(
-            {"max_tokens": 50}
+            {"max_tokens": 1000}
             if models[task_model_id]["owned_by"] == "ollama"
             else {
-                "max_completion_tokens": 50,
+                "max_completion_tokens": 1000,
             }
         ),
         "metadata": {
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index 815e69739..00f4dd402 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -892,16 +892,24 @@ async def process_chat_response(
 
                         if res and isinstance(res, dict):
                             if len(res.get("choices", [])) == 1:
-                                title = (
+                                title_string = (
                                     res.get("choices", [])[0]
                                     .get("message", {})
-                                    .get(
-                                        "content",
-                                        message.get("content", "New Chat"),
-                                    )
-                                ).strip()
+                                    .get("content", message.get("content", "New Chat"))
+                                )
                             else:
-                                title = None
+                                title_string = ""
+
+                            title_string = title_string[
+                                title_string.find("{") : title_string.rfind("}") + 1
+                            ]
+
+                            try:
+                                title = json.loads(title_string).get(
+                                    "title", "New Chat"
+                                )
+                            except Exception as e:
+                                pass
 
                             if not title:
                                 title = messages[0].get("content", "New Chat")