diff --git a/alignai/client.py b/alignai/client.py index 43d93cf..1f9d6c7 100644 --- a/alignai/client.py +++ b/alignai/client.py @@ -19,6 +19,12 @@ from alignai.ingestion.v1alpha.event_pb2 import Event, EventProperties from alignai.logger import get_logger from alignai.utils import CustomProperties, datetime_to_timestamp, serialize_custom_properties +from alignai.validate import ( + ValidationError, + validate_custom_properties_or_raise, + validate_session_id_or_raise, + validate_user_id_or_raise, +) from alignai.worker import Worker @@ -76,6 +82,13 @@ def open_session( assistant_id (str, optional): Assistant ID. Defaults to "DEFAULT". custom_properties (dict[str, str] | None, optional): Custom properties associated with the session. Defaults to None. """ + validate_session_id_or_raise(session_id) + validate_user_id_or_raise(user_id) + if len(assistant_id) > 64: + raise ValidationError("assistant_id must be at most 64 characters") + if custom_properties: + validate_custom_properties_or_raise(custom_properties) + session_properties_args = {"session_id": session_id, "user_id": user_id, "assistant_id": assistant_id} open_session_event = Event( id=uuid.uuid4().hex, @@ -95,6 +108,8 @@ def close_session(self, session_id: str) -> None: Args: session_id (str): Session ID. """ + validate_session_id_or_raise(session_id) + close_session_event = Event( id=uuid.uuid4().hex, type=EventTypes.SESSION_CLOSE, @@ -125,6 +140,16 @@ def identify_user( create_time (datetime | None, optional): User creation time. Defaults to None. custom_properties (dict[str, str] | None, optional): Custom properties associated with the user. Defaults to None. """ # noqa: E501 + validate_user_id_or_raise(user_id) + if display_name and len(display_name) > 64: + raise ValidationError("display_name must be at most 64 characters") + if email and len(email) > 256: + raise ValidationError("email must be at most 256 characters") + if country_code and len(country_code) != 2: + raise ValidationError("country_code must be ISO Alpha-2 code") + if custom_properties: + validate_custom_properties_or_raise(custom_properties) + user_properties_args = {"user_id": user_id} if display_name is not None: user_properties_args["user_display_name"] = display_name @@ -166,6 +191,14 @@ def create_message( content (str): Content of the message. custom_properties (dict[str, str] | None, optional): Custom properties associated with the message. Defaults to None. """ # noqa: E501 + validate_session_id_or_raise(session_id) + if message_index <= 0: + raise ValidationError("message_index must be greater than 0") + if not content: + raise ValidationError("content is required") + if custom_properties: + validate_custom_properties_or_raise(custom_properties) + if message_index <= 0: self.logger.error(f"Invalid message index '{message_index}': Message index must be a positive integer") return diff --git a/alignai/validate.py b/alignai/validate.py new file mode 100644 index 0000000..0ff3993 --- /dev/null +++ b/alignai/validate.py @@ -0,0 +1,32 @@ +class ValidationError(ValueError): + def __init__(self, message: str): + self.message = message + + def __str__(self): + return f"ValidationError: {self.message}" + + +def validate_session_id_or_raise(session_id: str): + if not session_id: + raise ValidationError("session_id is required") + if len(session_id) > 64: + raise ValidationError("session_id must be at most 64 characters") + + +def validate_user_id_or_raise(user_id: str): + if not user_id: + raise ValidationError("user_id is required") + if len(user_id) > 64: + raise ValidationError("user_id must be at most 64 characters") + + +def validate_custom_properties_or_raise(custom_properties: dict[str, str]): + if len(custom_properties) > 10: + raise ValidationError("custom_properties must have at most 10 keys") + for key, value in custom_properties.items(): + if not key: + raise ValidationError("key of custom_property is required") + if len(key) > 64: + raise ValidationError("key of custom_property must be at most 64 characters") + if len(value) > 256: + raise ValidationError("value of custom_property must be at most 256 characters")