Brand New Streamlit Component: Token Key Mastery

Author(s): Stavros Theocharis

Originally published on Towards AI.

Introduction

In the dynamic parts of web application development and AI, the Streamlit Python library stands out as an important resource, known for its promptness and effectiveness in the application, especially in software prototype development. One of the most important aspects of Streamlit is its ability to craft specialized components that elevate its functionality. Through this, the community is able to develop and contribute custom components that can broaden the possible use cases for which Streamlit can be used.

My new contribution is the ‘streamlit-token-craft’ component, a brand new custom component I made. It is intended to manage and display token keys within the user interface. This component represents a key segment of token key management solutions integral to enhancing the Streamlit experience.

Such a component (custom or not) was missing from the Streamlit library, and I had to enhance an existing Streamlit app to manage the user’s keys. So, after a bit of research, I saw that I had to build something from scratch. I hope that this component will help developers in the future to build such apps and that some of them will want to contribute to it and step in further.

What is the Streamlit Token Key Management Component?

This specific component is tailored for incorporation into Streamlit applications. It mainly focuses on the administration and visualization of existing token keys within Streamlit applications. If you have used services across the Internet to create your keys, you may remember tables displaying your created keys and the ability to create a new key, rename existing keys, or delete one. With these keys, you are actually able to get access to APIs and various other protected resources. This feature is often a standard one in numerous products that distribute API access tokens to their users.

For example, given the current hype of Large Language Models and their popular quick prototype apps within the Streamlit Community, the “streamlit-craft-token” component is going to be a valuable enhancement in such apps (for their user management dashboards).

For more information, you can access the GitHub repo.

A demo app showcasing the functionality of the component is also deployed in Streamlit community cloud!

Features U+1F389

This part comes directly from the documentation of the component.

• Inline Editing: Edit tokens directly in the table like a ninja! U+1F977
• Dynamic Column Visibility: Play hide and seek with your columns! U+1F648
• Action Handling: Manage token deletion with style. It’s like having a mini-command center. U+1F3AE
• Responsive Design: Looks great on screens of all sizes, even on your grandma’s old monitor! U+1F475U+1F4BB

Use Cases

This component finds its application in a specific use case, from small-scale projects to larger ones. Whether it’s for accessing third-party APIs, cloud services, or internal authentication systems, it offers a robust solution for managing access credentials.

Getting Started U+1F680

In order to get started you only need to install the "streamlit-token-craft" library into your project! This library is alraedy uploaded in pypi!

pip install streamlit-token-craft

Usage Guide

Here’s a very simple example of how to unleash the power of the Token Craft in your app:

import streamlit as st
from token_craft import st_token_table

mock_tokens = [
 {
 "id": "token98a1c077",
 "key": "token98-e316-49d9",
 "display_key": "token98a...e5437d75",
 "name": "Token 1",
 "dateCreated": "2023-12-20",
 "lastUsed": "Never",
 "is_active": True,
 },
]

rendered_tokens = st_token_table(
 tokens=mock_tokens,
 key="token_table",
)

Through this basic code, you will be able to visualize the Tokens table in your app. Of course, more complex functionality is needed to use the component in combination with your own token key creation and handling service.

Hands-On Examples

Example 1: Usage Scenario with token generation simulator

In this example, I am going to showcase how the component can be used in combination with a token-handling service. Additional functionality, such as visualizing success after creating a new key, displaying warning messages, and adding interactive buttons, will demonstrate an almost ‘real’ use case.

The token handling service, simulated through a set of functions, allows for creating, deleting, and renaming tokens within a database (represented as a list of dictionaries in this case). This simulated service can easily be substituted with real API calls for an actual token-handling service.

So first, let’s create our token handling service:

import uuid
from datetime import date

# Simulated database (in-memory for this example)
_token_db = []

uuid will be used to generate random unique "token keys". The temporary db will actually be the _token_db list.

We need to define the essential functions of our service, including retrieving, adding, removing, and renaming tokens in the database. Additionally, a function will be used to format the token keys. This function will modify the display of a token after its initial viewing. For instance, a token is initially shown as ‘dc22ac88-aab4–4109-bf03-d27fbaaed7a8’ will subsequently be displayed in a shortened form like ‘dc22xxxxxxxd7a8’ for enhanced security.

