ai-mirow-projects
/
swarms
mirror of https://github.com/kyegomez/swarms
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update agent_api_test.py

Browse Source
668^2
Kye Gomez 3 weeks ago committed by GitHub
parent 886ca11370
commit 95a7d59623
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 84 additions and 58 deletions
Show Stats Download Patch File Download Diff File

142
api/agent_api_test.py
Unescape View File

@ -3,11 +3,11 @@ from loguru import logger
import time
from typing import Dict, Optional, Tuple
from uuid import UUID
from datetime import datetime
import sys
BASE_URL = "http://localhost:8000/v1"
def check_api_server() -> bool:
"""Check if the API server is running and accessible."""
try:
@ -22,9 +22,10 @@ def check_api_server() -> bool:
logger.error(f"Error checking API server: {str(e)}")
return False
class TestSession:
"""Manages test session state and authentication."""
def __init__(self):
self.user_id: Optional[UUID] = None
self.api_key: Optional[str] = None
@ -35,16 +36,17 @@ class TestSession:
"""Get headers with authentication."""
return {"api-key": self.api_key} if self.api_key else {}
def create_test_user(session: TestSession) -> Tuple[bool, str]:
"""Create a test user and store credentials in session."""
logger.info("Creating test user")
try:
response = requests.post(
f"{BASE_URL}/users",
json={"username": f"test_user_{int(time.time())}"}
json={"username": f"test_user_{int(time.time())}"},
)
if response.status_code == 200:
data = response.json()
session.user_id = data["user_id"]
@ -58,17 +60,20 @@ def create_test_user(session: TestSession) -> Tuple[bool, str]:
logger.exception("Exception during user creation")
return False, str(e)
def create_additional_api_key(session: TestSession) -> Tuple[bool, str]:
def create_additional_api_key(
session: TestSession,
) -> Tuple[bool, str]:
"""Test creating an additional API key."""
logger.info("Creating additional API key")
try:
response = requests.post(
f"{BASE_URL}/users/{session.user_id}/api-keys",
headers=session.headers,
json={"name": "Test Key"}
json={"name": "Test Key"},
)
if response.status_code == 200:
logger.success("Created additional API key")
return True, response.json()["key"]
@ -79,7 +84,10 @@ def create_additional_api_key(session: TestSession) -> Tuple[bool, str]:
logger.exception("Exception during API key creation")
return False, str(e)
def test_create_agent(session: TestSession) -> Tuple[bool, Optional[UUID]]:
def test_create_agent(
session: TestSession,
) -> Tuple[bool, Optional[UUID]]:
"""Test creating a new agent."""
logger.info("Testing agent creation")
@ -88,16 +96,14 @@ def test_create_agent(session: TestSession) -> Tuple[bool, Optional[UUID]]:
"system_prompt": "You are a helpful assistant",
"model_name": "gpt-4",
"description": "Test agent",
"tags": ["test", "automated"]
"tags": ["test", "automated"],
}
try:
response = requests.post(
f"{BASE_URL}/agent",
headers=session.headers,
json=payload
f"{BASE_URL}/agent", headers=session.headers, json=payload
)
if response.status_code == 200:
agent_id = response.json()["agent_id"]
session.test_agents.append(agent_id)
@ -106,35 +112,40 @@ def test_create_agent(session: TestSession) -> Tuple[bool, Optional[UUID]]:
else:
logger.error(f"Failed to create agent: {response.text}")
return False, None
except Exception as e:
except Exception:
logger.exception("Exception during agent creation")
return False, None
def test_list_user_agents(session: TestSession) -> bool:
"""Test listing user's agents."""
logger.info("Testing user agent listing")
try:
response = requests.get(
f"{BASE_URL}/users/me/agents",
headers=session.headers
f"{BASE_URL}/users/me/agents", headers=session.headers
)
if response.status_code == 200:
agents = response.json()
logger.success(f"Found {len(agents)} user agents")
return True
else:
logger.error(f"Failed to list user agents: {response.text}")
logger.error(
f"Failed to list user agents: {response.text}"
)
return False
except Exception as e:
except Exception:
logger.exception("Exception during agent listing")
return False
def test_agent_operations(session: TestSession, agent_id: UUID) -> bool:
def test_agent_operations(
session: TestSession, agent_id: UUID
) -> bool:
"""Test various operations on an agent."""
logger.info(f"Testing operations for agent {agent_id}")
# Test update
try:
update_response = requests.patch(
@ -142,28 +153,33 @@ def test_agent_operations(session: TestSession, agent_id: UUID) -> bool:
headers=session.headers,
json={
"description": "Updated description",
"tags": ["test", "updated"]
}
"tags": ["test", "updated"],
},
)
if update_response.status_code != 200:
logger.error(f"Failed to update agent: {update_response.text}")
logger.error(
f"Failed to update agent: {update_response.text}"
)
return False
# Test metrics
metrics_response = requests.get(
f"{BASE_URL}/agent/{agent_id}/metrics",
headers=session.headers
headers=session.headers,
)
if metrics_response.status_code != 200:
logger.error(f"Failed to get agent metrics: {metrics_response.text}")
logger.error(
f"Failed to get agent metrics: {metrics_response.text}"
)
return False
logger.success("Successfully performed agent operations")
return True
except Exception as e:
except Exception:
logger.exception("Exception during agent operations")
return False
def test_completion(session: TestSession, agent_id: UUID) -> bool:
"""Test running a completion."""
logger.info("Testing completion")
@ -171,16 +187,16 @@ def test_completion(session: TestSession, agent_id: UUID) -> bool:
payload = {
"prompt": "What is the weather like today?",
"agent_id": agent_id,
"max_tokens": 100
"max_tokens": 100,
}
try:
response = requests.post(
f"{BASE_URL}/agent/completions",
headers=session.headers,
json=payload
json=payload,
)
if response.status_code == 200:
completion_data = response.json()
logger.success(
@ -190,26 +206,29 @@ def test_completion(session: TestSession, agent_id: UUID) -> bool:
else:
logger.error(f"Failed to get completion: {response.text}")
return False
except Exception as e:
except Exception:
logger.exception("Exception during completion")
return False
def cleanup_test_resources(session: TestSession):
"""Clean up all test resources."""
logger.info("Cleaning up test resources")
# Delete test agents
for agent_id in session.test_agents:
try:
response = requests.delete(
f"{BASE_URL}/agent/{agent_id}",
headers=session.headers
headers=session.headers,
)
if response.status_code == 200:
logger.debug(f"Deleted agent {agent_id}")
else:
logger.warning(f"Failed to delete agent {agent_id}: {response.text}")
except Exception as e:
logger.warning(
f"Failed to delete agent {agent_id}: {response.text}"
)
except Exception:
logger.exception(f"Exception deleting agent {agent_id}")
# Revoke API keys
@ -217,78 +236,85 @@ def cleanup_test_resources(session: TestSession):
try:
response = requests.get(
f"{BASE_URL}/users/{session.user_id}/api-keys",
headers=session.headers
headers=session.headers,
)
if response.status_code == 200:
for key in response.json():
try:
revoke_response = requests.delete(
f"{BASE_URL}/users/{session.user_id}/api-keys/{key['key']}",
headers=session.headers
headers=session.headers,
)
if revoke_response.status_code == 200:
logger.debug(f"Revoked API key {key['name']}")
logger.debug(
f"Revoked API key {key['name']}"
)
else:
logger.warning(f"Failed to revoke API key {key['name']}")
except Exception as e:
logger.exception(f"Exception revoking API key {key['name']}")
except Exception as e:
logger.warning(
f"Failed to revoke API key {key['name']}"
)
except Exception:
logger.exception(
f"Exception revoking API key {key['name']}"
)
except Exception:
logger.exception("Exception getting API keys for cleanup")
def run_test_workflow():
"""Run complete test workflow."""
logger.info("Starting API tests")
# Check if API server is running first
if not check_api_server():
return False
session = TestSession()
success = True
try:
# Create user
user_success, message = create_test_user(session)
if not user_success:
logger.error(f"User creation failed: {message}")
return False
# Create additional API key
key_success, key = create_additional_api_key(session)
if not key_success:
logger.error(f"API key creation failed: {key}")
return False
# Create agent
agent_success, agent_id = test_create_agent(session)
if not agent_success or not agent_id:
logger.error("Agent creation failed")
return False
# Test user agent listing
if not test_list_user_agents(session):
logger.error("Agent listing failed")
return False
# Test agent operations
if not test_agent_operations(session, agent_id):
logger.error("Agent operations failed")
return False
# Test completion
if not test_completion(session, agent_id):
logger.error("Completion test failed")
return False
logger.success("All tests completed successfully")
return True
except Exception as e:
except Exception:
logger.exception("Exception during test workflow")
return False
finally:
cleanup_test_resources(session)
if __name__ == "__main__":
success = run_test_workflow()
sys.exit(0 if success else 1)

Write Preview
Loading…
Cancel
Save