progress

Author: davidgomes

agentstack/templates/crewai/tools/neon_tool.py

@@ -75,22 +75,3 @@ def delete_database(project_id: str, database_name: str) -> str:
         return f"Database {database_name} deleted successfully"
     except Exception as e:
         return f"Failed to delete database: {str(e)}"
-
-@tool("Create Branch")
-def create_branch(project_id: str, branch_name: str) -> str:
-    """
-    Creates a new branch in specified Neon project.
-    Args:
-        project_id: ID of the project
-        branch_name: Name of the branch to create
-    Returns:
-        Status message about branch creation
-    """
-    try:
-        result = neon_client.branches.create(
-            project_id=project_id,
-            name=branch_name
-        )
-        return f"Branch {branch_name} created successfully"
-    except Exception as e:
-        return f"Failed to create branch: {str(e)}"

examples/web_researcher/poetry.lock

@@ -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/pyproject.toml

@@ -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/agents.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
+  expected_output: >-
+    the sql query to retrieve the contents of the website from the database
+  agent: >-
+    content_storer

examples/web_researcher/src/config/tasks.yaml

@@ -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/crew.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

examples/web_researcher/src/tools/__init__.py

@@ -0,0 +1,65 @@
+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.
+    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 command to execute
+    Returns:
+        the result of the DDL command
+    """
+    conn = psycopg2.connect(connection_uri)
+    cur = conn.cursor(cursor_factory=RealDictCursor)
+    cur.execute(command)
+    result = cur.fetchone()
+    cur.close()
+    conn.close()
+    return f"Command result: {result}"
+
+
+@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)
+    cur.query(query)
+    records = cur.fetchall()
+    cur.close()
+    conn.close()
+    return f"Query result: {records}"