andrewyng--aisuite
ERIK
37 行
1.5 KiB
Python
37 行
1.5 KiB
Python
import os
|
|
from aisuite.provider import Provider
|
|
from openai import OpenAI
|
|
|
|
|
|
class FeatherlessProvider(Provider):
|
|
def __init__(self, **config):
|
|
"""
|
|
Initialize the Featherless provider with the given configuration.
|
|
Pass the entire configuration dictionary to the OpenAI client constructor.
|
|
"""
|
|
# Ensure API key is provided either in config or via environment variable
|
|
config.setdefault("api_key", os.getenv("FEATHERLESS_API_KEY"))
|
|
if not config["api_key"]:
|
|
raise ValueError(
|
|
"Featherless API key is missing. Please provide it in the config or set the FEATHERLESS_API_KEY environment variable."
|
|
)
|
|
|
|
# NOTE: We could choose to remove above lines for api_key since OpenAI will automatically
|
|
# infer certain values from the environment variables.
|
|
# Eg: OPENAI_API_KEY, OPENAI_ORG_ID, OPENAI_PROJECT_ID, OPENAI_BASE_URL, etc.
|
|
|
|
# Pass the entire config to the OpenAI client constructor
|
|
self.client = OpenAI(
|
|
base_url="https://api.featherless.ai/v1/",
|
|
api_key=config["api_key"],
|
|
)
|
|
|
|
def chat_completions_create(self, model, messages, **kwargs):
|
|
# Any exception raised by OpenAI will be returned to the caller.
|
|
# Maybe we should catch them and raise a custom LLMError.
|
|
return self.client.chat.completions.create(
|
|
model=model,
|
|
messages=messages,
|
|
**kwargs # Pass any additional arguments to the OpenAI API
|
|
)
|