378 lines
14 KiB
Python
378 lines
14 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import uuid
|
|
import venv
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from dotenv import load_dotenv
|
|
from letta_client import Letta, McpTool, ToolCallMessage, ToolReturnMessage
|
|
|
|
from letta.functions.mcp_client.types import SSEServerConfig, StdioServerConfig, StreamableHTTPServerConfig
|
|
from letta.schemas.embedding_config import EmbeddingConfig
|
|
from letta.schemas.letta_message_content import TextContent
|
|
from letta.schemas.llm_config import LLMConfig
|
|
from letta.schemas.message import MessageCreate
|
|
from tests.utils import wait_for_server
|
|
|
|
|
|
def create_virtualenv_and_install_requirements(requirements_path: Path, name="venv", force_recreate=True) -> Path:
|
|
requirements_path = requirements_path.resolve()
|
|
|
|
if not requirements_path.exists():
|
|
raise FileNotFoundError(f"Requirements file not found: {requirements_path}")
|
|
if requirements_path.name != "requirements.txt":
|
|
raise ValueError(f"Expected file named 'requirements.txt', got: {requirements_path.name}")
|
|
|
|
venv_dir = requirements_path.parent / name
|
|
|
|
# Always clean up existing venv if force_recreate is True (default)
|
|
# This prevents corruption issues
|
|
if venv_dir.exists() and force_recreate:
|
|
try:
|
|
shutil.rmtree(venv_dir)
|
|
print(f"Cleaned up existing venv at {venv_dir}")
|
|
except Exception as e:
|
|
print(f"Warning: Failed to remove existing venv: {e}")
|
|
# Continue anyway, might still work
|
|
|
|
# Create fresh venv
|
|
if not venv_dir.exists():
|
|
venv.EnvBuilder(with_pip=True, clear=True).create(venv_dir)
|
|
|
|
pip_path = venv_dir / ("Scripts" if os.name == "nt" else "bin") / "pip"
|
|
|
|
# Wait a moment for venv creation to complete
|
|
for _ in range(10):
|
|
if pip_path.exists():
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
if not pip_path.exists():
|
|
raise FileNotFoundError(f"pip executable not found at: {pip_path}")
|
|
|
|
try:
|
|
# Upgrade pip first to avoid potential issues
|
|
subprocess.check_call([str(pip_path), "install", "--upgrade", "pip"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
# Install requirements
|
|
subprocess.check_call([str(pip_path), "install", "-r", str(requirements_path)])
|
|
except subprocess.CalledProcessError as exc:
|
|
# On failure, try to clean up and recreate once more
|
|
if not force_recreate: # Avoid infinite recursion
|
|
print(f"Initial pip install failed, attempting clean recreation...")
|
|
return create_virtualenv_and_install_requirements(requirements_path, name, force_recreate=False)
|
|
raise RuntimeError(f"pip install failed with exit code {exc.returncode}")
|
|
|
|
return venv_dir
|
|
|
|
|
|
# --- Server Management --- #
|
|
|
|
|
|
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
|
|
def empty_mcp_config():
|
|
path = Path(__file__).parent / "mcp_config.json"
|
|
path.write_text(json.dumps({})) # writes "{}"
|
|
|
|
return path
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def cleanup_test_venvs():
|
|
"""Fixture to clean up test virtual environments before and after tests."""
|
|
venv_path = Path(__file__).parent / "weather" / "venv"
|
|
|
|
# Clean before test (in case of previous failure)
|
|
if venv_path.exists():
|
|
try:
|
|
shutil.rmtree(venv_path)
|
|
except Exception:
|
|
pass # Ignore errors during cleanup
|
|
|
|
yield # Run the test
|
|
|
|
# Note: We don't clean after test to allow debugging if needed
|
|
# The next test run will clean it anyway
|
|
|
|
|
|
@pytest.fixture()
|
|
def server_url(empty_mcp_config):
|
|
"""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()
|
|
def client(server_url):
|
|
"""Creates a REST client for testing."""
|
|
client = Letta(base_url=server_url)
|
|
return client
|
|
|
|
|
|
@pytest.fixture()
|
|
def agent_state(client):
|
|
"""Creates an agent and ensures cleanup after tests."""
|
|
agent_state = client.agents.create(
|
|
name=f"test_compl_{str(uuid.uuid4())[5:]}",
|
|
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"),
|
|
)
|
|
yield agent_state
|
|
client.agents.delete(agent_state.id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_mcp_server(client, agent_state):
|
|
try:
|
|
mcp_server_name = "deepwiki"
|
|
server_url = "https://mcp.deepwiki.com/sse"
|
|
sse_mcp_config = SSEServerConfig(server_name=mcp_server_name, server_url=server_url)
|
|
client.tools.add_mcp_server(request=sse_mcp_config)
|
|
|
|
# Check that it's in the server mapping
|
|
mcp_server_mapping = client.tools.list_mcp_servers()
|
|
assert mcp_server_name in mcp_server_mapping
|
|
|
|
# Check tools
|
|
tools = client.tools.list_mcp_tools_by_server(mcp_server_name=mcp_server_name)
|
|
assert len(tools) > 0
|
|
assert isinstance(tools[0], McpTool)
|
|
|
|
# Test with the ask_question tool which is one of the available deepwiki tools
|
|
ask_question_tool = next((t for t in tools if t.name == "ask_question"), None)
|
|
assert ask_question_tool is not None, f"ask_question tool not found. Available tools: {[t.name for t in tools]}"
|
|
|
|
# Check that the tool is executable
|
|
letta_tool = client.tools.add_mcp_tool(mcp_server_name=mcp_server_name, mcp_tool_name=ask_question_tool.name)
|
|
|
|
tool_args = {"repoName": "facebook/react", "question": "What is React?"}
|
|
|
|
# Add to agent, have agent invoke tool
|
|
client.agents.tools.attach(agent_id=agent_state.id, tool_id=letta_tool.id)
|
|
response = client.agents.messages.create(
|
|
agent_id=agent_state.id,
|
|
messages=[
|
|
MessageCreate(
|
|
role="user",
|
|
content=[TextContent(text=f"Use the `{letta_tool.name}` tool with these arguments: {tool_args}.")],
|
|
)
|
|
],
|
|
)
|
|
seq = response.messages
|
|
calls = [m for m in seq if isinstance(m, ToolCallMessage)]
|
|
assert calls, "Expected a ToolCallMessage"
|
|
assert calls[0].tool_call.name == "ask_question"
|
|
|
|
returns = [m for m in seq if isinstance(m, ToolReturnMessage)]
|
|
assert returns, "Expected a ToolReturnMessage"
|
|
tr = returns[0]
|
|
# status field
|
|
assert tr.status == "success", f"Bad status: {tr.status}"
|
|
# Check that we got some content back
|
|
assert len(tr.tool_return.strip()) > 0, f"Expected non-empty tool return, got: {tr.tool_return}"
|
|
finally:
|
|
client.tools.delete_mcp_server(mcp_server_name=mcp_server_name)
|
|
assert mcp_server_name not in client.tools.list_mcp_servers()
|
|
|
|
|
|
def test_stdio_mcp_server(client, agent_state, server_url):
|
|
req_file = Path(__file__).parent / "weather" / "requirements.txt"
|
|
create_virtualenv_and_install_requirements(req_file, name="venv")
|
|
|
|
mcp_server_name = "weather"
|
|
command = str(Path(__file__).parent / "weather" / "venv" / "bin" / "python3")
|
|
args = [str(Path(__file__).parent / "weather" / "weather.py")]
|
|
|
|
stdio_config = StdioServerConfig(
|
|
server_name=mcp_server_name,
|
|
command=command,
|
|
args=args,
|
|
)
|
|
|
|
try:
|
|
client.tools.add_mcp_server(request=stdio_config)
|
|
|
|
servers = client.tools.list_mcp_servers()
|
|
assert mcp_server_name in servers
|
|
|
|
tools = client.tools.list_mcp_tools_by_server(mcp_server_name=mcp_server_name)
|
|
assert tools, "Expected at least one tool from the weather MCP server"
|
|
assert any(t.name == "get_alerts" for t in tools), f"Got: {[t.name for t in tools]}"
|
|
|
|
get_alerts = next(t for t in tools if t.name == "get_alerts")
|
|
|
|
letta_tool = client.tools.add_mcp_tool(
|
|
mcp_server_name=mcp_server_name,
|
|
mcp_tool_name=get_alerts.name,
|
|
)
|
|
|
|
client.agents.tools.attach(agent_id=agent_state.id, tool_id=letta_tool.id)
|
|
|
|
response = client.agents.messages.create(
|
|
agent_id=agent_state.id,
|
|
messages=[
|
|
MessageCreate(
|
|
role="user",
|
|
content=[TextContent(text=(f"Use the `{letta_tool.name}` tool with these arguments: {{'state': 'CA'}}."))],
|
|
)
|
|
],
|
|
)
|
|
|
|
calls = [m for m in response.messages if isinstance(m, ToolCallMessage) and m.tool_call.name == "get_alerts"]
|
|
assert calls, "Expected a get_alerts ToolCallMessage"
|
|
|
|
returns = [m for m in response.messages if isinstance(m, ToolReturnMessage) and m.tool_call_id == calls[0].tool_call.tool_call_id]
|
|
assert returns, "Expected a ToolReturnMessage for get_alerts"
|
|
ret = returns[0]
|
|
|
|
assert ret.status == "success", f"Unexpected status: {ret.status}"
|
|
# make sure there's at least some payload
|
|
assert len(ret.tool_return.strip()) >= 10, f"Expected at least 10 characters in tool_return, got {len(ret.tool_return.strip())}"
|
|
finally:
|
|
client.tools.delete_mcp_server(mcp_server_name=mcp_server_name)
|
|
assert mcp_server_name not in client.tools.list_mcp_servers()
|
|
|
|
|
|
# Optional OpenAI validation test for MCP-normalized schema
|
|
# Skips unless OPENAI_API_KEY is set to avoid network flakiness in CI
|
|
EXAMPLE_BAD_SCHEMA = {
|
|
"type": "object",
|
|
"properties": {
|
|
"conversation_type": {
|
|
"type": "string",
|
|
"const": "Group",
|
|
"description": "Specifies the type of conversation to be created. Must be 'Group' for this action.",
|
|
},
|
|
"message": {
|
|
"type": "object",
|
|
"additionalProperties": {}, # invalid for OpenAI: missing "type"
|
|
"description": "Initial message payload",
|
|
},
|
|
"participant_ids": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Participant IDs",
|
|
},
|
|
},
|
|
"required": ["conversation_type", "message", "participant_ids"],
|
|
"additionalProperties": False,
|
|
"$schema": "http://json-schema.org/draft-07/schema#",
|
|
}
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not os.getenv("OPENAI_API_KEY"),
|
|
reason="Requires OPENAI_API_KEY to call OpenAI for schema validation",
|
|
)
|
|
def test_openai_rejects_untyped_additional_properties_and_accepts_normalized_schema():
|
|
"""Test written to check if our extra schema validation works.
|
|
|
|
Some MCP servers will return faulty schemas that require correction, or they will brick the LLM client calls.
|
|
"""
|
|
import copy
|
|
|
|
try:
|
|
from openai import OpenAI
|
|
except Exception as e: # pragma: no cover
|
|
pytest.skip(f"openai package not available: {e}")
|
|
|
|
client = OpenAI()
|
|
|
|
def run_request_with_schema(schema: dict):
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "TWITTER_CREATE_A_NEW_DM_CONVERSATION",
|
|
"description": "Create a DM conversation",
|
|
"parameters": schema,
|
|
"strict": True,
|
|
},
|
|
}
|
|
]
|
|
|
|
return client.chat.completions.create(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
tools=tools,
|
|
)
|
|
|
|
# Bad schema should raise
|
|
with pytest.raises(Exception):
|
|
run_request_with_schema(EXAMPLE_BAD_SCHEMA)
|
|
|
|
# Normalized should succeed
|
|
normalized = copy.deepcopy(EXAMPLE_BAD_SCHEMA)
|
|
normalized["properties"]["message"]["additionalProperties"] = False
|
|
normalized["properties"]["message"]["properties"] = {"text": {"type": "string"}}
|
|
normalized["properties"]["message"]["required"] = ["text"]
|
|
resp = run_request_with_schema(normalized)
|
|
assert getattr(resp, "id", None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamable_http_mcp_server_update_schema_no_docstring_required(client, agent_state, server_url):
|
|
"""
|
|
Repro for schema-derivation-on-update error with MCP tools.
|
|
|
|
Without the fix, calling add_mcp_tool a second time for the same MCP tool
|
|
triggers a docstring-based schema derivation on a generated wrapper that has
|
|
no docstring, causing a 500. With the fix in place, updates should succeed.
|
|
"""
|
|
mcp_server_name = f"deepwiki_http_{uuid.uuid4().hex[:6]}"
|
|
mcp_url = "https://mcp.deepwiki.com/mcp"
|
|
|
|
http_mcp_config = StreamableHTTPServerConfig(server_name=mcp_server_name, server_url=mcp_url)
|
|
try:
|
|
client.tools.add_mcp_server(request=http_mcp_config)
|
|
|
|
# Ensure server is registered
|
|
servers = client.tools.list_mcp_servers()
|
|
assert mcp_server_name in servers
|
|
|
|
# Fetch available tools from server
|
|
tools = client.tools.list_mcp_tools_by_server(mcp_server_name=mcp_server_name)
|
|
assert tools, "Expected at least one tool from deepwiki streamable-http MCP server"
|
|
ask_question_tool = next((t for t in tools if t.name == "ask_question"), None)
|
|
assert ask_question_tool is not None, f"ask_question tool not found. Available: {[t.name for t in tools]}"
|
|
|
|
# Initial create
|
|
letta_tool_1 = client.tools.add_mcp_tool(mcp_server_name=mcp_server_name, mcp_tool_name=ask_question_tool.name)
|
|
assert letta_tool_1 is not None
|
|
|
|
# Update path (re-register same tool); should not attempt Python docstring schema derivation
|
|
letta_tool_2 = client.tools.add_mcp_tool(mcp_server_name=mcp_server_name, mcp_tool_name=ask_question_tool.name)
|
|
assert letta_tool_2 is not None
|
|
finally:
|
|
client.tools.delete_mcp_server(mcp_server_name=mcp_server_name)
|
|
assert mcp_server_name not in client.tools.list_mcp_servers()
|