MemGPT/tests/test_base_functions.py

import os
import threading

import pytest
from dotenv import load_dotenv
from letta_client import Letta

import letta.functions.function_sets.base as base_functions
from letta.config import LettaConfig
from letta.schemas.embedding_config import EmbeddingConfig
from letta.schemas.llm_config import LLMConfig
from letta.schemas.message import MessageCreate
from letta.server.server import SyncServer
from tests.test_tool_schema_parsing_files.expected_base_tool_schemas import (
    get_finish_rethinking_memory_schema,
    get_rethink_user_memory_schema,
    get_search_memory_schema,
    get_store_memories_schema,
)
from tests.utils import wait_for_server


def _run_server():
    """Starts the Letta server in a background thread."""
    load_dotenv()
    from letta.server.rest_api.app import start_server

    start_server(debug=True)


@pytest.fixture(scope="module")
def server():
    """
    Creates a SyncServer instance for testing.

    Loads and saves config to ensure proper initialization.
    """
    config = LettaConfig.load()

    config.save()

    server = SyncServer(init_with_default_org_and_user=True)
    yield server


@pytest.fixture(scope="session")
def server_url():
    """Ensures a server is running and returns its base URL."""
    url = 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()
        wait_for_server(url)

    return url


@pytest.fixture(scope="session")
def letta_client(server_url):
    """Creates a REST client for testing."""
    client = Letta(base_url=server_url)
    client.tools.upsert_base_tools()
    return client


@pytest.fixture(scope="function")
def agent_obj(letta_client, server):
    """Create a test agent that we can call functions on"""
    send_message_to_agent_and_wait_for_reply_tool_id = letta_client.tools.list(name="send_message_to_agent_and_wait_for_reply")[0].id
    agent_state = letta_client.agents.create(
        tool_ids=[send_message_to_agent_and_wait_for_reply_tool_id],
        include_base_tools=True,
        memory_blocks=[
            {
                "label": "human",
                "value": "Name: Matt",
            },
            {
                "label": "persona",
                "value": "Friendly agent",
            },
        ],
        llm_config=LLMConfig.default_config(model_name="gpt-4o-mini"),
        embedding_config=EmbeddingConfig.default_config(provider="openai"),
    )
    actor = server.user_manager.get_user_or_default()
    agent_obj = server.load_agent(agent_id=agent_state.id, actor=actor)
    yield agent_obj


def query_in_search_results(search_results, query):
    for result in search_results:
        if query.lower() in result["content"].lower():
            return True
    return False


def test_archival(agent_obj):
    """Test archival memory functions comprehensively."""
    # Test 1: Basic insertion and retrieval
    base_functions.archival_memory_insert(agent_obj, "The cat sleeps on the mat")
    base_functions.archival_memory_insert(agent_obj, "The dog plays in the park")
    base_functions.archival_memory_insert(agent_obj, "Python is a programming language")

    # Test exact text search
    results, _ = base_functions.archival_memory_search(agent_obj, "cat")
    assert query_in_search_results(results, "cat")

    # Test semantic search (should return animal-related content)
    results, _ = base_functions.archival_memory_search(agent_obj, "animal pets")
    assert query_in_search_results(results, "cat") or query_in_search_results(results, "dog")

    # Test unrelated search (should not return animal content)
    results, _ = base_functions.archival_memory_search(agent_obj, "programming computers")
    assert query_in_search_results(results, "python")

    # Test 2: Test pagination
    # Insert more items to test pagination
    for i in range(10):
        base_functions.archival_memory_insert(agent_obj, f"Test passage number {i}")

    # Get first page
    page0_results, next_page = base_functions.archival_memory_search(agent_obj, "Test passage", page=0)
    # Get second page
    page1_results, _ = base_functions.archival_memory_search(agent_obj, "Test passage", page=1, start=next_page)

    assert page0_results != page1_results
    assert query_in_search_results(page0_results, "Test passage")
    assert query_in_search_results(page1_results, "Test passage")

    # Test 3: Test complex text patterns
    base_functions.archival_memory_insert(agent_obj, "Important meeting on 2024-01-15 with John")
    base_functions.archival_memory_insert(agent_obj, "Follow-up meeting scheduled for next week")
    base_functions.archival_memory_insert(agent_obj, "Project deadline is approaching")

    # Search for meeting-related content
    results, _ = base_functions.archival_memory_search(agent_obj, "meeting schedule")
    assert query_in_search_results(results, "meeting")
    assert query_in_search_results(results, "2024-01-15") or query_in_search_results(results, "next week")

    # Test 4: Test error handling
    # Test invalid page number
    try:
        base_functions.archival_memory_search(agent_obj, "test", page="invalid")
        assert False, "Should have raised ValueError"
    except ValueError:
        pass


def test_recall(server, agent_obj, default_user):
    """Test that an agent can recall messages using a keyword via conversation search."""
    keyword = "banana"
    "".join(reversed(keyword))

    # Send messages
    for msg in ["hello", keyword, "tell me a fun fact"]:
        server.send_messages(
            actor=default_user,
            agent_id=agent_obj.agent_state.id,
            input_messages=[MessageCreate(role="user", content=msg)],
        )

    # Search memory
    result = base_functions.conversation_search(agent_obj, "banana")
    assert keyword in result


def test_get_rethink_user_memory_parsing(letta_client):
    tool = letta_client.tools.list(name="rethink_user_memory")[0]
    json_schema = tool.json_schema
    # Remove `request_heartbeat` from properties
    json_schema["parameters"]["properties"].pop("request_heartbeat", None)

    # Remove it from the required list if present
    required = json_schema["parameters"].get("required", [])
    if "request_heartbeat" in required:
        required.remove("request_heartbeat")

    assert json_schema == get_rethink_user_memory_schema()


def test_get_finish_rethinking_memory_parsing(letta_client):
    tool = letta_client.tools.list(name="finish_rethinking_memory")[0]
    json_schema = tool.json_schema
    # Remove `request_heartbeat` from properties
    json_schema["parameters"]["properties"].pop("request_heartbeat", None)

    # Remove it from the required list if present
    required = json_schema["parameters"].get("required", [])
    if "request_heartbeat" in required:
        required.remove("request_heartbeat")

    assert json_schema == get_finish_rethinking_memory_schema()


def test_store_memories_parsing(letta_client):
    tool = letta_client.tools.list(name="store_memories")[0]
    json_schema = tool.json_schema
    # Remove `request_heartbeat` from properties
    json_schema["parameters"]["properties"].pop("request_heartbeat", None)

    # Remove it from the required list if present
    required = json_schema["parameters"].get("required", [])
    if "request_heartbeat" in required:
        required.remove("request_heartbeat")
    assert json_schema == get_store_memories_schema()


def test_search_memory_parsing(letta_client):
    tool = letta_client.tools.list(name="search_memory")[0]
    json_schema = tool.json_schema
    # Remove `request_heartbeat` from properties
    json_schema["parameters"]["properties"].pop("request_heartbeat", None)

    # Remove it from the required list if present
    required = json_schema["parameters"].get("required", [])
    if "request_heartbeat" in required:
        required.remove("request_heartbeat")
    assert json_schema == get_search_memory_schema()