""" title: Cerebras description: https://cloud.cerebras.ai author: Olof Larsson author_url: https://olof.tech/cerebras-for-open-webui version: 1.8.0 license: MIT """ import os import requests from typing import List, Union, Generator, Iterator, Optional from pydantic import BaseModel, Field # Learn more here: # https://olof.tech/cerebras-for-open-webui class Pipe: class Valves(BaseModel): API_KEY: str = Field( default=os.getenv("CEREBRAS_API_KEY", ""), description="Get your key here: https://cloud.cerebras.ai", ) def __init__(self): self.type = "manifold" self.id = "cerebras" self.name = "cerebras/" self.valves = self.Valves() self.base_url = "https://api.cerebras.ai/v1" # Q: Why hard-code the model ids? # A: Because it yields the best performance. OpenWebUI is very sensitive to slowness here. It impacts page loads. def get_model_ids(self): # How to update: # curl "https://api.cerebras.ai/v1/models" -H "Authorization: Bearer ${CEREBRAS_API_KEY}" | jq '.data | sort_by(.id) | map(.id)' return ["llama-3.3-70b", "llama3.1-8b"] def pipes(self) -> List[dict]: model_ids = self.get_model_ids() return [{"id": model_id, "name": model_id} for model_id in model_ids] def pipe( self, body: dict, ) -> Union[str, Generator, Iterator]: body["model"] = body["model"].removeprefix(self.id + ".") try: r = requests.post( url=f"{self.base_url}/chat/completions", json=body, headers={ "Content-Type": "application/json", "Authorization": f"Bearer {self.valves.API_KEY}", }, stream=True, ) r.raise_for_status() if body["stream"]: return r.iter_lines() else: return r.json() except Exception as e: return f"Error: {e}"