diff --git a/agentstack/_tools/sql/__init__.py b/agentstack/_tools/sql/__init__.py new file mode 100644 index 00000000..292ff2af --- /dev/null +++ b/agentstack/_tools/sql/__init__.py @@ -0,0 +1,86 @@ +import os +import psycopg2 +from typing import Dict, Any + +connection = None + +def _get_connection(): + """Get PostgreSQL database connection""" + + 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]: + """ + 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/sql/config.json b/agentstack/_tools/sql/config.json new file mode 100644 index 00000000..67756cff --- /dev/null +++ b/agentstack/_tools/sql/config.json @@ -0,0 +1,20 @@ +{ + "name": "sql", + "url": "https://pypi.org/project/psycopg2/", + "category": "database", + "env": { + "POSTGRES_DB": null, + "POSTGRES_USER": null, + "POSTGRES_PASSWORD": null, + "POSTGRES_HOST": null, + "POSTGRES_PORT": null + }, + "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 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/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 diff --git a/tests/test_compile_llms.py b/tests/test_compile_llms.py new file mode 100644 index 00000000..1c15bebb --- /dev/null +++ b/tests/test_compile_llms.py @@ -0,0 +1,92 @@ +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): + 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) + + # Change to the temporary directory + 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