Spaces:

marvin-cusm-chatbot
/

champ-chatbot

Paused

Deploy from GitLab 3d688434

9c98929 about 1 month ago

5.39 kB

from enum import StrEnum
from typing import Literal, Set
from uuid import UUID

import nh3
from pydantic import BaseModel, Field, field_validator

from constants import (
    CHAMP_RAG_ENABLED,
    MAX_COMMENT_LENGTH,
    MAX_FILE_NAME_LENGTH,
    MAX_ID_LENGTH,
    MAX_MESSAGE_LENGTH,
    MAX_RESPONSE_LENGTH,
)


class AgeGroup(StrEnum):
    AGE_0_17 = "0-17"
    AGE_18_24 = "18-24"
    AGE_25_34 = "25-34"
    AGE_35_44 = "35-44"
    AGE_45_54 = "45-54"
    AGE_55_64 = "55-64"
    AGE_65_PLUS = "65+"


class Gender(StrEnum):
    MALE = "male"
    FEMALE = "female"
    NON_BINARY = "non-binary"
    OTHER = "other"
    PREFER_NOT = "prefer-not"


class Role(StrEnum):
    PATIENT = "patient"
    CLINICIAN = "clinician"
    COMPUTER_SCIENTIST = "computer-scientist"
    RESEARCHER = "researcher"
    OTHER = "other"


ModelType = Literal[
    "champ",
    "openai",
    "google-conservative",
    "google-creative",
    "fake",
    "skills_wiki",
]


class IdentifierBase(BaseModel):
    # Participant ID could be in ProfileBase instead. It doesn't really matter.
    participant_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )


class ProfileBase(BaseModel):
    consent: bool
    age_group: AgeGroup
    gender: Gender
    roles: Set[Role] = Field(min_length=1, max_length=5)


class ChatRequest(IdentifierBase, ProfileBase):
    conversation_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )
    lang: Literal["en", "fr"]
    human_message: str = Field(min_length=1, max_length=MAX_MESSAGE_LENGTH)

    @field_validator("human_message")
    def sanitize_human_message(cls, human_message: str):
        """Remove HTML tags to prevent XSS"""
        return nh3.clean(human_message)


class StartConversationRequest(BaseModel):
    model_type: ModelType

    @field_validator("model_type")
    def refuse_rag_champ_when_off(cls, model_type: str) -> str:
        """Refuse the RAG model type (champ) unless CHAMP_RAG_ENABLED turned it on.
        Args:
            model_type (str): The model type asked for.
        Raises:
            ValueError: When it is the RAG model type and that is off; answered with a 422.
        Returns:
            str: The model type.
        """
        if model_type == "champ" and not CHAMP_RAG_ENABLED:
            raise ValueError(
                "the RAG model type 'champ' is turned off (CHAMP_RAG_ENABLED is not true)"
            )
        return model_type


class ClearConversationRequest(BaseModel):
    conversation_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )


class RestoredTurn(BaseModel):
    """A turn as the user saw it, with the figures its reply showed."""

    human_message: str
    reply: str
    reply_id: str
    gwp_kgcoeq: float
    water_L: float
    electricity_kWh: float
    n_tokens: int


class CurrentConversation(BaseModel):
    """A conversation to show again: its turns as the user saw them, and its files' names."""

    conversation_id: str
    model_type: str
    turns: list[RestoredTurn]
    file_names: list[str]


class FeedbackRequest(IdentifierBase, ProfileBase):
    conversation_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )
    message_index: int = Field(ge=0, le=10_000)
    rating: Literal["like", "dislike", "mixed"]
    comment: str = Field(min_length=0, max_length=MAX_COMMENT_LENGTH)
    reply_content: str = Field(min_length=1, max_length=MAX_RESPONSE_LENGTH)
    reply_id: UUID

    @field_validator("comment")
    def sanitize_comment(cls, comment: str):
        """Remove HTML tags to prevent XSS"""
        return nh3.clean(comment)

    @field_validator("reply_content")
    def sanitize_reply_content(cls, reply_content: str):
        """Remove HTML tags to prevent XSS"""
        return nh3.clean(reply_content)


class CommentRequest(IdentifierBase, ProfileBase):
    comment: str = Field(min_length=1, max_length=MAX_COMMENT_LENGTH)
    # The conversation on screen; None when the user hasn't sent a message yet.
    conversation_id: str | None = Field(
        default=None, pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )

    @field_validator("comment")
    def sanitize_comment(cls, comment: str):
        """Remove HTML tags to prevent XSS"""
        return nh3.clean(comment)


