AI21 Studio Tutorial: Mastering the Chat Completions API

#!pip install -q ai21 python-dotenv

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

# Load environment variables from a .env file
load_dotenv()

# The client automatically uses the AI21_API_KEY environment variable
client = AI21Client()

def generate_text_with_chat(user_prompt):
    """Generates text using the Chat Completions API."""
    try:
        messages = [ChatMessage(role="user", content=user_prompt)]
        
        response = client.chat.completions.create(
            model="jamba-mini",
            messages=messages,
            max_tokens=150
        )
        
        return response.choices[0].message.content
    except Exception as e:
        return f"An error occurred: {e}"

# Example usage
prompt = "Write a short, optimistic paragraph about the future of renewable energy."
generated_text = generate_text_with_chat(prompt)
print(f"Generated Text:\n{generated_text}")

def paraphrase_text_with_chat(text_to_paraphrase):
    """Paraphrases text using an instruction-based chat call."""
    try:
        system_prompt = "You are an expert editor. Your task is to paraphrase the given text, ensuring the original meaning is perfectly preserved but the wording is different."
        user_prompt = f"Please paraphrase the following text: \"{text_to_paraphrase}\""
        
        messages = [
            ChatMessage(role="system", content=system_prompt),
            ChatMessage(role="user", content=user_prompt)
        ]
        
        response = client.chat.completions.create(
            model="jamba-mini",
            messages=messages,
            max_tokens=100
        )
        
        return response.choices[0].message.content
    except Exception as e:
        return f"An error occurred: {e}"

# Example usage
original_text = "The quick brown fox jumps over the lazy dog."
paraphrased_text = paraphrase_text_with_chat(original_text)
print(f"Original Text: {original_text}")
print(f"Paraphrased Text: {paraphrased_text}")

def summarize_text_with_chat(long_text):
    """Summarizes text using an instruction-based chat call."""
    try:
        system_prompt = "You are a helpful assistant that summarizes long texts into a single, concise paragraph."
        user_prompt = f"Please summarize the following text: \n\n{long_text}"
        
        messages = [
            ChatMessage(role="system", content=system_prompt),
            ChatMessage(role="user", content=user_prompt)
        ]
        
        response = client.chat.completions.create(
            model="jamba-mini",
            messages=messages,
            max_tokens=200
        )
        
        return response.choices[0].message.content
    except Exception as e:
        return f"An error occurred: {e}"

# Example usage
text_to_summarize = """AI21 Labs is an AI company specializing in Natural Language Processing. Founded in 2017, its goal is to build AI systems with an unprecedented capacity to understand and generate natural language. They have developed a family of large language models, including the Jurassic and Jamba series, which are among the largest and most sophisticated in the world. These models power applications for text generation, summarization, and paraphrasing."""
summary = summarize_text_with_chat(text_to_summarize)
print(f"Summary:\n{summary}")

def get_contextual_answer_with_chat(context, question):
    """Gets a contextual answer using an instruction-based chat call."""
    try:
        system_prompt = "You are a question-answering assistant. You must answer the user's question based ONLY on the provided context. If the answer is not in the context, say 'The answer is not available in the provided text.'"
        user_prompt = f"Context: {context}\n\nQuestion: {question}"
        
        messages = [
            ChatMessage(role="system", content=system_prompt),
            ChatMessage(role="user", content=user_prompt)
        ]
        
        response = client.chat.completions.create(
            model="jamba-mini",
            messages=messages,
            max_tokens=100
        )
        
        return response.choices[0].message.content
    except Exception as e:
        return f"An error occurred: {e}"

# Example usage
qa_context = "The capital of France is Paris. Paris is known for its art, fashion, and culture."
qa_question = "What is the primary industry in Paris?"
answer = get_contextual_answer_with_chat(qa_context, qa_question)
print(f"Context: {qa_context}")
print(f"Question: {qa_question}")
print(f"Answer: {answer}")