Merge branch 'main' into create_custom_tools

Author: bboynton97

agentstack/_tools/py_sql/config.json

@@ -2,23 +2,30 @@
 import psycopg2
 from typing import Dict, Any
 
-def get_connection():
+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]:
     """
     Initialize connection and get database schema.
     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
@@ -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:
@@ -63,15 +70,15 @@ def execute_query(query: str) -> list:
         List of query results
     """
     try:
-        conn = get_connection()
+        conn = _get_connection()
         cursor = conn.cursor()
 
         # Execute the query
         cursor.execute(query)
         results = cursor.fetchall()
 
         cursor.close()
-        conn.close()
+        # conn.close()
         return results
 
     except Exception as e:

agentstack/_tools/py_sql/__init__.py renamed to agentstack/_tools/sql/__init__.py

@@ -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."
+}

agentstack/_tools/sql/config.json

@@ -1,24 +1,26 @@
 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)
         if current_dir in excluded_names:
             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()

docs/compile_llms_txt.py

@@ -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)
+
 <CardGroup cols={1}>
   <Card
     title="Community Tools"

docs/tools/core.mdx

@@ -0,0 +1,26 @@
+---
+title: 'Perplexity'
+description: 'Agentic Search'
+---
+
+## Description
+Enable your agent to perform queries directly on a database
+
+<Note>
+    There is no built-in sandboxing using this tool. Agents may perform destructive queries that may be irreversible.
+</Note>
+
+## Installation
+
+```bash
+agentstack tools add sql
+```
+
+Set the API keys
+```env
+POSTGRES_DB=...
+POSTGRES_USER=...
+POSTGRES_PASSWORD=...
+POSTGRES_HOST=...
+POSTGRES_PORT=...
+```

docs/tools/tool/sql.mdx

@@ -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, "")