mirror of
https://github.com/cpacker/MemGPT.git
synced 2025-06-03 04:30:22 +00:00
chore: remove the admin client and tests (#1923)
This commit is contained in:
Sarah Wooders
GitHub
parent
ab3431f56a
commit
d84bf39c4b
@ -1,33 +1,19 @@
|
||||
import json
|
||||
|
||||
from letta import Admin, create_client
|
||||
from letta.memory import ChatMemory
|
||||
from letta import create_client
|
||||
from letta.schemas.memory import ChatMemory
|
||||
|
||||
"""
|
||||
Make sure you run the Letta server before running this example.
|
||||
```
|
||||
export MEMGPT_SERVER_PASS=your_token
|
||||
letta server
|
||||
```
|
||||
"""
|
||||
|
||||
|
||||
def main():
|
||||
# Create an admin client
|
||||
admin = Admin(base_url="http://localhost:8283", token="your_token")
|
||||
|
||||
# Create a user + token
|
||||
create_user_response = admin.create_user()
|
||||
user_id = create_user_response.user_id
|
||||
token = create_user_response.api_key
|
||||
print(f"Created user: {user_id} with token: {token}")
|
||||
|
||||
# List available keys
|
||||
get_keys_response = admin.get_keys(user_id=user_id)
|
||||
print(f"User {user_id} has keys: {get_keys_response}")
|
||||
|
||||
# Connect to the server as a user
|
||||
client = create_client(base_url="http://localhost:8283", token=token)
|
||||
client = create_client(base_url="http://localhost:8283")
|
||||
|
||||
# Create an agent
|
||||
agent_state = client.create_agent(name="my_agent", memory=ChatMemory(human="My name is Sarah.", persona="I am a friendly AI."))
|
||||
@ -42,10 +28,6 @@ def main():
|
||||
client.delete_agent(agent_id=agent_state.id)
|
||||
print(f"Deleted agent: {agent_state.name} with ID {str(agent_state.id)}")
|
||||
|
||||
# Delete user
|
||||
admin.delete_user(user_id=user_id)
|
||||
print(f"Deleted user: {user_id} with token: {token}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@ -11,6 +11,8 @@ from google_auth_oauthlib.flow import InstalledAppFlow
|
||||
from googleapiclient.discovery import build
|
||||
from googleapiclient.errors import HttpError
|
||||
|
||||
# NOTE: THIS file it out of date for >=0.5.0
|
||||
|
||||
# If modifying these scopes, delete the file token.json.
|
||||
SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]
|
||||
TOKEN_PATH = os.path.expanduser("~/.letta/gmail_token.json")
|
||||
|
||||
@ -1,7 +1,6 @@
|
||||
__version__ = "0.5.0"
|
||||
|
||||
# import clients
|
||||
from letta.client.admin import Admin
|
||||
from letta.client.client import LocalClient, RESTClient, create_client
|
||||
|
||||
# imports for easier access
|
||||
@ -13,7 +12,13 @@ from letta.schemas.file import FileMetadata
|
||||
from letta.schemas.job import Job
|
||||
from letta.schemas.letta_message import LettaMessage
|
||||
from letta.schemas.llm_config import LLMConfig
|
||||
from letta.schemas.memory import ArchivalMemorySummary, Memory, RecallMemorySummary
|
||||
from letta.schemas.memory import (
|
||||
ArchivalMemorySummary,
|
||||
BasicBlockMemory,
|
||||
ChatMemory,
|
||||
Memory,
|
||||
RecallMemorySummary,
|
||||
)
|
||||
from letta.schemas.message import Message
|
||||
from letta.schemas.openai.chat_completion_response import UsageStatistics
|
||||
from letta.schemas.organization import Organization
|
||||
|
||||
@ -105,7 +105,7 @@ def add_tool(
|
||||
"""Add or update a tool from a Python file."""
|
||||
from letta.client.client import create_client
|
||||
|
||||
client = create_client(base_url=os.getenv("MEMGPT_BASE_URL"), token=os.getenv("MEMGPT_SERVER_PASS"))
|
||||
client = create_client()
|
||||
|
||||
# 1. Parse the Python file
|
||||
with open(filename, "r", encoding="utf-8") as file:
|
||||
@ -145,7 +145,7 @@ def list_tools():
|
||||
"""List all available tools."""
|
||||
from letta.client.client import create_client
|
||||
|
||||
client = create_client(base_url=os.getenv("MEMGPT_BASE_URL"), token=os.getenv("MEMGPT_SERVER_PASS"))
|
||||
client = create_client()
|
||||
|
||||
tools = client.list_tools()
|
||||
for tool in tools:
|
||||
|
||||
@ -1,171 +0,0 @@
|
||||
from typing import List, Optional
|
||||
|
||||
import requests
|
||||
from requests import HTTPError
|
||||
|
||||
from letta.functions.functions import parse_source_code
|
||||
from letta.functions.schema_generator import generate_schema
|
||||
from letta.schemas.api_key import APIKey, APIKeyCreate
|
||||
from letta.schemas.organization import Organization, OrganizationCreate
|
||||
from letta.schemas.user import User, UserCreate
|
||||
|
||||
|
||||
class Admin:
|
||||
"""
|
||||
Admin client allows admin-level operations on the Letta server.
|
||||
- Creating users
|
||||
- Generating user keys
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
token: str,
|
||||
api_prefix: str = "v1",
|
||||
):
|
||||
self.base_url = base_url
|
||||
self.api_prefix = api_prefix
|
||||
self.token = token
|
||||
self.headers = {"accept": "application/json", "content-type": "application/json", "authorization": f"Bearer {token}"}
|
||||
|
||||
def get_users(self, cursor: Optional[str] = None, limit: Optional[int] = 50) -> List[User]:
|
||||
params = {}
|
||||
if cursor:
|
||||
params["cursor"] = str(cursor)
|
||||
if limit:
|
||||
params["limit"] = limit
|
||||
response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/users", params=params, headers=self.headers)
|
||||
if response.status_code != 200:
|
||||
raise HTTPError(response.json())
|
||||
return [User(**user) for user in response.json()]
|
||||
|
||||
def create_key(self, user_id: str, key_name: Optional[str] = None) -> APIKey:
|
||||
request = APIKeyCreate(user_id=user_id, name=key_name)
|
||||
response = requests.post(f"{self.base_url}/{self.api_prefix}/admin/users/keys", headers=self.headers, json=request.model_dump())
|
||||
if response.status_code != 200:
|
||||
raise HTTPError(response.json())
|
||||
return APIKey(**response.json())
|
||||
|
||||
def get_keys(self, user_id: str) -> List[APIKey]:
|
||||
params = {"user_id": str(user_id)}
|
||||
response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/users/keys", params=params, headers=self.headers)
|
||||
if response.status_code != 200:
|
||||
raise HTTPError(response.json())
|
||||
return [APIKey(**key) for key in response.json()]
|
||||
|
||||
def delete_key(self, api_key: str) -> APIKey:
|
||||
params = {"api_key": api_key}
|
||||
response = requests.delete(f"{self.base_url}/{self.api_prefix}/admin/users/keys", params=params, headers=self.headers)
|
||||
if response.status_code != 200:
|
||||
raise HTTPError(response.json())
|
||||
return APIKey(**response.json())
|
||||
|
||||
def create_user(self, name: Optional[str] = None, org_id: Optional[str] = None) -> User:
|
||||
request = UserCreate(name=name, org_id=org_id)
|
||||
response = requests.post(f"{self.base_url}/{self.api_prefix}/admin/users", headers=self.headers, json=request.model_dump())
|
||||
if response.status_code != 200:
|
||||
raise HTTPError(response.json())
|
||||
response_json = response.json()
|
||||
return User(**response_json)
|
||||
|
||||
def delete_user(self, user_id: str) -> User:
|
||||
params = {"user_id": str(user_id)}
|
||||
response = requests.delete(f"{self.base_url}/{self.api_prefix}/admin/users", params=params, headers=self.headers)
|
||||
if response.status_code != 200:
|
||||
raise HTTPError(response.json())
|
||||
return User(**response.json())
|
||||
|
||||
def create_organization(self, name: Optional[str] = None) -> Organization:
|
||||
request = OrganizationCreate(name=name)
|
||||
response = requests.post(f"{self.base_url}/{self.api_prefix}/admin/orgs", headers=self.headers, json=request.model_dump())
|
||||
if response.status_code != 200:
|
||||
raise HTTPError(response.json())
|
||||
response_json = response.json()
|
||||
return Organization(**response_json)
|
||||
|
||||
def delete_organization(self, org_id: str) -> Organization:
|
||||
params = {"org_id": str(org_id)}
|
||||
response = requests.delete(f"{self.base_url}/{self.api_prefix}/admin/orgs", params=params, headers=self.headers)
|
||||
if response.status_code != 200:
|
||||
raise HTTPError(response.json())
|
||||
return Organization(**response.json())
|
||||
|
||||
def get_organizations(self, cursor: Optional[str] = None, limit: Optional[int] = 50) -> List[Organization]:
|
||||
params = {}
|
||||
if cursor:
|
||||
params["cursor"] = str(cursor)
|
||||
if limit:
|
||||
params["limit"] = limit
|
||||
response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/orgs", params=params, headers=self.headers)
|
||||
if response.status_code != 200:
|
||||
raise HTTPError(response.json())
|
||||
return [Organization(**org) for org in response.json()]
|
||||
|
||||
def _reset_server(self):
|
||||
# DANGER: this will delete all users and keys
|
||||
# clear all state associated with users
|
||||
# TODO: clear out all agents, presets, etc.
|
||||
users = self.get_users()
|
||||
for user in users:
|
||||
keys = self.get_keys(user.id)
|
||||
for key in keys:
|
||||
self.delete_key(key.key)
|
||||
self.delete_user(user.id)
|
||||
|
||||
# tools
|
||||
def create_tool(
|
||||
self,
|
||||
func,
|
||||
name: Optional[str] = None,
|
||||
update: Optional[bool] = True, # TODO: actually use this
|
||||
tags: Optional[List[str]] = None,
|
||||
):
|
||||
"""Create a tool
|
||||
Args:
|
||||
func (callable): The function to create a tool for.
|
||||
tags (Optional[List[str]], optional): Tags for the tool. Defaults to None.
|
||||
update (bool, optional): Update the tool if it already exists. Defaults to True.
|
||||
Returns:
|
||||
Tool object
|
||||
"""
|
||||
|
||||
# TODO: check if tool already exists
|
||||
# TODO: how to load modules?
|
||||
# parse source code/schema
|
||||
source_code = parse_source_code(func)
|
||||
json_schema = generate_schema(func, name)
|
||||
source_type = "python"
|
||||
json_schema["name"]
|
||||
|
||||
if "memory" in tags:
|
||||
# special modifications to memory functions
|
||||
# self.memory -> self.memory.memory, since Agent.memory.memory needs to be modified (not BaseMemory.memory)
|
||||
source_code = source_code.replace("self.memory", "self.memory.memory")
|
||||
|
||||
# create data
|
||||
data = {"source_code": source_code, "source_type": source_type, "tags": tags, "json_schema": json_schema}
|
||||
CreateToolRequest(**data) # validate
|
||||
|
||||
# make REST request
|
||||
response = requests.post(f"{self.base_url}/{self.api_prefix}/admin/tools", json=data, headers=self.headers)
|
||||
if response.status_code != 200:
|
||||
raise ValueError(f"Failed to create tool: {response.text}")
|
||||
return ToolModel(**response.json())
|
||||
|
||||
def list_tools(self):
|
||||
response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/tools", headers=self.headers)
|
||||
return ListToolsResponse(**response.json()).tools
|
||||
|
||||
def delete_tool(self, name: str):
|
||||
response = requests.delete(f"{self.base_url}/{self.api_prefix}/admin/tools/{name}", headers=self.headers)
|
||||
if response.status_code != 200:
|
||||
raise ValueError(f"Failed to delete tool: {response.text}")
|
||||
return response.json()
|
||||
|
||||
def get_tool(self, name: str):
|
||||
response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/tools/{name}", headers=self.headers)
|
||||
if response.status_code == 404:
|
||||
return None
|
||||
elif response.status_code != 200:
|
||||
raise ValueError(f"Failed to get tool: {response.text}")
|
||||
return ToolModel(**response.json())
|
||||
@ -1777,6 +1777,19 @@ class LocalClient(AbstractClient):
|
||||
"""
|
||||
self.server.delete_agent(user_id=self.user_id, agent_id=agent_id)
|
||||
|
||||
def get_agent_by_name(self, agent_name: str, user_id: str) -> AgentState:
|
||||
"""
|
||||
Get an agent by its name
|
||||
|
||||
Args:
|
||||
agent_name (str): Name of the agent
|
||||
|
||||
Returns:
|
||||
agent_state (AgentState): State of the agent
|
||||
"""
|
||||
self.interface.clear()
|
||||
return self.server.get_agent(agent_name=agent_name, user_id=user_id, agent_id=None)
|
||||
|
||||
def get_agent(self, agent_id: str) -> AgentState:
|
||||
"""
|
||||
Get an agent's state by its ID.
|
||||
|
||||
@ -7,7 +7,7 @@ from typing import Union
|
||||
import pytest
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from letta import Admin, create_client
|
||||
from letta import create_client
|
||||
from letta.client.client import LocalClient, RESTClient
|
||||
from letta.constants import DEFAULT_PRESET
|
||||
from letta.schemas.agent import AgentState
|
||||
@ -31,10 +31,6 @@ client = None
|
||||
test_agent_state_post_message = None
|
||||
|
||||
|
||||
# admin credentials
|
||||
test_server_token = "test_server_token"
|
||||
|
||||
|
||||
def run_server():
|
||||
load_dotenv()
|
||||
|
||||
@ -58,7 +54,6 @@ def client(request):
|
||||
server_url = os.getenv("LETTA_SERVER_URL")
|
||||
if server_url is None:
|
||||
# run server in thread
|
||||
# NOTE: must set MEMGPT_SERVER_PASS enviornment variable
|
||||
server_url = "http://localhost:8283"
|
||||
print("Starting server thread")
|
||||
thread = threading.Thread(target=run_server, daemon=True)
|
||||
@ -66,10 +61,7 @@ def client(request):
|
||||
time.sleep(5)
|
||||
print("Running client tests with server:", server_url)
|
||||
# create user via admin client
|
||||
admin = Admin(server_url, test_server_token)
|
||||
user = admin.create_user() # Adjust as per your client's method
|
||||
api_key = admin.create_key(user.id)
|
||||
client = create_client(base_url=server_url, token=api_key.key) # This yields control back to the test function
|
||||
client = create_client(base_url=server_url, token=None) # This yields control back to the test function
|
||||
else:
|
||||
# use local client (no server)
|
||||
server_url = None
|
||||
@ -77,12 +69,7 @@ def client(request):
|
||||
|
||||
client.set_default_llm_config(LLMConfig.default_config("gpt-4"))
|
||||
client.set_default_embedding_config(EmbeddingConfig.default_config(provider="openai"))
|
||||
try:
|
||||
yield client
|
||||
finally:
|
||||
# cleanup user
|
||||
if server_url:
|
||||
admin.delete_user(user.id)
|
||||
yield client
|
||||
|
||||
|
||||
# Fixture for test agent
|
||||
|
||||
@ -7,7 +7,7 @@ from typing import Union
|
||||
import pytest
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from letta import Admin, create_client
|
||||
from letta import create_client
|
||||
from letta.agent import Agent
|
||||
from letta.client.client import LocalClient, RESTClient
|
||||
from letta.constants import DEFAULT_PRESET
|
||||
@ -25,10 +25,6 @@ test_agent_state_post_message = None
|
||||
test_user_id = uuid.uuid4()
|
||||
|
||||
|
||||
# admin credentials
|
||||
test_server_token = "test_server_token"
|
||||
|
||||
|
||||
def run_server():
|
||||
load_dotenv()
|
||||
|
||||
@ -53,33 +49,21 @@ def client(request):
|
||||
server_url = os.getenv("MEMGPT_SERVER_URL")
|
||||
if server_url is None:
|
||||
# run server in thread
|
||||
# NOTE: must set MEMGPT_SERVER_PASS enviornment variable
|
||||
server_url = "http://localhost:8283"
|
||||
print("Starting server thread")
|
||||
thread = threading.Thread(target=run_server, daemon=True)
|
||||
thread.start()
|
||||
time.sleep(5)
|
||||
print("Running client tests with server:", server_url)
|
||||
# create user via admin client
|
||||
admin = Admin(server_url, test_server_token)
|
||||
user = admin.create_user() # Adjust as per your client's method
|
||||
api_key = admin.create_key(user.id)
|
||||
else:
|
||||
# use local client (no server)
|
||||
assert False, "Local client not implemented"
|
||||
server_url = None
|
||||
assert False, "Local client not implemented"
|
||||
|
||||
assert server_url is not None
|
||||
assert api_key.key is not None
|
||||
client = create_client(base_url=server_url, token=api_key.key) # This yields control back to the test function
|
||||
client = create_client(base_url=server_url) # This yields control back to the test function
|
||||
client.set_default_llm_config(LLMConfig.default_config("gpt-4o-mini"))
|
||||
client.set_default_embedding_config(EmbeddingConfig.default_config(provider="openai"))
|
||||
try:
|
||||
yield client
|
||||
finally:
|
||||
# cleanup user
|
||||
if server_url:
|
||||
admin.delete_user(user.id)
|
||||
yield client
|
||||
|
||||
|
||||
# Fixture for test agent
|
||||
@ -127,37 +111,6 @@ def test_create_tool(client: Union[LocalClient, RESTClient]):
|
||||
response = client.user_message(agent_id=agent_state.id, message="hi")
|
||||
|
||||
|
||||
# TODO: add back once we fix admin client tool creation
|
||||
# def test_create_agent_tool_admin(admin_client):
|
||||
# if admin_client is None:
|
||||
# return
|
||||
#
|
||||
# def print_tool(message: str):
|
||||
# """
|
||||
# Args:
|
||||
# message (str): The message to print.
|
||||
#
|
||||
# Returns:
|
||||
# str: The message that was printed.
|
||||
#
|
||||
# """
|
||||
# print(message)
|
||||
# return message
|
||||
#
|
||||
# tools = admin_client.list_tools()
|
||||
# print(f"Original tools {[t.name for t in tools]}")
|
||||
#
|
||||
# tool = admin_client.create_tool(print_tool, tags=["extras"])
|
||||
#
|
||||
# tools = admin_client.list_tools()
|
||||
# assert tool in tools, f"Expected {tool.name} in {[t.name for t in tools]}"
|
||||
# print(f"Updated tools {[t.name for t in tools]}")
|
||||
#
|
||||
# # check tool id
|
||||
# tool = admin_client.get_tool(tool.name)
|
||||
# assert tool.user_id is None, f"Expected {tool.user_id} to be None"
|
||||
|
||||
|
||||
def test_create_agent_tool(client):
|
||||
"""Test creation of a agent tool"""
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user