Browse Source

refac: title generation

Timothy Jaeryang Baek 3 months ago
parent
commit
5420c165c6

+ 18 - 8
backend/open_webui/config.py

@@ -1094,16 +1094,26 @@ TITLE_GENERATION_PROMPT_TEMPLATE = PersistentConfig(
     os.environ.get("TITLE_GENERATION_PROMPT_TEMPLATE", ""),
     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:
+- 📉 Stock Market Trends
+- 🍪 Perfect Chocolate Chip Recipe
+- Evolution of Music Streaming
+- Remote Work Productivity Tips
+- Artificial Intelligence in Healthcare
+- 🎮 Video Game Development Insights
 
 
-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
+### Output:
+JSON format: { "title": "your concise title here" }
 
 
+### Chat History:
 <chat_history>
 <chat_history>
 {{MESSAGES:END:2}}
 {{MESSAGES:END:2}}
 </chat_history>"""
 </chat_history>"""

+ 2 - 2
backend/open_webui/routers/tasks.py

@@ -175,10 +175,10 @@ async def generate_title(
         "messages": [{"role": "user", "content": content}],
         "messages": [{"role": "user", "content": content}],
         "stream": False,
         "stream": False,
         **(
         **(
-            {"max_tokens": 50}
+            {"max_tokens": 1000}
             if models[task_model_id]["owned_by"] == "ollama"
             if models[task_model_id]["owned_by"] == "ollama"
             else {
             else {
-                "max_completion_tokens": 50,
+                "max_completion_tokens": 1000,
             }
             }
         ),
         ),
         "metadata": {
         "metadata": {

+ 15 - 7
backend/open_webui/utils/middleware.py

@@ -892,16 +892,24 @@ async def process_chat_response(
 
 
                         if res and isinstance(res, dict):
                         if res and isinstance(res, dict):
                             if len(res.get("choices", [])) == 1:
                             if len(res.get("choices", [])) == 1:
-                                title = (
+                                title_string = (
                                     res.get("choices", [])[0]
                                     res.get("choices", [])[0]
                                     .get("message", {})
                                     .get("message", {})
-                                    .get(
-                                        "content",
-                                        message.get("content", "New Chat"),
-                                    )
-                                ).strip()
+                                    .get("content", message.get("content", "New Chat"))
+                                )
                             else:
                             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:
                             if not title:
                                 title = messages[0].get("content", "New Chat")
                                 title = messages[0].get("content", "New Chat")