def get_tokens():
 """Retrieve all active tokens from the mock database, replacing 'key' with 'display_key'."""
 return [
 {**token, "key": token["display_key"]} # Replace 'key' with 'display_key'
 for token in _token_db
 if token["is_active"]
 ]

def format_text(text):
 if len(text) <= 8:
 return "xxxx" + text[-4:]
 return text[:4] + "xxxxxxx" + text[-4:]

def add_token(name="New Token"):
 """Add a new token to the mock database."""
 full_key = str(uuid.uuid4())
 # Create a simple hash of the key to display in the table, for example, take the first 8 characters.
 formated_key = format_text(full_key)
 new_token = {
 "id": "token" + str(uuid.uuid4()),
 "key": full_key, # Store the full key
 "display_key": formated_key, # Store the formated version for display
 "name": name,
 "dateCreated": str(date.today()),
 "lastUsed": "Never",
 "is_active": True,
 }
 _token_db.append(new_token)
 return new_token # Return the full new_token to show the full key once

def remove_token(token_key):
 """Deactivate a token in the mock database by its key."""
 token = next((token for token in _token_db if token["key"] == token_key), None)
 if token:
 token["is_active"] = False

def update_token_name(token_key, new_name):
 """Update the name of a specific token in the mock database."""
 token = next((token for token in _token_db if token["key"] == token_key), None)
 if token:
 token["name"] = new_name

Now, our token management service is ready to be used within our Streamlit app! (I suppose that we save it in a .py file named demo_mock_token_service.py)

Let’s initialize our Streamlit app now. To manage and integrate various components across different reruns of the app, we will utilize Streamlit’s session state class.

If you are not familiar with using streamlit’s session state class, I recommend checking out my two detailed articles: Unlocking the Power of Session State in Streamlit (1/2) and Unlocking the Power of Session State in Streamlit (2/2).

Import the needed libraries:

import streamlit as st
import demo_mock_token_service as mts
from token_craft import st_token_table

Initialization of the session state variables and a custom function that handles the changes in session state and reruns the app (it will be needed in several parts within the app code).

# Initialize session state keys if they don't exist
for key in [
 "show_basic_key",
 "show_add_key_form",
 "show_success_message",
 "success_message",
 "show_table",
]:
 if key not in st.session_state:
 st.session_state[key] = (
 True if key in ["show_basic_key", "show_table"] else False
 )


# Function to reset the success message and related state variables
def reset_state(
 show_basic_key=True,
 show_add_key_form=False,
 show_success_message=False,
 show_table=True,
):
 st.session_state.update(
 {
 "show_basic_key": show_basic_key,
 "show_add_key_form": show_add_key_form,
 "show_table": show_table,
 "show_success_message": show_success_message,
 }
 )
 st.rerun()

Now, the part for the main button for creating a new token key:

if st.session_state["show_basic_key"]:
 # Button to show the form for adding a new key
 if st.button("Create new secret key"):
 reset_state(show_basic_key=False, show_add_key_form=True, show_table=False)

When the button is pressed, a form should open to allow the user to input their preferred name. Upon the form’s opening, the initial button and the token table should be hidden.

Then, the form part and the key creation verification should take place. The form requires the user to input at least one character for the name, failing which a warning is displayed. Additionally, the name is limited to 18 characters. If the key creation is successful, a confirmation message is saved in the session state for subsequent display. For users opting not to create a key, a cancel button resets the app to its initial state.

It’s important to note how various parts of the token service are integrated within the code.

# The form for adding a new key
if st.session_state["show_add_key_form"]:
 with st.form(key="add_key_form"):
 new_key_name = st.text_input(
 "Name",
 placeholder="Enter the key name",
 disabled=st.session_state["show_success_message"],
 max_chars=18,
 )

 col1, col2 = st.columns([3, 10])
 with col1:
 submit_button = st.form_submit_button(
 label="Create secret key",
 )
 with col2:
 cancel_button = st.form_submit_button(label="Cancel", type="primary")

 if submit_button:
 if new_key_name: # Check if the name is not empty
 new_token = mts.add_token(new_key_name) # Call function to add token
 # Set the success message and flag
 st.session_state[
 "success_message"
 ] = f"New secret key created successfully! Key: {new_token['key']}"
 # st.session_state["show_success_message"] = True
 reset_state(
 show_basic_key=False,
 show_success_message=True,
 show_table=False,
 )
 else:
 st.warning("Please enter a name for the key.")
 elif cancel_button:
 reset_state()

