""" title: Cloudfare WorkerAI Manifold Pipe author: hnico21 funding_url: https://github.com/open-webui version: 1.0 license: MIT """ from pydantic import BaseModel, Field from typing import Optional, Union, Generator, Iterator from open_webui.utils.misc import get_last_user_message from open_webui.utils.misc import pop_system_message import os import requests class Pipe: class Valves(BaseModel): NAME_PREFIX: str = Field( default="WorkerAI/", description="The prefix applied before the model names.", ) ACCOUNT_ID: str = Field( default="AccoutID", description="Account_ID", ) GATEWAY_ID: str = Field( default="None", description="Gateway_ID, None if not used", ) API_KEY: str = Field( default="", description="Required API key to access your Worker.", ) def __init__(self): self.type = "manifold" self.valves = self.Valves() def pipes(self): return [ { "id": "@cf/meta/llama-3.1-8b-instruct", "name": f"{self.valves.NAME_PREFIX}Llama 3.1 8b", }, { "id": "@cf/meta/llama-2-7b-chat-fp16", "name": f"{self.valves.NAME_PREFIX}Llama 3.1 8b fp16", }, { "id": "@cf/meta/llama-3.2-11b-vision-instruct", "name": f"{self.valves.NAME_PREFIX}Llama 3.2 11b", }, { "id": "@cf/meta/llama-3.1-70b-instruct", "name": f"{self.valves.NAME_PREFIX}Llama 3.1 70b", }, ] def pipe(self, body: dict, __user__: dict) -> Union[str, Generator, Iterator]: print(f"pipe:{__name__}") if not self.valves.API_KEY: raise Exception("Please configure valves first.") headers = { "Authorization": f"Bearer {self.valves.API_KEY}", } system_message, messages = pop_system_message(body.get("messages", [])) system_prompt = "You are a helpful assistant." if system_message is not None: system_prompt = system_message["content"] model_id = body["model"] model_id = model_id[len(self.valves.NAME_PREFIX) :] base_url = f"https://api.cloudflare.com/client/v4/accounts/{self.valves.ACCOUNT_ID}/ai/run/{model_id}" if not self.valves.GATEWAY_ID == "None": base_url = f"https://gateway.ai.cloudflare.com/v1/{self.valves.ACCOUNT_ID}/{self.valves.GATEWAY_ID}/workers-ai/{model_id}" input = { "messages": [{"role": "system", "content": system_prompt}, *messages], "return_citations": True, "return_images": True, } try: r = requests.post( url=base_url, json=input, headers=headers, stream=True, ) r.raise_for_status() helper = r.json()["result"] response = helper["response"] return response except Exception as e: return f"Error: {e}"