From 13623442d2c6c51b561bee0b34b12f82cb806da1 Mon Sep 17 00:00:00 2001
From: David Gomes <davidrafagomes@gmail.com>
Date: Sun, 17 Nov 2024 12:05:22 +0100
Subject: [PATCH] get everything to work

---
 examples/web_researcher/.env.example          |  1 +
 examples/web_researcher/src/config/tasks.yaml |  4 +--
 .../web_researcher/src/tools/neon_tool.py     | 34 ++++++++++++++-----
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/examples/web_researcher/.env.example b/examples/web_researcher/.env.example
index 08789d5..149da23 100644
--- a/examples/web_researcher/.env.example
+++ b/examples/web_researcher/.env.example
@@ -4,3 +4,4 @@ OPENAI_API_KEY=...
 # Tools
 
 FIRECRAWL_API_KEY=...
+NEON_API_KEY=...
\ No newline at end of file
diff --git a/examples/web_researcher/src/config/tasks.yaml b/examples/web_researcher/src/config/tasks.yaml
index 3f3e3b3..8e8c6fa 100644
--- a/examples/web_researcher/src/config/tasks.yaml
+++ b/examples/web_researcher/src/config/tasks.yaml
@@ -14,8 +14,8 @@ summarize:
     content_summarizer
 store:
   description: >-
-    store the contents of the website in a postgres database
+    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 to retrieve the contents of the website from the database
+    the sql query that one could use to retrieve the data that was saved
   agent: >-
     content_storer
\ No newline at end of file
diff --git a/examples/web_researcher/src/tools/neon_tool.py b/examples/web_researcher/src/tools/neon_tool.py
index 540bee8..4c28b70 100644
--- a/examples/web_researcher/src/tools/neon_tool.py
+++ b/examples/web_researcher/src/tools/neon_tool.py
@@ -33,17 +33,21 @@ 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
+        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)
-    cur.execute(command)
-    result = cur.fetchone()
+    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 result: {result}"
+    return f"Command succeeded"
 
 
 @tool("Execute SQL DML")
@@ -58,8 +62,20 @@ def run_sql_query(connection_uri: str, query: str) -> str:
     """
     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}"
+    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()