LLM API

The LLM class is the core interface in LitAI for interacting with any large language model - OpenAI, Claude, Gemini, Mistral, or open-source. It abstracts away provider-specific setup and lets the user chat with models using a unified API. With built-in support for streaming, retries, fallbacks, memory, and usage tracking, the LLM class makes it easy to build production-grade LLM apps with just a few lines of code.

ImportImport LitAI for interacting with language modelsInitializeInitialize an LLM instance with any available modelAuthenticationProvide the Lightning AI account teamspace name, api key and user id.ChatSend a chat message to the LLM and print its response

1/4

2/4

3/4

4/4

from litai import LLM


llm = LLM(model="openai/gpt-4",
        # Find the api keys in the account's global setting
        teamspace="owner-name/teamspace-name",
        lightning_api_key="xxx-xxx-xxx",
        lightning_user_id="xxx-xxx-xxx"
    )


output = llm.chat("who are you?")
print(output)
# Example response: "I'm an AI developed by OpenAI."

Parameters

model

LLM provider to be used for generation

LLM(model="openai/gpt-4")

Example: This example sends user input to OpenAI GPT-4 LLM

from litai import LLM

llm = LLM(model="openai/gpt-4")
print(llm.chat("who are you?"))
# I'm an AI by OpenAI

fallback_models

A list of models which can be used in case of main model provider fails due to rate limits, or any kind of error.

LLM(model="openai/gpt-4", fallback_models=["mistral/mixtral", "anthropic/claude-3"])

Example: The following example will use mixtral, and then claude-3 in case gpt-4 fails for any reason.

from litai import LLM

llm = LLM(
    model="openai/gpt-4",
    fallback_models=["mistral/mixtral", "anthropic/claude-3"],
)

print(llm.chat("How do I fine-tune an LLM?"))

teamspace

Credits to run the model will be deducted from the provided teamspace.

llm = LLM(name="openai/gpt-4", teamspace="owner-name/teamspace-name")

Example: This example will use the user provided teamspace credits.

from litai import LLM

llm = LLM(name="openai/gpt-4", teamspace="owner-name/teamspace-name")
response = llm.chat("Hello, how are you?")
print(response)

lightning_api_key

To run LitAI outside of Lightning Studio, you'll need your Lightning AI account API key for authentication. You can find it by clicking on your profile, navigating to Global Settings → Keys, and copying the API key listed there.

llm = LLM(..., lightning_api_key="xxx-xxx-xxx")

Example: The following code authenticates with the api key.

from litai import LLM


llm = LLM(model="openai/gpt-4",
        # Find the api keys in the account's global setting
        teamspace="owner-name/teamspace-name",
        lightning_api_key="xxx-xxx-xxx",
        lightning_user_id="xxx-xxx-xxx"
    )

output = llm.chat("who are you?")
print(output)
# Example response: "I'm an AI developed by OpenAI."

lightning_user_id

To run LitAI outside of Lightning Studio, you'll need your Lightning AI account user id for identification. You can find it by clicking on your profile, navigating to Global Settings → Keys, and copying the USER ID listed there.

llm = LLM(model="openai/gpt-4", lightning_api_key="xxx-xxx-xxx", lightning_user_id="xxx-xxx-xxx")

Example: The following code authenticates with the api key and user id.

from litai import LLM


llm = LLM(model="openai/gpt-4",
        # Find the api keys in the account's global setting
        teamspace="owner-name/teamspace-name",
        lightning_api_key="xxx-xxx-xxx",
        lightning_user_id="xxx-xxx-xxx"
    )

output = llm.chat("who are you?")
print(output)
# Example response: "I'm an AI developed by OpenAI."

max_retries

LitAI retries the LLM API call in cases of failure. max_retries can be provided as a threshold. Defaults to 3.

llm = LLM(
    model="openai/gpt-4",
    max_retries=4,
)

Example: The following LLM will be called 4 times before returning the error to the user in case of failure.

from litai import LLM

llm = LLM(
    model="openai/gpt-4",
    max_retries=4,
)

print(llm.chat("How do I fine-tune an LLM?"))

enable_async

Run LLM generation concurrently by setting enable_async=True

await llm.chat("How do I fine-tune an LLM?")

Example: The following code makes two LLM calls concurrently reduce the total time to generate the output.

import asyncio
from litai import LLM

llm = LLM(
    model="openai/gpt-4",
    enable_async=True
)

async def main():
    result1 = llm.chat("How do I fine-tune an LLM?")
    result2 = llm.chat("where is Paris?")
    results = await asyncio.gather(result1, result2)
    print(results)

if __name__=="__main__":
    asyncio.run(main())

Was this page helpful?