class DeleteFileRequest(IdentifierBase, ProfileBase):
    conversation_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )
    file_name: str = Field(
        # Matches is_valid_filename in helpers/file_helper.py: block only
        # control characters, path separators, and HTML/attribute-injection
        # characters — everything else a real filename can contain (commas,
        # apostrophes, colons, etc.) is allowed, so anything accepted on
        # upload can also be deleted. The "." / ".." exclusion lives in the
        # validator below — pydantic-core's regex engine has no look-ahead.
        pattern=r'^[^\x00-\x1f\\/<>"]+$',
        min_length=1,
        max_length=MAX_FILE_NAME_LENGTH,
    )

    @field_validator("file_name")
    def reject_dot_names(cls, file_name: str):
        if file_name in {".", ".."}:
            raise ValueError('file_name cannot be "." or ".."')
        return file_name

	from enum import StrEnum
	from typing import Literal, Set
	from uuid import UUID

	import nh3
	from pydantic import BaseModel, Field, field_validator

	from constants import (
	CHAMP_RAG_ENABLED,
	MAX_COMMENT_LENGTH,
	MAX_FILE_NAME_LENGTH,
	MAX_ID_LENGTH,
	MAX_MESSAGE_LENGTH,
	MAX_RESPONSE_LENGTH,
	)


	class AgeGroup(StrEnum):
	AGE_0_17 = "0-17"
	AGE_18_24 = "18-24"
	AGE_25_34 = "25-34"
	AGE_35_44 = "35-44"
	AGE_45_54 = "45-54"
	AGE_55_64 = "55-64"
	AGE_65_PLUS = "65+"


	class Gender(StrEnum):
	MALE = "male"
	FEMALE = "female"
	NON_BINARY = "non-binary"
	OTHER = "other"
	PREFER_NOT = "prefer-not"


	class Role(StrEnum):
	PATIENT = "patient"
	CLINICIAN = "clinician"
	COMPUTER_SCIENTIST = "computer-scientist"
	RESEARCHER = "researcher"
	OTHER = "other"


	ModelType = Literal[
	"champ",
	"openai",
	"google-conservative",
	"google-creative",
	"fake",
	"skills_wiki",
	]


	class IdentifierBase(BaseModel):
	# Participant ID could be in ProfileBase instead. It doesn't really matter.
	participant_id: str = Field(
	pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
	)


	class ProfileBase(BaseModel):
	consent: bool
	age_group: AgeGroup
	gender: Gender
	roles: Set[Role] = Field(min_length=1, max_length=5)


	class ChatRequest(IdentifierBase, ProfileBase):
	conversation_id: str = Field(
	pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
	)
	lang: Literal["en", "fr"]
	human_message: str = Field(min_length=1, max_length=MAX_MESSAGE_LENGTH)

	@field_validator("human_message")
	def sanitize_human_message(cls, human_message: str):
	"""Remove HTML tags to prevent XSS"""
	return nh3.clean(human_message)


	class StartConversationRequest(BaseModel):
	model_type: ModelType

	@field_validator("model_type")
	def refuse_rag_champ_when_off(cls, model_type: str) -> str:
	"""Refuse the RAG model type (champ) unless CHAMP_RAG_ENABLED turned it on.

	Args:
	model_type (str): The model type asked for.

	Raises:
	ValueError: When it is the RAG model type and that is off; answered with a 422.

	Returns:
	str: The model type.
	"""
	if model_type == "champ" and not CHAMP_RAG_ENABLED:
	raise ValueError(
	"the RAG model type 'champ' is turned off (CHAMP_RAG_ENABLED is not true)"
	)
	return model_type


	class ClearConversationRequest(BaseModel):
	conversation_id: str = Field(
	pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
	)


	class RestoredTurn(BaseModel):
	"""A turn as the user saw it, with the figures its reply showed."""

	human_message: str
	reply: str
	reply_id: str
	gwp_kgcoeq: float
	water_L: float
	electricity_kWh: float
	n_tokens: int


	class CurrentConversation(BaseModel):
	"""A conversation to show again: its turns as the user saw them, and its files' names."""

	conversation_id: str
	model_type: str
	turns: list[RestoredTurn]
	file_names: list[str]


	class FeedbackRequest(IdentifierBase, ProfileBase):
	conversation_id: str = Field(
	pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
	)
	message_index: int = Field(ge=0, le=10_000)
	rating: Literal["like", "dislike", "mixed"]
	comment: str = Field(min_length=0, max_length=MAX_COMMENT_LENGTH)
	reply_content: str = Field(min_length=1, max_length=MAX_RESPONSE_LENGTH)
	reply_id: UUID

	@field_validator("comment")
	def sanitize_comment(cls, comment: str):
	"""Remove HTML tags to prevent XSS"""
	return nh3.clean(comment)

	@field_validator("reply_content")
	def sanitize_reply_content(cls, reply_content: str):
	"""Remove HTML tags to prevent XSS"""
	return nh3.clean(reply_content)


	class CommentRequest(IdentifierBase, ProfileBase):
	comment: str = Field(min_length=1, max_length=MAX_COMMENT_LENGTH)
	# The conversation on screen; None when the user hasn't sent a message yet.
	conversation_id: str \| None = Field(
	default=None, pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
	)

	@field_validator("comment")
	def sanitize_comment(cls, comment: str):
	"""Remove HTML tags to prevent XSS"""
	return nh3.clean(comment)


	class DeleteFileRequest(IdentifierBase, ProfileBase):
	conversation_id: str = Field(
	pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
	)
	file_name: str = Field(
	# Matches is_valid_filename in helpers/file_helper.py: block only
	# control characters, path separators, and HTML/attribute-injection
	# characters — everything else a real filename can contain (commas,
	# apostrophes, colons, etc.) is allowed, so anything accepted on
	# upload can also be deleted. The "." / ".." exclusion lives in the
	# validator below — pydantic-core's regex engine has no look-ahead.
	pattern=r'^[^\x00-\x1f\\/<>"]+$',
	min_length=1,
	max_length=MAX_FILE_NAME_LENGTH,
	)

	@field_validator("file_name")
	def reject_dot_names(cls, file_name: str):
	if file_name in {".", ".."}:
	raise ValueError('file_name cannot be "." or ".."')
	return file_name