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 typing import List, Union, Generator, Iterator
from schemas import OpenAIChatMessage from schemas import OpenAIChatMessage
from pydantic import BaseModel from pydantic import BaseModel
import os import os
import requests import requests
@ -10,40 +9,33 @@ class Pipeline:
class Valves(BaseModel): class Valves(BaseModel):
CLOUDFLARE_ACCOUNT_ID: str = "" CLOUDFLARE_ACCOUNT_ID: str = ""
CLOUDFLARE_API_KEY: str = "" CLOUDFLARE_API_KEY: str = ""
CLOUDFLARE_MODELS: str = "" CLOUDFLARE_MODEL: str = ""
pass pass
def __init__(self): def __init__(self):
self.type = "manifold"
# Optionally, you can set the id and name of the pipeline. # 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. # 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 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. # 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.id = "openai_pipeline"
self.name = "Cloudflare AI: " self.name = "Cloudlfare AI"
self.valves = self.Valves( self.valves = self.Valves(
**{ **{
"CLOUDFLARE_ACCOUNT_ID": os.getenv(
"CLOUDFLARE_ACCOUNT_ID",
"your-account-id",
),
"CLOUDFLARE_API_KEY": os.getenv( "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", "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 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): async def on_startup(self):
# This function is called when the server is started. # This function is called when the server is started.
print(f"on_startup:{__name__}") print(f"on_startup:{__name__}")
@ -54,26 +46,17 @@ class Pipeline:
print(f"on_shutdown:{__name__}") print(f"on_shutdown:{__name__}")
pass 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( def pipe(
self, user_message: str, model_id: str, messages: List[dict], body: dict self, user_message: str, model_id: str, messages: List[dict], body: dict
) -> Union[str, Generator, Iterator]: ) -> Union[str, Generator, Iterator]:
# This is where you can add your custom pipelines like RAG. # This is where you can add your custom pipelines like RAG.
print(f"pipe:{__name__}") print(f"pipe:{__name__}")
print(messages)
print(user_message)
headers = {} headers = {}
headers["Authorization"] = f"Bearer {self.valves.CLOUDFLARE_API_KEY}" headers["Authorization"] = f"Bearer {self.valves.CLOUDFLARE_API_KEY}"
headers["Content-Type"] = "application/json" headers["Content-Type"] = "application/json"
payload = {**body, "model": model_id} payload = {**body, "model": self.valves.CLOUDFLARE_MODEL}
if "user" in payload: if "user" in payload:
del payload["user"] del payload["user"]
@ -82,8 +65,6 @@ class Pipeline:
if "title" in payload: if "title" in payload:
del payload["title"] del payload["title"]
print(payload)
try: try:
r = requests.post( r = requests.post(
url=f"https://api.cloudflare.com/client/v4/accounts/{self.valves.CLOUDFLARE_ACCOUNT_ID}/ai/v1/chat/completions", url=f"https://api.cloudflare.com/client/v4/accounts/{self.valves.CLOUDFLARE_ACCOUNT_ID}/ai/v1/chat/completions",