cloudflare ai initial draft

examples/pipelines/providers/cloudflare_ai_pipeline.py

@@ -0,0 +1,102 @@
+from typing import List, Union, Generator, Iterator
+from schemas import OpenAIChatMessage
+from pydantic import BaseModel
+
+import os
+import requests
+
+
+class Pipeline:
+    class Valves(BaseModel):
+        CLOUDFLARE_ACCOUNT_ID: str = ""
+        CLOUDFLARE_API_KEY: str = ""
+        CLOUDFLARE_MODELS: 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.valves = self.Valves(
+            **{
+                "CLOUDFLARE_API_KEY": os.getenv(
+                    "CLOUDFLARE_API_KEY", "your-openai-api-key-here"
+                ),
+                "CLOUDFLARE_MODELS": os.getenv(
+                    "CLOUDFLARE_MODELS",
+                    "@cf/meta/llama-3.1-8,@cf/deepseek-ai/deepseek-math-7b-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__}")
+        pass
+
+    async def on_shutdown(self):
+        # This function is called when the server is stopped.
+        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}
+
+        if "user" in payload:
+            del payload["user"]
+        if "chat_id" in payload:
+            del payload["chat_id"]
+        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",
+                json=payload,
+                headers=headers,
+                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}"