Skip to content

Commit

Permalink
some code cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
davidgomes committed Nov 17, 2024
1 parent 1362344 commit 7172d97
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 68 deletions.
69 changes: 2 additions & 67 deletions agentstack/templates/crewai/tools/neon_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,70 +8,5 @@
NEON_API_KEY = os.getenv('NEON_API_KEY')
neon_client = Neon(api_key=NEON_API_KEY)

@tool("Create Project")
def create_project(project_name: str) -> str:
"""
Creates a new Neon project.
Args:
project_name: Name of the project to create
Returns:
Project ID and status message
"""
try:
project = neon_client.projects.create(name=project_name)
return f"Project created successfully. Project ID: {project.id}"
except Exception as e:
return f"Failed to create project: {str(e)}"

@tool("Create Database")
def create_database(project_id: str, database_name: str) -> str:
"""
Creates a new database in specified Neon project.
Args:
project_id: ID of the project
database_name: Name of the database to create
Returns:
Status message about database creation
"""
try:
result = neon_client.databases.create(
project_id=project_id,
name=database_name
)
return f"Database {database_name} created successfully"
except Exception as e:
return f"Failed to create database: {str(e)}"

@tool("List Databases")
def list_databases(project_id: str) -> str:
"""
Lists all databases in specified Neon project.
Args:
project_id: ID of the project
Returns:
List of database names
"""
try:
databases = neon_client.databases.list(project_id=project_id)
return "\n".join([db.name for db in databases])
except Exception as e:
return f"Failed to list databases: {str(e)}"

@tool("Delete Database")
def delete_database(project_id: str, database_name: str) -> str:
"""
Deletes a database from specified Neon project.
Args:
project_id: ID of the project
database_name: Name of the database to delete
Returns:
Status message about database deletion
"""
try:
neon_client.databases.delete(
project_id=project_id,
name=database_name
)
return f"Database {database_name} deleted successfully"
except Exception as e:
return f"Failed to delete database: {str(e)}"
# TODO: Should I duplicate this from the web_researcher example? Or is it
# possible to somehow have the web_researcher example import this?
2 changes: 1 addition & 1 deletion examples/web_researcher/src/tools/neon_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
@tool("Create Neon Project and Database")
def create_database(project_name: str) -> str:
"""
Creates a new Neon project.
Creates a new Neon project. (this takes less than 500ms)
Args:
project_name: Name of the project to create
Returns:
Expand Down

0 comments on commit 7172d97

Please sign in to comment.