OpenWebUI/pipelines
1
0
Fork 0
mirror of https://github.com/open-webui/pipelines synced 2025-06-26 18:15:58 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

enh: litellm manifold

Browse Source 
Co-Authored-By: Artur Zdolinski <15941777+azdolinski@users.noreply.github.com>
This commit is contained in:
Timothy J. Baek 2024-06-03 13:24:57 -07:00
parent 8b5e0a05e9
commit 17bc338c4a
1 changed files with 32 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
examples/pipelines/providers/litellm_manifold_pipeline.py
View File

@ -11,12 +11,15 @@ from typing import List, Union, Generator, Iterator
from schemas import OpenAIChatMessage
from pydantic import BaseModel
import requests
import os
class Pipeline:
class Valves(BaseModel):
LITELLM_BASE_URL: str
LITELLM_BASE_URL: str = ""
LITELLM_API_KEY: str = ""
LITELLM_PIPELINE_DEBUG: bool = False
def __init__(self):
# You can also set the pipelines that are available in this pipeline.
@ -34,7 +37,15 @@ class Pipeline:
self.name = "LiteLLM: "
# Initialize rate limits
self.valves = self.Valves(**{"LITELLM_BASE_URL": "http://localhost:4001"})
self.valves = self.Valves(
**{
"LITELLM_BASE_URL": os.getenv(
"LITELLM_BASE_URL", "http://localhost:4001"
),
"LITELLM_API_KEY": os.getenv("LITELLM_API_KEY", "your-api-key-here"),
"LITELLM_PIPELINE_DEBUG": os.getenv("LITELLM_PIPELINE_DEBUG", False),
}
)
self.pipelines = []
pass
@ -55,9 +66,16 @@ class Pipeline:
pass
def get_litellm_models(self):
headers = {}
if self.valves.LITELLM_API_KEY:
headers["Authorization"] = f"Bearer {self.valves.LITELLM_API_KEY}"
if self.valves.LITELLM_BASE_URL:
try:
r = requests.get(f"{self.valves.LITELLM_BASE_URL}/v1/models")
r = requests.get(
f"{self.valves.LITELLM_BASE_URL}/v1/models", headers=headers
)
models = r.json()
return [
{
@ -86,10 +104,20 @@ class Pipeline:
print(f"# Message: {user_message}")
print("######################################")
headers = {}
if self.valves.LITELLM_API_KEY:
headers["Authorization"] = f"Bearer {self.valves.LITELLM_API_KEY}"
try:
payload = {**body, "model": model_id, "user_id": body["user"]["id"]}
payload.pop("chat_id", None)
payload.pop("user", None)
payload.pop("title", None)
r = requests.post(
url=f"{self.valves.LITELLM_BASE_URL}/v1/chat/completions",
json={**body, "model": model_id, "user_id": body["user"]["id"]},
json=payload,
headers=headers,
stream=True,
)