llama-index-tools-database 0.0.1__tar.gz

llama_index_tools_database-0.0.1/PKG-INFO

@@ -0,0 +1,53 @@
+Metadata-Version: 2.1
+Name: llama-index-tools-database
+Version: 0.0.1
+Summary: llama-index tools database integration
+License: MIT
+Author: Your Name
+Author-email: you@example.com
+Requires-Python: >=3.8.1,<3.12
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: llama-index-core (>=0.9.32,<0.10.0)
+Description-Content-Type: text/markdown
+
+# Database Tool
+
+This tool connects to a database (using SQLAlchemy under the hood) and allows an Agent to query the database and get information about the tables.
+
+## Usage
+
+This tool has more extensive example usage documented in a Jupyter notebook [here](https://github.com/emptycrown/llama-hub/tree/main/llama_hub/tools/notebooks/database.ipynb) and [here](https://github.com/emptycrown/llama-hub/tree/main/llama_hub/tools/notebooks/intro_to_tools.ipynb)
+
+Here's an example usage of the DatabaseToolSpec.
+
+```python
+from llama_index.tools.database import DatabaseToolSpec
+from llama_index.agent import OpenAIAgent
+
+db_tools = DatabaseToolSpec(
+    scheme="postgresql",  # Database Scheme
+    host="localhost",  # Database Host
+    port="5432",  # Database Port
+    user="postgres",  # Database User
+    password="FakeExamplePassword",  # Database Password
+    dbname="postgres",  # Database Name
+)
+agent = OpenAIAgent.from_tools(db_tools.to_tool_list())
+
+agent.chat("What tables does this database contain")
+agent.chat("Describe the first table")
+agent.chat("Retrieve the first row of that table")
+```
+
+The tools available are:
+
+`list_tables`: A tool to list the tables in the database schema
+`describe_tables`: A tool to describe the schema of a table
+`load_data`: A tool that accepts an SQL query and returns the result
+
+This loader is designed to be used as a way to load data as a Tool in a Agent. See [here](https://github.com/emptycrown/llama-hub/tree/main) for examples.
+

llama_index_tools_database-0.0.1/README.md

@@ -0,0 +1,36 @@
+# Database Tool
+
+This tool connects to a database (using SQLAlchemy under the hood) and allows an Agent to query the database and get information about the tables.
+
+## Usage
+
+This tool has more extensive example usage documented in a Jupyter notebook [here](https://github.com/emptycrown/llama-hub/tree/main/llama_hub/tools/notebooks/database.ipynb) and [here](https://github.com/emptycrown/llama-hub/tree/main/llama_hub/tools/notebooks/intro_to_tools.ipynb)
+
+Here's an example usage of the DatabaseToolSpec.
+
+```python
+from llama_index.tools.database import DatabaseToolSpec
+from llama_index.agent import OpenAIAgent
+
+db_tools = DatabaseToolSpec(
+    scheme="postgresql",  # Database Scheme
+    host="localhost",  # Database Host
+    port="5432",  # Database Port
+    user="postgres",  # Database User
+    password="FakeExamplePassword",  # Database Password
+    dbname="postgres",  # Database Name
+)
+agent = OpenAIAgent.from_tools(db_tools.to_tool_list())
+
+agent.chat("What tables does this database contain")
+agent.chat("Describe the first table")
+agent.chat("Retrieve the first row of that table")
+```
+
+The tools available are:
+
+`list_tables`: A tool to list the tables in the database schema
+`describe_tables`: A tool to describe the schema of a table
+`load_data`: A tool that accepts an SQL query and returns the result
+
+This loader is designed to be used as a way to load data as a Tool in a Agent. See [here](https://github.com/emptycrown/llama-hub/tree/main) for examples.

llama_index_tools_database-0.0.1/llama_index/tools/database/__init__.py

@@ -0,0 +1,6 @@
+# __init__.py
+from llama_index.tools.database.base import (
+    DatabaseToolSpec,
+)
+
+__all__ = ["DatabaseToolSpec"]

llama_index_tools_database-0.0.1/llama_index/tools/database/base.py

@@ -0,0 +1,135 @@
+"""Database Tool."""
+
+from typing import Any, List, Optional
+
+from llama_index.core.utilities.sql_wrapper import SQLDatabase
+from llama_index.core.readers.base import BaseReader
+from llama_index.core.schema import Document
+from llama_index.core.tools.tool_spec.base import BaseToolSpec
+from sqlalchemy import MetaData, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import NoSuchTableError
+from sqlalchemy.schema import CreateTable
+
+
+class DatabaseToolSpec(BaseToolSpec, BaseReader):
+    """Simple Database tool.
+
+    Concatenates each row into Document used by LlamaIndex.
+
+    Args:
+        sql_database (Optional[SQLDatabase]): SQL database to use,
+            including table names to specify.
+            See :ref:`Ref-Struct-Store` for more details.
+
+        OR
+
+        engine (Optional[Engine]): SQLAlchemy Engine object of the database connection.
+
+        OR
+
+        uri (Optional[str]): uri of the database connection.
+
+        OR
+
+        scheme (Optional[str]): scheme of the database connection.
+        host (Optional[str]): host of the database connection.
+        port (Optional[int]): port of the database connection.
+        user (Optional[str]): user of the database connection.
+        password (Optional[str]): password of the database connection.
+        dbname (Optional[str]): dbname of the database connection.
+
+    """
+
+    spec_functions = ["load_data", "describe_tables", "list_tables"]
+
+    def __init__(
+        self,
+        sql_database: Optional[SQLDatabase] = None,
+        engine: Optional[Engine] = None,
+        uri: Optional[str] = None,
+        scheme: Optional[str] = None,
+        host: Optional[str] = None,
+        port: Optional[str] = None,
+        user: Optional[str] = None,
+        password: Optional[str] = None,
+        dbname: Optional[str] = None,
+        *args: Optional[Any],
+        **kwargs: Optional[Any],
+    ) -> None:
+        """Initialize with parameters."""
+        if sql_database:
+            self.sql_database = sql_database
+        elif engine:
+            self.sql_database = SQLDatabase(engine, *args, **kwargs)
+        elif uri:
+            self.uri = uri
+            self.sql_database = SQLDatabase.from_uri(uri, *args, **kwargs)
+        elif scheme and host and port and user and password and dbname:
+            uri = f"{scheme}://{user}:{password}@{host}:{port}/{dbname}"
+            self.uri = uri
+            self.sql_database = SQLDatabase.from_uri(uri, *args, **kwargs)
+        else:
+            raise ValueError(
+                "You must provide either a SQLDatabase, "
+                "a SQL Alchemy Engine, a valid connection URI, or a valid "
+                "set of credentials."
+            )
+        self._metadata = MetaData()
+        self._metadata.reflect(bind=self.sql_database.engine)
+
+    def load_data(self, query: str) -> List[Document]:
+        """Query and load data from the Database, returning a list of Documents.
+
+        Args:
+            query (str): an SQL query to filter tables and rows.
+
+        Returns:
+            List[Document]: A list of Document objects.
+        """
+        documents = []
+        with self.sql_database.engine.connect() as connection:
+            if query is None:
+                raise ValueError("A query parameter is necessary to filter the data")
+            else:
+                result = connection.execute(text(query))
+
+            for item in result.fetchall():
+                # fetch each item
+                doc_str = ", ".join([str(entry) for entry in item])
+                documents.append(Document(text=doc_str))
+        return documents
+
+    def list_tables(self) -> List[str]:
+        """
+        Returns a list of available tables in the database.
+        To retrieve details about the columns of specfic tables, use
+        the describe_tables endpoint
+        """
+        return list(map(lambda x: x.name, self._metadata.sorted_tables))
+
+    def describe_tables(self, tables: Optional[List[str]] = None) -> str:
+        """
+        Describes the specifed tables in the database
+
+        Args:
+            tables (List[str]): A list of table names to retrieve details about
+        """
+        table_names = tables or [table.name for table in self._metadata.sorted_tables]
+        table_schemas = []
+
+        for table_name in table_names:
+            table = next(
+                (
+                    table
+                    for table in self._metadata.sorted_tables
+                    if table.name == table_name
+                ),
+                None,
+            )
+            if table is None:
+                raise NoSuchTableError(f"Table '{table_name}' does not exist.")
+            schema = str(CreateTable(table).compile(self.sql_database._engine))
+            table_schemas.append(f"{schema}\n")
+
+        return "\n".join(table_schemas)

llama_index_tools_database-0.0.1/pyproject.toml

@@ -0,0 +1,49 @@
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.codespell]
+check-filenames = true
+check-hidden = true
+# Feel free to un-skip examples, and experimental, you will just need to
+# work through many typos (--write-changes and --interactive will help)
+skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
+
+[tool.mypy]
+disallow_untyped_defs = true
+# Remove venv skip when integrated with pre-commit
+exclude = ["_static", "build", "examples", "notebooks", "venv"]
+ignore_missing_imports = true
+python_version = "3.8"
+
+[tool.poetry]
+name = "llama-index-tools-database"
+version = "0.0.1"
+description = "llama-index tools database integration"
+authors = ["Your Name <you@example.com>"]
+license = "MIT"
+readme = "README.md"
+packages = [{include = "llama_index/"}]
+
+[tool.poetry.dependencies]
+python = ">=3.8.1,<3.12"
+llama-index-core = "^0.9.32"
+
+[tool.poetry.group.dev.dependencies]
+black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"}
+codespell = {extras = ["toml"], version = ">=v2.2.6"}
+ipython = "8.10.0"
+jupyter = "^1.0.0"
+mypy = "0.991"
+pre-commit = "3.2.0"
+pylint = "2.15.10"
+pytest = "7.2.1"
+pytest-mock = "3.11.1"
+ruff = "0.0.292"
+tree-sitter-languages = "^1.8.0"
+types-Deprecated = ">=0.1.0"
+types-PyYAML = "^6.0.12.12"
+types-protobuf = "^4.24.0.4"
+types-redis = "4.5.5.0"
+types-requests = "2.28.11.8" # TODO: unpin when mypy>0.991
+types-setuptools = "67.1.0.0"