Nat Taylor — Blog, AI, Product Management & Tinkering

Test Drive: litellm

Published on .

Today I’m test driving LiteLLM (Python SDK to call 100+ LLM APIs in OpenAI format). Although it’s possible to do extraordinary things with a single lab’s model, it’s also well known that different models perform differently on different tasks. I don’t want my code littered with slightly different incantations of completions, so LiteLLM’s unified approach seems great. I’ll continue with my companies dataset and use the language models to label the sector.

The process is extremely simple:

  1. Configure environment variables with keys
  2. Invoke LiteLLM

I’m not publishing the first step :) but I was pleased to learn that its only about a 1-minute process to get a Gemini API key at https://aistudio.google.com/app/apikey

Invoking LiteLLM is simple, and its easy to see why the unified approach is appealing.

desc = """Cummins Inc. designs, manufactures, distributes and services diesel and natural gas engines and engine-related component products. The Company's segments include Engine, Distribution, Components and Power Systems. The Engine segment manufactures and markets a range of diesel and natural gas powered engines under the Cummins brand name, as well as certain customer brand names, for the heavy and medium-duty truck, bus, recreational vehicle (RV), light-duty automotive and agricultural markets. The Distribution segment consists of the product lines, which service and/or distribute a range of products and services, including parts, engines, power generation and service. The Components segment supplies products, including aftertreatment systems, turbochargers, filtration products and fuel systems for commercial diesel applications. The Power Systems segment consists of businesses, including Power generation, Industrial and Generator technologies."""

from litellm import completion

models = [
    "claude-3-haiku-20240307",
    "gpt-4o-mini",
    "gemini/gemini-pro",
]

for model in models:
    print(model)
    print(completion(
      model=model,
      messages=[{ "content": f"Given the company description, what is the sector?  Answer in JSON no backticks with a key for sector.\n{desc}""","role": "user"}]
    ).choices[0].message.content)
    print()

This produces the following:

claude-3-haiku-20240307
{
  "sector": "Industrials"
}

gpt-4o-mini
{
  "sector": "Manufacturing"
}

gemini/gemini-pro
{"sector": "Industrials"}

Post Navigation

«
»