EE movement followup for Standard Answers (#2467)

* Move StandardAnswer to EE section of danswer/db/models

* Move StandardAnswer DB layer to EE

* Add EERequiredError for distinct error handling here

* Handle EE fallback for slack bot config

* Migrate all standard answer models to ee

* Flagging categories for removal

* Add missing versioned impl for update_slack_bot_config

---------

Co-authored-by: danswer-trial <danswer-trial@danswer-trials-MacBook-Pro.local>

backend/ee/danswer/danswerbot/slack/handlers/handle_standard_answers.py

@@ -21,11 +21,11 @@ from danswer.db.chat import get_or_create_root_message
 from danswer.db.models import Prompt
 from danswer.db.models import SlackBotConfig
 from danswer.db.models import StandardAnswer as StandardAnswerModel
-from danswer.server.manage.models import StandardAnswer as PydanticStandardAnswer
 from danswer.utils.logger import DanswerLoggingAdapter
 from danswer.utils.logger import setup_logger
 from ee.danswer.db.standard_answer import fetch_standard_answer_categories_by_names
 from ee.danswer.db.standard_answer import find_matching_standard_answers
+from ee.danswer.server.manage.models import StandardAnswer as PydanticStandardAnswer
 
 logger = setup_logger()
 

backend/ee/danswer/db/standard_answer.py

@@ -12,6 +12,198 @@ from danswer.utils.logger import setup_logger
 logger = setup_logger()
 
 
+def check_category_validity(category_name: str) -> bool:
+    """If a category name is too long, it should not be used (it will cause an error in Postgres
+    as the unique constraint can only apply to entries that are less than 2704 bytes).
+
+    Additionally, extremely long categories are not really usable / useful."""
+    if len(category_name) > 255:
+        logger.error(
+            f"Category with name '{category_name}' is too long, cannot be used"
+        )
+        return False
+
+    return True
+
+
+def insert_standard_answer_category(
+    category_name: str, db_session: Session
+) -> StandardAnswerCategory:
+    if not check_category_validity(category_name):
+        raise ValueError(f"Invalid category name: {category_name}")
+    standard_answer_category = StandardAnswerCategory(name=category_name)
+    db_session.add(standard_answer_category)
+    db_session.commit()
+
+    return standard_answer_category
+
+
+def insert_standard_answer(
+    keyword: str,
+    answer: str,
+    category_ids: list[int],
+    match_regex: bool,
+    match_any_keywords: bool,
+    db_session: Session,
+) -> StandardAnswer:
+    existing_categories = fetch_standard_answer_categories_by_ids(
+        standard_answer_category_ids=category_ids,
+        db_session=db_session,
+    )
+    if len(existing_categories) != len(category_ids):
+        raise ValueError(f"Some or all categories with ids {category_ids} do not exist")
+
+    standard_answer = StandardAnswer(
+        keyword=keyword,
+        answer=answer,
+        categories=existing_categories,
+        active=True,
+        match_regex=match_regex,
+        match_any_keywords=match_any_keywords,
+    )
+    db_session.add(standard_answer)
+    db_session.commit()
+    return standard_answer
+
+
+def update_standard_answer(
+    standard_answer_id: int,
+    keyword: str,
+    answer: str,
+    category_ids: list[int],
+    match_regex: bool,
+    match_any_keywords: bool,
+    db_session: Session,
+) -> StandardAnswer:
+    standard_answer = db_session.scalar(
+        select(StandardAnswer).where(StandardAnswer.id == standard_answer_id)
+    )
+    if standard_answer is None:
+        raise ValueError(f"No standard answer with id {standard_answer_id}")
+
+    existing_categories = fetch_standard_answer_categories_by_ids(
+        standard_answer_category_ids=category_ids,
+        db_session=db_session,
+    )
+    if len(existing_categories) != len(category_ids):
+        raise ValueError(f"Some or all categories with ids {category_ids} do not exist")
+
+    standard_answer.keyword = keyword
+    standard_answer.answer = answer
+    standard_answer.categories = list(existing_categories)
+    standard_answer.match_regex = match_regex
+    standard_answer.match_any_keywords = match_any_keywords
+
+    db_session.commit()
+
+    return standard_answer
+
+
+def remove_standard_answer(
+    standard_answer_id: int,
+    db_session: Session,
+) -> None:
+    standard_answer = db_session.scalar(
+        select(StandardAnswer).where(StandardAnswer.id == standard_answer_id)
+    )
+    if standard_answer is None:
+        raise ValueError(f"No standard answer with id {standard_answer_id}")
+
+    standard_answer.active = False
+    db_session.commit()
+
+
+def update_standard_answer_category(
+    standard_answer_category_id: int,
+    category_name: str,
+    db_session: Session,
+) -> StandardAnswerCategory:
+    standard_answer_category = db_session.scalar(
+        select(StandardAnswerCategory).where(
+            StandardAnswerCategory.id == standard_answer_category_id
+        )
+    )
+    if standard_answer_category is None:
+        raise ValueError(
+            f"No standard answer category with id {standard_answer_category_id}"
+        )
+
+    if not check_category_validity(category_name):
+        raise ValueError(f"Invalid category name: {category_name}")
+
+    standard_answer_category.name = category_name
+
+    db_session.commit()
+
+    return standard_answer_category
+
+
+def fetch_standard_answer_category(
+    standard_answer_category_id: int,
+    db_session: Session,
+) -> StandardAnswerCategory | None:
+    return db_session.scalar(
+        select(StandardAnswerCategory).where(
+            StandardAnswerCategory.id == standard_answer_category_id
+        )
+    )
+
+
+def fetch_standard_answer_categories_by_ids(
+    standard_answer_category_ids: list[int],
+    db_session: Session,
+) -> Sequence[StandardAnswerCategory]:
+    return db_session.scalars(
+        select(StandardAnswerCategory).where(
+            StandardAnswerCategory.id.in_(standard_answer_category_ids)
+        )
+    ).all()
+
+
+def fetch_standard_answer_categories(
+    db_session: Session,
+) -> Sequence[StandardAnswerCategory]:
+    return db_session.scalars(select(StandardAnswerCategory)).all()
+
+
+def fetch_standard_answer(
+    standard_answer_id: int,
+    db_session: Session,
+) -> StandardAnswer | None:
+    return db_session.scalar(
+        select(StandardAnswer).where(StandardAnswer.id == standard_answer_id)
+    )
+
+
+def fetch_standard_answers(db_session: Session) -> Sequence[StandardAnswer]:
+    return db_session.scalars(
+        select(StandardAnswer).where(StandardAnswer.active.is_(True))
+    ).all()
+
+
+def create_initial_default_standard_answer_category(db_session: Session) -> None:
+    default_category_id = 0
+    default_category_name = "General"
+    default_category = fetch_standard_answer_category(
+        standard_answer_category_id=default_category_id,
+        db_session=db_session,
+    )
+    if default_category is not None:
+        if default_category.name != default_category_name:
+            raise ValueError(
+                "DB is not in a valid initial state. "
+                "Default standard answer category does not have expected name."
+            )
+        return
+
+    standard_answer_category = StandardAnswerCategory(
+        id=default_category_id,
+        name=default_category_name,
+    )
+    db_session.add(standard_answer_category)
+    db_session.commit()
+
+
 def fetch_standard_answer_categories_by_names(
     standard_answer_category_names: list[str],
     db_session: Session,

backend/ee/danswer/server/manage/models.py

@@ -0,0 +1,98 @@
+import re
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic import field_validator
+from pydantic import model_validator
+
+from danswer.db.models import StandardAnswer as StandardAnswerModel
+from danswer.db.models import StandardAnswerCategory as StandardAnswerCategoryModel
+
+
+class StandardAnswerCategoryCreationRequest(BaseModel):
+    name: str
+
+
+class StandardAnswerCategory(BaseModel):
+    id: int
+    name: str
+
+    @classmethod
+    def from_model(
+        cls, standard_answer_category: StandardAnswerCategoryModel
+    ) -> "StandardAnswerCategory":
+        return cls(
+            id=standard_answer_category.id,
+            name=standard_answer_category.name,
+        )
+
+
+class StandardAnswer(BaseModel):
+    id: int
+    keyword: str
+    answer: str
+    categories: list[StandardAnswerCategory]
+    match_regex: bool
+    match_any_keywords: bool
+
+    @classmethod
+    def from_model(cls, standard_answer_model: StandardAnswerModel) -> "StandardAnswer":
+        return cls(
+            id=standard_answer_model.id,
+            keyword=standard_answer_model.keyword,
+            answer=standard_answer_model.answer,
+            match_regex=standard_answer_model.match_regex,
+            match_any_keywords=standard_answer_model.match_any_keywords,
+            categories=[
+                StandardAnswerCategory.from_model(standard_answer_category_model)
+                for standard_answer_category_model in standard_answer_model.categories
+            ],
+        )
+
+
+class StandardAnswerCreationRequest(BaseModel):
+    keyword: str
+    answer: str
+    categories: list[int]
+    match_regex: bool
+    match_any_keywords: bool
+
+    @field_validator("categories", mode="before")
+    @classmethod
+    def validate_categories(cls, value: list[int]) -> list[int]:
+        if len(value) < 1:
+            raise ValueError(
+                "At least one category must be attached to a standard answer"
+            )
+        return value
+
+    @model_validator(mode="after")
+    def validate_only_match_any_if_not_regex(self) -> Any:
+        if self.match_regex and self.match_any_keywords:
+            raise ValueError(
+                "Can only match any keywords in keyword mode, not regex mode"
+            )
+
+        return self
+
+    @model_validator(mode="after")
+    def validate_keyword_if_regex(self) -> Any:
+        if not self.match_regex:
+            # no validation for keywords
+            return self
+
+        try:
+            re.compile(self.keyword)
+            return self
+        except re.error as err:
+            if isinstance(err.pattern, bytes):
+                raise ValueError(
+                    f'invalid regex pattern r"{err.pattern.decode()}" in `keyword`: {err.msg}'
+                )
+            else:
+                pattern = f'r"{err.pattern}"' if err.pattern is not None else ""
+                raise ValueError(
+                    " ".join(
+                        ["invalid regex pattern", pattern, f"in `keyword`: {err.msg}"]
+                    )
+                )

backend/ee/danswer/server/manage/standard_answer.py

@@ -6,19 +6,19 @@ from sqlalchemy.orm import Session
 from danswer.auth.users import current_admin_user
 from danswer.db.engine import get_session
 from danswer.db.models import User
-from danswer.db.standard_answer import fetch_standard_answer
-from danswer.db.standard_answer import fetch_standard_answer_categories
-from danswer.db.standard_answer import fetch_standard_answer_category
-from danswer.db.standard_answer import fetch_standard_answers
-from danswer.db.standard_answer import insert_standard_answer
-from danswer.db.standard_answer import insert_standard_answer_category
-from danswer.db.standard_answer import remove_standard_answer
-from danswer.db.standard_answer import update_standard_answer
-from danswer.db.standard_answer import update_standard_answer_category
-from danswer.server.manage.models import StandardAnswer
-from danswer.server.manage.models import StandardAnswerCategory
-from danswer.server.manage.models import StandardAnswerCategoryCreationRequest
-from danswer.server.manage.models import StandardAnswerCreationRequest
+from ee.danswer.db.standard_answer import fetch_standard_answer
+from ee.danswer.db.standard_answer import fetch_standard_answer_categories
+from ee.danswer.db.standard_answer import fetch_standard_answer_category
+from ee.danswer.db.standard_answer import fetch_standard_answers
+from ee.danswer.db.standard_answer import insert_standard_answer
+from ee.danswer.db.standard_answer import insert_standard_answer_category
+from ee.danswer.db.standard_answer import remove_standard_answer
+from ee.danswer.db.standard_answer import update_standard_answer
+from ee.danswer.db.standard_answer import update_standard_answer_category
+from ee.danswer.server.manage.models import StandardAnswer
+from ee.danswer.server.manage.models import StandardAnswerCategory
+from ee.danswer.server.manage.models import StandardAnswerCategoryCreationRequest
+from ee.danswer.server.manage.models import StandardAnswerCreationRequest
 
 router = APIRouter(prefix="/manage")
 

backend/ee/danswer/server/query_and_chat/models.py

@@ -8,7 +8,7 @@ from danswer.search.enums import SearchType
 from danswer.search.models import ChunkContext
 from danswer.search.models import RerankingDetails
 from danswer.search.models import RetrievalDetails
-from danswer.server.manage.models import StandardAnswer
+from ee.danswer.server.manage.models import StandardAnswer
 
 
 class StandardAnswerRequest(BaseModel):

backend/ee/danswer/server/seeding.py

@@ -13,6 +13,9 @@ from danswer.server.manage.llm.models import LLMProviderUpsertRequest
 from danswer.server.settings.models import Settings
 from danswer.server.settings.store import store_settings as store_base_settings
 from danswer.utils.logger import setup_logger
+from ee.danswer.db.standard_answer import (
+    create_initial_default_standard_answer_category,
+)
 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 store_analytics_script
@@ -21,6 +24,7 @@ from ee.danswer.server.enterprise_settings.store import (
 )
 from ee.danswer.server.enterprise_settings.store import upload_logo
 
+
 logger = setup_logger()
 
 _SEED_CONFIG_ENV_VAR_NAME = "ENV_SEED_CONFIGURATION"
@@ -146,3 +150,6 @@ def seed_db() -> None:
         _seed_logo(db_session, seed_config.seeded_logo_path)
         _seed_enterprise_settings(seed_config)
         _seed_analytics_script(seed_config)
+
+        logger.notice("Verifying default standard answer category exists.")
+        create_initial_default_standard_answer_category(db_session)