diff --git a/alembic/versions/c7ac45f69849_add_timezone_to_agents_table.py b/alembic/versions/c7ac45f69849_add_timezone_to_agents_table.py new file mode 100644 index 00000000..3f11a102 --- /dev/null +++ b/alembic/versions/c7ac45f69849_add_timezone_to_agents_table.py @@ -0,0 +1,31 @@ +"""Add timezone to agents table + +Revision ID: c7ac45f69849 +Revises: 61ee53ec45a5 +Create Date: 2025-06-23 17:48:51.177458 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "c7ac45f69849" +down_revision: Union[str, None] = "61ee53ec45a5" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("agents", sa.Column("timezone", sa.String(), nullable=True, default="UTC")) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("agents", "timezone") + # ### end Alembic commands ### diff --git a/letta/agent.py b/letta/agent.py index ea009f2f..d6ce90fc 100644 --- a/letta/agent.py +++ b/letta/agent.py @@ -255,7 +255,7 @@ class Agent(BaseAgent): self.tool_rules_solver.register_tool_call(function_name) # Extend conversation with function response - function_response = package_function_response(False, error_msg) + function_response = package_function_response(False, error_msg, self.agent_state.timezone) new_message = Message( agent_id=self.agent_state.id, # Base info OpenAI-style @@ -640,7 +640,7 @@ class Agent(BaseAgent): function_response, return_char_limit=return_char_limit, truncate=truncate ) function_args.pop("self", None) - function_response = package_function_response(True, function_response_string) + function_response = package_function_response(True, function_response_string, self.agent_state.timezone) function_failed = False except Exception as e: function_args.pop("self", None) @@ -763,7 +763,7 @@ class Agent(BaseAgent): self.tool_rules_solver.clear_tool_history() # Convert MessageCreate objects to Message objects - next_input_messages = convert_message_creates_to_messages(input_messages, self.agent_state.id) + next_input_messages = convert_message_creates_to_messages(input_messages, self.agent_state.id, self.agent_state.timezone) counter = 0 total_usage = UsageStatistics() step_count = 0 @@ -823,7 +823,7 @@ class Agent(BaseAgent): model=self.model, openai_message_dict={ "role": "user", # TODO: change to system? - "content": get_heartbeat(FUNC_FAILED_HEARTBEAT_MESSAGE), + "content": get_heartbeat(self.agent_state.timezone, FUNC_FAILED_HEARTBEAT_MESSAGE), }, ) ] @@ -836,7 +836,7 @@ class Agent(BaseAgent): model=self.model, openai_message_dict={ "role": "user", # TODO: change to system? - "content": get_heartbeat(REQ_HEARTBEAT_MESSAGE), + "content": get_heartbeat(self.agent_state.timezone, REQ_HEARTBEAT_MESSAGE), }, ) ] @@ -1080,7 +1080,7 @@ class Agent(BaseAgent): assert user_message_str and isinstance( user_message_str, str ), f"user_message_str should be a non-empty string, got {type(user_message_str)}" - user_message_json_str = package_user_message(user_message_str) + user_message_json_str = package_user_message(user_message_str, self.agent_state.timezone) # Validate JSON via save/load user_message = validate_json(user_message_json_str) @@ -1143,7 +1143,9 @@ class Agent(BaseAgent): remaining_message_count = 1 + len(in_context_messages) - cutoff # System + remaining hidden_message_count = all_time_message_count - remaining_message_count summary_message_count = len(message_sequence_to_summarize) - summary_message = package_summarize_message(summary, summary_message_count, hidden_message_count, all_time_message_count) + summary_message = package_summarize_message( + summary, summary_message_count, hidden_message_count, all_time_message_count, self.agent_state.timezone + ) logger.info(f"Packaged into message: {summary_message}") prior_len = len(in_context_messages_openai) diff --git a/letta/agents/helpers.py b/letta/agents/helpers.py index de7c4d15..d98d6555 100644 --- a/letta/agents/helpers.py +++ b/letta/agents/helpers.py @@ -69,7 +69,8 @@ def _prepare_in_context_messages( # Create a new user message from the input and store it new_in_context_messages = message_manager.create_many_messages( - create_input_messages(input_messages=input_messages, agent_id=agent_state.id, actor=actor), actor=actor + create_input_messages(input_messages=input_messages, agent_id=agent_state.id, timezone=agent_state.timezone, actor=actor), + actor=actor, ) return current_in_context_messages, new_in_context_messages @@ -106,7 +107,8 @@ async def _prepare_in_context_messages_async( # Create a new user message from the input and store it new_in_context_messages = await message_manager.create_many_messages_async( - create_input_messages(input_messages=input_messages, agent_id=agent_state.id, actor=actor), actor=actor + create_input_messages(input_messages=input_messages, agent_id=agent_state.id, timezone=agent_state.timezone, actor=actor), + actor=actor, ) return current_in_context_messages, new_in_context_messages @@ -141,7 +143,9 @@ async def _prepare_in_context_messages_no_persist_async( current_in_context_messages = await message_manager.get_messages_by_ids_async(message_ids=agent_state.message_ids, actor=actor) # Create a new user message from the input but dont store it yet - new_in_context_messages = create_input_messages(input_messages=input_messages, agent_id=agent_state.id, actor=actor) + new_in_context_messages = create_input_messages( + input_messages=input_messages, agent_id=agent_state.id, timezone=agent_state.timezone, actor=actor + ) return current_in_context_messages, new_in_context_messages diff --git a/letta/agents/letta_agent.py b/letta/agents/letta_agent.py index 9b1bac28..2374e47e 100644 --- a/letta/agents/letta_agent.py +++ b/letta/agents/letta_agent.py @@ -1011,6 +1011,7 @@ class LettaAgent(BaseAgent): function_response = package_function_response( was_success=tool_execution_result.success_flag, response_string=function_response_string, + timezone=agent_state.timezone, ) # 4. Register tool call with tool rule solver @@ -1073,6 +1074,7 @@ class LettaAgent(BaseAgent): tool_call_id=tool_call_id, function_call_success=tool_execution_result.success_flag, function_response=function_response_string, + timezone=agent_state.timezone, actor=self.actor, add_heartbeat_request_system_message=continue_stepping, heartbeat_reason=heartbeat_reason, diff --git a/letta/agents/letta_agent_batch.py b/letta/agents/letta_agent_batch.py index 629714f2..d3343a03 100644 --- a/letta/agents/letta_agent_batch.py +++ b/letta/agents/letta_agent_batch.py @@ -548,6 +548,7 @@ class LettaAgentBatch(BaseAgent): function_call_success=success_flag, function_response=tool_exec_result, tool_execution_result=tool_exec_result_obj, + timezone=agent_state.timezone, actor=self.actor, add_heartbeat_request_system_message=False, reasoning_content=reasoning_content, diff --git a/letta/agents/voice_agent.py b/letta/agents/voice_agent.py index d49c7289..633f9925 100644 --- a/letta/agents/voice_agent.py +++ b/letta/agents/voice_agent.py @@ -153,7 +153,9 @@ class VoiceAgent(BaseAgent): previous_message_count=self.num_messages, archival_memory_size=self.num_archival_memories, ) - letta_message_db_queue = create_input_messages(input_messages=input_messages, agent_id=agent_state.id, actor=self.actor) + letta_message_db_queue = create_input_messages( + input_messages=input_messages, agent_id=agent_state.id, timezone=agent_state.timezone, actor=self.actor + ) in_memory_message_history = self.pre_process_input_message(input_messages) # TODO: Define max steps here @@ -212,6 +214,7 @@ class VoiceAgent(BaseAgent): agent_id=agent_state.id, model=agent_state.llm_config.model, actor=self.actor, + timezone=agent_state.timezone, ) letta_message_db_queue.extend(assistant_msgs) @@ -272,6 +275,7 @@ class VoiceAgent(BaseAgent): function_call_success=success_flag, function_response=tool_result, tool_execution_result=tool_execution_result, + timezone=agent_state.timezone, actor=self.actor, add_heartbeat_request_system_message=True, ) diff --git a/letta/helpers/message_helper.py b/letta/helpers/message_helper.py index 0c6a0584..3baf4897 100644 --- a/letta/helpers/message_helper.py +++ b/letta/helpers/message_helper.py @@ -12,6 +12,7 @@ from letta.schemas.message import Message, MessageCreate def convert_message_creates_to_messages( message_creates: list[MessageCreate], agent_id: str, + timezone: str, wrap_user_message: bool = True, wrap_system_message: bool = True, ) -> list[Message]: @@ -19,6 +20,7 @@ def convert_message_creates_to_messages( _convert_message_create_to_message( message_create=create, agent_id=agent_id, + timezone=timezone, wrap_user_message=wrap_user_message, wrap_system_message=wrap_system_message, ) @@ -29,6 +31,7 @@ def convert_message_creates_to_messages( def _convert_message_create_to_message( message_create: MessageCreate, agent_id: str, + timezone: str, wrap_user_message: bool = True, wrap_system_message: bool = True, ) -> Message: @@ -50,9 +53,9 @@ def _convert_message_create_to_message( if isinstance(content, TextContent): # Apply wrapping if needed if message_create.role == MessageRole.user and wrap_user_message: - content.text = system.package_user_message(user_message=content.text) + content.text = system.package_user_message(user_message=content.text, timezone=timezone) elif message_create.role == MessageRole.system and wrap_system_message: - content.text = system.package_system_message(system_message=content.text) + content.text = system.package_system_message(system_message=content.text, timezone=timezone) elif isinstance(content, ImageContent): if content.source.type == ImageSourceType.url: # Convert URL image to Base64Image if needed diff --git a/letta/orm/agent.py b/letta/orm/agent.py index a1ab5367..4a5dae96 100644 --- a/letta/orm/agent.py +++ b/letta/orm/agent.py @@ -89,6 +89,9 @@ class Agent(SqlalchemyBase, OrganizationMixin, AsyncAttrs): Integer, nullable=True, doc="The duration in milliseconds of the agent's last run." ) + # timezone + timezone: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The timezone of the agent (for the context window).") + # relationships organization: Mapped["Organization"] = relationship("Organization", back_populates="agents") tool_exec_environment_variables: Mapped[List["AgentEnvironmentVariable"]] = relationship( diff --git a/letta/schemas/agent.py b/letta/schemas/agent.py index 8b5d7dd4..e53a88b8 100644 --- a/letta/schemas/agent.py +++ b/letta/schemas/agent.py @@ -109,6 +109,9 @@ class AgentState(OrmMetadataBase, validate_assignment=True): last_run_completion: Optional[datetime] = Field(None, description="The timestamp when the agent last completed a run.") last_run_duration_ms: Optional[int] = Field(None, description="The duration in milliseconds of the agent's last run.") + # timezone + timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).") + def get_agent_env_vars_as_dict(self) -> Dict[str, str]: # Get environment variables for this agent specifically per_agent_env_vars = {} @@ -192,6 +195,7 @@ class CreateAgent(BaseModel, validate_assignment=True): # ) enable_sleeptime: Optional[bool] = Field(None, description="If set to True, memory management will move to a background agent thread.") response_format: Optional[ResponseFormatUnion] = Field(None, description="The response format for the agent.") + timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).") @field_validator("name") @classmethod @@ -285,6 +289,7 @@ class UpdateAgent(BaseModel): response_format: Optional[ResponseFormatUnion] = Field(None, description="The response format for the agent.") last_run_completion: Optional[datetime] = Field(None, description="The timestamp when the agent last completed a run.") last_run_duration_ms: Optional[int] = Field(None, description="The duration in milliseconds of the agent's last run.") + timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).") class Config: extra = "ignore" # Ignores extra fields diff --git a/letta/server/rest_api/utils.py b/letta/server/rest_api/utils.py index 271ae3bc..c6a4d902 100644 --- a/letta/server/rest_api/utils.py +++ b/letta/server/rest_api/utils.py @@ -169,7 +169,7 @@ def log_error_to_sentry(e): sentry_sdk.capture_exception(e) -def create_input_messages(input_messages: List[MessageCreate], agent_id: str, actor: User) -> List[Message]: +def create_input_messages(input_messages: List[MessageCreate], agent_id: str, timezone: str, actor: User) -> List[Message]: """ Converts a user input message into the internal structured format. @@ -177,7 +177,7 @@ def create_input_messages(input_messages: List[MessageCreate], agent_id: str, ac we should unify this when it's clear what message attributes we need. """ - messages = convert_message_creates_to_messages(input_messages, agent_id, wrap_user_message=False, wrap_system_message=False) + messages = convert_message_creates_to_messages(input_messages, agent_id, timezone, wrap_user_message=False, wrap_system_message=False) for message in messages: message.organization_id = actor.organization_id return messages @@ -192,6 +192,7 @@ def create_letta_messages_from_llm_response( tool_call_id: str, function_call_success: bool, function_response: Optional[str], + timezone: str, actor: User, add_heartbeat_request_system_message: bool = False, heartbeat_reason: Optional[str] = None, @@ -233,7 +234,7 @@ def create_letta_messages_from_llm_response( # TODO: This helps preserve ordering tool_message = Message( role=MessageRole.tool, - content=[TextContent(text=package_function_response(function_call_success, function_response))], + content=[TextContent(text=package_function_response(function_call_success, function_response, timezone))], organization_id=actor.organization_id, agent_id=agent_id, model=model, @@ -259,7 +260,7 @@ def create_letta_messages_from_llm_response( model=model, function_call_success=function_call_success, actor=actor, - llm_batch_item_id=llm_batch_item_id, + timezone=timezone, heartbeat_reason=heartbeat_reason, ) messages.append(heartbeat_system_message) @@ -274,6 +275,7 @@ def create_heartbeat_system_message( agent_id: str, model: str, function_call_success: bool, + timezone: str, actor: User, llm_batch_item_id: Optional[str] = None, heartbeat_reason: Optional[str] = None, @@ -285,7 +287,7 @@ def create_heartbeat_system_message( heartbeat_system_message = Message( role=MessageRole.user, - content=[TextContent(text=get_heartbeat(text_content))], + content=[TextContent(text=get_heartbeat(timezone, text_content))], organization_id=actor.organization_id, agent_id=agent_id, model=model, @@ -302,6 +304,7 @@ def create_assistant_messages_from_openai_response( agent_id: str, model: str, actor: User, + timezone: str, ) -> List[Message]: """ Converts an OpenAI response into Messages that follow the internal @@ -318,6 +321,7 @@ def create_assistant_messages_from_openai_response( tool_call_id=tool_call_id, function_call_success=True, function_response=None, + timezone=timezone, actor=actor, add_heartbeat_request_system_message=False, ) diff --git a/letta/server/server.py b/letta/server/server.py index b9d50a8c..ba3665f4 100644 --- a/letta/server/server.py +++ b/letta/server/server.py @@ -631,7 +631,7 @@ class SyncServer(Server): packaged_user_message = system.package_user_message( user_message=message, - time=timestamp.isoformat() if timestamp else None, + timezone=agent.timezone, ) # NOTE: eventually deprecate and only allow passing Message types diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py index eaa6851a..1e98c90f 100644 --- a/letta/services/agent_manager.py +++ b/letta/services/agent_manager.py @@ -317,6 +317,7 @@ class AgentManager: response_format=agent_create.response_format, created_by_id=actor.id, last_updated_by_id=actor.id, + timezone=agent_create.timezone, ) if _test_only_force_id: @@ -480,6 +481,7 @@ class AgentManager: response_format=agent_create.response_format, created_by_id=actor.id, last_updated_by_id=actor.id, + timezone=agent_create.timezone, ) if _test_only_force_id: @@ -539,6 +541,7 @@ class AgentManager: new_agent.message_ids = [msg.id for msg in init_messages] await session.refresh(new_agent) + result = await new_agent.to_pydantic_async() await self.message_manager.create_many_messages_async(pydantic_msgs=init_messages, actor=actor) @@ -561,7 +564,9 @@ class AgentManager: # Don't use anything else in the pregen sequence, instead use the provided sequence init_messages = [system_message_obj] init_messages.extend( - package_initial_message_sequence(agent_state.id, supplied_initial_message_sequence, agent_state.llm_config.model, actor) + package_initial_message_sequence( + agent_state.id, supplied_initial_message_sequence, agent_state.llm_config.model, agent_state.timezone, actor + ) ) else: init_messages = [ diff --git a/letta/services/helpers/agent_manager_helper.py b/letta/services/helpers/agent_manager_helper.py index c48dfd39..0f875ae8 100644 --- a/letta/services/helpers/agent_manager_helper.py +++ b/letta/services/helpers/agent_manager_helper.py @@ -309,15 +309,15 @@ def initialize_message_sequence( previous_message_count=previous_message_count, archival_memory_size=archival_memory_size, ) - first_user_message = get_login_event() # event letting Letta know the user just logged in + first_user_message = get_login_event(agent_state.timezone) # event letting Letta know the user just logged in if include_initial_boot_message: if agent_state.agent_type == AgentType.sleeptime_agent: initial_boot_messages = [] elif agent_state.llm_config.model is not None and "gpt-3.5" in agent_state.llm_config.model: - initial_boot_messages = get_initial_boot_messages("startup_with_send_message_gpt35") + initial_boot_messages = get_initial_boot_messages("startup_with_send_message_gpt35", agent_state.timezone) else: - initial_boot_messages = get_initial_boot_messages("startup_with_send_message") + initial_boot_messages = get_initial_boot_messages("startup_with_send_message", agent_state.timezone) messages = ( [ {"role": "system", "content": full_system_message}, @@ -338,7 +338,7 @@ def initialize_message_sequence( def package_initial_message_sequence( - agent_id: str, initial_message_sequence: List[MessageCreate], model: str, actor: User + agent_id: str, initial_message_sequence: List[MessageCreate], model: str, timezone: str, actor: User ) -> List[Message]: # create the agent object init_messages = [] @@ -347,6 +347,7 @@ def package_initial_message_sequence( if message_create.role == MessageRole.user: packed_message = system.package_user_message( user_message=message_create.content, + timezone=timezone, ) init_messages.append( Message( @@ -361,6 +362,7 @@ def package_initial_message_sequence( elif message_create.role == MessageRole.system: packed_message = system.package_system_message( system_message=message_create.content, + timezone=timezone, ) init_messages.append( Message( @@ -402,7 +404,7 @@ def package_initial_message_sequence( ) # add tool return - function_response = package_function_response(True, "None") + function_response = package_function_response(True, "None", timezone) init_messages.append( Message( role=MessageRole.tool, diff --git a/letta/system.py b/letta/system.py index 06acb8f9..33337569 100644 --- a/letta/system.py +++ b/letta/system.py @@ -13,7 +13,7 @@ from .helpers.datetime_helpers import get_local_time from .helpers.json_helpers import json_dumps -def get_initial_boot_messages(version="startup"): +def get_initial_boot_messages(version, timezone): if version == "startup": initial_boot_message = INITIAL_BOOT_MESSAGE messages = [ @@ -47,7 +47,7 @@ def get_initial_boot_messages(version="startup"): # "role": "function", "role": "tool", "name": "send_message", # NOTE: technically not up to spec, this is old functions style - "content": package_function_response(True, None), + "content": package_function_response(True, None, timezone), "tool_call_id": tool_call_id, }, ] @@ -76,7 +76,7 @@ def get_initial_boot_messages(version="startup"): # "role": "function", "role": "tool", "name": "send_message", - "content": package_function_response(True, None), + "content": package_function_response(True, None, timezone), "tool_call_id": tool_call_id, }, ] @@ -87,9 +87,9 @@ def get_initial_boot_messages(version="startup"): return messages -def get_heartbeat(reason: str = "Automated timer", include_location: bool = False, location_name: str = "San Francisco, CA, USA"): +def get_heartbeat(timezone, reason: str = "Automated timer", include_location: bool = False, location_name: str = "San Francisco, CA, USA"): # Package the message with time and location - formatted_time = get_local_time() + formatted_time = get_local_time(timezone=timezone) packaged_message = { "type": "heartbeat", "reason": reason, @@ -102,9 +102,9 @@ def get_heartbeat(reason: str = "Automated timer", include_location: bool = Fals return json_dumps(packaged_message) -def get_login_event(last_login="Never (first login)", include_location=False, location_name="San Francisco, CA, USA"): +def get_login_event(timezone, last_login="Never (first login)", include_location=False, location_name="San Francisco, CA, USA"): # Package the message with time and location - formatted_time = get_local_time() + formatted_time = get_local_time(timezone=timezone) packaged_message = { "type": "login", "last_login": last_login, @@ -119,13 +119,13 @@ def get_login_event(last_login="Never (first login)", include_location=False, lo def package_user_message( user_message: str, - time: Optional[str] = None, + timezone: str, include_location: bool = False, location_name: Optional[str] = "San Francisco, CA, USA", name: Optional[str] = None, ): # Package the message with time and location - formatted_time = time if time else get_local_time() + formatted_time = get_local_time(timezone=timezone) packaged_message = { "type": "user_message", "message": user_message, @@ -141,8 +141,8 @@ def package_user_message( return json_dumps(packaged_message) -def package_function_response(was_success, response_string, timestamp=None): - formatted_time = get_local_time() if timestamp is None else timestamp +def package_function_response(was_success, response_string, timezone): + formatted_time = get_local_time(timezone=timezone) packaged_message = { "status": "OK" if was_success else "Failed", "message": response_string, @@ -152,7 +152,7 @@ def package_function_response(was_success, response_string, timestamp=None): return json_dumps(packaged_message) -def package_system_message(system_message, message_type="system_alert", time=None): +def package_system_message(system_message, timezone, message_type="system_alert"): # error handling for recursive packaging try: message_json = json.loads(system_message) @@ -162,7 +162,7 @@ def package_system_message(system_message, message_type="system_alert", time=Non except: pass # do nothing, expected behavior that the message is not JSON - formatted_time = time if time else get_local_time() + formatted_time = get_local_time(timezone=timezone) packaged_message = { "type": message_type, "message": system_message, @@ -172,13 +172,13 @@ def package_system_message(system_message, message_type="system_alert", time=Non return json.dumps(packaged_message) -def package_summarize_message(summary, summary_message_count, hidden_message_count, total_message_count, timestamp=None): +def package_summarize_message(summary, summary_message_count, hidden_message_count, total_message_count, timezone): context_message = ( f"Note: prior messages ({hidden_message_count} of {total_message_count} total messages) have been hidden from view due to conversation memory constraints.\n" + f"The following is a summary of the previous {summary_message_count} messages:\n {summary}" ) - formatted_time = get_local_time() if timestamp is None else timestamp + formatted_time = get_local_time(timezone=timezone) packaged_message = { "type": "system_alert", "message": context_message, @@ -188,11 +188,11 @@ def package_summarize_message(summary, summary_message_count, hidden_message_cou return json_dumps(packaged_message) -def package_summarize_message_no_summary(hidden_message_count, timestamp=None, message=None): +def package_summarize_message_no_summary(hidden_message_count, message=None, timezone=None): """Add useful metadata to the summary message""" # Package the message with time and location - formatted_time = get_local_time() if timestamp is None else timestamp + formatted_time = get_local_time(timezone=timezone) context_message = ( message if message diff --git a/tests/test_timezone_formatting.py b/tests/test_timezone_formatting.py new file mode 100644 index 00000000..ca259e89 --- /dev/null +++ b/tests/test_timezone_formatting.py @@ -0,0 +1,221 @@ +import json +from datetime import datetime + +import pytest +import pytz + +from letta.helpers.datetime_helpers import get_local_time, get_local_time_timezone +from letta.system import ( + get_heartbeat, + get_login_event, + package_function_response, + package_summarize_message, + package_system_message, + package_user_message, +) + + +class TestTimezoneFormatting: + """Test suite for timezone formatting functions in system.py""" + + def _extract_time_from_json(self, json_str: str) -> str: + """Helper to extract time field from JSON string""" + data = json.loads(json_str) + return data["time"] + + def _validate_timezone_accuracy(self, formatted_time: str, expected_timezone: str, tolerance_minutes: int = 2): + """ + Validate that the formatted time is accurate for the given timezone within tolerance. + + Args: + formatted_time: The time string from the system functions + expected_timezone: The timezone string (e.g., "America/New_York") + tolerance_minutes: Acceptable difference in minutes + """ + # Parse the formatted time - handle the actual format produced + # Expected format: "2025-06-24 12:53:40 AM EDT-0400" + import re + from datetime import timedelta, timezone + + # Match pattern like "2025-06-24 12:53:40 AM EDT-0400" + pattern = r"(\d{4}-\d{2}-\d{2} \d{1,2}:\d{2}:\d{2} [AP]M) ([A-Z]{3,4})([-+]\d{4})" + match = re.match(pattern, formatted_time) + + if not match: + # Fallback: just check basic format without detailed parsing + assert len(formatted_time) > 20, f"Time string too short: {formatted_time}" + assert " AM " in formatted_time or " PM " in formatted_time, f"No AM/PM in time: {formatted_time}" + return + + time_part, tz_name, tz_offset = match.groups() + + # Parse the time part without timezone + time_without_tz = datetime.strptime(time_part, "%Y-%m-%d %I:%M:%S %p") + + # Create timezone offset + hours_offset = int(tz_offset[:3]) + minutes_offset = int(tz_offset[3:5]) if len(tz_offset) > 3 else 0 + if tz_offset[0] == "-" and hours_offset >= 0: + hours_offset = -hours_offset + total_offset = timedelta(hours=hours_offset, minutes=minutes_offset) + tz_info = timezone(total_offset) + + parsed_time = time_without_tz.replace(tzinfo=tz_info) + + # Get current time in the expected timezone + tz = pytz.timezone(expected_timezone) + current_time_in_tz = datetime.now(tz) + + # Check that times are within tolerance + time_diff = abs((parsed_time - current_time_in_tz).total_seconds()) + assert time_diff <= tolerance_minutes * 60, ( + f"Time difference too large: {time_diff}s. " f"Parsed: {parsed_time}, Expected timezone: {current_time_in_tz}" + ) + + # Verify timezone info exists and format looks reasonable + assert parsed_time.tzinfo is not None, "Parsed time should have timezone info" + assert tz_name in formatted_time, f"Timezone abbreviation {tz_name} should be in formatted time" + + def test_get_heartbeat_timezone_accuracy(self): + """Test that get_heartbeat produces accurate timestamps for different timezones""" + test_timezones = ["UTC", "America/New_York", "America/Los_Angeles", "Europe/London", "Asia/Tokyo"] + + for tz in test_timezones: + heartbeat = get_heartbeat(timezone=tz, reason="Test heartbeat") + time_str = self._extract_time_from_json(heartbeat) + self._validate_timezone_accuracy(time_str, tz) + + def test_get_login_event_timezone_accuracy(self): + """Test that get_login_event produces accurate timestamps for different timezones""" + test_timezones = ["UTC", "US/Eastern", "US/Pacific", "Australia/Sydney"] + + for tz in test_timezones: + login = get_login_event(timezone=tz, last_login="2024-01-01") + time_str = self._extract_time_from_json(login) + self._validate_timezone_accuracy(time_str, tz) + + def test_package_user_message_timezone_accuracy(self): + """Test that package_user_message produces accurate timestamps for different timezones""" + test_timezones = ["UTC", "America/Chicago", "Europe/Paris", "Asia/Shanghai"] + + for tz in test_timezones: + message = package_user_message("Test message", timezone=tz) + time_str = self._extract_time_from_json(message) + self._validate_timezone_accuracy(time_str, tz) + + def test_package_function_response_timezone_accuracy(self): + """Test that package_function_response produces accurate timestamps for different timezones""" + test_timezones = ["UTC", "America/Denver", "Europe/Berlin", "Pacific/Auckland"] + + for tz in test_timezones: + response = package_function_response(True, "Success", timezone=tz) + time_str = self._extract_time_from_json(response) + self._validate_timezone_accuracy(time_str, tz) + + def test_package_system_message_timezone_accuracy(self): + """Test that package_system_message produces accurate timestamps for different timezones""" + test_timezones = ["UTC", "America/Phoenix", "Europe/Rome", "Asia/Kolkata"] # Mumbai is now called Kolkata in pytz + + for tz in test_timezones: + message = package_system_message("System alert", timezone=tz) + time_str = self._extract_time_from_json(message) + self._validate_timezone_accuracy(time_str, tz) + + def test_package_summarize_message_timezone_accuracy(self): + """Test that package_summarize_message produces accurate timestamps for different timezones""" + test_timezones = ["UTC", "America/Anchorage", "Europe/Stockholm", "Asia/Seoul"] + + for tz in test_timezones: + summary = package_summarize_message( + summary="Test summary", summary_message_count=2, hidden_message_count=5, total_message_count=7, timezone=tz + ) + time_str = self._extract_time_from_json(summary) + self._validate_timezone_accuracy(time_str, tz) + + def test_get_local_time_timezone_direct(self): + """Test get_local_time_timezone directly for accuracy""" + test_timezones = ["UTC", "America/New_York", "Europe/London", "Asia/Tokyo", "Australia/Melbourne"] + + for tz in test_timezones: + time_str = get_local_time_timezone(timezone=tz) + self._validate_timezone_accuracy(time_str, tz) + + def test_get_local_time_with_timezone_param(self): + """Test get_local_time when timezone parameter is provided""" + test_timezones = ["UTC", "America/Los_Angeles", "Europe/Madrid", "Asia/Bangkok"] + + for tz in test_timezones: + time_str = get_local_time(timezone=tz) + self._validate_timezone_accuracy(time_str, tz) + + def test_timezone_offset_differences(self): + """Test that different timezones produce appropriately offset times""" + # Get times for different timezones at the same moment + utc_heartbeat = get_heartbeat(timezone="UTC") + utc_time_str = self._extract_time_from_json(utc_heartbeat) + + ny_heartbeat = get_heartbeat(timezone="America/New_York") + ny_time_str = self._extract_time_from_json(ny_heartbeat) + + tokyo_heartbeat = get_heartbeat(timezone="Asia/Tokyo") + tokyo_time_str = self._extract_time_from_json(tokyo_heartbeat) + + # Just validate that all times have the expected format + # UTC should have UTC in the string + assert "UTC" in utc_time_str, f"UTC timezone not found in: {utc_time_str}" + + # NY should have EST or EDT + assert any(tz in ny_time_str for tz in ["EST", "EDT"]), f"EST/EDT not found in: {ny_time_str}" + + # Tokyo should have JST + assert "JST" in tokyo_time_str, f"JST not found in: {tokyo_time_str}" + + def test_daylight_saving_time_handling(self): + """Test that DST transitions are handled correctly""" + # Test timezone that observes DST + eastern_tz = "America/New_York" + + # Get current time in Eastern timezone + message = package_user_message("DST test", timezone=eastern_tz) + time_str = self._extract_time_from_json(message) + + # Validate against current Eastern time + self._validate_timezone_accuracy(time_str, eastern_tz) + + # The timezone abbreviation should be either EST or EDT + assert any(tz in time_str for tz in ["EST", "EDT"]), f"EST/EDT not found in: {time_str}" + + @pytest.mark.parametrize( + "timezone_str,expected_format_parts", + [ + ("UTC", ["UTC", "+0000"]), + ("America/New_York", ["EST", "EDT"]), # Either EST or EDT depending on date + ("Europe/London", ["GMT", "BST"]), # Either GMT or BST depending on date + ("Asia/Tokyo", ["JST", "+0900"]), + ("Australia/Sydney", ["AEDT", "AEST"]), # Either AEDT or AEST depending on date + ], + ) + def test_timezone_format_components(self, timezone_str, expected_format_parts): + """Test that timezone formatting includes expected components""" + heartbeat = get_heartbeat(timezone=timezone_str) + time_str = self._extract_time_from_json(heartbeat) + + # Check that at least one expected format part is present + found_expected_part = any(part in time_str for part in expected_format_parts) + assert found_expected_part, f"None of expected format parts {expected_format_parts} found in time string: {time_str}" + + # Validate the time is accurate + self._validate_timezone_accuracy(time_str, timezone_str) + + def test_timezone_parameter_working(self): + """Test that timezone parameter correctly affects the output""" + # Test that different timezones produce different time formats + utc_message = package_user_message("Test", timezone="UTC") + utc_time = self._extract_time_from_json(utc_message) + + ny_message = package_user_message("Test", timezone="America/New_York") + ny_time = self._extract_time_from_json(ny_message) + + # Times should have different timezone indicators + assert "UTC" in utc_time, f"UTC not found in: {utc_time}" + assert any(tz in ny_time for tz in ["EST", "EDT"]), f"EST/EDT not found in: {ny_time}"