|
from __future__ import annotations |
|
|
|
import requests |
|
from aiohttp import ClientSession |
|
|
|
from .OpenaiAPI import OpenaiAPI |
|
from ...typing import AsyncResult, Messages, Cookies |
|
from ...requests.raise_for_status import raise_for_status |
|
from ...cookies import get_cookies |
|
|
|
class Cerebras(OpenaiAPI): |
|
label = "Cerebras Inference" |
|
url = "https://inference.cerebras.ai/" |
|
working = True |
|
default_model = "llama3.1-70b" |
|
fallback_models = [ |
|
"llama3.1-70b", |
|
"llama3.1-8b", |
|
] |
|
model_aliases = {"llama-3.1-70b": "llama3.1-70b", "llama-3.1-8b": "llama3.1-8b"} |
|
|
|
@classmethod |
|
def get_models(cls, api_key: str = None): |
|
if not cls.models: |
|
try: |
|
headers = {} |
|
if api_key: |
|
headers["authorization"] = f"Bearer ${api_key}" |
|
response = requests.get(f"https://api.cerebras.ai/v1/models", headers=headers) |
|
raise_for_status(response) |
|
data = response.json() |
|
cls.models = [model.get("model") for model in data.get("models")] |
|
except Exception: |
|
cls.models = cls.fallback_models |
|
return cls.models |
|
|
|
@classmethod |
|
async def create_async_generator( |
|
cls, |
|
model: str, |
|
messages: Messages, |
|
api_base: str = "https://api.cerebras.ai/v1", |
|
api_key: str = None, |
|
cookies: Cookies = None, |
|
**kwargs |
|
) -> AsyncResult: |
|
if api_key is None and cookies is None: |
|
cookies = get_cookies(".cerebras.ai") |
|
async with ClientSession(cookies=cookies) as session: |
|
async with session.get("https://inference.cerebras.ai/api/auth/session") as response: |
|
raise_for_status(response) |
|
data = await response.json() |
|
if data: |
|
api_key = data.get("user", {}).get("demoApiKey") |
|
async for chunk in super().create_async_generator( |
|
model, messages, |
|
api_base=api_base, |
|
impersonate="chrome", |
|
api_key=api_key, |
|
headers={ |
|
"User-Agent": "ex/JS 1.5.0", |
|
}, |
|
**kwargs |
|
): |
|
yield chunk |
|
|
