Use Athena SDK in Notebooks

Athena SDK enables Athena capabilities integrated in your pipelines and workflows

Most of Athena's capabilities exist in a Python SDK & API form factor as well as the UI platform. Using the SDK to run agentic workflows and data analysis tools allows you to integrate and scale workflows made in Athena UI platform so these would run in your pipelines and ad-hocs.

API & SDK Documentation

For full documentation refer to docs: Athena Documentation

API Key Access

Athena API and SDK is authenticated with your account using a personal Athena API Key. Reach out to Athena team to set up a personal token.

Open a Notebook and Set up the SDK Client

1. Once you opened a Athena Notebook, install the latest Athena SDK package:

! pip install -U athena-intelligence

DISCLAIMER

Your API KEY is auto-injected into your Athena Notebook Environment. Try the below code to see it!

import os

print(os.environ['ATHENA_API_KEY'])

2. Initialize the client with your personal Athena Token

import os
from athena import Model, Tools
from athena.client import AsyncAthena

async_client = AsyncAthena(
    api_key=os.environ['ATHENA_API_KEY']
)

3. Send a request to Athena, receive a response, and render it in Markdown for readability:

import json
from IPython.display import Markdown
user_name = 'Athena'

USER_MESSAGE_INPUT = f"""
What is the origin of this name: {user_name}?
"""

message = await async_client.message.submit_and_poll(
    content=USER_MESSAGE_INPUT,
    model=Model.GPT_4_TURBO_PREVIEW,
    tools=[Tools.SEARCH, Tools.BROWSE],
)

message_output= json.loads(message.json())["content"]
display(Markdown(message_output))

Learn more with Athena SDK documentation and Sample Notebooks

Refer to docs to learn about all capabilities and parameters of the SDK client.

Explore examples and sample notebooks of end-to-end workflows built with Athena SDK in Athena Notebooks module: Click Athena Notebooks -> File Viewer -> Sample Notebooks.

View image