Skip to content

Commit

Permalink
enh: litellm manifold
Browse files Browse the repository at this point in the history
Co-Authored-By: Artur Zdolinski <15941777+azdolinski@users.noreply.github.com>
  • Loading branch information
tjbck and azdolinski committed Jun 3, 2024
1 parent 8b5e0a0 commit 17bc338
Showing 1 changed file with 32 additions and 4 deletions.
36 changes: 32 additions & 4 deletions examples/pipelines/providers/litellm_manifold_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,15 @@
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.
Expand All @@ -34,7 +37,15 @@ def __init__(self):
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

Expand All @@ -55,9 +66,16 @@ async def on_valves_updated(self):
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 [
{
Expand Down Expand Up @@ -86,10 +104,20 @@ def pipe(
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,
)

Expand Down

0 comments on commit 17bc338

Please sign in to comment.