From 18402b7d78e9cee5b418e6ab76ef90dd206963d6 Mon Sep 17 00:00:00 2001 From: "Richard Kuo (Onyx)" Date: Thu, 13 Mar 2025 16:50:43 -0700 Subject: [PATCH] refactor, use inline attachment for image (base64 encoding doesn't work) --- backend/ee/onyx/file_store/onyx_file_store.py | 18 ---- .../ee/onyx/server/enterprise_settings/api.py | 12 ++- .../onyx/server/enterprise_settings/store.py | 15 ++- backend/onyx/auth/email_utils.py | 92 +++++++++--------- backend/onyx/configs/constants.py | 2 + backend/onyx/file_store/onyx_file_store.py | 27 ++---- backend/onyx/server/runtime/onyx_runtime.py | 95 +++++++++++++++++++ backend/onyx/utils/file.py | 19 ++++ backend/onyx/utils/variable_functionality.py | 19 ++++ 9 files changed, 206 insertions(+), 93 deletions(-) delete mode 100644 backend/ee/onyx/file_store/onyx_file_store.py create mode 100644 backend/onyx/server/runtime/onyx_runtime.py diff --git a/backend/ee/onyx/file_store/onyx_file_store.py b/backend/ee/onyx/file_store/onyx_file_store.py deleted file mode 100644 index 4f53bdf7a..000000000 --- a/backend/ee/onyx/file_store/onyx_file_store.py +++ /dev/null @@ -1,18 +0,0 @@ -from sqlalchemy.orm import Session - -from ee.onyx.server.enterprise_settings.store import _LOGO_FILENAME -from ee.onyx.server.enterprise_settings.store import _LOGOTYPE_FILENAME -from onyx.file_store.onyx_file_store import OnyxFileStore -from onyx.utils.file import OnyxFile - - -class OnyxEnterpriseFileStore(OnyxFileStore): - def get_logo(self) -> OnyxFile: - return self.get_db_image(_LOGO_FILENAME) - - def get_logotype(self) -> OnyxFile: - return self.get_db_image(_LOGOTYPE_FILENAME) - - -def get_onyx_file_store(db_session: Session) -> OnyxFileStore: - return OnyxEnterpriseFileStore(db_session=db_session) diff --git a/backend/ee/onyx/server/enterprise_settings/api.py b/backend/ee/onyx/server/enterprise_settings/api.py index 4a130a84a..98b05880c 100644 --- a/backend/ee/onyx/server/enterprise_settings/api.py +++ b/backend/ee/onyx/server/enterprise_settings/api.py @@ -13,9 +13,10 @@ from pydantic import BaseModel from pydantic import Field from sqlalchemy.orm import Session -from ee.onyx.file_store.onyx_file_store import get_onyx_file_store from ee.onyx.server.enterprise_settings.models import AnalyticsScriptUpload from ee.onyx.server.enterprise_settings.models import EnterpriseSettings +from ee.onyx.server.enterprise_settings.store import get_logo_filename +from ee.onyx.server.enterprise_settings.store import get_logotype_filename from ee.onyx.server.enterprise_settings.store import load_analytics_script from ee.onyx.server.enterprise_settings.store import load_settings from ee.onyx.server.enterprise_settings.store import store_analytics_script @@ -27,6 +28,7 @@ from onyx.auth.users import get_user_manager from onyx.auth.users import UserManager from onyx.db.engine import get_session from onyx.db.models import User +from onyx.file_store.onyx_file_store import OnyxFileStore from onyx.utils.logger import setup_logger admin_router = APIRouter(prefix="/admin/enterprise-settings") @@ -146,8 +148,8 @@ def put_logo( def fetch_logo_helper(db_session: Session) -> Response: try: - file_store = get_onyx_file_store(db_session) - onyx_file = file_store.get_logo() + file_store = OnyxFileStore(db_session) + onyx_file = file_store.get_onyx_file(get_logo_filename()) if not onyx_file: raise except Exception: @@ -161,8 +163,8 @@ def fetch_logo_helper(db_session: Session) -> Response: def fetch_logotype_helper(db_session: Session) -> Response: try: - file_store = get_onyx_file_store(db_session) - onyx_file = file_store.get_logotype() + file_store = OnyxFileStore(db_session) + onyx_file = file_store.get_onyx_file(get_logotype_filename()) if not onyx_file: raise except Exception: diff --git a/backend/ee/onyx/server/enterprise_settings/store.py b/backend/ee/onyx/server/enterprise_settings/store.py index c5ae4d2a9..71f562bcb 100644 --- a/backend/ee/onyx/server/enterprise_settings/store.py +++ b/backend/ee/onyx/server/enterprise_settings/store.py @@ -22,6 +22,9 @@ from onyx.utils.logger import setup_logger logger = setup_logger() +_LOGO_FILENAME = "__logo__" +_LOGOTYPE_FILENAME = "__logotype__" + def load_settings() -> EnterpriseSettings: """Loads settings data directly from DB. This should be used primarily @@ -83,10 +86,6 @@ def store_analytics_script(analytics_script_upload: AnalyticsScriptUpload) -> No get_kv_store().store(KV_CUSTOM_ANALYTICS_SCRIPT_KEY, analytics_script_upload.script) -_LOGO_FILENAME = "__logo__" -_LOGOTYPE_FILENAME = "__logotype__" - - def is_valid_file_type(filename: str) -> bool: valid_extensions = (".png", ".jpg", ".jpeg") return filename.endswith(valid_extensions) @@ -139,3 +138,11 @@ def upload_logo( file_type=file_type, ) return True + + +def get_logo_filename() -> str: + return _LOGO_FILENAME + + +def get_logotype_filename() -> str: + return _LOGOTYPE_FILENAME diff --git a/backend/onyx/auth/email_utils.py b/backend/onyx/auth/email_utils.py index b1eb54ca2..1650007e0 100644 --- a/backend/onyx/auth/email_utils.py +++ b/backend/onyx/auth/email_utils.py @@ -1,6 +1,6 @@ -import base64 import smtplib from datetime import datetime +from email.mime.image import MIMEImage from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText from email.utils import formatdate @@ -17,9 +17,8 @@ from onyx.configs.constants import AuthType from onyx.configs.constants import ONYX_DEFAULT_APPLICATION_NAME from onyx.configs.constants import ONYX_SLACK_URL from onyx.configs.constants import TENANT_ID_COOKIE_NAME -from onyx.db.engine import get_session_with_shared_schema from onyx.db.models import User -from onyx.file_store.onyx_file_store import OnyxFileStore +from onyx.server.runtime.onyx_runtime import OnyxRuntime from onyx.utils.file import OnyxFile from onyx.utils.variable_functionality import fetch_versioned_implementation from shared_configs.configs import MULTI_TENANT @@ -104,7 +103,7 @@ HTML_EMAIL_TEMPLATE = """\ {application_name} Logo @@ -134,13 +133,11 @@ def build_html_email( application_name: str | None, heading: str, message: str, - logo: bytes, - logo_mimetype: str, cta_text: str | None = None, cta_link: str | None = None, ) -> str: slack_fragment = "" - if application_name != ONYX_DEFAULT_APPLICATION_NAME: + if application_name == ONYX_DEFAULT_APPLICATION_NAME: slack_fragment = f'
Have questions? Join our Slack community here.' if cta_text and cta_link: @@ -152,8 +149,6 @@ def build_html_email( title=heading, heading=heading, message=message, - logo=base64.b64encode(logo), - logo_mimetype=logo_mimetype, cta_block=cta_block, slack_fragment=slack_fragment, year=datetime.now().year, @@ -166,6 +161,7 @@ def send_email( html_body: str, text_body: str, mail_from: str = EMAIL_FROM, + inline_png: tuple[str, bytes] | None = None, ) -> None: if not EMAIL_CONFIGURED: raise ValueError("Email is not configured.") @@ -184,6 +180,12 @@ def send_email( msg.attach(part_text) msg.attach(part_html) + if inline_png: + img = MIMEImage(inline_png[1], _subtype="png") + img.add_header("Content-ID", inline_png[0]) # CID reference + img.add_header("Content-Disposition", "inline", filename=inline_png[0]) + msg.attach(img) + try: with smtplib.SMTP(SMTP_SERVER, SMTP_PORT) as s: s.starttls() @@ -206,12 +208,7 @@ def send_subscription_cancellation_email(user_email: str) -> None: except ModuleNotFoundError: application_name = ONYX_DEFAULT_APPLICATION_NAME - with get_session_with_shared_schema() as db_session: - get_onyx_file_store_fn = fetch_versioned_implementation( - "onyx.file_store.onyx_file_store", "get_onyx_file_store" - ) - file_store: OnyxFileStore = get_onyx_file_store_fn(db_session) - onyx_file = file_store.get_logo() + onyx_file = OnyxRuntime.get_emailable_logo() subject = f"Your {application_name} Subscription Has Been Canceled" heading = "Subscription Canceled" @@ -226,8 +223,6 @@ def send_subscription_cancellation_email(user_email: str) -> None: application_name, heading, message, - onyx_file.data, - onyx_file.mime_type, cta_text, cta_link, ) @@ -236,7 +231,13 @@ def send_subscription_cancellation_email(user_email: str) -> None: "Your subscription has been canceled and will end on your next billing date.\n" "If you change your mind, visit https://www.onyx.app/pricing" ) - send_email(user_email, subject, html_content, text_content) + send_email( + user_email, + subject, + html_content, + text_content, + inline_png=("logo.png", onyx_file.data), + ) def send_user_email_invite( @@ -253,18 +254,13 @@ def send_user_email_invite( except ModuleNotFoundError: application_name = ONYX_DEFAULT_APPLICATION_NAME - with get_session_with_shared_schema() as db_session: - get_onyx_file_store_fn = fetch_versioned_implementation( - "onyx.file_store.onyx_file_store", "get_onyx_file_store" - ) - file_store: OnyxFileStore = get_onyx_file_store_fn(db_session) - onyx_file = file_store.get_logo() + onyx_file = OnyxRuntime.get_emailable_logo() subject = f"Invitation to Join {application_name} Organization" heading = "You've Been Invited!" # the exact action taken by the user, and thus the message, depends on the auth type - message = f"

You have been invited by {current_user.email} to join an organization on {application_name} .

" + message = f"

You have been invited by {current_user.email} to join an organization on {application_name}.

" if auth_type == AuthType.CLOUD: message += ( "

To join the organization, please click the button below to set a password " @@ -295,8 +291,6 @@ def send_user_email_invite( application_name, heading, message, - onyx_file.data, - onyx_file.mime_type, cta_text, cta_link, ) @@ -311,7 +305,13 @@ def send_user_email_invite( if auth_type == AuthType.CLOUD: text_content += "You'll be asked to set a password or login with Google to complete your registration." - send_email(user_email, subject, html_content, text_content) + send_email( + user_email, + subject, + html_content, + text_content, + inline_png=("logo.png", onyx_file.data), + ) def send_forgot_password_email( @@ -330,12 +330,7 @@ def send_forgot_password_email( except ModuleNotFoundError: application_name = ONYX_DEFAULT_APPLICATION_NAME - with get_session_with_shared_schema() as db_session: - get_onyx_file_store_fn = fetch_versioned_implementation( - "onyx.file_store.onyx_file_store", "get_onyx_file_store" - ) - file_store: OnyxFileStore = get_onyx_file_store_fn(db_session) - onyx_file = file_store.get_logo() + onyx_file = OnyxRuntime.get_emailable_logo() subject = f"{application_name} Forgot Password" link = f"{WEB_DOMAIN}/auth/reset-password?token={token}" @@ -346,11 +341,16 @@ def send_forgot_password_email( application_name, "Reset Your Password", message, - onyx_file.data, - onyx_file.mime_type, ) text_content = f"Click the following link to reset your password: {link}" - send_email(user_email, subject, html_content, text_content, mail_from) + send_email( + user_email, + subject, + html_content, + text_content, + mail_from, + inline_png=("logo.png", onyx_file.data), + ) def send_user_verification_email( @@ -368,12 +368,7 @@ def send_user_verification_email( except ModuleNotFoundError: application_name = ONYX_DEFAULT_APPLICATION_NAME - with get_session_with_shared_schema() as db_session: - get_onyx_file_store_fn = fetch_versioned_implementation( - "onyx.file_store.onyx_file_store", "get_onyx_file_store" - ) - file_store: OnyxFileStore = get_onyx_file_store_fn(db_session) - onyx_file = file_store.get_logo() + onyx_file = OnyxRuntime.get_emailable_logo() subject = f"{application_name} Email Verification" link = f"{WEB_DOMAIN}/auth/verify-email?token={token}" @@ -384,8 +379,13 @@ def send_user_verification_email( application_name, "Verify Your Email", message, - onyx_file.data, - onyx_file.mime_type, ) text_content = f"Click the following link to verify your email address: {link}" - send_email(user_email, subject, html_content, text_content, mail_from) + send_email( + user_email, + subject, + html_content, + text_content, + mail_from, + inline_png=("logo.png", onyx_file.data), + ) diff --git a/backend/onyx/configs/constants.py b/backend/onyx/configs/constants.py index dcd795e0e..4feec228f 100644 --- a/backend/onyx/configs/constants.py +++ b/backend/onyx/configs/constants.py @@ -5,6 +5,7 @@ from enum import Enum ONYX_DEFAULT_APPLICATION_NAME = "Onyx" ONYX_SLACK_URL = "https://join.slack.com/t/onyx-dot-app/shared_invite/zt-2twesxdr6-5iQitKZQpgq~hYIZ~dv3KA" +ONYX_EMAILABLE_LOGO_MAX_DIM = 512 SOURCE_TYPE = "source_type" # stored in the `metadata` of a chunk. Used to signify that this chunk should @@ -43,6 +44,7 @@ DISABLED_GEN_AI_MSG = ( "You can still use Onyx as a search engine." ) + DEFAULT_PERSONA_ID = 0 DEFAULT_CC_PAIR_ID = 1 diff --git a/backend/onyx/file_store/onyx_file_store.py b/backend/onyx/file_store/onyx_file_store.py index 01fcc707d..4a628668c 100644 --- a/backend/onyx/file_store/onyx_file_store.py +++ b/backend/onyx/file_store/onyx_file_store.py @@ -7,27 +7,14 @@ from onyx.utils.file import OnyxFile class OnyxFileStore(PostgresBackedFileStore): - def get_static_image(self, filename: str) -> OnyxFile: + def get_onyx_file(self, filename: str) -> OnyxFile | None: mime_type: str = "application/octet-stream" - with open(filename, "b") as f: - file_content = f.read() + try: + file_io = self.read_file(filename, mode="b") + file_content = file_io.read() matches = puremagic.magic_string(file_content) if matches: mime_type = cast(str, matches[0].mime_type) - - return OnyxFile(data=file_content, mime_type=mime_type) - - def get_db_image(self, filename: str) -> OnyxFile: - mime_type: str = "application/octet-stream" - file_io = self.read_file(filename, mode="b") - file_content = file_io.read() - matches = puremagic.magic_string(file_content) - if matches: - mime_type = cast(str, matches[0].mime_type) - return OnyxFile(data=file_content, mime_type=mime_type) - - def get_logo(self) -> OnyxFile: - return self.get_static_image("static/images/logo.png") - - def get_logotype(self) -> OnyxFile: - return self.get_static_image("static/images/logotype.png") + return OnyxFile(data=file_content, mime_type=mime_type) + except Exception: + return None diff --git a/backend/onyx/server/runtime/onyx_runtime.py b/backend/onyx/server/runtime/onyx_runtime.py new file mode 100644 index 000000000..736155bdc --- /dev/null +++ b/backend/onyx/server/runtime/onyx_runtime.py @@ -0,0 +1,95 @@ +import io + +from PIL import Image + +from onyx.configs.constants import ONYX_EMAILABLE_LOGO_MAX_DIM +from onyx.db.engine import get_session_with_shared_schema +from onyx.file_store.onyx_file_store import OnyxFileStore +from onyx.utils.file import OnyxFile +from onyx.utils.file import OnyxStaticFileManager +from onyx.utils.variable_functionality import ( + fetch_ee_implementation_or_none, +) + + +class OnyxRuntime: + @staticmethod + def _get_with_static_fallback( + db_filename: str | None, static_filename: str + ) -> OnyxFile: + onyx_file: OnyxFile | None = None + + while True: + if db_filename: + with get_session_with_shared_schema() as db_session: + file_store: OnyxFileStore = OnyxFileStore(db_session) + onyx_file = file_store.get_onyx_file(db_filename) + + if onyx_file: + break + + onyx_file = OnyxStaticFileManager.get_static(static_filename) + break + + if not onyx_file: + raise RuntimeError( + f"Resource not found: db={db_filename} static={static_filename}" + ) + + return onyx_file + + @staticmethod + def get_logo() -> OnyxFile: + STATIC_FILENAME = "static/images/logo.png" + + db_filename: str | None = None + get_logo_filename_fn = fetch_ee_implementation_or_none( + "onyx.server.enterprise_settings.store", "get_logo_filename" + ) + if get_logo_filename_fn: + db_filename = get_logo_filename_fn() + + return OnyxRuntime._get_with_static_fallback(db_filename, STATIC_FILENAME) + + @staticmethod + def get_emailable_logo() -> OnyxFile: + STATIC_FILENAME = "static/images/logo.png" + + db_filename: str | None = None + get_logo_filename_fn = fetch_ee_implementation_or_none( + "onyx.server.enterprise_settings.store", "get_logo_filename" + ) + if get_logo_filename_fn: + db_filename = get_logo_filename_fn() + + onyx_file = OnyxRuntime._get_with_static_fallback(db_filename, STATIC_FILENAME) + + # check dimensions and resize downwards if necessary or if not PNG + image = Image.open(io.BytesIO(onyx_file.data)) + if ( + image.size[0] > ONYX_EMAILABLE_LOGO_MAX_DIM + or image.size[1] > ONYX_EMAILABLE_LOGO_MAX_DIM + or image.format != "PNG" + ): + image.thumbnail( + (ONYX_EMAILABLE_LOGO_MAX_DIM, ONYX_EMAILABLE_LOGO_MAX_DIM), + Image.LANCZOS, + ) # maintains aspect ratio + output_buffer = io.BytesIO() + image.save(output_buffer, format="PNG") + onyx_file = OnyxFile(data=output_buffer.getvalue(), mime_type="image/png") + + return onyx_file + + @staticmethod + def get_logotype() -> OnyxFile: + STATIC_FILENAME = "static/images/logotype.png" + + db_filename: str | None = None + get_logotype_filename_fn = fetch_ee_implementation_or_none( + "onyx.server.enterprise_settings.store", "get_logotype_filename" + ) + if get_logotype_filename_fn: + db_filename = get_logotype_filename_fn() + + return OnyxRuntime._get_with_static_fallback(db_filename, STATIC_FILENAME) diff --git a/backend/onyx/utils/file.py b/backend/onyx/utils/file.py index 116948ea4..60a965b6b 100644 --- a/backend/onyx/utils/file.py +++ b/backend/onyx/utils/file.py @@ -1,6 +1,25 @@ +from typing import cast + +import puremagic from pydantic import BaseModel class OnyxFile(BaseModel): data: bytes mime_type: str + + +class OnyxStaticFileManager: + @staticmethod + def get_static(filename: str) -> OnyxFile | None: + try: + mime_type: str = "application/octet-stream" + with open(filename, "rb") as f: + file_content = f.read() + matches = puremagic.magic_string(file_content) + if matches: + mime_type = cast(str, matches[0].mime_type) + except (OSError, FileNotFoundError, PermissionError): + return None + + return OnyxFile(data=file_content, mime_type=mime_type) diff --git a/backend/onyx/utils/variable_functionality.py b/backend/onyx/utils/variable_functionality.py index 8d4f6be8d..57b9c1f35 100644 --- a/backend/onyx/utils/variable_functionality.py +++ b/backend/onyx/utils/variable_functionality.py @@ -108,6 +108,25 @@ def fetch_versioned_implementation_with_fallback( return fallback +def fetch_ee_implementation_or_none(module: str, attribute: str) -> T | None: + """ + Attempts to fetch a versioned implementation of a specified attribute from a given module. + If the attempt fails (e.g., due to an import error or missing attribute), the function + returns None + + Args: + module (str): The name of the module from which to fetch the attribute. + attribute (str): The name of the attribute to fetch from the module. + + Returns: + T: The fetched implementation if successful, otherwise None. + """ + try: + return fetch_versioned_implementation(module, attribute) + except Exception: + return None + + def noop_fallback(*args: Any, **kwargs: Any) -> None: """ A no-op (no operation) fallback function that accepts any arguments but does nothing.