From f00453c1528ab5c9721c169aa19656835e8a511e Mon Sep 17 00:00:00 2001 From: cthomas Date: Mon, 8 Sep 2025 22:25:01 -0700 Subject: [PATCH] feat: migrate sleeptime to new agent loop (#4485) * feat: migrate sleeptime to new agent loop * only load agent state when needed --- letta/agents/letta_agent_v2.py | 2 +- letta/groups/sleeptime_multi_agent_v3.py | 225 ++++++++++ tests/integration_test_sleeptime_agent.py | 476 ++++++++-------------- 3 files changed, 391 insertions(+), 312 deletions(-) create mode 100644 letta/groups/sleeptime_multi_agent_v3.py diff --git a/letta/agents/letta_agent_v2.py b/letta/agents/letta_agent_v2.py index 471af76b..5beaf44d 100644 --- a/letta/agents/letta_agent_v2.py +++ b/letta/agents/letta_agent_v2.py @@ -213,7 +213,7 @@ class LettaAgentV2(BaseAgentV2): self, input_messages: list[MessageCreate], max_steps: int = DEFAULT_MAX_STEPS, - stream_tokens: bool = True, + stream_tokens: bool = False, run_id: str | None = None, use_assistant_message: bool = True, include_return_message_types: list[MessageType] | None = None, diff --git a/letta/groups/sleeptime_multi_agent_v3.py b/letta/groups/sleeptime_multi_agent_v3.py new file mode 100644 index 00000000..2e2e0da2 --- /dev/null +++ b/letta/groups/sleeptime_multi_agent_v3.py @@ -0,0 +1,225 @@ +import asyncio +from collections.abc import AsyncGenerator +from datetime import datetime, timezone + +from letta.agents.letta_agent_v2 import LettaAgentV2 +from letta.constants import DEFAULT_MAX_STEPS +from letta.groups.helpers import stringify_message +from letta.schemas.agent import AgentState +from letta.schemas.enums import JobStatus +from letta.schemas.group import Group, ManagerType +from letta.schemas.job import JobUpdate +from letta.schemas.letta_message import MessageType +from letta.schemas.letta_message_content import TextContent +from letta.schemas.letta_response import LettaResponse +from letta.schemas.message import Message, MessageCreate +from letta.schemas.run import Run +from letta.schemas.user import User +from letta.services.group_manager import GroupManager + + +class SleeptimeMultiAgentV3(LettaAgentV2): + def __init__( + self, + agent_state: AgentState, + actor: User, + group: Group, + ): + super().__init__(agent_state, actor) + assert group.manager_type == ManagerType.sleeptime, f"Expected group type to be 'sleeptime', got {group.manager_type}" + self.group = group + self.run_ids = [] + + # Additional manager classes + self.group_manager = GroupManager() + + async def step( + self, + input_messages: list[MessageCreate], + max_steps: int = DEFAULT_MAX_STEPS, + run_id: str | None = None, + use_assistant_message: bool = False, + include_return_message_types: list[MessageType] | None = None, + request_start_timestamp_ns: int | None = None, + ) -> LettaResponse: + self.run_ids = [] + + for i in range(len(input_messages)): + input_messages[i].group_id = self.group.id + + response = await super().step( + input_messages=input_messages, + max_steps=max_steps, + run_id=run_id, + use_assistant_message=use_assistant_message, + include_return_message_types=include_return_message_types, + request_start_timestamp_ns=request_start_timestamp_ns, + ) + + await self.run_sleeptime_agents(use_assistant_message=use_assistant_message) + + response.usage.run_ids = self.run_ids + return response + + async def stream( + self, + input_messages: list[MessageCreate], + max_steps: int = DEFAULT_MAX_STEPS, + stream_tokens: bool = True, + run_id: str | None = None, + use_assistant_message: bool = True, + request_start_timestamp_ns: int | None = None, + include_return_message_types: list[MessageType] | None = None, + ) -> AsyncGenerator[str, None]: + self.run_ids = [] + + for i in range(len(input_messages)): + input_messages[i].group_id = self.group.id + + # Perform foreground agent step + async for chunk in super().stream( + input_messages=input_messages, + max_steps=max_steps, + stream_tokens=stream_tokens, + run_id=run_id, + use_assistant_message=use_assistant_message, + include_return_message_types=include_return_message_types, + request_start_timestamp_ns=request_start_timestamp_ns, + ): + yield chunk + + await self.run_sleeptime_agents(use_assistant_message=use_assistant_message) + + async def run_sleeptime_agents(self, use_assistant_message: bool = True): + # Get response messages + last_response_messages = self.response_messages + + # Update turns counter + turns_counter = None + if self.group.sleeptime_agent_frequency is not None and self.group.sleeptime_agent_frequency > 0: + turns_counter = await self.group_manager.bump_turns_counter_async(group_id=self.group.id, actor=self.actor) + + # Perform participant steps + if self.group.sleeptime_agent_frequency is None or ( + turns_counter is not None and turns_counter % self.group.sleeptime_agent_frequency == 0 + ): + last_processed_message_id = await self.group_manager.get_last_processed_message_id_and_update_async( + group_id=self.group.id, last_processed_message_id=last_response_messages[-1].id, actor=self.actor + ) + for sleeptime_agent_id in self.group.agent_ids: + try: + sleeptime_run_id = await self._issue_background_task( + sleeptime_agent_id, + last_response_messages, + last_processed_message_id, + use_assistant_message, + ) + self.run_ids.append(sleeptime_run_id) + except Exception as e: + # Individual task failures + print(f"Sleeptime agent processing failed: {e!s}") + raise e + + async def _issue_background_task( + self, + sleeptime_agent_id: str, + response_messages: list[Message], + last_processed_message_id: str, + use_assistant_message: bool = True, + ) -> str: + run = Run( + user_id=self.actor.id, + status=JobStatus.created, + metadata={ + "job_type": "sleeptime_agent_send_message_async", # is this right? + "agent_id": sleeptime_agent_id, + }, + ) + run = await self.job_manager.create_job_async(pydantic_job=run, actor=self.actor) + + asyncio.create_task( + self._participant_agent_step( + foreground_agent_id=self.agent_state.id, + sleeptime_agent_id=sleeptime_agent_id, + response_messages=response_messages, + last_processed_message_id=last_processed_message_id, + run_id=run.id, + use_assistant_message=use_assistant_message, + ) + ) + return run.id + + async def _participant_agent_step( + self, + foreground_agent_id: str, + sleeptime_agent_id: str, + response_messages: list[Message], + last_processed_message_id: str, + run_id: str, + use_assistant_message: bool = True, + ) -> LettaResponse: + try: + # Update job status + job_update = JobUpdate(status=JobStatus.running) + await self.job_manager.update_job_by_id_async(job_id=run_id, job_update=job_update, actor=self.actor) + + # Create conversation transcript + prior_messages = [] + if self.group.sleeptime_agent_frequency: + try: + prior_messages = await self.message_manager.list_messages_for_agent_async( + agent_id=foreground_agent_id, + actor=self.actor, + after=last_processed_message_id, + before=response_messages[0].id, + ) + except Exception: + pass # continue with just latest messages + + transcript_summary = [stringify_message(message) for message in prior_messages + response_messages] + transcript_summary = [summary for summary in transcript_summary if summary is not None] + message_text = "\n".join(transcript_summary) + + sleeptime_agent_messages = [ + MessageCreate( + role="user", + content=[TextContent(text=message_text)], + id=Message.generate_id(), + agent_id=sleeptime_agent_id, + group_id=self.group.id, + ) + ] + + # Load sleeptime agent + sleeptime_agent_state = await self.agent_manager.get_agent_by_id_async(agent_id=sleeptime_agent_id) + sleeptime_agent = LettaAgentV2( + agent_state=sleeptime_agent_state, + actor=self.actor, + ) + + # Perform sleeptime agent step + result = await sleeptime_agent.step( + input_messages=sleeptime_agent_messages, + run_id=run_id, + use_assistant_message=use_assistant_message, + ) + + # Update job status + job_update = JobUpdate( + status=JobStatus.completed, + completed_at=datetime.now(timezone.utc).replace(tzinfo=None), + metadata={ + "result": result.model_dump(mode="json"), + "agent_id": sleeptime_agent.id, + }, + ) + await self.job_manager.update_job_by_id_async(job_id=run_id, job_update=job_update, actor=self.actor) + return result + except Exception as e: + job_update = JobUpdate( + status=JobStatus.failed, + completed_at=datetime.now(timezone.utc).replace(tzinfo=None), + metadata={"error": str(e)}, + ) + await self.job_manager.update_job_by_id_async(job_id=run_id, job_update=job_update, actor=self.actor) + raise diff --git a/tests/integration_test_sleeptime_agent.py b/tests/integration_test_sleeptime_agent.py index 5f08452b..34e993c3 100644 --- a/tests/integration_test_sleeptime_agent.py +++ b/tests/integration_test_sleeptime_agent.py @@ -1,85 +1,91 @@ +import os +import threading import time import pytest -from sqlalchemy import delete +import requests +from dotenv import load_dotenv +from letta_client import Letta +from letta_client.core.api_error import ApiError -from letta.config import LettaConfig from letta.constants import DEFAULT_HUMAN -from letta.groups.sleeptime_multi_agent_v2 import SleeptimeMultiAgentV2 -from letta.orm import Provider, ProviderTrace, Step from letta.orm.errors import NoResultFound -from letta.schemas.agent import CreateAgent from letta.schemas.block import CreateBlock from letta.schemas.enums import JobStatus, JobType, ToolRuleType -from letta.schemas.group import GroupUpdate, ManagerType, SleeptimeManagerUpdate +from letta.schemas.group import ManagerType, SleeptimeManagerUpdate from letta.schemas.message import MessageCreate from letta.schemas.run import Run -from letta.server.db import db_registry -from letta.server.server import SyncServer from letta.utils import get_human_text, get_persona_text @pytest.fixture(scope="module") -def server(): - config = LettaConfig.load() - print("CONFIG PATH", config.config_path) +def server_url() -> str: + """ + Provides the URL for the Letta server. + If LETTA_SERVER_URL is not set, starts the server in a background thread + and polls until it's accepting connections. + """ - config.save() + def _run_server() -> None: + load_dotenv() + from letta.server.rest_api.app import start_server - server = SyncServer() - return server + start_server(debug=True) + + url: str = os.getenv("LETTA_SERVER_URL", "http://localhost:8283") + + if not os.getenv("LETTA_SERVER_URL"): + thread = threading.Thread(target=_run_server, daemon=True) + thread.start() + + # Poll until the server is up (or timeout) + timeout_seconds = 30 + deadline = time.time() + timeout_seconds + while time.time() < deadline: + try: + resp = requests.get(url + "/v1/health") + if resp.status_code < 500: + break + except requests.exceptions.RequestException: + pass + time.sleep(0.1) + else: + raise RuntimeError(f"Could not reach {url} within {timeout_seconds}s") + + return url @pytest.fixture(scope="module") -def org_id(server): - org = server.organization_manager.create_default_organization() - - yield org.id - - # cleanup - with db_registry.session() as session: - session.execute(delete(ProviderTrace)) - session.execute(delete(Step)) - session.execute(delete(Provider)) - session.commit() - server.organization_manager.delete_organization_by_id(org.id) - - -@pytest.fixture(scope="module") -def actor(server, org_id): - user = server.user_manager.create_default_user() - yield user - - # cleanup - server.user_manager.delete_user_by_id(user.id) +def client(server_url: str) -> Letta: + """ + Creates and returns a synchronous Letta REST client for testing. + """ + client_instance = Letta(base_url=server_url) + yield client_instance @pytest.mark.flaky(max_runs=3) @pytest.mark.asyncio(loop_scope="module") -async def test_sleeptime_group_chat(server, actor): +async def test_sleeptime_group_chat(client): # 0. Refresh base tools - server.tool_manager.upsert_base_tools(actor=actor) + client.tools.upsert_base_tools() # 1. Create sleeptime agent - main_agent = server.create_agent( - request=CreateAgent( - name="main_agent", - memory_blocks=[ - CreateBlock( - label="persona", - value="You are a personal assistant that helps users with requests.", - ), - CreateBlock( - label="human", - value="My favorite plant is the fiddle leaf\nMy favorite color is lavender", - ), - ], - # model="openai/gpt-4o-mini", - model="anthropic/claude-3-5-sonnet-20240620", - embedding="openai/text-embedding-3-small", - enable_sleeptime=True, - ), - actor=actor, + main_agent = client.agents.create( + name="main_agent", + memory_blocks=[ + CreateBlock( + label="persona", + value="You are a personal assistant that helps users with requests.", + ), + CreateBlock( + label="human", + value="My favorite plant is the fiddle leaf\nMy favorite color is lavender", + ), + ], + model="anthropic/claude-3-5-sonnet-20240620", + embedding="openai/text-embedding-3-small", + enable_sleeptime=True, ) assert main_agent.enable_sleeptime == True @@ -89,14 +95,11 @@ async def test_sleeptime_group_chat(server, actor): assert "archival_memory_insert" not in main_agent_tools # 2. Override frequency for test - group = await server.group_manager.modify_group_async( + group = client.groups.modify( group_id=main_agent.multi_agent_group.id, - group_update=GroupUpdate( - manager_config=SleeptimeManagerUpdate( - sleeptime_agent_frequency=2, - ), + manager_config=SleeptimeManagerUpdate( + sleeptime_agent_frequency=2, ), - actor=actor, ) assert group.manager_type == ManagerType.sleeptime @@ -105,14 +108,14 @@ async def test_sleeptime_group_chat(server, actor): # 3. Verify shared blocks sleeptime_agent_id = group.agent_ids[0] - shared_block = server.agent_manager.get_block_with_label(agent_id=main_agent.id, block_label="human", actor=actor) - agents = await server.block_manager.get_agents_for_block_async(block_id=shared_block.id, actor=actor) + shared_block = client.agents.blocks.retrieve(agent_id=main_agent.id, block_label="human") + agents = client.blocks.agents.list(block_id=shared_block.id) assert len(agents) == 2 assert sleeptime_agent_id in [agent.id for agent in agents] assert main_agent.id in [agent.id for agent in agents] # 4 Verify sleeptime agent tools - sleeptime_agent = server.agent_manager.get_agent_by_id(agent_id=sleeptime_agent_id, actor=actor) + sleeptime_agent = client.agents.retrieve(agent_id=sleeptime_agent_id) sleeptime_agent_tools = [tool.name for tool in sleeptime_agent.tools] assert "memory_rethink" in sleeptime_agent_tools assert "memory_finish_edits" in sleeptime_agent_tools @@ -132,136 +135,9 @@ async def test_sleeptime_group_chat(server, actor): ] run_ids = [] for i, text in enumerate(message_text): - response = await server.send_message_to_agent( + response = client.agents.messages.create( agent_id=main_agent.id, - actor=actor, - input_messages=[ - MessageCreate( - role="user", - content=text, - ), - ], - stream_steps=False, - stream_tokens=False, - ) - - assert len(response.messages) > 0 - assert len(response.usage.run_ids or []) == (i + 1) % 2 - run_ids.extend(response.usage.run_ids or []) - - jobs = server.job_manager.list_jobs(actor=actor, job_type=JobType.RUN) - runs = [Run.from_job(job) for job in jobs] - agent_runs = [run for run in runs if "agent_id" in run.metadata and run.metadata["agent_id"] == sleeptime_agent_id] - assert len(agent_runs) == len(run_ids) - - # 6. Verify run status after sleep - time.sleep(2) - for run_id in run_ids: - job = server.job_manager.get_job_by_id(job_id=run_id, actor=actor) - assert job.status == JobStatus.running or job.status == JobStatus.completed - - # 7. Delete agent - server.agent_manager.delete_agent(agent_id=main_agent.id, actor=actor) - - with pytest.raises(NoResultFound): - server.group_manager.retrieve_group(group_id=group.id, actor=actor) - with pytest.raises(NoResultFound): - server.agent_manager.get_agent_by_id(agent_id=sleeptime_agent_id, actor=actor) - - -@pytest.mark.asyncio(loop_scope="module") -async def test_sleeptime_group_chat_v2(server, actor): - # 0. Refresh base tools - server.tool_manager.upsert_base_tools(actor=actor) - - # 1. Create sleeptime agent - main_agent = server.create_agent( - request=CreateAgent( - name="main_agent", - memory_blocks=[ - CreateBlock( - label="persona", - value="You are a personal assistant that helps users with requests.", - ), - CreateBlock( - label="human", - value="My favorite plant is the fiddle leaf\nMy favorite color is lavender", - ), - ], - # model="openai/gpt-4o-mini", - model="anthropic/claude-3-5-sonnet-20240620", - embedding="openai/text-embedding-3-small", - enable_sleeptime=True, - include_base_tool_rules=True, - ), - actor=actor, - ) - - assert main_agent.enable_sleeptime == True - main_agent_tools = [tool.name for tool in main_agent.tools] - assert "core_memory_append" not in main_agent_tools - assert "core_memory_replace" not in main_agent_tools - assert "archival_memory_insert" not in main_agent_tools - - # 2. Override frequency for test - group = await server.group_manager.modify_group_async( - group_id=main_agent.multi_agent_group.id, - group_update=GroupUpdate( - manager_config=SleeptimeManagerUpdate( - sleeptime_agent_frequency=2, - ), - ), - actor=actor, - ) - - assert group.manager_type == ManagerType.sleeptime - assert group.sleeptime_agent_frequency == 2 - assert len(group.agent_ids) == 1 - - # 3. Verify shared blocks - sleeptime_agent_id = group.agent_ids[0] - shared_block = server.agent_manager.get_block_with_label(agent_id=main_agent.id, block_label="human", actor=actor) - agents = await server.block_manager.get_agents_for_block_async(block_id=shared_block.id, actor=actor) - assert len(agents) == 2 - assert sleeptime_agent_id in [agent.id for agent in agents] - assert main_agent.id in [agent.id for agent in agents] - - # 4 Verify sleeptime agent tools - sleeptime_agent = server.agent_manager.get_agent_by_id(agent_id=sleeptime_agent_id, actor=actor) - sleeptime_agent_tools = [tool.name for tool in sleeptime_agent.tools] - assert "memory_rethink" in sleeptime_agent_tools - assert "memory_finish_edits" in sleeptime_agent_tools - assert "memory_replace" in sleeptime_agent_tools - assert "memory_insert" in sleeptime_agent_tools - - assert len([rule for rule in sleeptime_agent.tool_rules if rule.type == ToolRuleType.exit_loop]) > 0 - - # 5. Send messages and verify run ids - message_text = [ - "my favorite color is orange", - "not particularly. today is a good day", - "actually my favorite color is coral", - "let's change the subject", - "actually my fav plant is the the african spear", - "indeed", - ] - run_ids = [] - for i, text in enumerate(message_text): - agent = SleeptimeMultiAgentV2( - agent_id=main_agent.id, - message_manager=server.message_manager, - agent_manager=server.agent_manager, - block_manager=server.block_manager, - passage_manager=server.passage_manager, - group_manager=server.group_manager, - job_manager=server.job_manager, - actor=actor, - group=main_agent.multi_agent_group, - step_manager=server.step_manager, - ) - - response = await agent.step( - input_messages=[ + messages=[ MessageCreate( role="user", content=text, @@ -273,172 +149,153 @@ async def test_sleeptime_group_chat_v2(server, actor): assert len(response.usage.run_ids or []) == (i + 1) % 2 run_ids.extend(response.usage.run_ids or []) - jobs = server.job_manager.list_jobs(actor=actor, job_type=JobType.RUN) - runs = [Run.from_job(job) for job in jobs] + runs = client.runs.list() agent_runs = [run for run in runs if "agent_id" in run.metadata and run.metadata["agent_id"] == sleeptime_agent_id] assert len(agent_runs) == len(run_ids) # 6. Verify run status after sleep time.sleep(2) for run_id in run_ids: - job = server.job_manager.get_job_by_id(job_id=run_id, actor=actor) + job = client.runs.retrieve(run_id=run_id) assert job.status == JobStatus.running or job.status == JobStatus.completed # 7. Delete agent - server.agent_manager.delete_agent(agent_id=main_agent.id, actor=actor) + client.agents.delete(agent_id=main_agent.id) - with pytest.raises(NoResultFound): - server.group_manager.retrieve_group(group_id=group.id, actor=actor) - with pytest.raises(NoResultFound): - server.agent_manager.get_agent_by_id(agent_id=sleeptime_agent_id, actor=actor) + with pytest.raises(ApiError): + client.groups.retrieve(group_id=group.id) + with pytest.raises(ApiError): + client.agents.retrieve(agent_id=sleeptime_agent_id) @pytest.mark.skip @pytest.mark.asyncio(loop_scope="module") -async def test_sleeptime_removes_redundant_information(server, actor): +async def test_sleeptime_removes_redundant_information(client): # 1. set up sleep-time agent as in test_sleeptime_group_chat - server.tool_manager.upsert_base_tools(actor=actor) - main_agent = server.create_agent( - request=CreateAgent( - name="main_agent", - memory_blocks=[ - CreateBlock( - label="persona", - value="You are a personal assistant that helps users with requests.", - ), - CreateBlock( - label="human", - value="My favorite plant is the fiddle leaf\nMy favorite dog is the husky\nMy favorite plant is the fiddle leaf\nMy favorite plant is the fiddle leaf", - ), - ], - model="anthropic/claude-3-5-sonnet-20240620", - embedding="openai/text-embedding-3-small", - enable_sleeptime=True, - ), - actor=actor, + client.tools.upsert_base_tools() + main_agent = client.agents.create( + name="main_agent", + memory_blocks=[ + CreateBlock( + label="persona", + value="You are a personal assistant that helps users with requests.", + ), + CreateBlock( + label="human", + value="My favorite plant is the fiddle leaf\nMy favorite dog is the husky\nMy favorite plant is the fiddle leaf\nMy favorite plant is the fiddle leaf", + ), + ], + model="anthropic/claude-3-5-sonnet-20240620", + embedding="openai/text-embedding-3-small", + enable_sleeptime=True, ) - group = await server.group_manager.modify_group_async( + group = client.groups.modify( group_id=main_agent.multi_agent_group.id, - group_update=GroupUpdate( - manager_config=SleeptimeManagerUpdate( - sleeptime_agent_frequency=1, - ), + manager_config=SleeptimeManagerUpdate( + sleeptime_agent_frequency=1, ), - actor=actor, ) sleeptime_agent_id = group.agent_ids[0] - shared_block = server.agent_manager.get_block_with_label(agent_id=main_agent.id, block_label="human", actor=actor) + shared_block = client.agents.blocks.retrieve(agent_id=main_agent.id, block_label="human") count_before_memory_edits = shared_block.value.count("fiddle leaf") test_messages = ["hello there", "my favorite bird is the sparrow"] for test_message in test_messages: - _ = await server.send_message_to_agent( + _ = client.agents.messages.create( agent_id=main_agent.id, - actor=actor, - input_messages=[ + messages=[ MessageCreate( role="user", content=test_message, ), ], - stream_steps=False, - stream_tokens=False, ) # 2. Allow memory blocks time to update time.sleep(5) # 3. Check that the memory blocks have been collapsed - shared_block = server.agent_manager.get_block_with_label(agent_id=main_agent.id, block_label="human", actor=actor) + shared_block = client.agents.blocks.retrieve(agent_id=main_agent.id, block_label="human") count_after_memory_edits = shared_block.value.count("fiddle leaf") assert count_after_memory_edits < count_before_memory_edits # 4. Delete agent - server.agent_manager.delete_agent(agent_id=main_agent.id, actor=actor) + client.agents.delete(agent_id=main_agent.id) - with pytest.raises(NoResultFound): - server.group_manager.retrieve_group(group_id=group.id, actor=actor) - with pytest.raises(NoResultFound): - server.agent_manager.get_agent_by_id(agent_id=sleeptime_agent_id, actor=actor) + with pytest.raises(ApiError): + client.groups.retrieve(group_id=group.id) + with pytest.raises(ApiError): + client.agents.retrieve(agent_id=sleeptime_agent_id) @pytest.mark.asyncio(loop_scope="module") -async def test_sleeptime_edit(server, actor): - sleeptime_agent = server.create_agent( - request=CreateAgent( - name="sleeptime_agent", - agent_type="sleeptime_agent", - memory_blocks=[ - CreateBlock( - label="human", - value=get_human_text(DEFAULT_HUMAN), - limit=2000, - ), - CreateBlock( - label="memory_persona", - value=get_persona_text("sleeptime_memory_persona"), - limit=2000, - ), - CreateBlock( - label="fact_block", - value="""Messi resides in the Paris. - Messi plays in the league Ligue 1. - Messi plays for the team Paris Saint-Germain. - The national team Messi plays for is the Argentina team. - Messi is also known as Leo Messi - Victor Ulloa plays for Inter Miami""", - limit=2000, - ), - ], - model="anthropic/claude-3-5-sonnet-20240620", - embedding="openai/text-embedding-3-small", - enable_sleeptime=True, - ), - actor=actor, +async def test_sleeptime_edit(client): + sleeptime_agent = client.agents.create( + name="sleeptime_agent", + agent_type="sleeptime_agent", + memory_blocks=[ + CreateBlock( + label="human", + value=get_human_text(DEFAULT_HUMAN), + limit=2000, + ), + CreateBlock( + label="memory_persona", + value=get_persona_text("sleeptime_memory_persona"), + limit=2000, + ), + CreateBlock( + label="fact_block", + value="""Messi resides in the Paris. + Messi plays in the league Ligue 1. + Messi plays for the team Paris Saint-Germain. + The national team Messi plays for is the Argentina team. + Messi is also known as Leo Messi + Victor Ulloa plays for Inter Miami""", + limit=2000, + ), + ], + model="anthropic/claude-3-5-sonnet-20240620", + embedding="openai/text-embedding-3-small", + enable_sleeptime=True, ) - _ = await server.send_message_to_agent( + _ = client.agents.messages.create( agent_id=sleeptime_agent.id, - actor=actor, - input_messages=[ + messages=[ MessageCreate( role="user", content="Messi has now moved to playing for Inter Miami", ), ], - stream_steps=False, - stream_tokens=False, ) - fact_block = server.agent_manager.get_block_with_label(agent_id=sleeptime_agent.id, block_label="fact_block", actor=actor) + fact_block = client.agents.blocks.retrieve(agent_id=sleeptime_agent.id, block_label="fact_block") print(fact_block.value) assert fact_block.value.count("Inter Miami") > 1 @pytest.mark.asyncio(loop_scope="module") -async def test_sleeptime_agent_new_block_attachment(server, actor): +async def test_sleeptime_agent_new_block_attachment(client): """Test that a new block created after agent creation is properly attached to both main and sleeptime agents.""" # 0. Refresh base tools - server.tool_manager.upsert_base_tools(actor=actor) + client.tools.upsert_base_tools() # 1. Create sleeptime agent - main_agent = server.create_agent( - request=CreateAgent( - name="main_agent", - memory_blocks=[ - CreateBlock( - label="persona", - value="You are a personal assistant that helps users with requests.", - ), - CreateBlock( - label="human", - value="My favorite plant is the fiddle leaf\nMy favorite color is lavender", - ), - ], - model="anthropic/claude-3-5-sonnet-20240620", - embedding="openai/text-embedding-3-small", - enable_sleeptime=True, - ), - actor=actor, + main_agent = client.agents.create( + name="main_agent", + memory_blocks=[ + CreateBlock( + label="persona", + value="You are a personal assistant that helps users with requests.", + ), + CreateBlock( + label="human", + value="My favorite plant is the fiddle leaf\nMy favorite color is lavender", + ), + ], + model="anthropic/claude-3-5-sonnet-20240620", + embedding="openai/text-embedding-3-small", + enable_sleeptime=True, ) assert main_agent.enable_sleeptime == True @@ -448,13 +305,13 @@ async def test_sleeptime_agent_new_block_attachment(server, actor): sleeptime_agent_id = group.agent_ids[0] # 3. Verify initial shared blocks - main_agent_refreshed = server.agent_manager.get_agent_by_id(agent_id=main_agent.id, actor=actor) + main_agent_refreshed = client.agents.retrieve(agent_id=main_agent.id) initial_blocks = main_agent_refreshed.memory.blocks initial_block_count = len(initial_blocks) # Verify both agents share the initial blocks for block in initial_blocks: - agents = await server.block_manager.get_agents_for_block_async(block_id=block.id, actor=actor) + agents = client.blocks.agents.list(block_id=block.id) assert len(agents) == 2 assert sleeptime_agent_id in [agent.id for agent in agents] assert main_agent.id in [agent.id for agent in agents] @@ -462,26 +319,23 @@ async def test_sleeptime_agent_new_block_attachment(server, actor): # 4. Create a new block after agent creation from letta.schemas.block import Block as PydanticBlock - new_block = server.block_manager.create_or_update_block( - PydanticBlock( - label="preferences", - value="My favorite season is autumn\nI prefer tea over coffee", - ), - actor=actor, + new_block = client.blocks.create( + label="preferences", + value="My favorite season is autumn\nI prefer tea over coffee", ) # 5. Attach the new block to the main agent - server.agent_manager.attach_block(agent_id=main_agent.id, block_id=new_block.id, actor=actor) + client.agents.blocks.attach(agent_id=main_agent.id, block_id=new_block.id) # 6. Verify the new block is attached to the main agent - main_agent_refreshed = server.agent_manager.get_agent_by_id(agent_id=main_agent.id, actor=actor) + main_agent_refreshed = client.agents.retrieve(agent_id=main_agent.id) main_agent_blocks = main_agent_refreshed.memory.blocks assert len(main_agent_blocks) == initial_block_count + 1 main_agent_block_ids = [block.id for block in main_agent_blocks] assert new_block.id in main_agent_block_ids # 7. Check if the new block is also attached to the sleeptime agent (this is where the bug might be) - sleeptime_agent = server.agent_manager.get_agent_by_id(agent_id=sleeptime_agent_id, actor=actor) + sleeptime_agent = client.agents.retrieve(agent_id=sleeptime_agent_id) sleeptime_agent_blocks = sleeptime_agent.memory.blocks sleeptime_agent_block_ids = [block.id for block in sleeptime_agent_blocks] @@ -489,7 +343,7 @@ async def test_sleeptime_agent_new_block_attachment(server, actor): assert new_block.id in sleeptime_agent_block_ids, f"New block {new_block.id} not attached to sleeptime agent {sleeptime_agent_id}" # 8. Verify that agents sharing the new block include both main and sleeptime agents - agents_with_new_block = await server.block_manager.get_agents_for_block_async(block_id=new_block.id, actor=actor) + agents_with_new_block = client.blocks.agents.list(block_id=new_block.id) agent_ids_with_new_block = [agent.id for agent in agents_with_new_block] assert main_agent.id in agent_ids_with_new_block, "Main agent should have access to the new block" @@ -497,4 +351,4 @@ async def test_sleeptime_agent_new_block_attachment(server, actor): assert len(agents_with_new_block) == 2, "Both main and sleeptime agents should share the new block" # 9. Clean up - server.agent_manager.delete_agent(agent_id=main_agent.id, actor=actor) + client.agents.delete(agent_id=main_agent.id)