AISuite: Simplifying AI Development with Unified LLM Integration

In the fast-paced world of Artificial Intelligence (AI), the challenge of integrating multiple large language models (LLMs) from different providers can be a real obstacle for developers. The constant need to switch between APIs and adjust code can slow down innovation, making it harder to experiment with various models and find the best solution for a given task.

This is where aisuite, a new open-source Python package from Andrew Ng’s team, comes in to revolutionize how developers build AI products.

What is AISuite?

aisuite is a Python library designed to simplify the process of interacting with different LLM providers such as Watsonx, OpenAI, HuggingFace, etc.

With aisuite, developers can easily switch between models just by changing a single string. For example, you can switch from aws:anthropic.claude-v2 to watsonx:ibm/granite-3–8b-instruct without having to rewrite any of your existing code. This flexibility allows for efficient testing and integration of different models, enabling robust AI application development while also allowing developers to analyze outputs from different models and compare them across various LLMs and providers.

Note: The API call uses the following format — model_provider:model_name.

Key Features

• Unified Interface: Interact with multiple LLMs using a single API, saving time on integration.
• Model-Agnostic Approach: Switch between different models like watsonx:mistralai/mistral-large or groq:mixtral-8x7b-32768 without altering code.
• Wide Provider Support: Compatible with a variety of top LLM providers such as OpenAI, Watsonx, Ollama, Anthropic, HuggingFace, AWS, Google, and more.
• Easy Setup: Install easily with or without provider-specific packages, and set API keys via environment variables or configuration.
• Open Source: Licensed under MIT, it encourages community contributions and ongoing development.

Code Walkthrough

Now, let’s look at a code example to understand how aisuite simplifies interactions with multiple LLMs. It provides a standardized interface that abstracts interactions, acting as a thin wrapper around Python client libraries. This enables seamless switching and comparison between different LLMs and providers without modifying the underlying code.

Load environment variables

To get started, you will need API Keys for the providers you intend to use

from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv())

Below is an example of what your .env file should look like:

# WatsonX Credentials
WATSONX_API_KEY=
WATSONX_SERVICE_URL=
WATSONX_PROJECT_ID=

# Groq Credentials
GROQ_API_KEY=

# AWS Credentials
AWS_ACCESS_KEY_ID=
AWS_SECRET_ACCESS_KEY=
AWS_REGION=

# Ollama API URL
OLLAMA_API_URL=

Initialize the AI client

import aisuite as ai
client = ai.Client()

Define a Function to interact with different LLMs

The below function will be used to send a query to an LLM and retrieve its response.

def query_llm(message, model, sys_message="You are a helpful assistant. Briefly answer the user query."):

    messages = [
        {"role": "system", "content": sys_message},
        {"role": "user", "content": message}
    ]

    response = client.chat.completions.create(model = model, messages = messages)
    return response.choices[0].message.content.strip()

Query different LLM providers for a set of questions

Let’s specify some sample questions and different LLMs from different providers we intend to compare.

questions = [
    "How does a DNS server work?",
    "What is multi-factor authentication (MFA)?",
    "What is CI/CD in DevOps?"
]

Specify the LLMs in the format — <model_provider>:<model_name>

llms = [
    "groq:gemma2-9b-it",
    "groq:mixtral-8x7b-32768",
    "ollama:llama3.1:8b",
    "aws:anthropic.claude-v2",
    "watsonx:ibm/granite-3-8b-instruct",
    "watsonx:mistralai/mistral-large"
]

Now, loop through questions and LLMs to query responses and measure execution times. Log the provider, model, response time, and result for each query.

execution_times = []
responses = []

for ques in questions:
    print(f"{Fore.BLUE}Question: {ques}{Style.RESET_ALL}\n")
    for llm in llms:
        provider, model_id = llm.split(":", 1)  # Split provider and model_id
        start_time = time.time()
        response = query_llm(message=ques, model=llm)
        end_time = time.time()
        execution_time = end_time - start_time
        responses.append(response)
        execution_times.append(execution_time)
        print(f"{Fore.GREEN}Provider: {provider}{Style.RESET_ALL}, "
              f"{Fore.CYAN}Model ID: {model_id}{Style.RESET_ALL} - "
              f"{Fore.YELLOW}{execution_time:.2f} seconds{Style.RESET_ALL}: {response}")
    print("\n")

Create a DataFrame to store questions, LLM responses, execution times, and model details for better analysis.

expanded_questions = []
for i in range(len(questions)):
    expanded_questions.extend([questions[i]] * len(llms))

data = {
    'Question': expanded_questions,
    'Provider:Model_ID': llms * len(questions),
    'Execution Time (in seconds)': execution_times,
    'LLM Response': responses
}


df = pd.DataFrame(data)
df.index = df.index + 1

styled_df = df.style.set_table_styles(
    [{'selector': 'th', 'props': [('text-align', 'center')]},  
     {'selector': 'td', 'props': [('text-align', 'center')]}]   
).set_properties(**{'text-align': 'center'}) 
styled_df

LLM Comparison DataFrame

Why Use AISuite?

AISuite simplifies AI workflows by allowing you to compare models and switch between them effortlessly, saving both time and effort. It allows you to focus on optimizing outputs rather than managing model integration. It’s scalable for both small projects and large-scale deployments, making it a valuable tool for developers and researchers.

Conclusion

AISuite makes AI development simpler and more flexible by centralizing access to multiple LLMs. Its model-agnostic approach and easy integration offer a streamlined experience, freeing you from provider-specific constraints. Currently, aisuite is focused on chat completions, but the team has plans to expand its functionality. But with plans for further expansion, AISuite is poised to be an essential tool for the AI community.

🚀 Explore AISuite on GitHub and start building smarter AI solutions today!