fbd89c9360
* fix: replace all 'PRODUCTION' references with 'prod' for consistency
Problem: Codebase had 11 references to 'PRODUCTION' (uppercase) that should
use 'prod' (lowercase) for consistency with the deployment workflows and
environment normalization.
Changes across 8 files:
1. Source files (using settings.environment):
- letta/functions/function_sets/multi_agent.py
- letta/services/tool_manager.py
- letta/services/tool_executor/multi_agent_tool_executor.py
- letta/services/helpers/agent_manager_helper.py
All checks changed from: settings.environment == "PRODUCTION"
To: settings.environment == "prod"
2. OTEL resource configuration:
- letta/otel/resource.py
- Updated _normalize_environment_tag() to handle 'prod' directly
- Removed 'PRODUCTION' -> 'prod' mapping (no longer needed)
- Updated device.id check from _env != "PRODUCTION" to _env != "prod"
3. Test files:
- tests/managers/conftest.py
- Fixture parameter changed from "PRODUCTION" to "prod"
- tests/managers/test_agent_manager.py (3 occurrences)
- tests/managers/test_tool_manager.py (2 occurrences)
All test checks changed to use "prod"
Result: Complete consistency across the codebase:
- All environment checks use "prod" instead of "PRODUCTION"
- Normalization function simplified (no special case for PRODUCTION)
- Tests use correct "prod" value
- Matches deployment workflow configuration from PR #6626
This completes the environment naming standardization effort.
* fix: update settings.py environment description to use 'prod' instead of 'PRODUCTION'
The field description still referenced PRODUCTION as an example value.
Updated to use lowercase 'prod' for consistency with actual usage.
Before: "Application environment (PRODUCTION, DEV, CANARY, etc. - normalized to lowercase for OTEL tags)"
After: "Application environment (prod, dev, canary, etc. - lowercase values used for OTEL tags)"
1784 lines
73 KiB
Python
1784 lines
73 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import re
|
|
import string
|
|
import time
|
|
import uuid
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import List
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
import pytest
|
|
from _pytest.python_api import approx
|
|
from anthropic.types.beta import BetaMessage
|
|
from anthropic.types.beta.messages import BetaMessageBatchIndividualResponse, BetaMessageBatchSucceededResult
|
|
|
|
# Import shared fixtures and constants from conftest
|
|
from conftest import (
|
|
CREATE_DELAY_SQLITE,
|
|
DEFAULT_EMBEDDING_CONFIG,
|
|
USING_SQLITE,
|
|
)
|
|
from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall as OpenAIToolCall, Function as OpenAIFunction
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.exc import IntegrityError, InvalidRequestError
|
|
from sqlalchemy.orm.exc import StaleDataError
|
|
|
|
from letta.config import LettaConfig
|
|
from letta.constants import (
|
|
BASE_MEMORY_TOOLS,
|
|
BASE_SLEEPTIME_TOOLS,
|
|
BASE_TOOLS,
|
|
BASE_VOICE_SLEEPTIME_CHAT_TOOLS,
|
|
BASE_VOICE_SLEEPTIME_TOOLS,
|
|
BUILTIN_TOOLS,
|
|
DEFAULT_ORG_ID,
|
|
DEFAULT_ORG_NAME,
|
|
FILES_TOOLS,
|
|
LETTA_TOOL_EXECUTION_DIR,
|
|
LETTA_TOOL_SET,
|
|
LOCAL_ONLY_MULTI_AGENT_TOOLS,
|
|
MCP_TOOL_TAG_NAME_PREFIX,
|
|
MULTI_AGENT_TOOLS,
|
|
)
|
|
from letta.data_sources.redis_client import NoopAsyncRedisClient, get_redis_client
|
|
from letta.errors import LettaAgentNotFoundError
|
|
from letta.functions.functions import derive_openai_json_schema, parse_source_code
|
|
from letta.functions.mcp_client.types import MCPTool
|
|
from letta.helpers import ToolRulesSolver
|
|
from letta.helpers.datetime_helpers import AsyncTimer
|
|
from letta.jobs.types import ItemUpdateInfo, RequestStatusUpdateInfo, StepStatusUpdateInfo
|
|
from letta.orm import Base, Block
|
|
from letta.orm.block_history import BlockHistory
|
|
from letta.orm.errors import NoResultFound, UniqueConstraintViolationError
|
|
from letta.orm.file import FileContent as FileContentModel, FileMetadata as FileMetadataModel
|
|
from letta.schemas.agent import CreateAgent, InternalTemplateAgentCreate, UpdateAgent
|
|
from letta.schemas.block import Block as PydanticBlock, BlockUpdate, CreateBlock
|
|
from letta.schemas.embedding_config import EmbeddingConfig
|
|
from letta.schemas.enums import (
|
|
ActorType,
|
|
AgentStepStatus,
|
|
FileProcessingStatus,
|
|
JobStatus,
|
|
JobType,
|
|
MessageRole,
|
|
ProviderType,
|
|
SandboxType,
|
|
StepStatus,
|
|
TagMatchMode,
|
|
ToolType,
|
|
VectorDBProvider,
|
|
)
|
|
from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate, SandboxEnvironmentVariableUpdate
|
|
from letta.schemas.file import FileMetadata, FileMetadata as PydanticFileMetadata
|
|
from letta.schemas.identity import IdentityCreate, IdentityProperty, IdentityPropertyType, IdentityType, IdentityUpdate, IdentityUpsert
|
|
from letta.schemas.job import BatchJob, Job, Job as PydanticJob, JobUpdate, LettaRequestConfig
|
|
from letta.schemas.letta_message import UpdateAssistantMessage, UpdateReasoningMessage, UpdateSystemMessage, UpdateUserMessage
|
|
from letta.schemas.letta_message_content import TextContent
|
|
from letta.schemas.letta_stop_reason import LettaStopReason, StopReasonType
|
|
from letta.schemas.llm_batch_job import AgentStepState, LLMBatchItem
|
|
from letta.schemas.llm_config import LLMConfig
|
|
from letta.schemas.message import Message as PydanticMessage, MessageCreate, MessageUpdate
|
|
from letta.schemas.openai.chat_completion_response import UsageStatistics
|
|
from letta.schemas.organization import Organization, Organization as PydanticOrganization, OrganizationUpdate
|
|
from letta.schemas.passage import Passage as PydanticPassage
|
|
from letta.schemas.pip_requirement import PipRequirement
|
|
from letta.schemas.run import Run as PydanticRun
|
|
from letta.schemas.sandbox_config import E2BSandboxConfig, LocalSandboxConfig, SandboxConfigCreate, SandboxConfigUpdate
|
|
from letta.schemas.source import Source as PydanticSource, SourceUpdate
|
|
from letta.schemas.tool import Tool as PydanticTool, ToolCreate, ToolUpdate
|
|
from letta.schemas.tool_rule import InitToolRule
|
|
from letta.schemas.user import User as PydanticUser, UserUpdate
|
|
from letta.server.db import db_registry
|
|
from letta.server.server import SyncServer
|
|
from letta.services.block_manager import BlockManager
|
|
from letta.services.helpers.agent_manager_helper import calculate_base_tools, calculate_multi_agent_tools, validate_agent_exists_async
|
|
from letta.services.step_manager import FeedbackType
|
|
from letta.settings import settings, tool_settings
|
|
from letta.utils import calculate_file_defaults_based_on_context_window
|
|
from tests.helpers.utils import comprehensive_agent_checks, validate_context_window_overview
|
|
from tests.utils import random_string
|
|
|
|
# ======================================================================================================================
|
|
# Helper Functions
|
|
# ======================================================================================================================
|
|
|
|
|
|
async def _count_file_content_rows(session, file_id: str) -> int:
|
|
q = select(func.count()).select_from(FileContentModel).where(FileContentModel.file_id == file_id)
|
|
result = await session.execute(q)
|
|
return result.scalar_one()
|
|
|
|
|
|
# ======================================================================================================================
|
|
# AgentManager Tests - Basic
|
|
# ======================================================================================================================
|
|
async def test_validate_agent_exists_async(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
"""Test the validate_agent_exists_async helper function"""
|
|
created_agent, _ = comprehensive_test_agent_fixture
|
|
|
|
# test with valid agent
|
|
async with db_registry.async_session() as session:
|
|
# should not raise exception
|
|
await validate_agent_exists_async(session, created_agent.id, default_user)
|
|
|
|
# test with non-existent agent
|
|
async with db_registry.async_session() as session:
|
|
with pytest.raises(LettaAgentNotFoundError):
|
|
await validate_agent_exists_async(session, "non-existent-id", default_user)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_get_list_agent(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
# Test agent creation
|
|
created_agent, create_agent_request = comprehensive_test_agent_fixture
|
|
comprehensive_agent_checks(created_agent, create_agent_request, actor=default_user)
|
|
|
|
# Test get agent
|
|
get_agent = await server.agent_manager.get_agent_by_id_async(agent_id=created_agent.id, actor=default_user)
|
|
comprehensive_agent_checks(get_agent, create_agent_request, actor=default_user)
|
|
|
|
# Test get agent name
|
|
agents = await server.agent_manager.list_agents_async(name=created_agent.name, actor=default_user)
|
|
get_agent_name = agents[0]
|
|
comprehensive_agent_checks(get_agent_name, create_agent_request, actor=default_user)
|
|
|
|
# Test list agent
|
|
list_agents = await server.agent_manager.list_agents_async(actor=default_user)
|
|
assert len(list_agents) == 1
|
|
comprehensive_agent_checks(list_agents[0], create_agent_request, actor=default_user)
|
|
|
|
# Test deleting the agent
|
|
await server.agent_manager.delete_agent_async(get_agent.id, default_user)
|
|
list_agents = await server.agent_manager.list_agents_async(actor=default_user)
|
|
assert len(list_agents) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_agent_include_base_tools(server: SyncServer, default_user):
|
|
"""Test agent creation with include_default_source=True"""
|
|
# Upsert base tools
|
|
await server.tool_manager.upsert_base_tools_async(actor=default_user)
|
|
|
|
memory_blocks = [CreateBlock(label="human", value="TestUser"), CreateBlock(label="persona", value="I am a test assistant")]
|
|
|
|
create_agent_request = CreateAgent(
|
|
name="test_default_source_agent",
|
|
agent_type="memgpt_v2_agent",
|
|
system="test system",
|
|
memory_blocks=memory_blocks,
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
include_base_tools=True,
|
|
)
|
|
|
|
# Create the agent
|
|
created_agent = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
|
|
# Assert the tools exist
|
|
tool_names = [t.name for t in created_agent.tools]
|
|
expected_tools = calculate_base_tools(is_v2=True)
|
|
assert sorted(tool_names) == sorted(expected_tools)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_agent_base_tool_rules_excluded_providers(server: SyncServer, default_user):
|
|
"""Test that include_base_tool_rules is overridden to False for excluded providers"""
|
|
# Upsert base tools
|
|
await server.tool_manager.upsert_base_tools_async(actor=default_user)
|
|
|
|
memory_blocks = [CreateBlock(label="human", value="TestUser"), CreateBlock(label="persona", value="I am a test assistant")]
|
|
|
|
# Test with excluded provider (openai)
|
|
create_agent_request = CreateAgent(
|
|
name="test_excluded_provider_agent",
|
|
agent_type="memgpt_v2_agent",
|
|
system="test system",
|
|
memory_blocks=memory_blocks,
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"), # This has model_endpoint_type="openai"
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
include_base_tool_rules=False,
|
|
)
|
|
|
|
# Create the agent
|
|
created_agent = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
|
|
# Assert that no base tool rules were added (since include_base_tool_rules was overridden to False)
|
|
print(created_agent.tool_rules)
|
|
assert created_agent.tool_rules is None or len(created_agent.tool_rules) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_agent_base_tool_rules_non_excluded_providers(server: SyncServer, default_user):
|
|
"""Test that include_base_tool_rules is NOT overridden for non-excluded providers"""
|
|
# Upsert base tools
|
|
await server.tool_manager.upsert_base_tools_async(actor=default_user)
|
|
|
|
memory_blocks = [CreateBlock(label="human", value="TestUser"), CreateBlock(label="persona", value="I am a test assistant")]
|
|
|
|
# Test with non-excluded provider (together)
|
|
create_agent_request = CreateAgent(
|
|
name="test_non_excluded_provider_agent",
|
|
agent_type="memgpt_v2_agent",
|
|
system="test system",
|
|
memory_blocks=memory_blocks,
|
|
llm_config=LLMConfig(
|
|
model="llama-3.1-8b-instruct",
|
|
model_endpoint_type="together", # Model doesn't match EXCLUDE_MODEL_KEYWORDS_FROM_BASE_TOOL_RULES
|
|
model_endpoint="https://api.together.xyz",
|
|
context_window=8192,
|
|
),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
include_base_tool_rules=True, # Should remain True
|
|
)
|
|
|
|
# Create the agent
|
|
created_agent = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
|
|
# Assert that base tool rules were added (since include_base_tool_rules remained True)
|
|
assert created_agent.tool_rules is not None
|
|
assert len(created_agent.tool_rules) > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_calculate_multi_agent_tools(set_letta_environment):
|
|
"""Test that calculate_multi_agent_tools excludes local-only tools in production."""
|
|
result = calculate_multi_agent_tools()
|
|
|
|
if settings.environment == "prod":
|
|
# Production environment should exclude local-only tools
|
|
expected_tools = set(MULTI_AGENT_TOOLS) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS)
|
|
assert result == expected_tools, "Production should exclude local-only multi-agent tools"
|
|
assert not set(LOCAL_ONLY_MULTI_AGENT_TOOLS).intersection(result), "Production should not include local-only tools"
|
|
|
|
# Verify specific tools
|
|
assert "send_message_to_agent_and_wait_for_reply" in result, "Standard multi-agent tools should be in production"
|
|
assert "send_message_to_agents_matching_tags" in result, "Standard multi-agent tools should be in production"
|
|
assert "send_message_to_agent_async" not in result, "Local-only tools should not be in production"
|
|
else:
|
|
# Non-production environment should include all multi-agent tools
|
|
assert result == set(MULTI_AGENT_TOOLS), "Non-production should include all multi-agent tools"
|
|
assert set(LOCAL_ONLY_MULTI_AGENT_TOOLS).issubset(result), "Non-production should include local-only tools"
|
|
|
|
# Verify specific tools
|
|
assert "send_message_to_agent_and_wait_for_reply" in result, "All multi-agent tools should be in non-production"
|
|
assert "send_message_to_agents_matching_tags" in result, "All multi-agent tools should be in non-production"
|
|
assert "send_message_to_agent_async" in result, "Local-only tools should be in non-production"
|
|
|
|
|
|
async def test_upsert_base_tools_excludes_local_only_in_production(server: SyncServer, default_user, set_letta_environment):
|
|
"""Test that upsert_base_tools excludes local-only multi-agent tools in production."""
|
|
# Upsert all base tools
|
|
tools = await server.tool_manager.upsert_base_tools_async(actor=default_user)
|
|
tool_names = {tool.name for tool in tools}
|
|
|
|
if settings.environment == "prod":
|
|
# Production environment should exclude local-only multi-agent tools
|
|
for local_only_tool in LOCAL_ONLY_MULTI_AGENT_TOOLS:
|
|
assert local_only_tool not in tool_names, f"Local-only tool '{local_only_tool}' should not be upserted in production"
|
|
|
|
# But should include standard multi-agent tools
|
|
standard_multi_agent_tools = set(MULTI_AGENT_TOOLS) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS)
|
|
for standard_tool in standard_multi_agent_tools:
|
|
assert standard_tool in tool_names, f"Standard multi-agent tool '{standard_tool}' should be upserted in production"
|
|
else:
|
|
# Non-production environment should include all multi-agent tools
|
|
for tool in MULTI_AGENT_TOOLS:
|
|
assert tool in tool_names, f"Multi-agent tool '{tool}' should be upserted in non-production"
|
|
|
|
|
|
async def test_upsert_multi_agent_tools_only(server: SyncServer, default_user, set_letta_environment):
|
|
"""Test that upserting only multi-agent tools respects production filtering."""
|
|
from letta.schemas.enums import ToolType
|
|
|
|
# Upsert only multi-agent tools
|
|
tools = await server.tool_manager.upsert_base_tools_async(actor=default_user, allowed_types={ToolType.LETTA_MULTI_AGENT_CORE})
|
|
tool_names = {tool.name for tool in tools}
|
|
|
|
if settings.environment == "prod":
|
|
# Should only have non-local multi-agent tools
|
|
expected_tools = set(MULTI_AGENT_TOOLS) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS)
|
|
assert tool_names == expected_tools, "Production multi-agent upsert should exclude local-only tools"
|
|
assert "send_message_to_agent_async" not in tool_names, "Local-only async tool should not be upserted in production"
|
|
else:
|
|
# Should have all multi-agent tools
|
|
assert tool_names == set(MULTI_AGENT_TOOLS), "Non-production multi-agent upsert should include all tools"
|
|
assert "send_message_to_agent_async" in tool_names, "Local-only async tool should be upserted in non-production"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_agent_with_default_source(server: SyncServer, default_user, print_tool, default_block):
|
|
"""Test agent creation with include_default_source=True"""
|
|
memory_blocks = [CreateBlock(label="human", value="TestUser"), CreateBlock(label="persona", value="I am a test assistant")]
|
|
|
|
create_agent_request = CreateAgent(
|
|
name="test_default_source_agent",
|
|
agent_type="memgpt_v2_agent",
|
|
system="test system",
|
|
memory_blocks=memory_blocks,
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
tool_ids=[print_tool.id],
|
|
include_default_source=True, # This is the key field we're testing
|
|
include_base_tools=False,
|
|
)
|
|
|
|
# Create the agent
|
|
created_agent = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
|
|
# Verify agent was created
|
|
assert created_agent is not None
|
|
assert created_agent.name == "test_default_source_agent"
|
|
|
|
# Verify that a default source was created and attached
|
|
attached_sources = await server.agent_manager.list_attached_sources_async(agent_id=created_agent.id, actor=default_user)
|
|
|
|
# Should have exactly one source (the default one)
|
|
assert len(attached_sources) == 1
|
|
auto_default_source = attached_sources[0]
|
|
|
|
# Verify the default source properties
|
|
assert created_agent.name in auto_default_source.name
|
|
assert auto_default_source.embedding_config.embedding_endpoint_type == "openai"
|
|
|
|
# Test with include_default_source=False
|
|
create_agent_request_no_source = CreateAgent(
|
|
name="test_no_default_source_agent",
|
|
agent_type="memgpt_v2_agent",
|
|
system="test system",
|
|
memory_blocks=memory_blocks,
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
tool_ids=[print_tool.id],
|
|
include_default_source=False, # Explicitly set to False
|
|
include_base_tools=False,
|
|
)
|
|
|
|
created_agent_no_source = await server.agent_manager.create_agent_async(
|
|
create_agent_request_no_source,
|
|
actor=default_user,
|
|
)
|
|
|
|
# Verify no sources are attached
|
|
attached_sources_no_source = await server.agent_manager.list_attached_sources_async(
|
|
agent_id=created_agent_no_source.id, actor=default_user
|
|
)
|
|
|
|
assert len(attached_sources_no_source) == 0
|
|
|
|
# Clean up
|
|
await server.agent_manager.delete_agent_async(created_agent.id, default_user)
|
|
await server.agent_manager.delete_agent_async(created_agent_no_source.id, default_user)
|
|
|
|
|
|
async def test_get_context_window_basic(
|
|
server: SyncServer, comprehensive_test_agent_fixture, default_user, default_file, set_letta_environment
|
|
):
|
|
# Test agent creation
|
|
created_agent, create_agent_request = comprehensive_test_agent_fixture
|
|
|
|
# Attach a file
|
|
assoc, closed_files = await server.file_agent_manager.attach_file(
|
|
agent_id=created_agent.id,
|
|
file_id=default_file.id,
|
|
file_name=default_file.file_name,
|
|
source_id=default_file.source_id,
|
|
actor=default_user,
|
|
visible_content="hello",
|
|
max_files_open=created_agent.max_files_open,
|
|
)
|
|
|
|
# Get context window and check for basic appearances
|
|
context_window_overview = await server.agent_manager.get_context_window(agent_id=created_agent.id, actor=default_user)
|
|
validate_context_window_overview(created_agent, context_window_overview, assoc)
|
|
|
|
# Test deleting the agent
|
|
await server.agent_manager.delete_agent_async(created_agent.id, default_user)
|
|
list_agents = await server.agent_manager.list_agents_async(actor=default_user)
|
|
assert len(list_agents) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_agent_passed_in_initial_messages(server: SyncServer, default_user, default_block):
|
|
memory_blocks = [CreateBlock(label="human", value="BananaBoy"), CreateBlock(label="persona", value="I am a helpful assistant")]
|
|
create_agent_request = CreateAgent(
|
|
agent_type="memgpt_v2_agent",
|
|
system="test system",
|
|
memory_blocks=memory_blocks,
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
tags=["a", "b"],
|
|
description="test_description",
|
|
initial_message_sequence=[MessageCreate(role=MessageRole.user, content="hello world")],
|
|
include_base_tools=False,
|
|
)
|
|
agent_state = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
assert await server.message_manager.size_async(agent_id=agent_state.id, actor=default_user) == 2
|
|
init_messages = await server.message_manager.get_messages_by_ids_async(message_ids=agent_state.message_ids, actor=default_user)
|
|
|
|
# Check that the system appears in the first initial message
|
|
assert create_agent_request.system in init_messages[0].content[0].text
|
|
assert create_agent_request.memory_blocks[0].value in init_messages[0].content[0].text
|
|
# Check that the second message is the passed in initial message seq
|
|
assert create_agent_request.initial_message_sequence[0].role == init_messages[1].role
|
|
assert create_agent_request.initial_message_sequence[0].content in init_messages[1].content[0].text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_agent_default_initial_message(server: SyncServer, default_user, default_block):
|
|
memory_blocks = [CreateBlock(label="human", value="BananaBoy"), CreateBlock(label="persona", value="I am a helpful assistant")]
|
|
create_agent_request = CreateAgent(
|
|
agent_type="memgpt_v2_agent",
|
|
system="test system",
|
|
memory_blocks=memory_blocks,
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
tags=["a", "b"],
|
|
description="test_description",
|
|
include_base_tools=False,
|
|
)
|
|
agent_state = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
assert await server.message_manager.size_async(agent_id=agent_state.id, actor=default_user) == 4
|
|
init_messages = await server.message_manager.get_messages_by_ids_async(message_ids=agent_state.message_ids, actor=default_user)
|
|
# Check that the system appears in the first initial message
|
|
assert create_agent_request.system in init_messages[0].content[0].text
|
|
assert create_agent_request.memory_blocks[0].value in init_messages[0].content[0].text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_agent_with_json_in_system_message(server: SyncServer, default_user, default_block):
|
|
system_prompt = (
|
|
"You are an expert teaching agent with encyclopedic knowledge. "
|
|
"When you receive a topic, query the external database for more "
|
|
"information. Format the queries as a JSON list of queries making "
|
|
"sure to include your reasoning for that query, e.g. "
|
|
"{'query1' : 'reason1', 'query2' : 'reason2'}"
|
|
)
|
|
create_agent_request = CreateAgent(
|
|
agent_type="memgpt_v2_agent",
|
|
system=system_prompt,
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
tags=["a", "b"],
|
|
description="test_description",
|
|
include_base_tools=False,
|
|
)
|
|
agent_state = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
assert agent_state is not None
|
|
system_message_id = agent_state.message_ids[0]
|
|
system_message = await server.message_manager.get_message_by_id_async(message_id=system_message_id, actor=default_user)
|
|
assert system_prompt in system_message.content[0].text
|
|
assert default_block.value in system_message.content[0].text
|
|
await server.agent_manager.delete_agent_async(agent_id=agent_state.id, actor=default_user)
|
|
|
|
|
|
async def test_update_agent(server: SyncServer, comprehensive_test_agent_fixture, other_tool, other_source, other_block, default_user):
|
|
agent, _ = comprehensive_test_agent_fixture
|
|
update_agent_request = UpdateAgent(
|
|
name="train_agent",
|
|
description="train description",
|
|
tool_ids=[other_tool.id],
|
|
source_ids=[other_source.id],
|
|
block_ids=[other_block.id],
|
|
tool_rules=[InitToolRule(tool_name=other_tool.name)],
|
|
tags=["c", "d"],
|
|
system="train system",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(model_name="letta"),
|
|
message_ids=["10", "20"],
|
|
metadata={"train_key": "train_value"},
|
|
tool_exec_environment_variables={"test_env_var_key_a": "a", "new_tool_exec_key": "n"},
|
|
message_buffer_autoclear=False,
|
|
)
|
|
|
|
last_updated_timestamp = agent.updated_at
|
|
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_agent_request, actor=default_user)
|
|
comprehensive_agent_checks(updated_agent, update_agent_request, actor=default_user)
|
|
assert updated_agent.message_ids == update_agent_request.message_ids
|
|
assert updated_agent.updated_at > last_updated_timestamp
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_file_defaults_based_on_context_window(server: SyncServer, default_user, default_block):
|
|
"""Test that file-related defaults are set based on the model's context window size"""
|
|
|
|
# test with small context window model (8k)
|
|
llm_config_small = LLMConfig.default_config("gpt-4o-mini")
|
|
llm_config_small.context_window = 8000
|
|
create_agent_request = CreateAgent(
|
|
name="test_agent_small_context",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=llm_config_small,
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
include_base_tools=False,
|
|
)
|
|
agent_state = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
assert agent_state.max_files_open == 3
|
|
assert (
|
|
agent_state.per_file_view_window_char_limit == calculate_file_defaults_based_on_context_window(llm_config_small.context_window)[1]
|
|
)
|
|
await server.agent_manager.delete_agent_async(agent_id=agent_state.id, actor=default_user)
|
|
|
|
# test with medium context window model (32k)
|
|
llm_config_medium = LLMConfig.default_config("gpt-4o-mini")
|
|
llm_config_medium.context_window = 32000
|
|
create_agent_request = CreateAgent(
|
|
name="test_agent_medium_context",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=llm_config_medium,
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
include_base_tools=False,
|
|
)
|
|
agent_state = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
assert agent_state.max_files_open == 5
|
|
assert (
|
|
agent_state.per_file_view_window_char_limit == calculate_file_defaults_based_on_context_window(llm_config_medium.context_window)[1]
|
|
)
|
|
await server.agent_manager.delete_agent_async(agent_id=agent_state.id, actor=default_user)
|
|
|
|
# test with large context window model (128k)
|
|
llm_config_large = LLMConfig.default_config("gpt-4o-mini")
|
|
llm_config_large.context_window = 128000
|
|
create_agent_request = CreateAgent(
|
|
name="test_agent_large_context",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=llm_config_large,
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
include_base_tools=False,
|
|
)
|
|
agent_state = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
assert agent_state.max_files_open == 10
|
|
assert (
|
|
agent_state.per_file_view_window_char_limit == calculate_file_defaults_based_on_context_window(llm_config_large.context_window)[1]
|
|
)
|
|
await server.agent_manager.delete_agent_async(agent_id=agent_state.id, actor=default_user)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_file_defaults_explicit_values(server: SyncServer, default_user, default_block):
|
|
"""Test that explicitly set file-related values are respected"""
|
|
|
|
llm_config_explicit = LLMConfig.default_config("gpt-4o-mini")
|
|
llm_config_explicit.context_window = 32000 # would normally get defaults of 5 and 30k
|
|
create_agent_request = CreateAgent(
|
|
name="test_agent_explicit_values",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=llm_config_explicit,
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
include_base_tools=False,
|
|
max_files_open=20, # explicit value
|
|
per_file_view_window_char_limit=500_000, # explicit value
|
|
)
|
|
agent_state = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
# verify explicit values are used instead of defaults
|
|
assert agent_state.max_files_open == 20
|
|
assert agent_state.per_file_view_window_char_limit == 500_000
|
|
await server.agent_manager.delete_agent_async(agent_id=agent_state.id, actor=default_user)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_file_fields(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
"""Test updating file-related fields on an existing agent"""
|
|
|
|
agent, _ = comprehensive_test_agent_fixture
|
|
|
|
# update file-related fields
|
|
update_request = UpdateAgent(
|
|
max_files_open=15,
|
|
per_file_view_window_char_limit=150_000,
|
|
)
|
|
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
|
|
|
|
assert updated_agent.max_files_open == 15
|
|
assert updated_agent.per_file_view_window_char_limit == 150_000
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_agent_last_stop_reason(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
"""Test updating last_stop_reason field on an existing agent"""
|
|
|
|
agent, _ = comprehensive_test_agent_fixture
|
|
|
|
assert agent.last_stop_reason is None
|
|
|
|
# Update with end_turn stop reason
|
|
update_request = UpdateAgent(
|
|
last_stop_reason=StopReasonType.end_turn,
|
|
last_run_completion=datetime.now(timezone.utc),
|
|
last_run_duration_ms=1500,
|
|
)
|
|
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
|
|
|
|
assert updated_agent.last_stop_reason == StopReasonType.end_turn
|
|
assert updated_agent.last_run_completion is not None
|
|
assert updated_agent.last_run_duration_ms == 1500
|
|
|
|
# Update with error stop reason
|
|
update_request = UpdateAgent(
|
|
last_stop_reason=StopReasonType.error,
|
|
last_run_completion=datetime.now(timezone.utc),
|
|
last_run_duration_ms=2500,
|
|
)
|
|
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
|
|
|
|
assert updated_agent.last_stop_reason == StopReasonType.error
|
|
assert updated_agent.last_run_duration_ms == 2500
|
|
|
|
# Update with requires_approval stop reason
|
|
update_request = UpdateAgent(
|
|
last_stop_reason=StopReasonType.requires_approval,
|
|
)
|
|
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
|
|
|
|
assert updated_agent.last_stop_reason == StopReasonType.requires_approval
|
|
|
|
|
|
# ======================================================================================================================
|
|
# AgentManager Tests - Listing
|
|
# ======================================================================================================================
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_select_fields_empty(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
# Create an agent using the comprehensive fixture.
|
|
created_agent, create_agent_request = comprehensive_test_agent_fixture
|
|
|
|
# List agents using an empty list for select_fields.
|
|
agents = await server.agent_manager.list_agents_async(actor=default_user, include_relationships=[])
|
|
# Assert that the agent is returned and basic fields are present.
|
|
assert len(agents) >= 1
|
|
agent = agents[0]
|
|
assert agent.id is not None
|
|
assert agent.name is not None
|
|
|
|
# Assert no relationships were loaded
|
|
assert len(agent.tools) == 0
|
|
assert len(agent.tags) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_select_fields_none(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
# Create an agent using the comprehensive fixture.
|
|
created_agent, create_agent_request = comprehensive_test_agent_fixture
|
|
|
|
# List agents using an empty list for select_fields.
|
|
agents = await server.agent_manager.list_agents_async(actor=default_user, include_relationships=None)
|
|
# Assert that the agent is returned and basic fields are present.
|
|
assert len(agents) >= 1
|
|
agent = agents[0]
|
|
assert agent.id is not None
|
|
assert agent.name is not None
|
|
|
|
# Assert no relationships were loaded
|
|
assert len(agent.tools) > 0
|
|
assert len(agent.tags) > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_select_fields_specific(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
created_agent, create_agent_request = comprehensive_test_agent_fixture
|
|
|
|
# Choose a subset of valid relationship fields.
|
|
valid_fields = ["tools", "tags"]
|
|
agents = await server.agent_manager.list_agents_async(actor=default_user, include_relationships=valid_fields)
|
|
assert len(agents) >= 1
|
|
agent = agents[0]
|
|
# Depending on your to_pydantic() implementation,
|
|
# verify that the fields exist in the returned pydantic model.
|
|
# (Note: These assertions may require that your CreateAgent fixture sets up these relationships.)
|
|
assert agent.tools
|
|
assert sorted(agent.tags) == ["a", "b"]
|
|
assert not agent.memory.blocks
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_select_fields_invalid(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
created_agent, create_agent_request = comprehensive_test_agent_fixture
|
|
|
|
# Provide field names that are not recognized.
|
|
invalid_fields = ["foobar", "nonexistent_field"]
|
|
# The expectation is that these fields are simply ignored.
|
|
agents = await server.agent_manager.list_agents_async(actor=default_user, include_relationships=invalid_fields)
|
|
assert len(agents) >= 1
|
|
agent = agents[0]
|
|
# Verify that standard fields are still present.c
|
|
assert agent.id is not None
|
|
assert agent.name is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_select_fields_duplicates(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
created_agent, create_agent_request = comprehensive_test_agent_fixture
|
|
|
|
# Provide duplicate valid field names.
|
|
duplicate_fields = ["tools", "tools", "tags", "tags"]
|
|
agents = await server.agent_manager.list_agents_async(actor=default_user, include_relationships=duplicate_fields)
|
|
assert len(agents) >= 1
|
|
agent = agents[0]
|
|
# Verify that the agent pydantic representation includes the relationships.
|
|
# Even if duplicates were provided, the query should not break.
|
|
assert isinstance(agent.tools, list)
|
|
assert isinstance(agent.tags, list)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_select_fields_mixed(server: SyncServer, comprehensive_test_agent_fixture, default_user):
|
|
created_agent, create_agent_request = comprehensive_test_agent_fixture
|
|
|
|
# Mix valid fields with an invalid one.
|
|
mixed_fields = ["tools", "invalid_field"]
|
|
agents = await server.agent_manager.list_agents_async(actor=default_user, include_relationships=mixed_fields)
|
|
assert len(agents) >= 1
|
|
agent = agents[0]
|
|
# Valid fields should be loaded and accessible.
|
|
assert agent.tools
|
|
# Since "invalid_field" is not recognized, it should have no adverse effect.
|
|
# You might optionally check that no extra attribute is created on the pydantic model.
|
|
assert not hasattr(agent, "invalid_field")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_ascending(server: SyncServer, default_user):
|
|
# Create two agents with known names
|
|
agent1 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_oldest",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
),
|
|
actor=default_user,
|
|
)
|
|
|
|
if USING_SQLITE:
|
|
time.sleep(CREATE_DELAY_SQLITE)
|
|
|
|
agent2 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_newest",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
),
|
|
actor=default_user,
|
|
)
|
|
|
|
agents = await server.agent_manager.list_agents_async(actor=default_user, ascending=True)
|
|
names = [agent.name for agent in agents]
|
|
assert names.index("agent_oldest") < names.index("agent_newest")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_descending(server: SyncServer, default_user):
|
|
# Create two agents with known names
|
|
agent1 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_oldest",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
),
|
|
actor=default_user,
|
|
)
|
|
|
|
if USING_SQLITE:
|
|
time.sleep(CREATE_DELAY_SQLITE)
|
|
|
|
agent2 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_newest",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
),
|
|
actor=default_user,
|
|
)
|
|
|
|
agents = await server.agent_manager.list_agents_async(actor=default_user, ascending=False)
|
|
names = [agent.name for agent in agents]
|
|
assert names.index("agent_newest") < names.index("agent_oldest")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_by_last_stop_reason(server: SyncServer, default_user):
|
|
# Create agent with requires_approval stop reason
|
|
agent1 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_requires_approval",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
),
|
|
actor=default_user,
|
|
)
|
|
await server.agent_manager.update_agent_async(
|
|
agent_id=agent1.id,
|
|
agent_update=UpdateAgent(last_stop_reason=StopReasonType.requires_approval),
|
|
actor=default_user,
|
|
)
|
|
|
|
# Create agent with error stop reason
|
|
agent2 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_error",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
),
|
|
actor=default_user,
|
|
)
|
|
await server.agent_manager.update_agent_async(
|
|
agent_id=agent2.id,
|
|
agent_update=UpdateAgent(last_stop_reason=StopReasonType.error),
|
|
actor=default_user,
|
|
)
|
|
|
|
# Create agent with no stop reason
|
|
agent3 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_no_stop_reason",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
),
|
|
actor=default_user,
|
|
)
|
|
|
|
# Filter by requires_approval
|
|
approval_agents = await server.agent_manager.list_agents_async(
|
|
actor=default_user, last_stop_reason=StopReasonType.requires_approval.value
|
|
)
|
|
approval_names = {agent.name for agent in approval_agents}
|
|
assert approval_names == {"agent_requires_approval"}
|
|
|
|
# Filter by error
|
|
error_agents = await server.agent_manager.list_agents_async(actor=default_user, last_stop_reason=StopReasonType.error.value)
|
|
error_names = {agent.name for agent in error_agents}
|
|
assert error_names == {"agent_error"}
|
|
|
|
# No filter - should return all agents
|
|
all_agents = await server.agent_manager.list_agents_async(actor=default_user)
|
|
all_names = {agent.name for agent in all_agents}
|
|
assert {"agent_requires_approval", "agent_error", "agent_no_stop_reason"}.issubset(all_names)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_count_agents_with_filters(server: SyncServer, default_user):
|
|
"""Test count_agents_async with various filters"""
|
|
# Create agents with different attributes
|
|
agent1 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_requires_approval",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
tags=["inbox", "test"],
|
|
),
|
|
actor=default_user,
|
|
)
|
|
await server.agent_manager.update_agent_async(
|
|
agent_id=agent1.id,
|
|
agent_update=UpdateAgent(last_stop_reason=StopReasonType.requires_approval),
|
|
actor=default_user,
|
|
)
|
|
|
|
agent2 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_error",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
tags=["error", "test"],
|
|
),
|
|
actor=default_user,
|
|
)
|
|
await server.agent_manager.update_agent_async(
|
|
agent_id=agent2.id,
|
|
agent_update=UpdateAgent(last_stop_reason=StopReasonType.error),
|
|
actor=default_user,
|
|
)
|
|
|
|
agent3 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_completed",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
tags=["completed"],
|
|
),
|
|
actor=default_user,
|
|
)
|
|
await server.agent_manager.update_agent_async(
|
|
agent_id=agent3.id,
|
|
agent_update=UpdateAgent(last_stop_reason=StopReasonType.end_turn),
|
|
actor=default_user,
|
|
)
|
|
|
|
agent4 = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name="agent_no_stop_reason",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
memory_blocks=[],
|
|
include_base_tools=False,
|
|
tags=["test"],
|
|
),
|
|
actor=default_user,
|
|
)
|
|
|
|
# Test count with no filters - should return total count
|
|
total_count = await server.agent_manager.count_agents_async(actor=default_user)
|
|
assert total_count >= 4
|
|
|
|
# Test count by last_stop_reason - requires_approval (inbox use case)
|
|
approval_count = await server.agent_manager.count_agents_async(
|
|
actor=default_user, last_stop_reason=StopReasonType.requires_approval.value
|
|
)
|
|
assert approval_count == 1
|
|
|
|
# Test count by last_stop_reason - error
|
|
error_count = await server.agent_manager.count_agents_async(actor=default_user, last_stop_reason=StopReasonType.error.value)
|
|
assert error_count == 1
|
|
|
|
# Test count by last_stop_reason - end_turn
|
|
completed_count = await server.agent_manager.count_agents_async(actor=default_user, last_stop_reason=StopReasonType.end_turn.value)
|
|
assert completed_count == 1
|
|
|
|
# Test count by tags
|
|
test_tag_count = await server.agent_manager.count_agents_async(actor=default_user, tags=["test"])
|
|
assert test_tag_count == 3
|
|
|
|
# Test count by tags with match_all_tags
|
|
inbox_test_count = await server.agent_manager.count_agents_async(actor=default_user, tags=["inbox", "test"], match_all_tags=True)
|
|
assert inbox_test_count == 1
|
|
|
|
# Test count by name
|
|
name_count = await server.agent_manager.count_agents_async(actor=default_user, name="agent_requires_approval")
|
|
assert name_count == 1
|
|
|
|
# Test count by query_text
|
|
query_count = await server.agent_manager.count_agents_async(actor=default_user, query_text="error")
|
|
assert query_count >= 1
|
|
|
|
# Test combined filters: last_stop_reason + tags
|
|
combined_count = await server.agent_manager.count_agents_async(
|
|
actor=default_user, last_stop_reason=StopReasonType.requires_approval.value, tags=["inbox"]
|
|
)
|
|
assert combined_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_agents_ordering_and_pagination(server: SyncServer, default_user):
|
|
names = ["alpha_agent", "beta_agent", "gamma_agent"]
|
|
created_agents = []
|
|
|
|
# Create agents in known order
|
|
for name in names:
|
|
agent = await server.agent_manager.create_agent_async(
|
|
agent_create=CreateAgent(
|
|
name=name,
|
|
agent_type="memgpt_v2_agent",
|
|
memory_blocks=[],
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
include_base_tools=False,
|
|
),
|
|
actor=default_user,
|
|
)
|
|
created_agents.append(agent)
|
|
if USING_SQLITE:
|
|
time.sleep(CREATE_DELAY_SQLITE)
|
|
|
|
agent_ids = {agent.name: agent.id for agent in created_agents}
|
|
|
|
# Ascending (oldest to newest)
|
|
agents_asc = await server.agent_manager.list_agents_async(actor=default_user, ascending=True)
|
|
asc_names = [agent.name for agent in agents_asc]
|
|
assert asc_names.index("alpha_agent") < asc_names.index("beta_agent") < asc_names.index("gamma_agent")
|
|
|
|
# Descending (newest to oldest)
|
|
agents_desc = await server.agent_manager.list_agents_async(actor=default_user, ascending=False)
|
|
desc_names = [agent.name for agent in agents_desc]
|
|
assert desc_names.index("gamma_agent") < desc_names.index("beta_agent") < desc_names.index("alpha_agent")
|
|
|
|
# After: Get agents after alpha_agent in ascending order (should exclude alpha)
|
|
after_alpha = await server.agent_manager.list_agents_async(actor=default_user, after=agent_ids["alpha_agent"], ascending=True)
|
|
after_names = [a.name for a in after_alpha]
|
|
assert "alpha_agent" not in after_names
|
|
assert "beta_agent" in after_names
|
|
assert "gamma_agent" in after_names
|
|
assert after_names == ["beta_agent", "gamma_agent"]
|
|
|
|
# Before: Get agents before gamma_agent in ascending order (should exclude gamma)
|
|
before_gamma = await server.agent_manager.list_agents_async(actor=default_user, before=agent_ids["gamma_agent"], ascending=True)
|
|
before_names = [a.name for a in before_gamma]
|
|
assert "gamma_agent" not in before_names
|
|
assert "alpha_agent" in before_names
|
|
assert "beta_agent" in before_names
|
|
assert before_names == ["alpha_agent", "beta_agent"]
|
|
|
|
# After: Get agents after gamma_agent in descending order (should exclude gamma, return beta then alpha)
|
|
after_gamma_desc = await server.agent_manager.list_agents_async(actor=default_user, after=agent_ids["gamma_agent"], ascending=False)
|
|
after_names_desc = [a.name for a in after_gamma_desc]
|
|
assert after_names_desc == ["beta_agent", "alpha_agent"]
|
|
|
|
# Before: Get agents before alpha_agent in descending order (should exclude alpha)
|
|
before_alpha_desc = await server.agent_manager.list_agents_async(actor=default_user, before=agent_ids["alpha_agent"], ascending=False)
|
|
before_names_desc = [a.name for a in before_alpha_desc]
|
|
assert before_names_desc == ["gamma_agent", "beta_agent"]
|
|
|
|
|
|
# ======================================================================================================================
|
|
# AgentManager Tests - Environment Variable Encryption
|
|
# ======================================================================================================================
|
|
|
|
|
|
@pytest.fixture
|
|
def encryption_key():
|
|
"""Fixture to ensure encryption key is set for tests."""
|
|
original_key = settings.encryption_key
|
|
# Set a test encryption key if not already set
|
|
if not settings.encryption_key:
|
|
settings.encryption_key = "test-encryption-key-32-bytes!!"
|
|
yield settings.encryption_key
|
|
# Restore original
|
|
settings.encryption_key = original_key
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_environment_variables_encrypt_on_create(server: SyncServer, default_user, encryption_key):
|
|
"""Test that creating an agent with secrets encrypts the values in the database."""
|
|
from letta.orm.sandbox_config import AgentEnvironmentVariable as AgentEnvironmentVariableModel
|
|
from letta.schemas.secret import Secret
|
|
|
|
# Create agent with secrets
|
|
agent_create = CreateAgent(
|
|
name="test-agent-with-secrets",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=DEFAULT_EMBEDDING_CONFIG,
|
|
include_base_tools=False,
|
|
secrets={
|
|
"API_KEY": "sk-test-secret-12345",
|
|
"DATABASE_URL": "postgres://user:pass@localhost/db",
|
|
},
|
|
)
|
|
|
|
created_agent = await server.agent_manager.create_agent_async(agent_create, actor=default_user)
|
|
|
|
# Verify agent has secrets
|
|
assert created_agent.secrets is not None
|
|
assert len(created_agent.secrets) == 2
|
|
|
|
# Verify secrets are AgentEnvironmentVariable objects with Secret fields
|
|
for secret_obj in created_agent.secrets:
|
|
assert secret_obj.key in ["API_KEY", "DATABASE_URL"]
|
|
assert secret_obj.value_enc is not None
|
|
assert isinstance(secret_obj.value_enc, Secret)
|
|
|
|
# Verify values are encrypted in the database
|
|
async with db_registry.async_session() as session:
|
|
env_vars = await session.execute(
|
|
select(AgentEnvironmentVariableModel).where(AgentEnvironmentVariableModel.agent_id == created_agent.id)
|
|
)
|
|
env_var_list = list(env_vars.scalars().all())
|
|
|
|
assert len(env_var_list) == 2
|
|
for env_var in env_var_list:
|
|
# Check that value_enc is not None and is encrypted
|
|
assert env_var.value_enc is not None
|
|
assert isinstance(env_var.value_enc, str)
|
|
|
|
# Decrypt and verify
|
|
decrypted = Secret.from_encrypted(env_var.value_enc).get_plaintext()
|
|
if env_var.key == "API_KEY":
|
|
assert decrypted == "sk-test-secret-12345"
|
|
elif env_var.key == "DATABASE_URL":
|
|
assert decrypted == "postgres://user:pass@localhost/db"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_environment_variables_decrypt_on_read(server: SyncServer, default_user, encryption_key):
|
|
"""Test that reading an agent deserializes secrets correctly to AgentEnvironmentVariable objects."""
|
|
from letta.schemas.environment_variables import AgentEnvironmentVariable
|
|
from letta.schemas.secret import Secret
|
|
|
|
# Create agent with secrets
|
|
agent_create = CreateAgent(
|
|
name="test-agent-read-secrets",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=DEFAULT_EMBEDDING_CONFIG,
|
|
include_base_tools=False,
|
|
secrets={
|
|
"TEST_KEY": "test-value-67890",
|
|
},
|
|
)
|
|
|
|
created_agent = await server.agent_manager.create_agent_async(agent_create, actor=default_user)
|
|
agent_id = created_agent.id
|
|
|
|
# Read the agent back
|
|
retrieved_agent = await server.agent_manager.get_agent_by_id_async(agent_id=agent_id, actor=default_user)
|
|
|
|
# Verify secrets are properly deserialized
|
|
assert retrieved_agent.secrets is not None
|
|
assert len(retrieved_agent.secrets) == 1
|
|
|
|
secret_obj = retrieved_agent.secrets[0]
|
|
assert isinstance(secret_obj, AgentEnvironmentVariable)
|
|
assert secret_obj.key == "TEST_KEY"
|
|
assert secret_obj.value == "test-value-67890"
|
|
|
|
# Verify value_enc is a Secret object (not a string)
|
|
assert secret_obj.value_enc is not None
|
|
assert isinstance(secret_obj.value_enc, Secret)
|
|
|
|
# Verify we can decrypt through the Secret object
|
|
decrypted = secret_obj.value_enc.get_plaintext()
|
|
assert decrypted == "test-value-67890"
|
|
|
|
# Verify get_value_secret() method works
|
|
value_secret = secret_obj.get_value_secret()
|
|
assert isinstance(value_secret, Secret)
|
|
assert value_secret.get_plaintext() == "test-value-67890"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_environment_variables_update_encryption(server: SyncServer, default_user, encryption_key):
|
|
"""Test that updating agent secrets encrypts new values."""
|
|
from letta.orm.sandbox_config import AgentEnvironmentVariable as AgentEnvironmentVariableModel
|
|
from letta.schemas.secret import Secret
|
|
|
|
# Create agent with initial secrets
|
|
agent_create = CreateAgent(
|
|
name="test-agent-update-secrets",
|
|
agent_type="memgpt_v2_agent",
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=DEFAULT_EMBEDDING_CONFIG,
|
|
include_base_tools=False,
|
|
secrets={
|
|
"INITIAL_KEY": "initial-value",
|
|
},
|
|
)
|
|
|
|
created_agent = await server.agent_manager.create_agent_async(agent_create, actor=default_user)
|
|
agent_id = created_agent.id
|
|
|
|
# Update with new secrets
|
|
agent_update = UpdateAgent(
|
|
secrets={
|
|
"UPDATED_KEY": "updated-value-abc",
|
|
"NEW_KEY": "new-value-xyz",
|
|
},
|
|
)
|
|
|
|
updated_agent = await server.agent_manager.update_agent_async(agent_id=agent_id, agent_update=agent_update, actor=default_user)
|
|
|
|
# Verify updated secrets
|
|
assert updated_agent.secrets is not None
|
|
assert len(updated_agent.secrets) == 2
|
|
|
|
# Verify in database
|
|
async with db_registry.async_session() as session:
|
|
env_vars = await session.execute(select(AgentEnvironmentVariableModel).where(AgentEnvironmentVariableModel.agent_id == agent_id))
|
|
env_var_list = list(env_vars.scalars().all())
|
|
|
|
assert len(env_var_list) == 2
|
|
for env_var in env_var_list:
|
|
assert env_var.value_enc is not None
|
|
|
|
# Decrypt and verify
|
|
decrypted = Secret.from_encrypted(env_var.value_enc).get_plaintext()
|
|
if env_var.key == "UPDATED_KEY":
|
|
assert decrypted == "updated-value-abc"
|
|
elif env_var.key == "NEW_KEY":
|
|
assert decrypted == "new-value-xyz"
|
|
else:
|
|
pytest.fail(f"Unexpected key: {env_var.key}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_state_schema_unchanged(server: SyncServer):
|
|
"""
|
|
Test that the AgentState pydantic schema structure has not changed.
|
|
This test validates all fields including nested pydantic objects to ensure
|
|
the schema remains stable across changes.
|
|
"""
|
|
from letta.schemas.agent import AgentState, AgentType
|
|
from letta.schemas.block import Block
|
|
from letta.schemas.embedding_config import EmbeddingConfig
|
|
from letta.schemas.environment_variables import AgentEnvironmentVariable
|
|
from letta.schemas.group import Group
|
|
from letta.schemas.llm_config import LLMConfig
|
|
from letta.schemas.memory import Memory
|
|
from letta.schemas.model import ModelSettingsUnion
|
|
from letta.schemas.response_format import ResponseFormatUnion
|
|
from letta.schemas.source import Source
|
|
from letta.schemas.tool import Tool
|
|
from letta.schemas.tool_rule import ToolRule
|
|
|
|
# Define the expected schema structure
|
|
expected_schema = {
|
|
# Core identification
|
|
"id": str,
|
|
"name": str,
|
|
# Tool rules
|
|
"tool_rules": (list, type(None)),
|
|
# In-context memory
|
|
"message_ids": (list, type(None)),
|
|
# System prompt
|
|
"system": str,
|
|
# Agent configuration
|
|
"agent_type": AgentType,
|
|
# LLM information
|
|
"llm_config": LLMConfig,
|
|
"model": str,
|
|
"embedding": str,
|
|
"embedding_config": EmbeddingConfig,
|
|
"model_settings": (ModelSettingsUnion, type(None)),
|
|
"response_format": (ResponseFormatUnion, type(None)),
|
|
# State fields
|
|
"description": (str, type(None)),
|
|
"metadata": (dict, type(None)),
|
|
# Memory and tools
|
|
"memory": Memory, # deprecated
|
|
"blocks": list,
|
|
"tools": list,
|
|
"sources": list,
|
|
"tags": list,
|
|
"tool_exec_environment_variables": list, # deprecated
|
|
"secrets": list,
|
|
# Project and template fields
|
|
"project_id": (str, type(None)),
|
|
"template_id": (str, type(None)),
|
|
"base_template_id": (str, type(None)),
|
|
"deployment_id": (str, type(None)),
|
|
"entity_id": (str, type(None)),
|
|
"identity_ids": list,
|
|
"identities": list,
|
|
# Advanced configuration
|
|
"message_buffer_autoclear": bool,
|
|
"enable_sleeptime": (bool, type(None)),
|
|
# Multi-agent
|
|
"multi_agent_group": (Group, type(None)), # deprecated
|
|
"managed_group": (Group, type(None)),
|
|
# Run metrics
|
|
"last_run_completion": (datetime, type(None)),
|
|
"last_run_duration_ms": (int, type(None)),
|
|
"last_stop_reason": (StopReasonType, type(None)),
|
|
# Timezone
|
|
"timezone": (str, type(None)),
|
|
# File controls
|
|
"max_files_open": (int, type(None)),
|
|
"per_file_view_window_char_limit": (int, type(None)),
|
|
# Indexing controls
|
|
"hidden": (bool, type(None)),
|
|
# Metadata fields (from OrmMetadataBase)
|
|
"created_by_id": (str, type(None)),
|
|
"last_updated_by_id": (str, type(None)),
|
|
"created_at": (datetime, type(None)),
|
|
"updated_at": (datetime, type(None)),
|
|
}
|
|
|
|
# Get the actual schema fields from AgentState
|
|
agent_state_fields = AgentState.model_fields
|
|
actual_field_names = set(agent_state_fields.keys())
|
|
expected_field_names = set(expected_schema.keys())
|
|
|
|
# Check for added fields
|
|
added_fields = actual_field_names - expected_field_names
|
|
if added_fields:
|
|
pytest.fail(
|
|
f"New fields detected in AgentState schema: {sorted(added_fields)}. "
|
|
"This test must be updated to include these fields, and the schema change must be intentional."
|
|
)
|
|
|
|
# Check for removed fields
|
|
removed_fields = expected_field_names - actual_field_names
|
|
if removed_fields:
|
|
pytest.fail(
|
|
f"Fields removed from AgentState schema: {sorted(removed_fields)}. "
|
|
"This test must be updated to remove these fields, and the schema change must be intentional."
|
|
)
|
|
|
|
# Validate field types
|
|
import typing
|
|
|
|
for field_name, expected_type in expected_schema.items():
|
|
field = agent_state_fields[field_name]
|
|
annotation = field.annotation
|
|
|
|
# Helper function to check if annotation matches expected type
|
|
def check_type_match(annotation, expected):
|
|
origin = typing.get_origin(annotation)
|
|
args = typing.get_args(annotation)
|
|
|
|
# Direct match
|
|
if annotation == expected:
|
|
return True
|
|
|
|
# Handle list type (List[X] should match list)
|
|
if expected is list and origin is list:
|
|
return True
|
|
|
|
# Handle dict type (Dict[X, Y] should match dict)
|
|
if expected is dict and origin is dict:
|
|
return True
|
|
|
|
# Handle Optional types
|
|
if origin is typing.Union:
|
|
# Check if expected type is in the union
|
|
if expected in args:
|
|
return True
|
|
# Handle list case within Union (e.g., Union[List[X], None])
|
|
if expected is list:
|
|
for arg in args:
|
|
if typing.get_origin(arg) is list:
|
|
return True
|
|
# Handle dict case within Union
|
|
if expected is dict:
|
|
for arg in args:
|
|
if typing.get_origin(arg) is dict:
|
|
return True
|
|
# Handle Annotated types within Union (e.g., Union[Annotated[...], None])
|
|
# This checks if any of the union args is an Annotated type that matches expected
|
|
for arg in args:
|
|
if typing.get_origin(arg) is typing.Annotated:
|
|
# For Annotated types, compare the first argument (the actual type)
|
|
annotated_args = typing.get_args(arg)
|
|
if annotated_args and annotated_args[0] == expected:
|
|
return True
|
|
|
|
return False
|
|
|
|
# Handle tuple of expected types (Optional)
|
|
if isinstance(expected_type, tuple):
|
|
valid = any(check_type_match(annotation, exp_t) for exp_t in expected_type)
|
|
if not valid:
|
|
pytest.fail(
|
|
f"Field '{field_name}' type changed. Expected one of {expected_type}, "
|
|
f"but got {annotation}. Schema changes must be intentional."
|
|
)
|
|
else:
|
|
# Single expected type
|
|
valid = check_type_match(annotation, expected_type)
|
|
if not valid:
|
|
pytest.fail(
|
|
f"Field '{field_name}' type changed. Expected {expected_type}, "
|
|
f"but got {annotation}. Schema changes must be intentional."
|
|
)
|
|
|
|
# Validate nested object schemas
|
|
# Memory schema
|
|
memory_fields = Memory.model_fields
|
|
expected_memory_fields = {"agent_type", "blocks", "file_blocks", "prompt_template"}
|
|
actual_memory_fields = set(memory_fields.keys())
|
|
if actual_memory_fields != expected_memory_fields:
|
|
pytest.fail(
|
|
f"Memory schema changed. Expected fields: {expected_memory_fields}, "
|
|
f"Got: {actual_memory_fields}. Schema changes must be intentional."
|
|
)
|
|
|
|
# Block schema
|
|
block_fields = Block.model_fields
|
|
expected_block_fields = {
|
|
"id",
|
|
"value",
|
|
"limit",
|
|
"project_id",
|
|
"template_name",
|
|
"is_template",
|
|
"template_id",
|
|
"base_template_id",
|
|
"deployment_id",
|
|
"entity_id",
|
|
"preserve_on_migration",
|
|
"label",
|
|
"read_only",
|
|
"description",
|
|
"metadata",
|
|
"hidden",
|
|
"created_by_id",
|
|
"last_updated_by_id",
|
|
}
|
|
actual_block_fields = set(block_fields.keys())
|
|
if actual_block_fields != expected_block_fields:
|
|
pytest.fail(
|
|
f"Block schema changed. Expected fields: {expected_block_fields}, "
|
|
f"Got: {actual_block_fields}. Schema changes must be intentional."
|
|
)
|
|
|
|
# Tool schema
|
|
tool_fields = Tool.model_fields
|
|
expected_tool_fields = {
|
|
"id",
|
|
"tool_type",
|
|
"description",
|
|
"source_type",
|
|
"name",
|
|
"tags",
|
|
"source_code",
|
|
"json_schema",
|
|
"args_json_schema",
|
|
"return_char_limit",
|
|
"pip_requirements",
|
|
"npm_requirements",
|
|
"default_requires_approval",
|
|
"enable_parallel_execution",
|
|
"created_by_id",
|
|
"last_updated_by_id",
|
|
"metadata_",
|
|
"project_id",
|
|
}
|
|
actual_tool_fields = set(tool_fields.keys())
|
|
if actual_tool_fields != expected_tool_fields:
|
|
pytest.fail(
|
|
f"Tool schema changed. Expected fields: {expected_tool_fields}, Got: {actual_tool_fields}. Schema changes must be intentional."
|
|
)
|
|
|
|
# Source schema
|
|
source_fields = Source.model_fields
|
|
expected_source_fields = {
|
|
"id",
|
|
"name",
|
|
"description",
|
|
"instructions",
|
|
"metadata",
|
|
"embedding_config",
|
|
"organization_id",
|
|
"vector_db_provider",
|
|
"created_by_id",
|
|
"last_updated_by_id",
|
|
"created_at",
|
|
"updated_at",
|
|
}
|
|
actual_source_fields = set(source_fields.keys())
|
|
if actual_source_fields != expected_source_fields:
|
|
pytest.fail(
|
|
f"Source schema changed. Expected fields: {expected_source_fields}, "
|
|
f"Got: {actual_source_fields}. Schema changes must be intentional."
|
|
)
|
|
|
|
# LLMConfig schema
|
|
llm_config_fields = LLMConfig.model_fields
|
|
expected_llm_config_fields = {
|
|
"model",
|
|
"display_name",
|
|
"model_endpoint_type",
|
|
"model_endpoint",
|
|
"provider_name",
|
|
"provider_category",
|
|
"model_wrapper",
|
|
"context_window",
|
|
"put_inner_thoughts_in_kwargs",
|
|
"handle",
|
|
"temperature",
|
|
"max_tokens",
|
|
"enable_reasoner",
|
|
"reasoning_effort",
|
|
"effort",
|
|
"response_format",
|
|
"max_reasoning_tokens",
|
|
"frequency_penalty",
|
|
"compatibility_type",
|
|
"verbosity",
|
|
"tier",
|
|
"parallel_tool_calls",
|
|
}
|
|
actual_llm_config_fields = set(llm_config_fields.keys())
|
|
if actual_llm_config_fields != expected_llm_config_fields:
|
|
pytest.fail(
|
|
f"LLMConfig schema changed. Expected fields: {expected_llm_config_fields}, "
|
|
f"Got: {actual_llm_config_fields}. Schema changes must be intentional."
|
|
)
|
|
|
|
# EmbeddingConfig schema
|
|
embedding_config_fields = EmbeddingConfig.model_fields
|
|
expected_embedding_config_fields = {
|
|
"embedding_endpoint_type",
|
|
"embedding_endpoint",
|
|
"embedding_model",
|
|
"embedding_dim",
|
|
"embedding_chunk_size",
|
|
"handle",
|
|
"batch_size",
|
|
"azure_endpoint",
|
|
"azure_version",
|
|
"azure_deployment",
|
|
}
|
|
actual_embedding_config_fields = set(embedding_config_fields.keys())
|
|
if actual_embedding_config_fields != expected_embedding_config_fields:
|
|
pytest.fail(
|
|
f"EmbeddingConfig schema changed. Expected fields: {expected_embedding_config_fields}, "
|
|
f"Got: {actual_embedding_config_fields}. Schema changes must be intentional."
|
|
)
|
|
|
|
# AgentEnvironmentVariable schema
|
|
agent_env_var_fields = AgentEnvironmentVariable.model_fields
|
|
expected_agent_env_var_fields = {
|
|
"id",
|
|
"key",
|
|
"value",
|
|
"description",
|
|
"organization_id",
|
|
"value_enc",
|
|
"agent_id",
|
|
# From OrmMetadataBase
|
|
"created_by_id",
|
|
"last_updated_by_id",
|
|
"created_at",
|
|
"updated_at",
|
|
}
|
|
actual_agent_env_var_fields = set(agent_env_var_fields.keys())
|
|
if actual_agent_env_var_fields != expected_agent_env_var_fields:
|
|
pytest.fail(
|
|
f"AgentEnvironmentVariable schema changed. Expected fields: {expected_agent_env_var_fields}, "
|
|
f"Got: {actual_agent_env_var_fields}. Schema changes must be intentional."
|
|
)
|
|
|
|
# Group schema
|
|
group_fields = Group.model_fields
|
|
expected_group_fields = {
|
|
"id",
|
|
"manager_type",
|
|
"agent_ids",
|
|
"description",
|
|
"project_id",
|
|
"template_id",
|
|
"base_template_id",
|
|
"deployment_id",
|
|
"shared_block_ids",
|
|
"manager_agent_id",
|
|
"termination_token",
|
|
"max_turns",
|
|
"sleeptime_agent_frequency",
|
|
"turns_counter",
|
|
"last_processed_message_id",
|
|
"max_message_buffer_length",
|
|
"min_message_buffer_length",
|
|
"hidden",
|
|
}
|
|
actual_group_fields = set(group_fields.keys())
|
|
if actual_group_fields != expected_group_fields:
|
|
pytest.fail(
|
|
f"Group schema changed. Expected fields: {expected_group_fields}, "
|
|
f"Got: {actual_group_fields}. Schema changes must be intentional."
|
|
)
|
|
|
|
|
|
async def test_agent_state_relationship_loads(server: SyncServer, default_user, print_tool, default_block):
|
|
memory_blocks = [CreateBlock(label="human", value="TestUser"), CreateBlock(label="persona", value="I am a test assistant")]
|
|
|
|
create_agent_request = CreateAgent(
|
|
name="test_default_source_agent",
|
|
agent_type="memgpt_v2_agent",
|
|
system="test system",
|
|
memory_blocks=memory_blocks,
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
tool_ids=[print_tool.id],
|
|
include_default_source=True,
|
|
include_base_tools=False,
|
|
tags=["test_tag"],
|
|
)
|
|
|
|
# Create the agent
|
|
created_agent = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
|
|
# Test legacy default include_relationships
|
|
agent_state = await server.agent_manager.get_agent_by_id_async(
|
|
agent_id=created_agent.id,
|
|
actor=default_user,
|
|
)
|
|
assert agent_state.blocks
|
|
assert agent_state.sources
|
|
assert agent_state.tags
|
|
assert agent_state.tools
|
|
|
|
# Test include_relationships override
|
|
agent_state = await server.agent_manager.get_agent_by_id_async(
|
|
agent_id=created_agent.id,
|
|
actor=default_user,
|
|
include_relationships=[],
|
|
)
|
|
assert not agent_state.blocks
|
|
assert not agent_state.sources
|
|
assert not agent_state.tags
|
|
assert not agent_state.tools
|
|
|
|
# Test include_relationships override with specific relationships
|
|
agent_state = await server.agent_manager.get_agent_by_id_async(
|
|
agent_id=created_agent.id,
|
|
actor=default_user,
|
|
include_relationships=["memory", "sources"],
|
|
)
|
|
assert agent_state.blocks
|
|
assert agent_state.sources
|
|
assert not agent_state.tags
|
|
assert not agent_state.tools
|
|
|
|
# Test include override with specific relationships
|
|
agent_state = await server.agent_manager.get_agent_by_id_async(
|
|
agent_id=created_agent.id,
|
|
actor=default_user,
|
|
include_relationships=[],
|
|
include=["agent.blocks", "agent.sources"],
|
|
)
|
|
assert agent_state.blocks
|
|
assert agent_state.sources
|
|
assert not agent_state.tags
|
|
assert not agent_state.tools
|
|
|
|
|
|
async def test_create_template_agent_with_files_from_sources(server: SyncServer, default_user, print_tool, default_block):
|
|
"""Test that agents created from templates properly attach files from their sources"""
|
|
from letta.schemas.file import FileMetadata as PydanticFileMetadata
|
|
|
|
memory_blocks = [CreateBlock(label="human", value="TestUser"), CreateBlock(label="persona", value="I am a test assistant")]
|
|
|
|
# Create a source with files
|
|
source = await server.source_manager.create_source(
|
|
source=PydanticSource(
|
|
name="test_template_source",
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
),
|
|
actor=default_user,
|
|
)
|
|
|
|
# Create files in the source
|
|
file1_metadata = PydanticFileMetadata(
|
|
file_name="template_file_1.txt",
|
|
organization_id=default_user.organization_id,
|
|
source_id=source.id,
|
|
)
|
|
file1 = await server.file_manager.create_file(file_metadata=file1_metadata, actor=default_user, text="content for file 1")
|
|
|
|
file2_metadata = PydanticFileMetadata(
|
|
file_name="template_file_2.txt",
|
|
organization_id=default_user.organization_id,
|
|
source_id=source.id,
|
|
)
|
|
file2 = await server.file_manager.create_file(file_metadata=file2_metadata, actor=default_user, text="content for file 2")
|
|
|
|
# Create agent using InternalTemplateAgentCreate with the source
|
|
create_agent_request = InternalTemplateAgentCreate(
|
|
name="test_template_agent_with_files",
|
|
agent_type="memgpt_v2_agent",
|
|
system="test system",
|
|
memory_blocks=memory_blocks,
|
|
llm_config=LLMConfig.default_config("gpt-4o-mini"),
|
|
embedding_config=EmbeddingConfig.default_config(provider="openai"),
|
|
block_ids=[default_block.id],
|
|
tool_ids=[print_tool.id],
|
|
source_ids=[source.id], # Attach the source with files
|
|
include_base_tools=False,
|
|
base_template_id="base_template_123",
|
|
template_id="template_456",
|
|
deployment_id="deployment_789",
|
|
entity_id="entity_012",
|
|
)
|
|
|
|
# Create the agent
|
|
created_agent = await server.agent_manager.create_agent_async(
|
|
create_agent_request,
|
|
actor=default_user,
|
|
)
|
|
|
|
# Verify agent was created
|
|
assert created_agent is not None
|
|
assert created_agent.name == "test_template_agent_with_files"
|
|
|
|
# Verify that the source is attached
|
|
attached_sources = await server.agent_manager.list_attached_sources_async(agent_id=created_agent.id, actor=default_user)
|
|
assert len(attached_sources) == 1
|
|
assert attached_sources[0].id == source.id
|
|
|
|
# Verify that files from the source are attached to the agent
|
|
attached_files = await server.file_agent_manager.list_files_for_agent(
|
|
created_agent.id, per_file_view_window_char_limit=created_agent.per_file_view_window_char_limit, actor=default_user
|
|
)
|
|
|
|
# Should have both files attached
|
|
assert len(attached_files) == 2
|
|
attached_file_names = {f.file_name for f in attached_files}
|
|
assert "template_file_1.txt" in attached_file_names
|
|
assert "template_file_2.txt" in attached_file_names
|
|
|
|
# Verify files are properly linked to the source
|
|
for attached_file in attached_files:
|
|
assert attached_file.source_id == source.id
|
|
|
|
# Clean up
|
|
await server.agent_manager.delete_agent_async(created_agent.id, default_user)
|
|
await server.source_manager.delete_source(source.id, default_user)
|