After successfully creating the token, we display it to the user as an ‘one-time view’ key:

# Show a success message if a new key was added
if st.session_state["show_success_message"]:
 container = st.container(border=True)
 container.write(
 "Stash this secret key in a super safe yet reachable hidey-hole! It's like a one-off magic ticket – once gone from here, it's gone for good. Lose it, and you're off to the wizarding world of making a new one! U+1F5DD️U+2728U+1F9D9‍U+2642️"
 )
 container.success(st.session_state["success_message"])
 # 'OK' button to reset the success state
 if container.button("OK"):
 reset_state(show_success_message=False)

Now, let’s delve into the code for the token table, which is the most important aspect of this tutorial:

if st.session_state["show_table"]:
 # Display the keys in the table with the hashed version
 rendered_tokens = st_token_table(
 tokens=mts.get_tokens(),
 key="token_table",
 )

 # Check for updates from the React component
 if rendered_tokens is not None:
 needs_rerun = False

 for rendered_token in rendered_tokens:
 current_token = next(
 (t for t in mts._token_db if t["display_key"] == rendered_token["key"]),
 None,
 )

 if current_token:
 # If the token should be deactivated
 if not rendered_token["is_active"] and current_token["is_active"]:
 mts.remove_token(current_token["key"]) # Use the full key here
 needs_rerun = True
 # If the token's name should be changed
 elif rendered_token["name"] != current_token["name"]:
 mts.update_token_name(
 current_token["key"], rendered_token["name"]
 ) # Use the full key here
 needs_rerun = True

 if needs_rerun:
 st.rerun()

It’s important to remember that Streamlit components typically operate in the background on React or another JavaScript framework. Consequently, we need a way to render and communicate possible changes between the Streamlit Interface and the React component. Fortunately, most of these processes occur automatically.

In the above blocks, we actually compare the token values based as they are changed by the user, between the component's state and the db.

If a token is active in the service but marked as inactive in our app, it should be converted to inactive (removed for the user's view). With the same logic, if a name is different between the app and the db, the db should be updated accordingly.

The function st.rerun() is used to directly update the token table with the possible changes made. In some apps, you might have noticed that changes made by the user are only reflected after a subsequent action triggers a rerun. The app actually shows the change when another action happens in a second phase (when something changes, we have a rerun in the app). In order to avoid this delayed update, we directly trigger the rerun, ensuring that the changes are visible in the app in “real-time”.

For the detailed code, check out the example app in GitHub.

Example 2: Advanced Implementation with Custom Settings for Columns Choice

Through this component, you also have the ability to hide specific columns on demand.

To accomplish this, we first need to create a sidebar with checkboxes for each column:

column_visibility_settings = {
 "name": st.sidebar.checkbox("Show Name", True),
 "key": st.sidebar.checkbox("Show Key", True),
 "dateCreated": st.sidebar.checkbox("Show Date Created", True),
 "lastUsed": st.sidebar.checkbox("Show Last Used", True),
 "actions": st.sidebar.checkbox("Show Actions", True),
}

And the token table needs to get the column_visibility argument:

rendered_tokens = st_token_table(
 tokens=mts.get_tokens(),
 key="token_table",
 column_visibility=column_visibility_settings,
)

For the detailed code, check out the example app on GitHub.

Conclusion

In conclusion, the "streamlit-token-craft" stands as an important component in streamlit web applications that target the user key management parts. This component introduces a simplified approach to managing and displaying token keys in the app, filling a gap not addressed by previous components.

By providing hands-on examples, the practicality and ease of use are being demonstrated. I encourage you to experiment with this new component for managing and displaying keys in your apps.

Possible feedback is more than welcome! If you want to contribute to the project, please read the contribution guide.

U+1F31F If you enjoyed this article, follow me to read more!

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI