File size: 4,635 Bytes
c7426d8 af61c79 c7426d8 af61c79 a111c12 af61c79 a111c12 af61c79 a111c12 af61c79 a111c12 c7426d8 a171580 a111c12 a171580 c7426d8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
"""Module containing the data models for the application."""
from datetime import datetime, timedelta
from typing import List, Optional
from pydantic import BaseModel, Field
class EmailQuery(BaseModel):
"""
EmailQuery model representing the structure of an email query.
Attributes:
subject (Optional[str]): The subject of the email to search for.
from_email (Optional[str]): The sender's email address.
to_email (Optional[str]): The recipient's email address.
cc_email (Optional[str]): The CC email address.
has_words (Optional[str]): Words that the email must contain.
not_has_words (Optional[str]): Words that the email must not contain.
size (Optional[int]): The size of the email in bytes.
date_within (Optional[str]): The date within which to search for emails.
after (Optional[str]): The date after which to search for emails.
max_results (Optional[int]): The maximum number of results to return.
"""
subject: Optional[str] = None
from_email: Optional[str] = Field(None, alias="from")
to_email: Optional[str] = Field(None, alias="to")
cc_email: Optional[str] = Field(None, alias="cc")
has_words: Optional[str] = None
not_has_words: Optional[str] = None
size: Optional[int] = None
before: Optional[str] = None
after: Optional[str] = None
@classmethod
def validate_before_after(
cls, before: Optional[str], after: Optional[str]) -> tuple[Optional[str], Optional[str]]:
"""
Validates and adjusts the 'before' and 'after' date parameters.
This method ensures that the 'before' date is greater than the 'after' date.
If 'before' is not provided, it defaults to six months prior to the current date.
Args:
before (Optional[str]): The 'before' date in the format "YYYY/MM/DD". Defaults to None.
after (Optional[str]): The 'after' date in the format "YYYY/MM/DD". Defaults to None.
Returns:
tuple[Optional[str], Optional[str]]:
A tuple containing the validated 'before' and 'after' dates.
Raises:
ValueError: If the 'before' date is not greater than the 'after' date.
"""
if after is None:
after = (datetime.now() - timedelta(days=6 * 30)).strftime("%Y/%m/%d")
if before and before >= after:
raise ValueError("The 'before' date must be greater than the 'after' date.")
return before, after
def __init__(self, **data):
super().__init__(**data)
self.before, self.after = self.validate_before_after(self.before, self.after)
max_results: Optional[int] = 10
class ReqData(BaseModel):
"""
RequestData is a Pydantic model that represents the data structure for a request.
Attributes:
query (str): The query string provided by the user.
chat_id (str): The unique identifier for the chat session.
user_id (str): The unique identifier for the user.
web (Optional[bool]): A flag indicating if the request is from the web. Defaults to False.
"""
query: str
id: Optional[List[str]] = []
site: Optional[List[str]] = []
chat_id: str
user_id: str
web: Optional[bool] = False
class MailReqData(BaseModel):
"""
MailReqData is a data model representing the structure of a mail request.
Attributes:
query (str): The query or message content sent by the user.
"""
query: EmailQuery
class ReqFollowUp(BaseModel):
"""
RequestFollowUp is a Pydantic model that represents a request for follow-up.
Attributes:
query (str): The query string that needs follow-up.
contexts (list[str]): A list of context strings related to the query.
"""
query: str
contexts: list[str]
class FollowUpQ(BaseModel):
"""
FollowUpQ model to represent a follow-up question based on context information.
Attributes:
question (list[str]): A list of follow-up questions based on context information.
"""
questions: list[str] = Field(..., description="3 Follow up questions based on context.")
class ChatHistory(BaseModel):
"""
ChatHistory model representing a chat session.
Attributes:
chat_id (str): The unique identifier for the chat session.
user_id (str): The unique identifier for the user.
"""
chat_id: str
user_id: str
class ChatSession(BaseModel):
"""
ChatSession model representing a chat session.
Attributes:
user_id (str): The unique identifier for the user.
"""
user_id: str
|