diff --git a/examples/pipelines/providers/litellm_manifold_pipeline.py b/examples/pipelines/providers/litellm_manifold_pipeline.py index 0b6771f..80fdf11 100644 --- a/examples/pipelines/providers/litellm_manifold_pipeline.py +++ b/examples/pipelines/providers/litellm_manifold_pipeline.py @@ -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, )