Guilherme34's picture
Upload folder using huggingface_hub
aa15bce verified
"""Schemas for the email search task tools."""
from __future__ import annotations
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, ConfigDict, Field
TASK_TOOL_NAME = "task_email_search"
SEARCH_TOOL_NAME = "gmail_fetch_emails"
COMPLETE_TOOL_NAME = "return_search_results"
_SCHEMAS: List[Dict[str, Any]] = [
{
"type": "function",
"function": {
"name": TASK_TOOL_NAME,
"description": "Expand a raw Gmail search request into multiple targeted queries and return relevant emails.",
"parameters": {
"type": "object",
"properties": {
"search_query": {
"type": "string",
"description": "Raw search request describing the emails to find.",
},
},
"required": ["search_query"],
"additionalProperties": False,
},
},
}
]
class GmailSearchEmail(BaseModel):
"""Clean email representation with enhanced content processing."""
model_config = ConfigDict(extra="ignore", frozen=True)
# Core identifiers
id: str # message_id from Gmail API
thread_id: Optional[str] = None
query: str # The search query that found this email
# Email metadata
subject: str
sender: str
recipient: str # to field
timestamp: datetime
label_ids: List[str] = Field(default_factory=list)
# Clean content (primary field for LLM consumption)
clean_text: str # Processed, readable email content
# Attachment information
has_attachments: bool = False
attachment_count: int = 0
attachment_filenames: List[str] = Field(default_factory=list)
class EmailSearchToolResult(BaseModel):
"""Structured payload for each tool-call response."""
status: Literal["success", "error"]
query: Optional[str] = None
result_count: Optional[int] = None
next_page_token: Optional[str] = None
messages: List[GmailSearchEmail] = Field(default_factory=list)
error: Optional[str] = None
class TaskEmailSearchPayload(BaseModel):
"""Envelope for the final email selection."""
model_config = ConfigDict(extra="forbid", frozen=True)
emails: List[GmailSearchEmail]
_COMPLETION_SCHEMAS: List[Dict[str, Any]] = [
{
"type": "function",
"function": {
"name": COMPLETE_TOOL_NAME,
"description": "Return the final list of relevant Gmail message ids that match the search criteria.",
"parameters": {
"type": "object",
"properties": {
"message_ids": {
"type": "array",
"description": "List of Gmail message ids deemed relevant.",
"items": {"type": "string"},
},
},
"required": ["message_ids"],
"additionalProperties": False,
},
},
}
]
def get_completion_schema() -> Dict[str, Any]:
return _COMPLETION_SCHEMAS[0]
def get_schemas() -> List[Dict[str, Any]]:
"""Return the JSON schema for the email search task."""
return _SCHEMAS
__all__ = [
"GmailSearchEmail",
"EmailSearchToolResult",
"TaskEmailSearchPayload",
"SEARCH_TOOL_NAME",
"COMPLETE_TOOL_NAME",
"TASK_TOOL_NAME",
"get_completion_schema",
"get_schemas",
]