/aita

Common text to {sql, dataframe APIs} AI agent.

Primary LanguagePythonApache License 2.0Apache-2.0

Aita

AI Powered Data Agent, a comprehensive solution for data analysis, engineering, and visualization.

Build status Python Version Dependencies Status

Code style: black Security: bandit Pre-commit License Coverage Report

Tech Stack

Our platform leverages a combination of cutting-edge technologies and frameworks:

  • Langchain: Facilitates the seamless integration of language models into application workflows, significantly enhancing AI interaction capabilities.
  • Reflex: An open-source framework for quickly building beautiful, interactive web applications in pure Python
  • Apache Arrow: A cross-language development platform for in-memory data that specifies a standardized language-independent columnar memory format for flat and hierarchical data, organized for efficient analytic operations on modern hardware like CPUs and GPUs.
  • Jupyter Ai Magics: A JupyterLab extension that provides a set of magics for working with AI models.

Installation

Simply install the package using pip:

pip install aita

Or with extra dependencies:

pip install aita[snowflake]

Example

Below is an example of using SqlAgent to connect to a PostgreSQL database and ask a question. You can read more details in the documentation.

from aita.agent.sql import SqlAgent
from aita.datasource.postgresql import PostgreSqlDataSource

# Connect to a PostgreSQL database
datasource = PostgreSqlDataSource("postgresql://user:password@localhost:5432/dbname")

# Create a SQL agent
sql_agent = SqlAgent("gpt-3.5-turbo").add_datasource(datasource)

# ask question to the agent
sql_agent.stream("What is the total number of orders?")

Supported Models

Model provider are supported by jupyter ai magics. Ensure the corresponding environment variables are set before using the Aita agent.

Provider Provider ID Environment variable(s) Python package(s)
AI21 ai21 AI21_API_KEY ai21
Anthropic anthropic ANTHROPIC_API_KEY langchain-anthropic
Anthropic (chat) anthropic-chat ANTHROPIC_API_KEY langchain-anthropic
Bedrock bedrock N/A boto3
Bedrock (chat) bedrock-chat N/A boto3
Cohere cohere COHERE_API_KEY cohere
ERNIE-Bot qianfan QIANFAN_AK, QIANFAN_SK qianfan
Gemini gemini GOOGLE_API_KEY langchain-google-genai
GPT4All gpt4all N/A gpt4all
Hugging Face Hub huggingface_hub HUGGINGFACEHUB_API_TOKEN huggingface_hub, ipywidgets, pillow
NVIDIA nvidia-chat NVIDIA_API_KEY langchain_nvidia_ai_endpoints
OpenAI openai OPENAI_API_KEY langchain-openai
OpenAI (chat) openai-chat OPENAI_API_KEY langchain-openai
SageMaker sagemaker-endpoint N/A boto3

Supported Data Sources

  • Snowflake
  • Sqlite
  • BigQuery
  • Postgres
  • MySQL
  • Redshift
  • DynamoDB
  • File (CSV, Excel, Parquet, etc.)

Supported Engines

  • Apache Spark
  • Apache Flink
  • Presto
  • Ray.io: A distributed computing framework that efficiently scales AI tasks and data processing across clusters, improving performance and resource utilization.

🛡 License

License

This project is licensed under the terms of the Apache Software License 2.0 license. See LICENSE for more details.