cloudlfare ai pipeline

This commit is contained in:
Jonas 2024-08-16 04:04:25 +00:00
parent d86ce893fd
commit c50d4eb8f8

View File

@ -1,7 +1,6 @@
from typing import List, Union, Generator, Iterator
from schemas import OpenAIChatMessage
from pydantic import BaseModel
import os
import requests
@ -10,40 +9,33 @@ class Pipeline:
class Valves(BaseModel):
CLOUDFLARE_ACCOUNT_ID: str = ""
CLOUDFLARE_API_KEY: str = ""
CLOUDFLARE_MODELS: str = ""
CLOUDFLARE_MODEL: str = ""
pass
def __init__(self):
self.type = "manifold"
# Optionally, you can set the id and name of the pipeline.
# Best practice is to not specify the id so that it can be automatically inferred from the filename, so that users can install multiple versions of the same pipeline.
# The identifier must be unique across all pipelines.
# The identifier must be an alphanumeric string that can include underscores or hyphens. It cannot contain spaces, special characters, slashes, or backslashes.
# self.id = "openai_pipeline"
self.name = "Cloudflare AI: "
self.name = "Cloudlfare AI"
self.valves = self.Valves(
**{
"CLOUDFLARE_ACCOUNT_ID": os.getenv(
"CLOUDFLARE_ACCOUNT_ID",
"your-account-id",
),
"CLOUDFLARE_API_KEY": os.getenv(
"CLOUDFLARE_API_KEY", "your-openai-api-key-here"
"CLOUDFLARE_API_KEY", "your-cloudflare-api-key"
),
"CLOUDFLARE_MODELS": os.getenv(
"CLOUDFLARE_MODEL": os.getenv(
"CLOUDFLARE_MODELS",
"@cf/meta/llama-3.1-8,@cf/deepseek-ai/deepseek-math-7b-instruct",
"@cf/meta/llama-3.1-8b-instruct",
),
},
}
)
self.pipelines = self.get_cloudflare_models()
pass
def get_cloudflare_models(self):
models = [
{"id": model, "name": model}
for model in self.valves.CLOUDFLARE_MODELS.split(",")
]
return models
async def on_startup(self):
# This function is called when the server is started.
print(f"on_startup:{__name__}")
@ -54,26 +46,17 @@ class Pipeline:
print(f"on_shutdown:{__name__}")
pass
async def on_valves_updated(self):
# This function is called when the valves are updated.
print(f"on_valves_updated:{__name__}")
self.pipelines = self.get_cloudflare_models()
pass
def pipe(
self, user_message: str, model_id: str, messages: List[dict], body: dict
) -> Union[str, Generator, Iterator]:
# This is where you can add your custom pipelines like RAG.
print(f"pipe:{__name__}")
print(messages)
print(user_message)
headers = {}
headers["Authorization"] = f"Bearer {self.valves.CLOUDFLARE_API_KEY}"
headers["Content-Type"] = "application/json"
payload = {**body, "model": model_id}
payload = {**body, "model": self.valves.CLOUDFLARE_MODEL}
if "user" in payload:
del payload["user"]
@ -82,8 +65,6 @@ class Pipeline:
if "title" in payload:
del payload["title"]
print(payload)
try:
r = requests.post(
url=f"https://api.cloudflare.com/client/v4/accounts/{self.valves.CLOUDFLARE_ACCOUNT_ID}/ai/v1/chat/completions",