From 8bc0967d5b236ee69e2adb61c59968b03115e1f9 Mon Sep 17 00:00:00 2001 From: Sri Laasya Nutheti Date: Fri, 24 Jan 2025 12:38:43 -0800 Subject: [PATCH 1/5] add tool to run sql queries over postgres --- agentstack/_tools/py_sql/__init__.py | 79 ++++++++++++++++++++++++++++ agentstack/_tools/py_sql/config.json | 37 +++++++++++++ 2 files changed, 116 insertions(+) create mode 100644 agentstack/_tools/py_sql/__init__.py create mode 100644 agentstack/_tools/py_sql/config.json diff --git a/agentstack/_tools/py_sql/__init__.py b/agentstack/_tools/py_sql/__init__.py new file mode 100644 index 00000000..4e387b23 --- /dev/null +++ b/agentstack/_tools/py_sql/__init__.py @@ -0,0 +1,79 @@ +import os +import psycopg2 +from typing import Dict, Any + +def get_connection(): + """Get PostgreSQL database connection""" + return psycopg2.connect( + dbname=os.getenv('POSTGRES_DB'), + user=os.getenv('POSTGRES_USER'), + password=os.getenv('POSTGRES_PASSWORD'), + host=os.getenv('POSTGRES_HOST', 'localhost'), + port=os.getenv('POSTGRES_PORT', '5432') + ) + +def get_schema() -> Dict[str, Any]: + """ + Initialize connection and get database schema. + Returns a dictionary containing the database schema. + """ + try: + conn = get_connection() + cursor = conn.cursor() + + # Query to get all tables in the current schema + schema_query = """ + SELECT table_name + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_type = 'BASE TABLE'; + """ + + cursor.execute(schema_query) + tables = cursor.fetchall() + + # Create schema dictionary + schema = {} + for (table_name,) in tables: + # Get column information for each table + column_query = """ + SELECT column_name + FROM information_schema.columns + WHERE table_schema = 'public' + AND table_name = %s; + """ + cursor.execute(column_query, (table_name,)) + columns = [col[0] for col in cursor.fetchall()] + schema[table_name] = columns + + cursor.close() + conn.close() + return schema + + except Exception as e: + print(f"Error getting database schema: {str(e)}") + return {} + +def execute_query(query: str) -> list: + """ + Execute a SQL query on the database. + Args: + query: SQL query to execute + Returns: + List of query results + """ + try: + conn = get_connection() + cursor = conn.cursor() + + # Execute the query + cursor.execute(query) + results = cursor.fetchall() + + cursor.close() + conn.close() + return results + + except Exception as e: + print(f"Error executing query: {str(e)}") + return [] diff --git a/agentstack/_tools/py_sql/config.json b/agentstack/_tools/py_sql/config.json new file mode 100644 index 00000000..751f95a0 --- /dev/null +++ b/agentstack/_tools/py_sql/config.json @@ -0,0 +1,37 @@ +{ + "name": "sql-queries", + "url": "https://pypi.org/project/psycopg2/", + "category": "database", + "env": { + "POSTGRES_DB": { + "description": "PostgreSQL database name", + "required": true + }, + "POSTGRES_USER": { + "description": "PostgreSQL username", + "required": true + }, + "POSTGRES_PASSWORD": { + "description": "PostgreSQL password", + "required": true + }, + "POSTGRES_HOST": { + "description": "PostgreSQL host address", + "required": true, + "default": "localhost" + }, + "POSTGRES_PORT": { + "description": "PostgreSQL port number", + "required": true, + "default": "5432" + } + }, + "dependencies": [ + "psycopg2-binary>=2.9.9" + ], + "tools": [ + "get_schema", + "execute_query" + ], + "cta": "Set up your PostgreSQL connection variables in the environment file." +} \ No newline at end of file From 5be3f584b3ac5c6973b8dc6885e2ea5d1bee3ba6 Mon Sep 17 00:00:00 2001 From: Sri Laasya Nutheti Date: Fri, 24 Jan 2025 13:04:04 -0800 Subject: [PATCH 2/5] fixed var names, assertion err --- agentstack/_tools/py_sql/config.json | 2 +- agentstack/_tools/py_sql/test.db | Bin 0 -> 12288 bytes agentstack/_tools/py_sql/test.py | 66 +++++++++++++++++++++++++++ 3 files changed, 67 insertions(+), 1 deletion(-) create mode 100644 agentstack/_tools/py_sql/test.db create mode 100644 agentstack/_tools/py_sql/test.py diff --git a/agentstack/_tools/py_sql/config.json b/agentstack/_tools/py_sql/config.json index 751f95a0..0e6c387b 100644 --- a/agentstack/_tools/py_sql/config.json +++ b/agentstack/_tools/py_sql/config.json @@ -1,5 +1,5 @@ { - "name": "sql-queries", + "name": "py_sql", "url": "https://pypi.org/project/psycopg2/", "category": "database", "env": { diff --git a/agentstack/_tools/py_sql/test.db b/agentstack/_tools/py_sql/test.db new file mode 100644 index 0000000000000000000000000000000000000000..fba01d7ddf92eefbcd3ad1b57e4258483ee3fa25 GIT binary patch literal 12288 zcmeI&O-jQ+6bJB`L@E_Dx(Y7T*IlJm5HDb=qf}$6ZHjgwlGf2i@}VY*TXE?PJd0=W z3ZBBflR%~E#)aTQ{*O%NBk#??Z#R(6vF)dlF5|nFcTn(m1 z6V+$)pF(5n4==MSOPhOV2B{YU0uX=z1Rwwb2tWV=5P$##An 18 + ) + print("Select users over 18:") + print(f"Query: {query}") + print(f"Params: {params}") + + # Test select with multiple conditions + query, params = construct_sql_query( + "select", + "users", + columns=[users.name, users.email], + where=(users.age > 18) & (users.active == True) + ) + print("\nSelect active users over 18:") + print(f"Query: {query}") + print(f"Params: {params}") + + print("\n=== Testing INSERT queries ===") + query, params = construct_sql_query( + "insert", + "users", + values=[["John Doe", 25, "john@example.com", True]] + ) + print("Insert new user:") + print(f"Query: {query}") + print(f"Params: {params}") + + print("\n=== Testing UPDATE queries ===") + query, params = construct_sql_query( + "update", + "users", + columns=[users.active], + values=[False], + where=users.age < 18 + ) + print("Deactivate users under 18:") + print(f"Query: {query}") + print(f"Params: {params}") + + print("\n=== Testing DELETE queries ===") + query, params = construct_sql_query( + "delete", + "users", + where=users.active == False + ) + print("Delete inactive users:") + print(f"Query: {query}") + print(f"Params: {params}") + +if __name__ == "__main__": + test_query_construction() \ No newline at end of file From 2d18ae3ed6a88f7e9ce81a3f8aa63b80dd2e1a99 Mon Sep 17 00:00:00 2001 From: Sri Laasya Nutheti Date: Fri, 7 Feb 2025 15:01:55 -0800 Subject: [PATCH 3/5] fixed workflow and added tests for compile_llms_txt.py, converted connection to a private function --- .github/workflows/compile-llms-txt.yml | 2 +- agentstack/_tools/py_sql/__init__.py | 6 +- agentstack/_tools/py_sql/config.json | 27 ++------ agentstack/_tools/py_sql/test.db | Bin 12288 -> 0 bytes agentstack/_tools/py_sql/test.py | 66 ------------------ docs/compile_llms_txt.py | 21 +++--- tests/test_compile_llms.py | 91 +++++++++++++++++++++++++ 7 files changed, 111 insertions(+), 102 deletions(-) delete mode 100644 agentstack/_tools/py_sql/test.db delete mode 100644 agentstack/_tools/py_sql/test.py create mode 100644 tests/test_compile_llms.py diff --git a/.github/workflows/compile-llms-txt.yml b/.github/workflows/compile-llms-txt.yml index 6859366e..d2043861 100644 --- a/.github/workflows/compile-llms-txt.yml +++ b/.github/workflows/compile-llms-txt.yml @@ -19,7 +19,7 @@ jobs: - name: Compile llms.txt run: | cd docs - python compile_llms.py + python compile_llms_txt.py - name: Commit and push changes run: | git config --local user.email "action@github.com" diff --git a/agentstack/_tools/py_sql/__init__.py b/agentstack/_tools/py_sql/__init__.py index 4e387b23..6634cd84 100644 --- a/agentstack/_tools/py_sql/__init__.py +++ b/agentstack/_tools/py_sql/__init__.py @@ -2,7 +2,7 @@ import psycopg2 from typing import Dict, Any -def get_connection(): +def _get_connection(): """Get PostgreSQL database connection""" return psycopg2.connect( dbname=os.getenv('POSTGRES_DB'), @@ -18,7 +18,7 @@ def get_schema() -> Dict[str, Any]: Returns a dictionary containing the database schema. """ try: - conn = get_connection() + conn = _get_connection() cursor = conn.cursor() # Query to get all tables in the current schema @@ -63,7 +63,7 @@ def execute_query(query: str) -> list: List of query results """ try: - conn = get_connection() + conn = _get_connection() cursor = conn.cursor() # Execute the query diff --git a/agentstack/_tools/py_sql/config.json b/agentstack/_tools/py_sql/config.json index 0e6c387b..69429eb2 100644 --- a/agentstack/_tools/py_sql/config.json +++ b/agentstack/_tools/py_sql/config.json @@ -3,28 +3,11 @@ "url": "https://pypi.org/project/psycopg2/", "category": "database", "env": { - "POSTGRES_DB": { - "description": "PostgreSQL database name", - "required": true - }, - "POSTGRES_USER": { - "description": "PostgreSQL username", - "required": true - }, - "POSTGRES_PASSWORD": { - "description": "PostgreSQL password", - "required": true - }, - "POSTGRES_HOST": { - "description": "PostgreSQL host address", - "required": true, - "default": "localhost" - }, - "POSTGRES_PORT": { - "description": "PostgreSQL port number", - "required": true, - "default": "5432" - } + "POSTGRES_DB": "...", + "POSTGRES_USER": "...", + "POSTGRES_PASSWORD": "...", + "POSTGRES_HOST": "...", + "POSTGRES_PORT": "..." }, "dependencies": [ "psycopg2-binary>=2.9.9" diff --git a/agentstack/_tools/py_sql/test.db b/agentstack/_tools/py_sql/test.db deleted file mode 100644 index fba01d7ddf92eefbcd3ad1b57e4258483ee3fa25..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 12288 zcmeI&O-jQ+6bJB`L@E_Dx(Y7T*IlJm5HDb=qf}$6ZHjgwlGf2i@}VY*TXE?PJd0=W z3ZBBflR%~E#)aTQ{*O%NBk#??Z#R(6vF)dlF5|nFcTn(m1 z6V+$)pF(5n4==MSOPhOV2B{YU0uX=z1Rwwb2tWV=5P$##An 18 - ) - print("Select users over 18:") - print(f"Query: {query}") - print(f"Params: {params}") - - # Test select with multiple conditions - query, params = construct_sql_query( - "select", - "users", - columns=[users.name, users.email], - where=(users.age > 18) & (users.active == True) - ) - print("\nSelect active users over 18:") - print(f"Query: {query}") - print(f"Params: {params}") - - print("\n=== Testing INSERT queries ===") - query, params = construct_sql_query( - "insert", - "users", - values=[["John Doe", 25, "john@example.com", True]] - ) - print("Insert new user:") - print(f"Query: {query}") - print(f"Params: {params}") - - print("\n=== Testing UPDATE queries ===") - query, params = construct_sql_query( - "update", - "users", - columns=[users.active], - values=[False], - where=users.age < 18 - ) - print("Deactivate users under 18:") - print(f"Query: {query}") - print(f"Params: {params}") - - print("\n=== Testing DELETE queries ===") - query, params = construct_sql_query( - "delete", - "users", - where=users.active == False - ) - print("Delete inactive users:") - print(f"Query: {query}") - print(f"Params: {params}") - -if __name__ == "__main__": - test_query_construction() \ No newline at end of file diff --git a/docs/compile_llms_txt.py b/docs/compile_llms_txt.py index 503f311b..87f722ca 100644 --- a/docs/compile_llms_txt.py +++ b/docs/compile_llms_txt.py @@ -1,15 +1,14 @@ import os +from pathlib import Path def compile_llms_txt(): - # Get the docs directory path (where this script is located) - docs_dir = os.path.dirname(os.path.abspath(__file__)) + # Get the current working directory + current_dir = Path(os.getcwd()) content = '' + # Define names of directories and files to exclude excluded_names = {'tool'} - # Change to docs directory - os.chdir(docs_dir) - for root, _, files in os.walk('.'): # Get the last part of the current directory current_dir = os.path.basename(root) @@ -17,8 +16,11 @@ def compile_llms_txt(): continue for file in files: + # Check if the file is an MDX file and not in excluded names if file.endswith('.mdx'): - if file in excluded_names: + # Extract the base name without extension for exclusion check + base_name = os.path.splitext(file)[0] + if base_name in excluded_names: continue file_path = os.path.join(root, file) @@ -28,10 +30,9 @@ def compile_llms_txt(): file_content = f.read() content += f"## {relative_path}\n\n{file_content}\n\n" - # Write the complete content, replacing the existing file - output_path = os.path.join(docs_dir, 'llms.txt') - with open(output_path, 'w', encoding='utf-8') as f: - f.write(content) + # Write the complete content to llms.txt in the current directory + output_path = Path('llms.txt') + output_path.write_text(content, encoding='utf-8') if __name__ == "__main__": compile_llms_txt() diff --git a/tests/test_compile_llms.py b/tests/test_compile_llms.py new file mode 100644 index 00000000..3ef1b08a --- /dev/null +++ b/tests/test_compile_llms.py @@ -0,0 +1,91 @@ +import os +import shutil +import tempfile +import unittest +from pathlib import Path + +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from docs.compile_llms_txt import compile_llms_txt + +class TestCompileLLMsTxt(unittest.TestCase): + def setUp(self): + # Create a temporary directory for test files + self.test_dir = tempfile.mkdtemp() + self.docs_dir = Path(self.test_dir) + + # Save current working directory + self.original_cwd = os.getcwd() + os.chdir(self.docs_dir) + + def tearDown(self): + os.chdir(self.original_cwd) + shutil.rmtree(self.test_dir) + + def create_test_mdx_file(self, path: str, content: str): + """Helper to create test MDX files""" + file_path = self.docs_dir / path + file_path.parent.mkdir(parents=True, exist_ok=True) + file_path.write_text(content) + + def test_basic_compilation(self): + """Test basic MDX file compilation""" + # Create test MDX files + self.create_test_mdx_file("test1.mdx", "Test content 1") + self.create_test_mdx_file("test2.mdx", "Test content 2") + + # Run compilation + compile_llms_txt() + + # Check output file exists and contains expected content + output_path = self.docs_dir / "llms.txt" + self.assertTrue(output_path.exists()) + + content = output_path.read_text() + self.assertIn("## test1.mdx", content) + self.assertIn("Test content 1", content) + self.assertIn("## test2.mdx", content) + self.assertIn("Test content 2", content) + + def test_excluded_directories(self): + """Test that files in excluded directories are skipped""" + # Create files in both regular and excluded directories + self.create_test_mdx_file("regular/file.mdx", "Regular content") + self.create_test_mdx_file("tool/file.mdx", "Tool content") + + compile_llms_txt() + + content = (self.docs_dir / "llms.txt").read_text() + self.assertIn("Regular content", content) + self.assertNotIn("Tool content", content) + + def test_excluded_files(self): + """Test that excluded files are skipped""" + self.create_test_mdx_file("regular.mdx", "Regular content") + self.create_test_mdx_file("tool.mdx", "Tool content") + + compile_llms_txt() + + content = (self.docs_dir / "llms.txt").read_text() + self.assertIn("Regular content", content) + self.assertNotIn("Tool content", content) + + def test_nested_directories(self): + """Test compilation from nested directory structure""" + self.create_test_mdx_file("dir1/test1.mdx", "Content 1") + self.create_test_mdx_file("dir1/dir2/test2.mdx", "Content 2") + + compile_llms_txt() + + content = (self.docs_dir / "llms.txt").read_text() + self.assertIn("## dir1/test1.mdx", content) + self.assertIn("## dir1/dir2/test2.mdx", content) + self.assertIn("Content 1", content) + self.assertIn("Content 2", content) + + def test_empty_directory(self): + """Test compilation with no MDX files""" + compile_llms_txt() + + content = (self.docs_dir / "llms.txt").read_text() + self.assertEqual(content, "") \ No newline at end of file From da14a29d324fe5be9e3c25e39ad00123f6acf95b Mon Sep 17 00:00:00 2001 From: Sri Laasya Nutheti Date: Fri, 7 Feb 2025 15:40:45 -0800 Subject: [PATCH 4/5] Err with accessing current working dir, fix updated --- tests/test_compile_llms.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_compile_llms.py b/tests/test_compile_llms.py index 3ef1b08a..1c15bebb 100644 --- a/tests/test_compile_llms.py +++ b/tests/test_compile_llms.py @@ -10,12 +10,13 @@ class TestCompileLLMsTxt(unittest.TestCase): def setUp(self): + self.original_cwd = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + # Create a temporary directory for test files self.test_dir = tempfile.mkdtemp() self.docs_dir = Path(self.test_dir) - # Save current working directory - self.original_cwd = os.getcwd() + # Change to the temporary directory os.chdir(self.docs_dir) def tearDown(self): From 3b799b7efafccf0e79dc8da7378026b2f6becc79 Mon Sep 17 00:00:00 2001 From: Braelyn Boynton Date: Thu, 13 Feb 2025 16:43:54 -0800 Subject: [PATCH 5/5] global connection and docs --- agentstack/_tools/{py_sql => sql}/__init__.py | 25 +++++++++++------- agentstack/_tools/{py_sql => sql}/config.json | 12 ++++----- docs/tools/core.mdx | 5 +++- docs/tools/tool/sql.mdx | 26 +++++++++++++++++++ 4 files changed, 52 insertions(+), 16 deletions(-) rename agentstack/_tools/{py_sql => sql}/__init__.py (81%) rename agentstack/_tools/{py_sql => sql}/config.json (63%) create mode 100644 docs/tools/tool/sql.mdx diff --git a/agentstack/_tools/py_sql/__init__.py b/agentstack/_tools/sql/__init__.py similarity index 81% rename from agentstack/_tools/py_sql/__init__.py rename to agentstack/_tools/sql/__init__.py index 6634cd84..292ff2af 100644 --- a/agentstack/_tools/py_sql/__init__.py +++ b/agentstack/_tools/sql/__init__.py @@ -2,15 +2,22 @@ import psycopg2 from typing import Dict, Any +connection = None + def _get_connection(): """Get PostgreSQL database connection""" - return psycopg2.connect( - dbname=os.getenv('POSTGRES_DB'), - user=os.getenv('POSTGRES_USER'), - password=os.getenv('POSTGRES_PASSWORD'), - host=os.getenv('POSTGRES_HOST', 'localhost'), - port=os.getenv('POSTGRES_PORT', '5432') - ) + + global connection + if connection is None: + connection = psycopg2.connect( + dbname=os.getenv('POSTGRES_DB'), + user=os.getenv('POSTGRES_USER'), + password=os.getenv('POSTGRES_PASSWORD'), + host=os.getenv('POSTGRES_HOST', 'localhost'), + port=os.getenv('POSTGRES_PORT', '5432') + ) + + return connection def get_schema() -> Dict[str, Any]: """ @@ -47,7 +54,7 @@ def get_schema() -> Dict[str, Any]: schema[table_name] = columns cursor.close() - conn.close() + # conn.close() return schema except Exception as e: @@ -71,7 +78,7 @@ def execute_query(query: str) -> list: results = cursor.fetchall() cursor.close() - conn.close() + # conn.close() return results except Exception as e: diff --git a/agentstack/_tools/py_sql/config.json b/agentstack/_tools/sql/config.json similarity index 63% rename from agentstack/_tools/py_sql/config.json rename to agentstack/_tools/sql/config.json index 69429eb2..67756cff 100644 --- a/agentstack/_tools/py_sql/config.json +++ b/agentstack/_tools/sql/config.json @@ -1,13 +1,13 @@ { - "name": "py_sql", + "name": "sql", "url": "https://pypi.org/project/psycopg2/", "category": "database", "env": { - "POSTGRES_DB": "...", - "POSTGRES_USER": "...", - "POSTGRES_PASSWORD": "...", - "POSTGRES_HOST": "...", - "POSTGRES_PORT": "..." + "POSTGRES_DB": null, + "POSTGRES_USER": null, + "POSTGRES_PASSWORD": null, + "POSTGRES_HOST": null, + "POSTGRES_PORT": null }, "dependencies": [ "psycopg2-binary>=2.9.9" diff --git a/docs/tools/core.mdx b/docs/tools/core.mdx index 464728ab..4879f81c 100644 --- a/docs/tools/core.mdx +++ b/docs/tools/core.mdx @@ -13,9 +13,12 @@ description: 'AgentStack tools that are not third-party integrations' - [Code Interpreter](/tools/tool/code-interpreter) -## Data Input +## Input - [Vision](/tools/tool/vision) +## Data +- [SQL](/tools/tool/sql) + + There is no built-in sandboxing using this tool. Agents may perform destructive queries that may be irreversible. + + +## Installation + +```bash +agentstack tools add sql +``` + +Set the API keys +```env +POSTGRES_DB=... +POSTGRES_USER=... +POSTGRES_PASSWORD=... +POSTGRES_HOST=... +POSTGRES_PORT=... +``` \ No newline at end of file