updates template file

Author: davidgomes

agentstack/templates/crewai/tools/neon_tool.py

@@ -1,12 +1,84 @@
 import os
 from crewai_tools import tool
 from dotenv import load_dotenv
-from neon import Neon
+from neon_api import NeonAPI
+import psycopg2
+from psycopg2.extras import RealDictCursor
 
 load_dotenv()
 
-NEON_API_KEY = os.getenv('NEON_API_KEY')
-neon_client = Neon(api_key=NEON_API_KEY)
+NEON_API_KEY = os.getenv("NEON_API_KEY")
+neon_client = NeonAPI(api_key=NEON_API_KEY)
 
-# TODO: Should I duplicate this from the web_researcher example? Or is it
-# possible to somehow have the web_researcher example import this?
+
+@tool("Create Neon Project and Database")
+def create_database(project_name: str) -> str:
+    """
+    Creates a new Neon project. (this takes less than 500ms)
+    Args:
+        project_name: Name of the project to create
+    Returns:
+        the connection URI for the new project
+    """
+    try:
+        project = neon_client.project_create(project={"name": project_name}).project
+        connection_uri = neon_client.connection_uri(
+            project_id=project.id, database_name="neondb", role_name="neondb_owner"
+        ).uri
+        return f"Project/database created, connection URI: {connection_uri}"
+    except Exception as e:
+        return f"Failed to create project: {str(e)}"
+
+
+@tool("Execute SQL DDL")
+def execute_sql_ddl(connection_uri: str, command: str) -> str:
+    """
+    Inserts data into a specified Neon database.
+    Args:
+        connection_uri: The connection URI for the Neon database
+        command: The DDL SQL command to execute
+    Returns:
+        the result of the DDL command
+    """
+    conn = psycopg2.connect(connection_uri)
+    cur = conn.cursor(cursor_factory=RealDictCursor)
+    try:
+        cur.execute(command)
+        conn.commit()
+    except Exception as e:
+        conn.rollback()
+        return f"Failed to execute DDL command: {str(e)}"
+    cur.close()
+    conn.close()
+    return f"Command succeeded"
+
+
+@tool("Execute SQL DML")
+def run_sql_query(connection_uri: str, query: str) -> str:
+    """
+    Inserts data into a specified Neon database.
+    Args:
+        connection_uri: The connection URI for the Neon database
+        query: The SQL query to execute
+    Returns:
+        the result of the SQL query
+    """
+    conn = psycopg2.connect(connection_uri)
+    cur = conn.cursor(cursor_factory=RealDictCursor)
+    try:
+        cur.execute(query)
+        conn.commit()
+
+        # Try to fetch results (for SELECT queries)
+        try:
+            records = cur.fetchall()
+            return f"Query result: {records}"
+        except psycopg2.ProgrammingError:
+            # For INSERT/UPDATE/DELETE operations
+            return f"Query executed successfully"
+    except Exception as e:
+        conn.rollback()
+        return f"Failed to execute SQL query: {str(e)}"
+    finally:
+        cur.close()
+        conn.close()

examples/web_researcher/src/tools/neon_tool.py

@@ -7,9 +7,10 @@
 
 load_dotenv()
 
-NEON_API_KEY = os.getenv('NEON_API_KEY')
+NEON_API_KEY = os.getenv("NEON_API_KEY")
 neon_client = NeonAPI(api_key=NEON_API_KEY)
 
+
 @tool("Create Neon Project and Database")
 def create_database(project_name: str) -> str:
     """
@@ -21,7 +22,9 @@ def create_database(project_name: str) -> str:
     """
     try:
         project = neon_client.project_create(project={"name": project_name}).project
-        connection_uri = neon_client.connection_uri(project_id=project.id, database_name="neondb", role_name="neondb_owner").uri
+        connection_uri = neon_client.connection_uri(
+            project_id=project.id, database_name="neondb", role_name="neondb_owner"
+        ).uri
         return f"Project/database created, connection URI: {connection_uri}"
     except Exception as e:
         return f"Failed to create project: {str(e)}"
@@ -65,7 +68,7 @@ def run_sql_query(connection_uri: str, query: str) -> str:
     try:
         cur.execute(query)
         conn.commit()
-        
+
         # Try to fetch results (for SELECT queries)
         try:
             records = cur.fetchall()