AutoGPT/autogpt_platform/backend/test/test_data_creator.py

437 lines
15 KiB
Python

import asyncio
import random
from datetime import datetime
import prisma.enums
from faker import Faker
from prisma import Prisma
faker = Faker()
# Constants for data generation limits
# Base entities
NUM_USERS = 100 # Creates 100 user records
NUM_AGENT_BLOCKS = 100 # Creates 100 agent block templates
# Per-user entities
MIN_GRAPHS_PER_USER = 1 # Each user will have between 1-5 graphs
MAX_GRAPHS_PER_USER = 5 # Total graphs: 500-2500 (NUM_USERS * MIN/MAX_GRAPHS)
# Per-graph entities
MIN_NODES_PER_GRAPH = 2 # Each graph will have between 2-5 nodes
MAX_NODES_PER_GRAPH = (
5 # Total nodes: 1000-2500 (GRAPHS_PER_USER * NUM_USERS * MIN/MAX_NODES)
)
# Additional per-user entities
MIN_PRESETS_PER_USER = 1 # Each user will have between 1-2 presets
MAX_PRESETS_PER_USER = 5 # Total presets: 500-2500 (NUM_USERS * MIN/MAX_PRESETS)
MIN_AGENTS_PER_USER = 1 # Each user will have between 1-2 agents
MAX_AGENTS_PER_USER = 10 # Total agents: 500-5000 (NUM_USERS * MIN/MAX_AGENTS)
# Execution and review records
MIN_EXECUTIONS_PER_GRAPH = 1 # Each graph will have between 1-5 execution records
MAX_EXECUTIONS_PER_GRAPH = (
20 # Total executions: 1000-5000 (TOTAL_GRAPHS * MIN/MAX_EXECUTIONS)
)
MIN_REVIEWS_PER_VERSION = 1 # Each version will have between 1-3 reviews
MAX_REVIEWS_PER_VERSION = 5 # Total reviews depends on number of versions created
def get_image():
url = faker.image_url()
while "placekitten.com" in url:
url = faker.image_url()
return url
async def main():
db = Prisma()
await db.connect()
# Insert Users
print(f"Inserting {NUM_USERS} users")
users = []
for _ in range(NUM_USERS):
user = await db.user.create(
data={
"id": str(faker.uuid4()),
"email": faker.unique.email(),
"name": faker.name(),
"metadata": prisma.Json({}),
"integrations": "",
}
)
users.append(user)
# Insert AgentBlocks
agent_blocks = []
print(f"Inserting {NUM_AGENT_BLOCKS} agent blocks")
for _ in range(NUM_AGENT_BLOCKS):
block = await db.agentblock.create(
data={
"name": f"{faker.word()}_{str(faker.uuid4())[:8]}",
"inputSchema": "{}",
"outputSchema": "{}",
}
)
agent_blocks.append(block)
# Insert AgentGraphs
agent_graphs = []
print(f"Inserting {NUM_USERS * MAX_GRAPHS_PER_USER} agent graphs")
for user in users:
for _ in range(
random.randint(MIN_GRAPHS_PER_USER, MAX_GRAPHS_PER_USER)
): # Adjust the range to create more graphs per user if desired
graph = await db.agentgraph.create(
data={
"name": faker.sentence(nb_words=3),
"description": faker.text(max_nb_chars=200),
"userId": user.id,
"isActive": True,
"isTemplate": False,
}
)
agent_graphs.append(graph)
# Insert AgentNodes
agent_nodes = []
print(
f"Inserting {NUM_USERS * MAX_GRAPHS_PER_USER * MAX_NODES_PER_GRAPH} agent nodes"
)
for graph in agent_graphs:
num_nodes = random.randint(MIN_NODES_PER_GRAPH, MAX_NODES_PER_GRAPH)
for _ in range(num_nodes): # Create 5 AgentNodes per graph
block = random.choice(agent_blocks)
node = await db.agentnode.create(
data={
"agentBlockId": block.id,
"agentGraphId": graph.id,
"agentGraphVersion": graph.version,
"constantInput": "{}",
"metadata": "{}",
}
)
agent_nodes.append(node)
# Insert AgentPresets
agent_presets = []
print(f"Inserting {NUM_USERS * MAX_PRESETS_PER_USER} agent presets")
for user in users:
num_presets = random.randint(MIN_PRESETS_PER_USER, MAX_PRESETS_PER_USER)
for _ in range(num_presets): # Create 1 AgentPreset per user
graph = random.choice(agent_graphs)
preset = await db.agentpreset.create(
data={
"name": faker.sentence(nb_words=3),
"description": faker.text(max_nb_chars=200),
"userId": user.id,
"agentId": graph.id,
"agentVersion": graph.version,
"isActive": True,
}
)
agent_presets.append(preset)
# Insert UserAgents
user_agents = []
print(f"Inserting {NUM_USERS * MAX_AGENTS_PER_USER} user agents")
for user in users:
num_agents = random.randint(MIN_AGENTS_PER_USER, MAX_AGENTS_PER_USER)
for _ in range(num_agents): # Create 1 LibraryAgent per user
graph = random.choice(agent_graphs)
preset = random.choice(agent_presets)
user_agent = await db.libraryagent.create(
data={
"userId": user.id,
"agentId": graph.id,
"agentVersion": graph.version,
"agentPresetId": preset.id,
"isFavorite": random.choice([True, False]),
"isCreatedByUser": random.choice([True, False]),
"isArchived": random.choice([True, False]),
"isDeleted": random.choice([True, False]),
}
)
user_agents.append(user_agent)
# Insert AgentGraphExecutions
# Insert AgentGraphExecutions
agent_graph_executions = []
print(
f"Inserting {NUM_USERS * MAX_GRAPHS_PER_USER * MAX_EXECUTIONS_PER_GRAPH} agent graph executions"
)
graph_execution_data = []
for graph in agent_graphs:
user = random.choice(users)
num_executions = random.randint(
MIN_EXECUTIONS_PER_GRAPH, MAX_EXECUTIONS_PER_GRAPH
)
for _ in range(num_executions):
matching_presets = [p for p in agent_presets if p.agentId == graph.id]
preset = (
random.choice(matching_presets)
if matching_presets and random.random() < 0.5
else None
)
graph_execution_data.append(
{
"agentGraphId": graph.id,
"agentGraphVersion": graph.version,
"userId": user.id,
"executionStatus": prisma.enums.AgentExecutionStatus.COMPLETED,
"startedAt": faker.date_time_this_year(),
"agentPresetId": preset.id if preset else None,
}
)
agent_graph_executions = await db.agentgraphexecution.create_many(
data=graph_execution_data
)
# Need to fetch the created records since create_many doesn't return them
agent_graph_executions = await db.agentgraphexecution.find_many()
# Insert AgentNodeExecutions
print(
f"Inserting {NUM_USERS * MAX_GRAPHS_PER_USER * MAX_EXECUTIONS_PER_GRAPH} agent node executions"
)
node_execution_data = []
for execution in agent_graph_executions:
nodes = [
node for node in agent_nodes if node.agentGraphId == execution.agentGraphId
]
for node in nodes:
node_execution_data.append(
{
"agentGraphExecutionId": execution.id,
"agentNodeId": node.id,
"executionStatus": prisma.enums.AgentExecutionStatus.COMPLETED,
"addedTime": datetime.now(),
}
)
agent_node_executions = await db.agentnodeexecution.create_many(
data=node_execution_data
)
# Need to fetch the created records since create_many doesn't return them
agent_node_executions = await db.agentnodeexecution.find_many()
# Insert AgentNodeExecutionInputOutput
print(
f"Inserting {NUM_USERS * MAX_GRAPHS_PER_USER * MAX_EXECUTIONS_PER_GRAPH} agent node execution input/outputs"
)
input_output_data = []
for node_execution in agent_node_executions:
# Input data
input_output_data.append(
{
"name": "input1",
"data": "{}",
"time": datetime.now(),
"referencedByInputExecId": node_execution.id,
}
)
# Output data
input_output_data.append(
{
"name": "output1",
"data": "{}",
"time": datetime.now(),
"referencedByOutputExecId": node_execution.id,
}
)
await db.agentnodeexecutioninputoutput.create_many(data=input_output_data)
# Insert AgentNodeLinks
print(f"Inserting {NUM_USERS * MAX_GRAPHS_PER_USER} agent node links")
for graph in agent_graphs:
nodes = [node for node in agent_nodes if node.agentGraphId == graph.id]
if len(nodes) >= 2:
source_node = nodes[0]
sink_node = nodes[1]
await db.agentnodelink.create(
data={
"agentNodeSourceId": source_node.id,
"sourceName": "output1",
"agentNodeSinkId": sink_node.id,
"sinkName": "input1",
"isStatic": False,
}
)
# Insert AnalyticsDetails
print(f"Inserting {NUM_USERS} analytics details")
for user in users:
for _ in range(1):
await db.analyticsdetails.create(
data={
"userId": user.id,
"type": faker.word(),
"data": prisma.Json({}),
"dataIndex": faker.word(),
}
)
# Insert AnalyticsMetrics
print(f"Inserting {NUM_USERS} analytics metrics")
for user in users:
for _ in range(1):
await db.analyticsmetrics.create(
data={
"userId": user.id,
"analyticMetric": faker.word(),
"value": random.uniform(0, 100),
"dataString": faker.word(),
}
)
# Insert CreditTransaction (formerly UserBlockCredit)
print(f"Inserting {NUM_USERS} credit transactions")
for user in users:
for _ in range(1):
block = random.choice(agent_blocks)
await db.credittransaction.create(
data={
"transactionKey": str(faker.uuid4()),
"userId": user.id,
"blockId": block.id,
"amount": random.randint(1, 100),
"type": (
prisma.enums.CreditTransactionType.TOP_UP
if random.random() < 0.5
else prisma.enums.CreditTransactionType.USAGE
),
"metadata": prisma.Json({}),
}
)
# Insert Profiles
profiles = []
print(f"Inserting {NUM_USERS} profiles")
for user in users:
profile = await db.profile.create(
data={
"userId": user.id,
"name": user.name or faker.name(),
"username": faker.unique.user_name(),
"description": faker.text(),
"links": [faker.url() for _ in range(3)],
"avatarUrl": get_image(),
}
)
profiles.append(profile)
# Insert StoreListings
store_listings = []
print(f"Inserting {NUM_USERS} store listings")
for graph in agent_graphs:
user = random.choice(users)
listing = await db.storelisting.create(
data={
"agentId": graph.id,
"agentVersion": graph.version,
"owningUserId": user.id,
"isApproved": random.choice([True, False]),
}
)
store_listings.append(listing)
# Insert StoreListingVersions
store_listing_versions = []
print(f"Inserting {NUM_USERS} store listing versions")
for listing in store_listings:
graph = [g for g in agent_graphs if g.id == listing.agentId][0]
version = await db.storelistingversion.create(
data={
"agentId": graph.id,
"agentVersion": graph.version,
"slug": faker.slug(),
"name": graph.name or faker.sentence(nb_words=3),
"subHeading": faker.sentence(),
"videoUrl": faker.url(),
"imageUrls": [get_image() for _ in range(3)],
"description": faker.text(),
"categories": [faker.word() for _ in range(3)],
"isFeatured": random.choice([True, False]),
"isAvailable": True,
"isApproved": random.choice([True, False]),
"storeListingId": listing.id,
}
)
store_listing_versions.append(version)
# Insert StoreListingReviews
print(f"Inserting {NUM_USERS * MAX_REVIEWS_PER_VERSION} store listing reviews")
for version in store_listing_versions:
# Create a copy of users list and shuffle it to avoid duplicates
available_reviewers = users.copy()
random.shuffle(available_reviewers)
# Limit number of reviews to available unique reviewers
num_reviews = min(
random.randint(MIN_REVIEWS_PER_VERSION, MAX_REVIEWS_PER_VERSION),
len(available_reviewers),
)
# Take only the first num_reviews reviewers
for reviewer in available_reviewers[:num_reviews]:
await db.storelistingreview.create(
data={
"storeListingVersionId": version.id,
"reviewByUserId": reviewer.id,
"score": random.randint(1, 5),
"comments": faker.text(),
}
)
# Insert StoreListingSubmissions
print(f"Inserting {NUM_USERS} store listing submissions")
for listing in store_listings:
version = random.choice(store_listing_versions)
reviewer = random.choice(users)
status: prisma.enums.SubmissionStatus = random.choice(
[
prisma.enums.SubmissionStatus.PENDING,
prisma.enums.SubmissionStatus.APPROVED,
prisma.enums.SubmissionStatus.REJECTED,
]
)
await db.storelistingsubmission.create(
data={
"storeListingId": listing.id,
"storeListingVersionId": version.id,
"reviewerId": reviewer.id,
"Status": status,
"reviewComments": faker.text(),
}
)
# Insert APIKeys
print(f"Inserting {NUM_USERS} api keys")
for user in users:
await db.apikey.create(
data={
"name": faker.word(),
"prefix": str(faker.uuid4())[:8],
"postfix": str(faker.uuid4())[-8:],
"key": str(faker.sha256()),
"status": prisma.enums.APIKeyStatus.ACTIVE,
"permissions": [
prisma.enums.APIKeyPermission.EXECUTE_GRAPH,
prisma.enums.APIKeyPermission.READ_GRAPH,
],
"description": faker.text(),
"userId": user.id,
}
)
await db.disconnect()
if __name__ == "__main__":
asyncio.run(main())