highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-03-17 21:32:36 +01:00
Code Issues Packages Projects Releases Wiki Activity

refactor, use inline attachment for image (base64 encoding doesn't work)

Browse Source
This commit is contained in:
Richard Kuo (Onyx) 2025-03-13 16:50:43 -07:00
parent f6c1f1cd80
commit 18402b7d78
9 changed files with 206 additions and 93 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
backend/ee/onyx/file_store/onyx_file_store.py
View File

@ -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)

12
backend/ee/onyx/server/enterprise_settings/api.py
View File

@ -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:

15
backend/ee/onyx/server/enterprise_settings/store.py
View File

@ -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

92
backend/onyx/auth/email_utils.py
View File

@ -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 = """\
<td class="header">
<img
style="background-color: #ffffff; border-radius: 8px;"
src="data:{logo_mimetype};base64,{logo_b64}"
src="cid:logo.png"
alt="{application_name} Logo"
>
</td>
@ -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'<br>Have questions? Join our Slack community <a href="{ONYX_SLACK_URL}">here</a>.'
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"<p>You have been invited by {current_user.email} to join an organization on {application_name} .</p>"
message = f"<p>You have been invited by {current_user.email} to join an organization on {application_name}.</p>"
if auth_type == AuthType.CLOUD:
message += (
"<p>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),
)

2
backend/onyx/configs/constants.py
View File

@ -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

27
backend/onyx/file_store/onyx_file_store.py
View File

@ -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

95
backend/onyx/server/runtime/onyx_runtime.py Normal file
View File

@ -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)

19
backend/onyx/utils/file.py
View File

@ -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)

19
backend/onyx/utils/variable_functionality.py
View File

@ -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.