Sql tool

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 []

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

docs/compile_llms_txt.py

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

docs/tools/tool/sql.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=...
+```

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