From dcb689cbfa5da7d5d3f297065956083adee2730c Mon Sep 17 00:00:00 2001 From: Matthew Zhou Date: Thu, 21 Aug 2025 11:08:06 -0700 Subject: [PATCH] feat: Add basic test serialization tests (#4076) --- letta/schemas/agent_file.py | 17 +- letta/schemas/message.py | 2 + letta/services/agent_serialization_manager.py | 8 +- ...sic_agent_with_blocks_tools_messages_v2.af | 623 ++++++++++++++++++ tests/test_sdk_client.py | 328 ++++++++- 5 files changed, 975 insertions(+), 3 deletions(-) create mode 100644 tests/test_agent_files/test_basic_agent_with_blocks_tools_messages_v2.af diff --git a/letta/schemas/agent_file.py b/letta/schemas/agent_file.py index 1d87fa20..02d3988a 100644 --- a/letta/schemas/agent_file.py +++ b/letta/schemas/agent_file.py @@ -1,15 +1,17 @@ from datetime import datetime from typing import Any, Dict, List, Optional +from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall as OpenAIToolCall from pydantic import BaseModel, Field +from letta.helpers.datetime_helpers import get_utc_time from letta.schemas.agent import AgentState, CreateAgent from letta.schemas.block import Block, CreateBlock from letta.schemas.enums import MessageRole from letta.schemas.file import FileAgent, FileAgentBase, FileMetadata, FileMetadataBase from letta.schemas.group import Group, GroupCreate from letta.schemas.mcp import MCPServer -from letta.schemas.message import Message, MessageCreate +from letta.schemas.message import Message, MessageCreate, ToolReturn from letta.schemas.source import Source, SourceCreate from letta.schemas.tool import Tool from letta.schemas.user import User @@ -46,6 +48,15 @@ class MessageSchema(MessageCreate): role: MessageRole = Field(..., description="The role of the participant.") model: Optional[str] = Field(None, description="The model used to make the function call") agent_id: Optional[str] = Field(None, description="The unique identifier of the agent") + tool_calls: Optional[List[OpenAIToolCall]] = Field( + default=None, description="The list of tool calls requested. Only applicable for role assistant." + ) + tool_call_id: Optional[str] = Field(default=None, description="The ID of the tool call. Only applicable for role tool.") + tool_returns: Optional[List[ToolReturn]] = Field(default=None, description="Tool execution return information for prior tool calls") + created_at: datetime = Field(default_factory=get_utc_time, description="The timestamp when the object was created.") + + # TODO: Should we also duplicate the steps here? + # TODO: What about tool_return? @classmethod def from_message(cls, message: Message) -> "MessageSchema": @@ -64,6 +75,10 @@ class MessageSchema(MessageCreate): group_id=message.group_id, model=message.model, agent_id=message.agent_id, + tool_calls=message.tool_calls, + tool_call_id=message.tool_call_id, + tool_returns=message.tool_returns, + created_at=message.created_at, ) diff --git a/letta/schemas/message.py b/letta/schemas/message.py index f49146dd..4fd0a8ea 100644 --- a/letta/schemas/message.py +++ b/letta/schemas/message.py @@ -414,6 +414,8 @@ class Message(BaseMessage): except json.JSONDecodeError: raise ValueError(f"Failed to decode function return: {text_content}") + # if self.tool_call_id is None: + # import pdb;pdb.set_trace() assert self.tool_call_id is not None return ToolReturnMessage( diff --git a/letta/services/agent_serialization_manager.py b/letta/services/agent_serialization_manager.py index edccd7b6..56f324d7 100644 --- a/letta/services/agent_serialization_manager.py +++ b/letta/services/agent_serialization_manager.py @@ -2,7 +2,13 @@ from datetime import datetime, timezone from typing import Any, Dict, List, Optional from letta.constants import MCP_TOOL_TAG_NAME_PREFIX -from letta.errors import AgentExportIdMappingError, AgentExportProcessingError, AgentFileImportError, AgentNotFoundForExportError +from letta.errors import ( + AgentExportIdMappingError, + AgentExportProcessingError, + AgentFileExportError, + AgentFileImportError, + AgentNotFoundForExportError, +) from letta.helpers.pinecone_utils import should_use_pinecone from letta.log import get_logger from letta.schemas.agent import AgentState, CreateAgent diff --git a/tests/test_agent_files/test_basic_agent_with_blocks_tools_messages_v2.af b/tests/test_agent_files/test_basic_agent_with_blocks_tools_messages_v2.af new file mode 100644 index 00000000..2971bcf0 --- /dev/null +++ b/tests/test_agent_files/test_basic_agent_with_blocks_tools_messages_v2.af @@ -0,0 +1,623 @@ +{ + "agents": [ + { + "name": "test_export_import_5297564a-d35a-4b42-9572-1ae313041430", + "memory_blocks": [], + "tools": [], + "tool_ids": [ + "tool-0", + "tool-1", + "tool-2", + "tool-3", + "tool-4", + "tool-5", + "tool-6" + ], + "source_ids": [], + "block_ids": [ + "block-0", + "block-1", + "block-2" + ], + "tool_rules": [ + { + "tool_name": "memory_insert", + "type": "continue_loop", + "prompt_template": "\n{{ tool_name }} requires continuing your response when called\n" + }, + { + "tool_name": "memory_replace", + "type": "continue_loop", + "prompt_template": "\n{{ tool_name }} requires continuing your response when called\n" + }, + { + "tool_name": "send_message", + "type": "exit_loop", + "prompt_template": "\n{{ tool_name }} ends your response (yields control) when called\n" + }, + { + "tool_name": "conversation_search", + "type": "continue_loop", + "prompt_template": "\n{{ tool_name }} requires continuing your response when called\n" + } + ], + "tags": [ + "test", + "export", + "import" + ], + "system": "You are a helpful assistant specializing in data analysis and mathematical computations.", + "agent_type": "memgpt_v2_agent", + "llm_config": { + "model": "gpt-4.1-mini", + "model_endpoint_type": "openai", + "model_endpoint": "https://api.openai.com/v1", + "provider_name": "openai", + "provider_category": "base", + "model_wrapper": null, + "context_window": 32000, + "put_inner_thoughts_in_kwargs": true, + "handle": "openai/gpt-4.1-mini", + "temperature": 0.7, + "max_tokens": 4096, + "enable_reasoner": true, + "reasoning_effort": null, + "max_reasoning_tokens": 0, + "frequency_penalty": 1.0, + "compatibility_type": null, + "verbosity": "medium" + }, + "embedding_config": { + "embedding_endpoint_type": "openai", + "embedding_endpoint": "https://api.openai.com/v1", + "embedding_model": "text-embedding-3-small", + "embedding_dim": 2000, + "embedding_chunk_size": 300, + "handle": "openai/text-embedding-3-small", + "batch_size": 1024, + "azure_endpoint": null, + "azure_version": null, + "azure_deployment": null + }, + "initial_message_sequence": null, + "include_base_tools": false, + "include_multi_agent_tools": false, + "include_base_tool_rules": false, + "include_default_source": false, + "description": null, + "metadata": null, + "model": null, + "embedding": null, + "context_window_limit": null, + "embedding_chunk_size": null, + "max_tokens": null, + "max_reasoning_tokens": null, + "enable_reasoner": false, + "reasoning": null, + "from_template": null, + "template": false, + "project": null, + "tool_exec_environment_variables": {}, + "memory_variables": null, + "project_id": null, + "template_id": null, + "base_template_id": null, + "identity_ids": null, + "message_buffer_autoclear": false, + "enable_sleeptime": false, + "response_format": null, + "timezone": "UTC", + "max_files_open": 5, + "per_file_view_window_char_limit": 15000, + "hidden": null, + "id": "agent-0", + "in_context_message_ids": [ + "message-0", + "message-1", + "message-2", + "message-3", + "message-4", + "message-5", + "message-6" + ], + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a helpful assistant specializing in data analysis and mathematical computations.\n\n\nThe following memory blocks are currently engaged in your core memory unit:\n\n\n\nThe persona block: Stores details about your current persona, guiding how you behave and respond. This helps you to maintain consistency and personality in your interactions.\n\n\n- chars_current=195\n- chars_limit=8000\n\n\n# NOTE: Line numbers shown below are to help during editing. Do NOT include line number prefixes in your memory edit tool calls.\nLine 1: You are Alex, a data analyst and mathematician who helps users with calculations and insights. You have extensive experience in statistical analysis and prefer to provide clear, accurate results.\n\n\n\n\n\nThe human block: Stores key details about the person you are conversing with, allowing for more personalized and friend-like conversation.\n\n\n- chars_current=175\n- chars_limit=4000\n\n\n# NOTE: Line numbers shown below are to help during editing. Do NOT include line number prefixes in your memory edit tool calls.\nLine 1: username: sarah_researcher\nLine 2: occupation: data scientist\nLine 3: interests: machine learning, statistics, fibonacci sequences\nLine 4: preferred_communication: detailed explanations with examples\n\n\n\n\n\nNone\n\n\n- chars_current=210\n- chars_limit=6000\n\n\n# NOTE: Line numbers shown below are to help during editing. Do NOT include line number prefixes in your memory edit tool calls.\nLine 1: Current project: Building predictive models for financial markets. Sarah is working on sequence analysis and pattern recognition. Recently interested in mathematical sequences like Fibonacci for trend analysis.\n\n\n\n\n\n\nThe following constraints define rules for tool usage and guide desired behavior. These rules must be followed to ensure proper tool execution and workflow. A single response may contain multiple tool calls.\n\n\nmemory_insert requires continuing your response when called\n\n\nmemory_replace requires continuing your response when called\n\n\nconversation_search requires continuing your response when called\n\n\nsend_message ends your response (yields control) when called\n\n\n\n\n\n\n- The current time is: 2025-08-21 05:49:47 PM UTC+0000\n- Memory blocks were last modified: 2025-08-21 05:49:47 PM UTC+0000\n- -1 previous messages between you and the user are stored in recall memory (use tools to access them)\n- 2 total memories you created are stored in archival memory (use tools to access them)\n" + } + ], + "name": null, + "otid": null, + "sender_id": null, + "batch_item_id": null, + "group_id": null, + "id": "message-0", + "model": "gpt-4.1-mini", + "agent_id": "agent-0", + "tool_calls": null, + "tool_call_id": null, + "tool_returns": [], + "created_at": "2025-08-21T17:49:45.737357+00:00" + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Bootup sequence complete. Persona activated. Testing messaging functionality." + } + ], + "name": null, + "otid": null, + "sender_id": null, + "batch_item_id": null, + "group_id": null, + "id": "message-1", + "model": "gpt-4.1-mini", + "agent_id": "agent-0", + "tool_calls": [ + { + "id": "3691c2df-84e5-479c-a871-3d5da7f0a7ee", + "function": { + "arguments": "{\n \"message\": \"More human than human is our motto.\"\n}", + "name": "send_message" + }, + "type": "function" + } + ], + "tool_call_id": null, + "tool_returns": [], + "created_at": "2025-08-21T17:49:45.739574+00:00" + }, + { + "role": "tool", + "content": [ + { + "type": "text", + "text": "{\n \"status\": \"OK\",\n \"message\": null,\n \"time\": \"2025-08-21 05:49:45 PM UTC+0000\"\n}" + } + ], + "name": "send_message", + "otid": null, + "sender_id": null, + "batch_item_id": null, + "group_id": null, + "id": "message-2", + "model": "gpt-4.1-mini", + "agent_id": "agent-0", + "tool_calls": null, + "tool_call_id": "3691c2df-84e5-479c-a871-3d5da7f0a7ee", + "tool_returns": [], + "created_at": "2025-08-21T17:49:45.739814+00:00" + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "{\n \"type\": \"login\",\n \"last_login\": \"Never (first login)\",\n \"time\": \"2025-08-21 05:49:45 PM UTC+0000\"\n}" + } + ], + "name": null, + "otid": null, + "sender_id": null, + "batch_item_id": null, + "group_id": null, + "id": "message-3", + "model": "gpt-4.1-mini", + "agent_id": "agent-0", + "tool_calls": null, + "tool_call_id": null, + "tool_returns": [], + "created_at": "2025-08-21T17:49:45.739827+00:00" + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Test message for export" + } + ], + "name": null, + "otid": null, + "sender_id": null, + "batch_item_id": null, + "group_id": null, + "id": "message-4", + "model": null, + "agent_id": "agent-0", + "tool_calls": null, + "tool_call_id": null, + "tool_returns": [], + "created_at": "2025-08-21T17:49:46.969096+00:00" + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Responding to test message for export." + } + ], + "name": null, + "otid": null, + "sender_id": null, + "batch_item_id": null, + "group_id": null, + "id": "message-5", + "model": "gpt-4.1-mini", + "agent_id": "agent-0", + "tool_calls": [ + { + "id": "call_CR0Pvrt3B7fUk4vz3b1MDQdR", + "function": { + "arguments": "{\"message\": \"Test message received successfully. Let me know how I can assist you further!\", \"request_heartbeat\": false}", + "name": "send_message" + }, + "type": "function" + } + ], + "tool_call_id": "call_CR0Pvrt3B7fUk4vz3b1MDQdR", + "tool_returns": [], + "created_at": "2025-08-21T17:49:47.926707+00:00" + }, + { + "role": "tool", + "content": [ + { + "type": "text", + "text": "{\n \"status\": \"OK\",\n \"message\": \"Sent message successfully.\",\n \"time\": \"2025-08-21 05:49:47 PM UTC+0000\"\n}" + } + ], + "name": "send_message", + "otid": null, + "sender_id": null, + "batch_item_id": null, + "group_id": null, + "id": "message-6", + "model": "gpt-4.1-mini", + "agent_id": "agent-0", + "tool_calls": null, + "tool_call_id": "call_CR0Pvrt3B7fUk4vz3b1MDQdR", + "tool_returns": [ + { + "status": "success", + "stdout": null, + "stderr": null + } + ], + "created_at": "2025-08-21T17:49:47.926951+00:00" + } + ], + "files_agents": [], + "group_ids": [] + } + ], + "groups": [], + "blocks": [ + { + "value": "username: sarah_researcher\noccupation: data scientist\ninterests: machine learning, statistics, fibonacci sequences\npreferred_communication: detailed explanations with examples", + "limit": 4000, + "project_id": null, + "template_name": null, + "is_template": false, + "preserve_on_migration": false, + "label": "human", + "read_only": false, + "description": "The human block: Stores key details about the person you are conversing with, allowing for more personalized and friend-like conversation.", + "metadata": {}, + "id": "block-1" + }, + { + "value": "You are Alex, a data analyst and mathematician who helps users with calculations and insights. You have extensive experience in statistical analysis and prefer to provide clear, accurate results.", + "limit": 8000, + "project_id": null, + "template_name": null, + "is_template": false, + "preserve_on_migration": false, + "label": "persona", + "read_only": false, + "description": "The persona block: Stores details about your current persona, guiding how you behave and respond. This helps you to maintain consistency and personality in your interactions.", + "metadata": {}, + "id": "block-0" + }, + { + "value": "Current project: Building predictive models for financial markets. Sarah is working on sequence analysis and pattern recognition. Recently interested in mathematical sequences like Fibonacci for trend analysis.", + "limit": 6000, + "project_id": null, + "template_name": null, + "is_template": false, + "preserve_on_migration": false, + "label": "project_context", + "read_only": false, + "description": null, + "metadata": {}, + "id": "block-2" + } + ], + "files": [], + "sources": [], + "tools": [ + { + "id": "tool-2", + "tool_type": "custom", + "description": "Analyze data and provide insights.", + "source_type": "python", + "name": "analyze_data", + "tags": [ + "analysis", + "data" + ], + "source_code": "def analyze_data(data_type: str, values: List[float]) -> str:\n \"\"\"Analyze data and provide insights.\n\n Args:\n data_type: Type of data to analyze.\n values: Numerical values to analyze.\n\n Returns:\n Analysis results including average, max, and min values.\n \"\"\"\n if not values:\n return \"No data provided\"\n avg = sum(values) / len(values)\n max_val = max(values)\n min_val = min(values)\n return f\"Analysis of {data_type}: avg={avg:.2f}, max={max_val}, min={min_val}\"\n", + "json_schema": { + "name": "analyze_data", + "description": "Analyze data and provide insights.", + "parameters": { + "type": "object", + "properties": { + "data_type": { + "type": "string", + "description": "Type of data to analyze." + }, + "values": { + "type": "array", + "items": { + "type": "number" + }, + "description": "Numerical values to analyze." + } + }, + "required": [ + "data_type", + "values" + ] + } + }, + "args_json_schema": null, + "return_char_limit": 50000, + "pip_requirements": null, + "npm_requirements": null, + "created_by_id": "user-00000000-0000-4000-8000-000000000000", + "last_updated_by_id": "user-00000000-0000-4000-8000-000000000000", + "metadata_": {} + }, + { + "id": "tool-5", + "tool_type": "custom", + "description": "Calculate the nth Fibonacci number.", + "source_type": "python", + "name": "calculate_fibonacci", + "tags": [ + "math", + "utility" + ], + "source_code": "def calculate_fibonacci(n: int) -> int:\n \"\"\"Calculate the nth Fibonacci number.\n\n Args:\n n: The position in the Fibonacci sequence to calculate.\n\n Returns:\n The nth Fibonacci number.\n \"\"\"\n if n <= 0:\n return 0\n elif n == 1:\n return 1\n else:\n a, b = 0, 1\n for _ in range(2, n + 1):\n a, b = b, a + b\n return b\n", + "json_schema": { + "name": "calculate_fibonacci", + "description": "Calculate the nth Fibonacci number.", + "parameters": { + "type": "object", + "properties": { + "n": { + "type": "integer", + "description": "The position in the Fibonacci sequence to calculate." + } + }, + "required": [ + "n" + ] + } + }, + "args_json_schema": null, + "return_char_limit": 50000, + "pip_requirements": null, + "npm_requirements": null, + "created_by_id": "user-00000000-0000-4000-8000-000000000000", + "last_updated_by_id": "user-00000000-0000-4000-8000-000000000000", + "metadata_": {} + }, + { + "id": "tool-4", + "tool_type": "letta_core", + "description": "Search prior conversation history using case-insensitive string matching.", + "source_type": "python", + "name": "conversation_search", + "tags": [ + "letta_core" + ], + "source_code": null, + "json_schema": { + "name": "conversation_search", + "description": "Search prior conversation history using case-insensitive string matching.", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "String to search for." + }, + "page": { + "type": "integer", + "description": "Allows you to page through results. Only use on a follow-up query. Defaults to 0 (first page)." + } + }, + "required": [ + "query" + ] + } + }, + "args_json_schema": null, + "return_char_limit": 1000000, + "pip_requirements": null, + "npm_requirements": null, + "created_by_id": "user-00000000-0000-4000-8000-000000000000", + "last_updated_by_id": "user-00000000-0000-4000-8000-000000000000", + "metadata_": {} + }, + { + "id": "tool-0", + "tool_type": "custom", + "description": "Get user preferences for a specific category.", + "source_type": "python", + "name": "get_user_preferences", + "tags": [ + "user", + "preferences" + ], + "source_code": "def get_user_preferences(category: str) -> str:\n \"\"\"Get user preferences for a specific category.\n\n Args:\n category: The preference category to retrieve (notification, theme, language).\n\n Returns:\n The user's preference for the specified category, or \"not specified\" if unknown.\n \"\"\"\n preferences = {\"notification\": \"email only\", \"theme\": \"dark mode\", \"language\": \"english\"}\n return preferences.get(category, \"not specified\")\n", + "json_schema": { + "name": "get_user_preferences", + "description": "Get user preferences for a specific category.", + "parameters": { + "type": "object", + "properties": { + "category": { + "type": "string", + "description": "The preference category to retrieve (notification, theme, language)." + } + }, + "required": [ + "category" + ] + } + }, + "args_json_schema": null, + "return_char_limit": 50000, + "pip_requirements": null, + "npm_requirements": null, + "created_by_id": "user-00000000-0000-4000-8000-000000000000", + "last_updated_by_id": "user-00000000-0000-4000-8000-000000000000", + "metadata_": {} + }, + { + "id": "tool-6", + "tool_type": "letta_sleeptime_core", + "description": "The memory_insert command allows you to insert text at a specific location in a memory block.", + "source_type": "python", + "name": "memory_insert", + "tags": [ + "letta_sleeptime_core" + ], + "source_code": null, + "json_schema": { + "name": "memory_insert", + "description": "The memory_insert command allows you to insert text at a specific location in a memory block.", + "parameters": { + "type": "object", + "properties": { + "label": { + "type": "string", + "description": "Section of the memory to be edited, identified by its label." + }, + "new_str": { + "type": "string", + "description": "The text to insert. Do not include line number prefixes." + }, + "insert_line": { + "type": "integer", + "description": "The line number after which to insert the text (0 for beginning of file). Defaults to -1 (end of the file)." + } + }, + "required": [ + "label", + "new_str" + ] + } + }, + "args_json_schema": null, + "return_char_limit": 1000000, + "pip_requirements": null, + "npm_requirements": null, + "created_by_id": "user-00000000-0000-4000-8000-000000000000", + "last_updated_by_id": "user-00000000-0000-4000-8000-000000000000", + "metadata_": {} + }, + { + "id": "tool-3", + "tool_type": "letta_sleeptime_core", + "description": "The memory_replace command allows you to replace a specific string in a memory block with a new string. This is used for making precise edits.", + "source_type": "python", + "name": "memory_replace", + "tags": [ + "letta_sleeptime_core" + ], + "source_code": null, + "json_schema": { + "name": "memory_replace", + "description": "The memory_replace command allows you to replace a specific string in a memory block with a new string. This is used for making precise edits.", + "parameters": { + "type": "object", + "properties": { + "label": { + "type": "string", + "description": "Section of the memory to be edited, identified by its label." + }, + "old_str": { + "type": "string", + "description": "The text to replace (must match exactly, including whitespace and indentation)." + }, + "new_str": { + "type": "string", + "description": "The new text to insert in place of the old text. Do not include line number prefixes." + } + }, + "required": [ + "label", + "old_str", + "new_str" + ] + } + }, + "args_json_schema": null, + "return_char_limit": 1000000, + "pip_requirements": null, + "npm_requirements": null, + "created_by_id": "user-00000000-0000-4000-8000-000000000000", + "last_updated_by_id": "user-00000000-0000-4000-8000-000000000000", + "metadata_": {} + }, + { + "id": "tool-1", + "tool_type": "letta_core", + "description": "Sends a message to the human user.", + "source_type": "python", + "name": "send_message", + "tags": [ + "letta_core" + ], + "source_code": null, + "json_schema": { + "name": "send_message", + "description": "Sends a message to the human user.", + "parameters": { + "type": "object", + "properties": { + "message": { + "type": "string", + "description": "Message contents. All unicode (including emojis) are supported." + } + }, + "required": [ + "message" + ] + } + }, + "args_json_schema": null, + "return_char_limit": 1000000, + "pip_requirements": null, + "npm_requirements": null, + "created_by_id": "user-00000000-0000-4000-8000-000000000000", + "last_updated_by_id": "user-00000000-0000-4000-8000-000000000000", + "metadata_": {} + } + ], + "mcp_servers": [], + "metadata": { + "revision_id": "ffb17eb241fc" + }, + "created_at": "2025-08-21T17:49:48.078363+00:00" +} diff --git a/tests/test_sdk_client.py b/tests/test_sdk_client.py index 250d3c55..b06f6934 100644 --- a/tests/test_sdk_client.py +++ b/tests/test_sdk_client.py @@ -1,5 +1,7 @@ +import io import json import os +import textwrap import threading import time import uuid @@ -60,6 +62,117 @@ def agent(client: LettaSDKClient): client.agents.delete(agent_id=agent_state.id) +@pytest.fixture(scope="function") +def fibonacci_tool(client: LettaSDKClient): + """Fixture providing Fibonacci calculation tool.""" + + def calculate_fibonacci(n: int) -> int: + """Calculate the nth Fibonacci number. + + Args: + n: The position in the Fibonacci sequence to calculate. + + Returns: + The nth Fibonacci number. + """ + if n <= 0: + return 0 + elif n == 1: + return 1 + else: + a, b = 0, 1 + for _ in range(2, n + 1): + a, b = b, a + b + return b + + tool = client.tools.upsert_from_function(func=calculate_fibonacci, tags=["math", "utility"]) + yield tool + client.tools.delete(tool.id) + + +@pytest.fixture(scope="function") +def preferences_tool(client: LettaSDKClient): + """Fixture providing user preferences tool.""" + + def get_user_preferences(category: str) -> str: + """Get user preferences for a specific category. + + Args: + category: The preference category to retrieve (notification, theme, language). + + Returns: + The user's preference for the specified category, or "not specified" if unknown. + """ + preferences = {"notification": "email only", "theme": "dark mode", "language": "english"} + return preferences.get(category, "not specified") + + tool = client.tools.upsert_from_function(func=get_user_preferences, tags=["user", "preferences"]) + yield tool + client.tools.delete(tool.id) + + +@pytest.fixture(scope="function") +def data_analysis_tool(client: LettaSDKClient): + """Fixture providing data analysis tool.""" + + def analyze_data(data_type: str, values: List[float]) -> str: + """Analyze data and provide insights. + + Args: + data_type: Type of data to analyze. + values: Numerical values to analyze. + + Returns: + Analysis results including average, max, and min values. + """ + if not values: + return "No data provided" + avg = sum(values) / len(values) + max_val = max(values) + min_val = min(values) + return f"Analysis of {data_type}: avg={avg:.2f}, max={max_val}, min={min_val}" + + tool = client.tools.upsert_from_function(func=analyze_data, tags=["analysis", "data"]) + yield tool + client.tools.delete(tool.id) + + +@pytest.fixture(scope="function") +def persona_block(client: LettaSDKClient): + """Fixture providing persona memory block.""" + block = client.blocks.create( + label="persona", + value="You are Alex, a data analyst and mathematician who helps users with calculations and insights. You have extensive experience in statistical analysis and prefer to provide clear, accurate results.", + limit=8000, + ) + yield block + client.blocks.delete(block.id) + + +@pytest.fixture(scope="function") +def human_block(client: LettaSDKClient): + """Fixture providing human memory block.""" + block = client.blocks.create( + label="human", + value="username: sarah_researcher\noccupation: data scientist\ninterests: machine learning, statistics, fibonacci sequences\npreferred_communication: detailed explanations with examples", + limit=4000, + ) + yield block + client.blocks.delete(block.id) + + +@pytest.fixture(scope="function") +def context_block(client: LettaSDKClient): + """Fixture providing project context memory block.""" + block = client.blocks.create( + label="project_context", + value="Current project: Building predictive models for financial markets. Sarah is working on sequence analysis and pattern recognition. Recently interested in mathematical sequences like Fibonacci for trend analysis.", + limit=6000, + ) + yield block + client.blocks.delete(block.id) + + def test_shared_blocks(client: LettaSDKClient): # create a block block = client.blocks.create( @@ -1465,7 +1578,6 @@ def test_tool_name_auto_update_with_multiple_functions(client: LettaSDKClient): def test_tool_rename_with_json_schema_and_source_code(client: LettaSDKClient): """Test that passing both new JSON schema AND source code still renames the tool based on source code""" - import textwrap # Create initial tool def initial_tool(x: int) -> int: @@ -1543,3 +1655,217 @@ def test_tool_rename_with_json_schema_and_source_code(client: LettaSDKClient): finally: # Clean up client.tools.delete(tool_id=tool.id) + + +def test_import_agent_file_from_disk( + client: LettaSDKClient, fibonacci_tool, preferences_tool, data_analysis_tool, persona_block, human_block, context_block +): + """Test exporting an agent to file and importing it back from disk.""" + # Create a comprehensive agent (similar to test_agent_serialization_v2) + name = f"test_export_import_{str(uuid.uuid4())}" + temp_agent = client.agents.create( + name=name, + memory_blocks=[persona_block, human_block, context_block], + model="openai/gpt-4.1-mini", + embedding="openai/text-embedding-3-small", + tool_ids=[fibonacci_tool.id, preferences_tool.id, data_analysis_tool.id], + include_base_tools=True, + tags=["test", "export", "import"], + system="You are a helpful assistant specializing in data analysis and mathematical computations.", + ) + + # Add archival memory + archival_passages = ["Test archival passage for export/import testing.", "Another passage with data about testing procedures."] + + for passage_text in archival_passages: + client.agents.passages.create(agent_id=temp_agent.id, text=passage_text) + + # Send a test message + client.agents.messages.create( + agent_id=temp_agent.id, + messages=[ + MessageCreate( + role="user", + content="Test message for export", + ), + ], + ) + + # Export the agent + serialized_v2 = client.agents.export_file(agent_id=temp_agent.id, use_legacy_format=False) + + # Save to file + file_path = os.path.join(os.path.dirname(__file__), "test_agent_files", "test_basic_agent_with_blocks_tools_messages_v2.af") + os.makedirs(os.path.dirname(file_path), exist_ok=True) + + with open(file_path, "w") as f: + json.dump(serialized_v2, f, indent=2) + + # Now import from the file + with open(file_path, "rb") as f: + import_result = client.agents.import_file( + file=f, append_copy_suffix=True, override_existing_tools=True # Use suffix to avoid name conflict + ) + + # Basic verification + assert import_result is not None, "Import result should not be None" + assert len(import_result.agent_ids) > 0, "Should have imported at least one agent" + + # Get the imported agent + imported_agent_id = import_result.agent_ids[0] + imported_agent = client.agents.retrieve(agent_id=imported_agent_id) + + # Basic checks + assert imported_agent is not None, "Should be able to retrieve imported agent" + assert imported_agent.name is not None, "Imported agent should have a name" + assert imported_agent.memory is not None, "Agent should have memory" + assert len(imported_agent.tools) > 0, "Agent should have tools" + assert imported_agent.system is not None, "Agent should have a system prompt" + + +def test_agent_serialization_v2( + client: LettaSDKClient, fibonacci_tool, preferences_tool, data_analysis_tool, persona_block, human_block, context_block +): + """Test agent serialization with comprehensive setup including custom tools, blocks, messages, and archival memory.""" + name = f"comprehensive_test_agent_{str(uuid.uuid4())}" + temp_agent = client.agents.create( + name=name, + memory_blocks=[persona_block, human_block, context_block], + model="openai/gpt-4.1-mini", + embedding="openai/text-embedding-3-small", + tool_ids=[fibonacci_tool.id, preferences_tool.id, data_analysis_tool.id], + include_base_tools=True, + tags=["test", "comprehensive", "serialization"], + system="You are a helpful assistant specializing in data analysis and mathematical computations.", + ) + + # Add archival memory + archival_passages = [ + "Project background: Sarah is working on a financial prediction model that uses Fibonacci retracements for technical analysis.", + "Research notes: Golden ratio (1.618) derived from Fibonacci sequence is often used in financial markets for support/resistance levels.", + ] + + for passage_text in archival_passages: + client.agents.passages.create(agent_id=temp_agent.id, text=passage_text) + + # Send some messages + client.agents.messages.create( + agent_id=temp_agent.id, + messages=[ + MessageCreate( + role="user", + content="Test message", + ), + ], + ) + + # Serialize using v2 + serialized_v2 = client.agents.export_file(agent_id=temp_agent.id, use_legacy_format=False) + # Convert dict to JSON bytes for import + json_str = json.dumps(serialized_v2) + file_obj = io.BytesIO(json_str.encode("utf-8")) + + # Import again + import_result = client.agents.import_file(file=file_obj, append_copy_suffix=False, override_existing_tools=True) + + # Verify import was successful + assert len(import_result.agent_ids) == 1, "Should have imported exactly one agent" + imported_agent_id = import_result.agent_ids[0] + imported_agent = client.agents.retrieve(agent_id=imported_agent_id) + + # ========== BASIC AGENT PROPERTIES ========== + # Name should be the same (if append_copy_suffix=False) or have suffix + assert imported_agent.name == name, f"Agent name mismatch: {imported_agent.name} != {name}" + + # LLM and embedding configs should be preserved + assert ( + imported_agent.llm_config.model == temp_agent.llm_config.model + ), f"LLM model mismatch: {imported_agent.llm_config.model} != {temp_agent.llm_config.model}" + assert imported_agent.embedding_config.embedding_model == temp_agent.embedding_config.embedding_model, "Embedding model mismatch" + + # System prompt should be preserved + assert imported_agent.system == temp_agent.system, "System prompt was not preserved" + + # Tags should be preserved + assert set(imported_agent.tags) == set(temp_agent.tags), f"Tags mismatch: {imported_agent.tags} != {temp_agent.tags}" + + # Agent type should be preserved + assert ( + imported_agent.agent_type == temp_agent.agent_type + ), f"Agent type mismatch: {imported_agent.agent_type} != {temp_agent.agent_type}" + + # ========== MEMORY BLOCKS ========== + # Compare memory blocks directly from AgentState objects + original_blocks = temp_agent.memory.blocks + imported_blocks = imported_agent.memory.blocks + + # Should have same number of blocks + assert len(imported_blocks) == len(original_blocks), f"Block count mismatch: {len(imported_blocks)} != {len(original_blocks)}" + + # Verify each block by label + original_blocks_by_label = {block.label: block for block in original_blocks} + imported_blocks_by_label = {block.label: block for block in imported_blocks} + + # Check persona block + assert "persona" in imported_blocks_by_label, "Persona block missing in imported agent" + assert "Alex" in imported_blocks_by_label["persona"].value, "Persona block content not preserved" + assert imported_blocks_by_label["persona"].limit == original_blocks_by_label["persona"].limit, "Persona block limit mismatch" + + # Check human block + assert "human" in imported_blocks_by_label, "Human block missing in imported agent" + assert "sarah_researcher" in imported_blocks_by_label["human"].value, "Human block content not preserved" + assert imported_blocks_by_label["human"].limit == original_blocks_by_label["human"].limit, "Human block limit mismatch" + + # Check context block + assert "project_context" in imported_blocks_by_label, "Context block missing in imported agent" + assert "financial markets" in imported_blocks_by_label["project_context"].value, "Context block content not preserved" + assert ( + imported_blocks_by_label["project_context"].limit == original_blocks_by_label["project_context"].limit + ), "Context block limit mismatch" + + # ========== TOOLS ========== + # Compare tools directly from AgentState objects + original_tools = temp_agent.tools + imported_tools = imported_agent.tools + + # Should have same number of tools + assert len(imported_tools) == len(original_tools), f"Tool count mismatch: {len(imported_tools)} != {len(original_tools)}" + + original_tool_names = {tool.name for tool in original_tools} + imported_tool_names = {tool.name for tool in imported_tools} + + # Check custom tools are present + assert "calculate_fibonacci" in imported_tool_names, "Fibonacci tool missing in imported agent" + assert "get_user_preferences" in imported_tool_names, "Preferences tool missing in imported agent" + assert "analyze_data" in imported_tool_names, "Data analysis tool missing in imported agent" + + # Check for base tools (since we set include_base_tools=True when creating the agent) + # Base tools should also be present (at least some core ones) + base_tool_names = {"send_message", "conversation_search"} + missing_base_tools = base_tool_names - imported_tool_names + assert len(missing_base_tools) == 0, f"Missing base tools: {missing_base_tools}" + + # Verify tool names match exactly + assert original_tool_names == imported_tool_names, f"Tool names don't match: {original_tool_names} != {imported_tool_names}" + + # ========== MESSAGE HISTORY ========== + # Get messages for both agents + original_messages = client.agents.messages.list(agent_id=temp_agent.id, limit=100) + imported_messages = client.agents.messages.list(agent_id=imported_agent_id, limit=100) + + # Should have same number of messages + assert len(imported_messages) >= 1, "Imported agent should have messages" + + # Filter for user messages (excluding system-generated login messages) + original_user_msgs = [msg for msg in original_messages if msg.message_type == "user_message" and "Test message" in msg.content] + imported_user_msgs = [msg for msg in imported_messages if msg.message_type == "user_message" and "Test message" in msg.content] + + # Should have the same number of test messages + assert len(imported_user_msgs) == len( + original_user_msgs + ), f"User message count mismatch: {len(imported_user_msgs)} != {len(original_user_msgs)}" + + # Verify test message content is preserved + if len(original_user_msgs) > 0 and len(imported_user_msgs) > 0: + assert imported_user_msgs[0].content == original_user_msgs[0].content, "User message content not preserved" + assert "Test message" in imported_user_msgs[0].content, "Test message content not found"