""" title: Auto Features description: Automatically enable Web Search and Image Generation features if the LLM requires it. author: Bastien Vidé version: 0.1.0 """ from pydantic import BaseModel, Field from typing import Callable, Awaitable, Any, Optional from open_webui.utils.middleware import ( chat_web_search_handler, chat_image_generation_handler, ) from open_webui.models.users import Users from open_webui.models.models import Models from open_webui.utils.chat import generate_chat_completion from open_webui.utils.misc import get_last_user_message class Filter: class Valves(BaseModel): pass class UserValves(BaseModel): auto_websearch: bool = Field( default=True, description="Automatically search the web" ) auto_image: bool = Field(default=True, description="Generate image when asked") pass def __init__(self): self.valves = self.Valves() self.user_valves = self.UserValves() async def inlet( self, body: dict, __event_emitter__: Callable[[Any], Awaitable[None]], __request__: Any, __user__: Optional[dict] = None, __model__: Optional[dict] = None, ) -> dict: if ( not body.get("features", {}).get("web_search", False) and self.user_valves.auto_websearch ): messages = body["messages"] user_message = get_last_user_message(messages) user = Users.get_user_by_id(__user__["id"]) prompt = ( "History:\n" + "\n".join( [ f"{message['role'].upper()}: \"\"\"{message['content']}\"\"\"" for message in messages[::-1][:4] ] ) + f"\nUser Query: {user_message}" ) try: system_prompt = """Your only goal is to determine if the user query requires a web search to be answered. A web search is required in at least one of this case: - If you do not know the answer to the user query. - If the user asks for up-to-date information. - If the user asks for for unknown knowledge, news, info, public contact info, weather. - If the user explicitly asks you to search something, or to perform a web search. Do not answer the user query. Only answer with the boolean, `true` if the web search is required to answer the user query or `false` otherwise. Do NOT add any explanations. Only answer with the boolean.""" payload = { "model": body["model"], "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, ], "stream": False, } response = await generate_chat_completion( request=__request__, form_data=payload, user=user ) content = response["choices"][0]["message"]["content"] if "true" in content.lower(): body = await chat_web_search_handler( __request__, body, {"__event_emitter__": __event_emitter__}, user, ) except Exception as e: print(e) if ( not body.get("features", {}).get("image_generation", False) and self.user_valves.auto_image ): try: system_prompt = """Your only goal is to determine if the user wants you to generate an image or a drawing. User must be explicit about generating an image, or modifying/adding/removing features from a previously generated image. Do not answer the user query. Only answer with the boolean, `true` if the user wants image/drawing generation or `false` otherwise. Do NOT add any explainations. Only answer with the boolean.""" payload = { "model": body["model"], "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, ], "stream": False, } response = await generate_chat_completion( request=__request__, form_data=payload, user=user ) content = response["choices"][0]["message"]["content"] if "true" in content.lower(): body = await chat_image_generation_handler( __request__, body, {"__event_emitter__": __event_emitter__}, user, ) except Exception as e: print(e) return body