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


[![Build status](https://github.com/project-aita/aita/workflows/build/badge.svg)](https://github.com/project-aita/aita/actions/workflows/build.yml?query=workflow%3Abuild)
[![Python Version](https://img.shields.io/pypi/pyversions/aita.svg)](https://pypi.org/project/aita/)
[![Dependencies Status](https://img.shields.io/badge/dependencies-up%20to%20date-brightgreen.svg)](https://github.com/project-aita/aita/pulls?utf8=%E2%9C%93&q=is%3Apr%20author%3Aapp%2Fdependabot)

[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
[![Security: bandit](https://img.shields.io/badge/security-bandit-green.svg)](https://github.com/PyCQA/bandit)
[![Pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white)](https://github.com/project-aita/aita/blob/main/.pre-commit-config.yaml)
[![License](https://img.shields.io/github/license/project-aita/aita)](https://github.com/project-aita/aita/blob/main/LICENSE)
[![Coverage Report](assets/images/coverage.svg)](https://github.com/project-aita/aita/blob/main/assets/images/coverage.svg)


## Tech Stack

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

- **[Langchain](https://www.langchain.com/)**: Facilitates the seamless integration of language models into application workflows, significantly enhancing AI interaction capabilities.
- **[Reflex](https://reflex.dev/)**: An open-source framework for quickly building beautiful, interactive web applications in pure Python
- **[Apache Arrow](https://arrow.apache.org/)**: 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](https://github.com/jupyterlab/jupyter-ai)**: A JupyterLab extension that provides a set of magics for working with AI models.

## Installation
Simply install the package using pip:

```shell
pip install aita

```
Or with extra dependencies:

```shell
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](https://aita-1.gitbook.io/aita/).

```python
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
- [x] Snowflake
- [x] Sqlite
- [ ] BigQuery
- [x] Postgres
- [ ] MySQL
- [ ] Redshift
- [ ] DynamoDB
- [x] File (CSV, Excel, Parquet, etc.)

## Supported Engines
- [x] Apache Spark
- [ ] Apache Flink
- [ ] Presto
- [ ] **[Ray.io](https://ray.io/)**: A distributed computing framework that efficiently scales AI tasks and data processing across clusters, improving performance and resource utilization.

## 🛡 License

[![License](https://img.shields.io/github/license/project-aita/aita)](https://github.com/project-aita/aita/blob/main/LICENSE)

This project is licensed under the terms of the `Apache Software License 2.0` license. See [LICENSE](https://github.com/aita/aita/blob/master/LICENSE) for more details.
