Adds Neon as a tool for storing structured data

agentstack/templates/crewai/tools/neon_tool.py

@@ -0,0 +1,84 @@
+import os
+from crewai_tools import tool
+from dotenv import load_dotenv
+from neon_api import NeonAPI
+import psycopg2
+from psycopg2.extras import RealDictCursor
+
+load_dotenv()
+
+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:
+    """
+    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()

agentstack/tools/neon.json

@@ -0,0 +1,7 @@
+{
+    "name": "neon",
+    "packages": ["neon-api", "psycopg2-binary"],
+    "env": "NEON_API_KEY=...",
+    "tools": ["create_database", "execute_sql_ddl", "run_sql_query"],
+    "cta": "Create an API key at https://www.neon.tech"
+  }

agentstack/tools/tools.json

@@ -9,6 +9,9 @@
   "storage": [{
     "name": "mem0",
     "url": "https://github.com/mem0ai/mem0"
+  }, {
+    "name": "neon",
+    "url": "https://github.com/neondatabase/neon"
   }],
   "code-execution": [{
     "name": "open_interpreter",

examples/howards_agent/.env.example

@@ -27,3 +27,5 @@ NEO4J_PASSWORD=...
 
 BROWSERBASE_API_KEY=...
 BROWSERBASE_PROJECT_ID=...
+
+NEON_API_KEY=...

examples/web_researcher/.env.example

@@ -4,3 +4,4 @@ OPENAI_API_KEY=...
 # Tools
 
 FIRECRAWL_API_KEY=...
+NEON_API_KEY=...

examples/web_researcher/agentstack.json

@@ -1,6 +1,7 @@
 {
     "framework": "crewai",
     "tools": [
-        "firecrawl"
+        "firecrawl",
+        "neon"
     ]
-}
+}

examples/web_researcher/pyproject.toml

@@ -12,6 +12,8 @@ crewai = "^0.63.6"
 crewai-tools= "0.12.1"
 python-dotenv="1.0.1"
 firecrawl-py = "^1.4.0"
+neon-api = "^0.1.4"
+psycopg2-binary = "^2.9.9"
 
 [project.scripts]
 web_researcher = "web_researcher.main:run"

examples/web_researcher/src/config/agents.yaml

@@ -13,4 +13,12 @@ content_summarizer:
     analyze the contents of a web page and summarize
   backstory: >-
     you are an expert content writer who concisely summarizes technical writing
+  llm: openai/gpt-4o
+content_storer:
+  role: >-
+    content storer
+  goal: >-
+    store the contents of a web page in a database
+  backstory: >-
+    you are an expert database engineer who knows how to use neon postgres to create tables, insert data, and query data
   llm: openai/gpt-4o

examples/web_researcher/src/config/tasks.yaml

@@ -12,3 +12,10 @@ summarize:
     one paragraph text output
   agent: >-
     content_summarizer
+store:
+  description: >-
+    store the contents of the website in a postgres database, then come up with the SQL query that one could use to retrieve the data that was just inserted (and test it)
+  expected_output: >-
+    the sql query that one could use to retrieve the data that was saved
+  agent: >-
+    content_storer

examples/web_researcher/src/crew.py

@@ -24,6 +24,14 @@ def web_scraper(self) -> Agent:
             verbose=True
         )
 
+    @agent
+    def content_storer(self) -> Agent:
+        return Agent(
+            config=self.agents_config['content_storer'],
+            tools=[tools.create_database, tools.execute_sql_ddl, tools.run_sql_query],  # add tools here or use `agentstack tools add <tool_name>
+            verbose=True
+        )
+
     # Task definitions
     @task
     def scrape_site(self) -> Task:
@@ -37,6 +45,12 @@ def summarize(self) -> Task:
             config=self.tasks_config['summarize'],
         )
 
+    @task
+    def store(self) -> Task:
+        return Task(
+            config=self.tasks_config['store'],
+        )
+
     @crew
     def crew(self) -> Crew:
         """Creates the Test crew"""
@@ -46,4 +60,4 @@ def crew(self) -> Crew:
             process=Process.sequential,
             verbose=True,
             # process=Process.hierarchical, # In case you wanna use that instead https://docs.crewai.com/how-to/Hierarchical/
-        )
+        )

examples/web_researcher/src/tools/__init__.py

@@ -2,3 +2,5 @@
 # tool import
 
 from .firecrawl_tool import web_scrape, web_crawl, retrieve_web_crawl
+
+from .neon_tool import create_database, execute_sql_ddl, run_sql_query