AI21 API Starter Sample for Gen AI Beginners

ai21-starter.py

"""
AI21 API Example: Get an answer to a question
"""

from ai21 import AI21Client
from ai21.models.chat import ChatMessage
import os

# Import necessary modules to work with the AI21 API
# AI21Client: Used to interact with the AI21 service
# ChatMessage: Represents a message in a chat conversation with the model
# os: Used to interact with the operating system's environment variables

# Set the AI21 API key as an environment variable.
# This way, the AI21Client can automatically find and use it.
os.environ["AI21_STUDIO_API_KEY"] = "XXX"  # Replace with your actual key

# Create an AI21Client instance, which is the main object to communicate with the AI21 API.
# It uses the API key from the environment variable for authentication.
client = AI21Client(api_key=os.environ["AI21_STUDIO_API_KEY"])

# Constants for configuration, adjust as needed
MAX_TOKENS = 30
TEMPERATURE = 0.7

def get_answer(prompt, max_tokens=MAX_TOKENS, temperature=TEMPERATURE):
    """
    Define a function called 'get_answer' that takes a 'prompt' (your question) as input.

    Parameters:
    - prompt (str): The question to ask the AI.
    - max_tokens (int): Limits the length of the generated response.
    - temperature (float): Controls the randomness of the output (higher values mean more creative responses).

    Returns:
    - str: The answer from the AI.
    """
    
    response = client.chat.completions.create(
        model="jamba-instruct",
        messages=[ChatMessage(role="user", content=prompt)],
        max_tokens=MAX_TOKENS,
        temperature=TEMPERATURE
    )
    # Send the prompt to the AI21 model and get a response.
    # - model: Specifies the AI21 model to use ("jamba-instruct" is an instruction-following model).
    # - messages: A list of ChatMessage objects, representing the conversation history. Here, only one message is used, which is the user's prompt.
    # - max_tokens: Limits the length of the generated response.
    # - temperature: Controls the randomness of the output (higher values mean more creative responses).
    
    return response.choices[0].message.content  # Extract the answer directly
    # Return the actual answer text from the response object.

# Get an answer
question = "Who was the first emperor of Rome?"
# Define the question you want to ask the AI.

answer = get_answer(question)
# Call the 'get_answer' function with your question and store the answer.

print(answer)
# Print the answer to the console.