import json import logging import uuid from typing import Callable, List, Optional, Sequence, Union from letta.llm_api.helpers import unpack_inner_thoughts_from_kwargs from letta.schemas.tool_rule import BaseToolRule logging.basicConfig(level=logging.DEBUG) logger = logging.getLogger(__name__) from letta import LocalClient, RESTClient, create_client from letta.agent import Agent from letta.config import LettaConfig from letta.constants import DEFAULT_HUMAN, DEFAULT_PERSONA from letta.embeddings import embedding_model from letta.errors import InvalidInnerMonologueError, InvalidToolCallError, MissingInnerMonologueError, MissingToolCallError from letta.helpers.json_helpers import json_dumps from letta.llm_api.llm_api_tools import create from letta.llm_api.llm_client import LLMClient from letta.local_llm.constants import INNER_THOUGHTS_KWARG from letta.schemas.agent import AgentState from letta.schemas.embedding_config import EmbeddingConfig from letta.schemas.letta_message import LettaMessage, ReasoningMessage, ToolCallMessage from letta.schemas.letta_response import LettaResponse from letta.schemas.llm_config import LLMConfig from letta.schemas.memory import ChatMemory from letta.schemas.openai.chat_completion_response import ChatCompletionResponse, Choice, FunctionCall, Message from letta.utils import get_human_text, get_persona_text from tests.helpers.utils import cleanup # Generate uuid for agent name for this example namespace = uuid.NAMESPACE_DNS agent_uuid = str(uuid.uuid5(namespace, "test-endpoints-agent")) # defaults (letta hosted) EMBEDDING_CONFIG_PATH = "tests/configs/embedding_model_configs/letta-hosted.json" LLM_CONFIG_PATH = "tests/configs/llm_model_configs/letta-hosted.json" # ====================================================================================================================== # Section: Test Setup # These functions help setup the test # ====================================================================================================================== def setup_agent( client: Union[LocalClient, RESTClient], filename: str, memory_human_str: str = get_human_text(DEFAULT_HUMAN), memory_persona_str: str = get_persona_text(DEFAULT_PERSONA), tool_ids: Optional[List[str]] = None, tool_rules: Optional[List[BaseToolRule]] = None, agent_uuid: str = agent_uuid, include_base_tools: bool = True, include_base_tool_rules: bool = True, ) -> AgentState: config_data = json.load(open(filename, "r")) llm_config = LLMConfig(**config_data) embedding_config = EmbeddingConfig(**json.load(open(EMBEDDING_CONFIG_PATH))) # setup config config = LettaConfig() config.default_llm_config = llm_config config.default_embedding_config = embedding_config config.save() memory = ChatMemory(human=memory_human_str, persona=memory_persona_str) agent_state = client.create_agent( name=agent_uuid, llm_config=llm_config, embedding_config=embedding_config, memory=memory, tool_ids=tool_ids, tool_rules=tool_rules, include_base_tools=include_base_tools, include_base_tool_rules=include_base_tool_rules, ) return agent_state # ====================================================================================================================== # Section: Complex E2E Tests # These functions describe individual testing scenarios. # ====================================================================================================================== def check_first_response_is_valid_for_llm_endpoint(filename: str, validate_inner_monologue_contents: bool = True) -> ChatCompletionResponse: """ Checks that the first response is valid: 1. Contains either send_message or archival_memory_search 2. Contains valid usage of the function 3. Contains inner monologue Note: This is acting on the raw LLM response, note the usage of `create` """ client = create_client() cleanup(client=client, agent_uuid=agent_uuid) agent_state = setup_agent(client, filename) full_agent_state = client.get_agent(agent_state.id) messages = client.server.agent_manager.get_in_context_messages(agent_id=full_agent_state.id, actor=client.user) agent = Agent(agent_state=full_agent_state, interface=None, user=client.user) llm_client = LLMClient.create(llm_config=agent_state.llm_config) if llm_client: response = llm_client.send_llm_request( messages=messages, tools=[t.json_schema for t in agent.agent_state.tools], ) else: response = create( llm_config=agent_state.llm_config, user_id=str(uuid.UUID(int=1)), # dummy user_id messages=messages, functions=[t.json_schema for t in agent.agent_state.tools], ) # Basic check assert response is not None, response assert response.choices is not None, response assert len(response.choices) > 0, response assert response.choices[0] is not None, response # Select first choice choice = response.choices[0] # Ensure that the first message returns a "send_message" validator_func = ( lambda function_call: function_call.name == "send_message" or function_call.name == "archival_memory_search" or function_call.name == "core_memory_append" ) assert_contains_valid_function_call(choice.message, validator_func) # Assert that the message has an inner monologue assert_contains_correct_inner_monologue( choice, agent_state.llm_config.put_inner_thoughts_in_kwargs, validate_inner_monologue_contents=validate_inner_monologue_contents, ) return response def check_response_contains_keyword(filename: str, keyword="banana") -> LettaResponse: """ Checks that the prompted response from the LLM contains a chosen keyword Note: This is acting on the Letta response, note the usage of `user_message` """ client = create_client() cleanup(client=client, agent_uuid=agent_uuid) agent_state = setup_agent(client, filename) keyword_message = f'This is a test to see if you can see my message. If you can see my message, please respond by calling send_message using a message that includes the word "{keyword}"' response = client.user_message(agent_id=agent_state.id, message=keyword_message) # Basic checks assert_sanity_checks(response) # Make sure the message was sent assert_invoked_send_message_with_keyword(response.messages, keyword) # Make sure some inner monologue is present assert_inner_monologue_is_present_and_valid(response.messages) return response def check_agent_uses_external_tool(filename: str) -> LettaResponse: """ Checks that the LLM will use external tools if instructed Note: This is acting on the Letta response, note the usage of `user_message` """ from composio_langchain import Action # Set up client client = create_client() cleanup(client=client, agent_uuid=agent_uuid) tool = client.load_composio_tool(action=Action.GITHUB_STAR_A_REPOSITORY_FOR_THE_AUTHENTICATED_USER) # Set up persona for tool usage persona = f""" My name is Letta. I am a personal assistant who uses a tool called {tool.name} to star a desired github repo. Don’t forget - inner monologue / inner thoughts should always be different than the contents of send_message! send_message is how you communicate with the user, whereas inner thoughts are your own personal inner thoughts. """ agent_state = setup_agent(client, filename, memory_persona_str=persona, tool_ids=[tool.id]) response = client.user_message(agent_id=agent_state.id, message="Please star the repo with owner=letta-ai and repo=letta") # Basic checks assert_sanity_checks(response) # Make sure the tool was called assert_invoked_function_call(response.messages, tool.name) # Make sure some inner monologue is present assert_inner_monologue_is_present_and_valid(response.messages) return response def check_agent_recall_chat_memory(filename: str) -> LettaResponse: """ Checks that the LLM will recall the chat memory, specifically the human persona. Note: This is acting on the Letta response, note the usage of `user_message` """ # Set up client client = create_client() cleanup(client=client, agent_uuid=agent_uuid) human_name = "BananaBoy" agent_state = setup_agent(client, filename, memory_human_str=f"My name is {human_name}.") response = client.user_message( agent_id=agent_state.id, message="Repeat my name back to me. You should search in your human memory block." ) # Basic checks assert_sanity_checks(response) # Make sure my name was repeated back to me assert_invoked_send_message_with_keyword(response.messages, human_name) # Make sure some inner monologue is present assert_inner_monologue_is_present_and_valid(response.messages) return response def check_agent_archival_memory_insert(filename: str) -> LettaResponse: """ Checks that the LLM will execute an archival memory insert. Note: This is acting on the Letta response, note the usage of `user_message` """ # Set up client client = create_client() cleanup(client=client, agent_uuid=agent_uuid) agent_state = setup_agent(client, filename) secret_word = "banana" response = client.user_message( agent_id=agent_state.id, message=f"Please insert the secret word '{secret_word}' into archival memory.", ) # Basic checks assert_sanity_checks(response) # Make sure archival_memory_search was called assert_invoked_function_call(response.messages, "archival_memory_insert") # Make sure some inner monologue is present assert_inner_monologue_is_present_and_valid(response.messages) return response def check_agent_archival_memory_retrieval(filename: str) -> LettaResponse: """ Checks that the LLM will execute an archival memory retrieval. Note: This is acting on the Letta response, note the usage of `user_message` """ # Set up client client = create_client() cleanup(client=client, agent_uuid=agent_uuid) agent_state = setup_agent(client, filename) secret_word = "banana" client.insert_archival_memory(agent_state.id, f"The secret word is {secret_word}!") response = client.user_message( agent_id=agent_state.id, message="Search archival memory for the secret word. If you find it successfully, you MUST respond by using the `send_message` function with a message that includes the secret word so I know you found it.", ) # Basic checks assert_sanity_checks(response) # Make sure archival_memory_search was called assert_invoked_function_call(response.messages, "archival_memory_search") # Make sure secret was repeated back to me assert_invoked_send_message_with_keyword(response.messages, secret_word) # Make sure some inner monologue is present assert_inner_monologue_is_present_and_valid(response.messages) return response def check_agent_edit_core_memory(filename: str) -> LettaResponse: """ Checks that the LLM is able to edit its core memories Note: This is acting on the Letta response, note the usage of `user_message` """ # Set up client client = create_client() cleanup(client=client, agent_uuid=agent_uuid) human_name_a = "AngryAardvark" human_name_b = "BananaBoy" agent_state = setup_agent(client, filename, memory_human_str=f"My name is {human_name_a}") client.user_message(agent_id=agent_state.id, message=f"Actually, my name changed. It is now {human_name_b}") response = client.user_message(agent_id=agent_state.id, message="Repeat my name back to me.") # Basic checks assert_sanity_checks(response) # Make sure my name was repeated back to me assert_invoked_send_message_with_keyword(response.messages, human_name_b) # Make sure some inner monologue is present assert_inner_monologue_is_present_and_valid(response.messages) return response def check_agent_summarize_memory_simple(filename: str) -> LettaResponse: """ Checks that the LLM is able to summarize its memory """ # Set up client client = create_client() cleanup(client=client, agent_uuid=agent_uuid) agent_state = setup_agent(client, filename) # Send a couple messages friend_name = "Shub" client.user_message(agent_id=agent_state.id, message="Hey, how's it going? What do you think about this whole shindig") client.user_message(agent_id=agent_state.id, message=f"By the way, my friend's name is {friend_name}!") client.user_message(agent_id=agent_state.id, message="Does the number 42 ring a bell?") # Summarize agent = client.server.load_agent(agent_id=agent_state.id, actor=client.user) agent.summarize_messages_inplace() print(f"Summarization succeeded: messages[1] = \n\n{json_dumps(agent.messages[1])}\n") response = client.user_message(agent_id=agent_state.id, message="What is my friend's name?") # Basic checks assert_sanity_checks(response) # Make sure my name was repeated back to me assert_invoked_send_message_with_keyword(response.messages, friend_name) # Make sure some inner monologue is present assert_inner_monologue_is_present_and_valid(response.messages) return response def run_embedding_endpoint(filename): # load JSON file config_data = json.load(open(filename, "r")) print(config_data) embedding_config = EmbeddingConfig(**config_data) model = embedding_model(embedding_config) query_text = "hello" query_vec = model.get_text_embedding(query_text) print("vector dim", len(query_vec)) assert query_vec is not None # ====================================================================================================================== # Section: Letta Message Assertions # These functions are validating elements of parsed Letta Messsage # ====================================================================================================================== def assert_sanity_checks(response: LettaResponse): assert response is not None, response assert response.messages is not None, response assert len(response.messages) > 0, response def assert_invoked_send_message_with_keyword(messages: Sequence[LettaMessage], keyword: str, case_sensitive: bool = False) -> None: # Find first instance of send_message target_message = None for message in messages: if isinstance(message, ToolCallMessage) and message.tool_call.name == "send_message": target_message = message break # No messages found with `send_messages` if target_message is None: raise MissingToolCallError(messages=messages, explanation="Missing `send_message` function call") send_message_function_call = target_message.tool_call try: arguments = json.loads(send_message_function_call.arguments) except: raise InvalidToolCallError(messages=[target_message], explanation="Function call arguments could not be loaded into JSON") # Message field not in send_message if "message" not in arguments: raise InvalidToolCallError( messages=[target_message], explanation=f"send_message function call does not have required field `message`" ) # Check that the keyword is in the message arguments if not case_sensitive: keyword = keyword.lower() arguments["message"] = arguments["message"].lower() if not keyword in arguments["message"]: raise InvalidToolCallError(messages=[target_message], explanation=f"Message argument did not contain keyword={keyword}") def assert_invoked_function_call(messages: Sequence[LettaMessage], function_name: str) -> None: for message in messages: if isinstance(message, ToolCallMessage) and message.tool_call.name == function_name: # Found it, do nothing return raise MissingToolCallError(messages=messages, explanation=f"No messages were found invoking function call with name: {function_name}") def assert_inner_monologue_is_present_and_valid(messages: List[LettaMessage]) -> None: for message in messages: if isinstance(message, ReasoningMessage): # Found it, do nothing return raise MissingInnerMonologueError(messages=messages) # ====================================================================================================================== # Section: Raw API Assertions # These functions are validating elements of the (close to) raw LLM API's response # ====================================================================================================================== def assert_contains_valid_function_call( message: Message, function_call_validator: Optional[Callable[[FunctionCall], bool]] = None, validation_failure_summary: Optional[str] = None, ) -> None: """ Helper function to check that a message contains a valid function call. There is an Optional parameter `function_call_validator` that specifies a validator function. This function gets called on the resulting function_call to validate the function is what we expect. """ if (hasattr(message, "function_call") and message.function_call is not None) and ( hasattr(message, "tool_calls") and message.tool_calls is not None ): raise InvalidToolCallError(messages=[message], explanation="Both function_call and tool_calls is present in the message") elif hasattr(message, "function_call") and message.function_call is not None: function_call = message.function_call elif hasattr(message, "tool_calls") and message.tool_calls is not None: # Note: We only take the first one for now. Is this a problem? @charles # This seems to be standard across the repo function_call = message.tool_calls[0].function else: # Throw a missing function call error raise MissingToolCallError(messages=[message]) if function_call_validator and not function_call_validator(function_call): raise InvalidToolCallError(messages=[message], explanation=validation_failure_summary) def assert_inner_monologue_is_valid(message: Message) -> None: """ Helper function to check that the inner monologue is valid. """ # Sometimes the syntax won't be correct and internal syntax will leak into message invalid_phrases = ["functions", "send_message", "arguments"] monologue = message.content for phrase in invalid_phrases: if phrase in monologue: raise InvalidInnerMonologueError(messages=[message], explanation=f"{phrase} is in monologue") def assert_contains_correct_inner_monologue( choice: Choice, inner_thoughts_in_kwargs: bool, validate_inner_monologue_contents: bool = True, ) -> None: """ Helper function to check that the inner monologue exists and is valid. """ # Unpack inner thoughts out of function kwargs, and repackage into choice if inner_thoughts_in_kwargs: choice = unpack_inner_thoughts_from_kwargs(choice, INNER_THOUGHTS_KWARG) message = choice.message monologue = message.content if not monologue or monologue is None or monologue == "": raise MissingInnerMonologueError(messages=[message]) if validate_inner_monologue_contents: assert_inner_monologue_is_valid(message)