From e34bcbbd06164e3484851d839537253766a9a084 Mon Sep 17 00:00:00 2001 From: pablodanswer Date: Fri, 21 Jun 2024 12:34:26 -0700 Subject: [PATCH] Add persistent name and logo seeding (#107) --- backend/Dockerfile | 6 +- backend/danswer/configs/app_configs.py | 1 + .../danswer/server/enterprise_settings/api.py | 29 ++------ .../server/enterprise_settings/store.py | 69 +++++++++++++++++++ backend/ee/danswer/server/seeding.py | 23 ++++++- .../docker_compose/docker-compose.dev.yml | 2 + .../docker_compose/docker-compose.prod.yml | 1 - 7 files changed, 103 insertions(+), 28 deletions(-) diff --git a/backend/Dockerfile b/backend/Dockerfile index 6d13053ca..370bc14d0 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -10,14 +10,15 @@ https://github.com/danswer-ai/danswer" # Default DANSWER_VERSION, typically overriden during builds by GitHub Actions. ARG DANSWER_VERSION=0.3-dev ENV DANSWER_VERSION=${DANSWER_VERSION} -RUN echo "DANSWER_VERSION: ${DANSWER_VERSION}" +RUN echo "DANSWER_VERSION: ${DANSWER_VERSION}" # Install system dependencies # cmake needed for psycopg (postgres) # libpq-dev needed for psycopg (postgres) # curl included just for users' convenience # zip for Vespa step futher down # ca-certificates for HTTPS + RUN apt-get update && \ apt-get install -y cmake curl zip ca-certificates libgnutls30=3.7.9-2+deb12u2 \ libblkid1=2.38.1-5+deb12u1 libmount1=2.38.1-5+deb12u1 libsmartcols1=2.38.1-5+deb12u1 \ @@ -80,6 +81,9 @@ COPY supervisord.conf /usr/etc/supervisord.conf # Escape hatch COPY ./scripts/force_delete_connector_by_id.py /app/scripts/force_delete_connector_by_id.py +# Put logo in assets +COPY ./assets /app/assets + ENV PYTHONPATH /app # Default command which does nothing diff --git a/backend/danswer/configs/app_configs.py b/backend/danswer/configs/app_configs.py index 01ce4c052..a87689f98 100644 --- a/backend/danswer/configs/app_configs.py +++ b/backend/danswer/configs/app_configs.py @@ -53,6 +53,7 @@ MASK_CREDENTIAL_PREFIX = ( os.environ.get("MASK_CREDENTIAL_PREFIX", "True").lower() != "false" ) + SESSION_EXPIRE_TIME_SECONDS = int( os.environ.get("SESSION_EXPIRE_TIME_SECONDS") or 86400 * 7 ) # 7 days diff --git a/backend/ee/danswer/server/enterprise_settings/api.py b/backend/ee/danswer/server/enterprise_settings/api.py index 4ca625c97..85f43ca55 100644 --- a/backend/ee/danswer/server/enterprise_settings/api.py +++ b/backend/ee/danswer/server/enterprise_settings/api.py @@ -6,17 +6,17 @@ from fastapi import UploadFile from sqlalchemy.orm import Session from danswer.auth.users import current_admin_user -from danswer.configs.constants import FileOrigin from danswer.db.engine import get_session from danswer.db.models import User from danswer.file_store.file_store import get_default_file_store from ee.danswer.server.enterprise_settings.models import AnalyticsScriptUpload from ee.danswer.server.enterprise_settings.models import EnterpriseSettings +from ee.danswer.server.enterprise_settings.store import _LOGO_FILENAME from ee.danswer.server.enterprise_settings.store import load_analytics_script from ee.danswer.server.enterprise_settings.store import load_settings from ee.danswer.server.enterprise_settings.store import store_analytics_script from ee.danswer.server.enterprise_settings.store import store_settings - +from ee.danswer.server.enterprise_settings.store import upload_logo admin_router = APIRouter(prefix="/admin/enterprise-settings") basic_router = APIRouter(prefix="/enterprise-settings") @@ -38,34 +38,13 @@ def fetch_settings() -> EnterpriseSettings: return load_settings() -_LOGO_FILENAME = "__logo__" - - @admin_router.put("/logo") -def upload_logo( +def put_logo( file: UploadFile, db_session: Session = Depends(get_session), _: User | None = Depends(current_admin_user), ) -> None: - if not file.filename or ( - not file.filename.endswith(".png") - and not file.filename.endswith(".jpg") - and not file.filename.endswith(".jpeg") - ): - raise HTTPException( - status_code=400, - detail="Invalid file type - only .png, .jpg, and .jpeg files are allowed", - ) - - file_store = get_default_file_store(db_session) - file_store.save_file( - file_name=_LOGO_FILENAME, - content=file.file, - # The rest aren't really used for anything - display_name=file.filename or "DanswerReplacementLogo", - file_origin=FileOrigin.OTHER, - file_type=file.content_type or "image/jpeg", - ) + upload_logo(file=file, db_session=db_session) @basic_router.get("/logo") diff --git a/backend/ee/danswer/server/enterprise_settings/store.py b/backend/ee/danswer/server/enterprise_settings/store.py index 023d36ef0..99fb1cc90 100644 --- a/backend/ee/danswer/server/enterprise_settings/store.py +++ b/backend/ee/danswer/server/enterprise_settings/store.py @@ -1,13 +1,24 @@ import os +from io import BytesIO +from typing import Any from typing import cast +from typing import IO +from fastapi import HTTPException +from fastapi import UploadFile +from sqlalchemy.orm import Session + +from danswer.configs.constants import FileOrigin from danswer.dynamic_configs.factory import get_dynamic_config_store from danswer.dynamic_configs.interface import ConfigNotFoundError +from danswer.file_store.file_store import get_default_file_store +from danswer.utils.logger import setup_logger from ee.danswer.server.enterprise_settings.models import AnalyticsScriptUpload from ee.danswer.server.enterprise_settings.models import EnterpriseSettings _ENTERPRISE_SETTINGS_KEY = "danswer_enterprise_settings" +logger = setup_logger() def load_settings() -> EnterpriseSettings: @@ -49,3 +60,61 @@ def store_analytics_script(analytics_script_upload: AnalyticsScriptUpload) -> No get_dynamic_config_store().store( _CUSTOM_ANALYTICS_SCRIPT_KEY, analytics_script_upload.script ) + + +_LOGO_FILENAME = "__logo__" + + +def is_valid_file_type(filename: str) -> bool: + valid_extensions = (".png", ".jpg", ".jpeg") + return filename.endswith(valid_extensions) + + +def guess_file_type(filename: str) -> str: + if filename.lower().endswith(".png"): + return "image/png" + elif filename.lower().endswith(".jpg") or filename.lower().endswith(".jpeg"): + return "image/jpeg" + return "application/octet-stream" + + +def upload_logo( + db_session: Session, + file: UploadFile | str, +) -> bool: + content: IO[Any] + + if isinstance(file, str): + logger.info(f"Uploading logo from local path {file}") + if not os.path.isfile(file) or not is_valid_file_type(file): + logger.error( + "Invalid file type- only .png, .jpg, and .jpeg files are allowed" + ) + return False + + with open(file, "rb") as file_handle: + file_content = file_handle.read() + content = BytesIO(file_content) + display_name = file + file_type = guess_file_type(file) + + else: + logger.info("Uploading logo from uploaded file") + if not file.filename or not is_valid_file_type(file.filename): + raise HTTPException( + status_code=400, + detail="Invalid file type- only .png, .jpg, and .jpeg files are allowed", + ) + content = file.file + display_name = file.filename + file_type = file.content_type or "image/jpeg" + + file_store = get_default_file_store(db_session) + file_store.save_file( + file_name=_LOGO_FILENAME, + content=content, + display_name=display_name, + file_origin=FileOrigin.OTHER, + file_type=file_type, + ) + return True diff --git a/backend/ee/danswer/server/seeding.py b/backend/ee/danswer/server/seeding.py index 225ffdc90..cd0f9e634 100644 --- a/backend/ee/danswer/server/seeding.py +++ b/backend/ee/danswer/server/seeding.py @@ -8,6 +8,10 @@ from danswer.db.llm import fetch_existing_llm_providers from danswer.db.llm import upsert_llm_provider from danswer.server.manage.llm.models import LLMProviderUpsertRequest from danswer.utils.logger import setup_logger +from ee.danswer.server.enterprise_settings.models import EnterpriseSettings +from ee.danswer.server.enterprise_settings.store import store_settings +from ee.danswer.server.enterprise_settings.store import upload_logo + logger = setup_logger() @@ -17,13 +21,14 @@ _SEED_CONFIG_ENV_VAR_NAME = "ENV_SEED_CONFIGURATION" class SeedConfiguration(BaseModel): llms: list[LLMProviderUpsertRequest] | None = None admin_user_emails: list[str] | None = None + seeded_name: str | None = None + seeded_logo_path: str | None = None def _parse_env() -> SeedConfiguration | None: seed_config_str = os.getenv(_SEED_CONFIG_ENV_VAR_NAME) if seed_config_str is None: return None - seed_config = SeedConfiguration.parse_raw(seed_config_str) return seed_config @@ -47,9 +52,25 @@ def get_seed_config() -> SeedConfiguration | None: def seed_db() -> None: seed_config = _parse_env() + if seed_config is None: + logger.info("No seeding configuration file passed") return with get_session_context_manager() as db_session: if seed_config.llms is not None: _seed_llms(db_session, seed_config.llms) + + is_seeded_logo = ( + upload_logo(db_session=db_session, file=seed_config.seeded_logo_path) + if seed_config.seeded_logo_path + else False + ) + seeded_name = seed_config.seeded_name + + if is_seeded_logo or seeded_name: + logger.info("Seeding enterprise settings") + seeded_settings = EnterpriseSettings( + application_name=seeded_name, use_custom_logo=is_seeded_logo + ) + store_settings(seeded_settings) diff --git a/deployment/docker_compose/docker-compose.dev.yml b/deployment/docker_compose/docker-compose.dev.yml index 4e751e30f..def99a72d 100644 --- a/deployment/docker_compose/docker-compose.dev.yml +++ b/deployment/docker_compose/docker-compose.dev.yml @@ -88,6 +88,8 @@ services: - LOG_ENDPOINT_LATENCY=${LOG_ENDPOINT_LATENCY:-} # Enterprise Edition only - API_KEY_HASH_ROUNDS=${API_KEY_HASH_ROUNDS:-} + # Seeding configuration + - ENV_SEED_CONFIGURATION=${ENV_SEED_CONFIGURATION:-} extra_hosts: - "host.docker.internal:host-gateway" logging: diff --git a/deployment/docker_compose/docker-compose.prod.yml b/deployment/docker_compose/docker-compose.prod.yml index 762a76917..3a8af61c6 100644 --- a/deployment/docker_compose/docker-compose.prod.yml +++ b/deployment/docker_compose/docker-compose.prod.yml @@ -58,7 +58,6 @@ services: max-size: "50m" max-file: "6" - web_server: image: danswer/danswer-ee-web-server:latest build: