mirror of
https://github.com/open-webui/pipelines
synced 2025-06-26 18:15:58 +00:00
chore: mv examples out of pipelines
This commit is contained in:
121
examples/anthropic_manifold_pipeline.py
Normal file
121
examples/anthropic_manifold_pipeline.py
Normal file
@@ -0,0 +1,121 @@
|
||||
"""
|
||||
title: Anthropic Manifold Pipeline
|
||||
author: justinh-rahb
|
||||
date: 2024-05-27
|
||||
version: 1.0
|
||||
license: MIT
|
||||
description: A pipeline for generating text using the Anthropic API.
|
||||
requirements: requests, anthropic
|
||||
environment_variables: ANTHROPIC_API_KEY
|
||||
"""
|
||||
|
||||
import os
|
||||
from anthropic import Anthropic, RateLimitError, APIStatusError, APIConnectionError
|
||||
|
||||
from schemas import OpenAIChatMessage
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from pydantic import BaseModel
|
||||
import requests
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
self.type = "manifold"
|
||||
self.id = "anthropic"
|
||||
self.name = "anthropic/"
|
||||
|
||||
class Valves(BaseModel):
|
||||
ANTHROPIC_API_KEY: str
|
||||
|
||||
self.valves = Valves(**{"ANTHROPIC_API_KEY": os.getenv("ANTHROPIC_API_KEY")})
|
||||
self.client = Anthropic(api_key=self.valves.ANTHROPIC_API_KEY)
|
||||
|
||||
def get_anthropic_models(self):
|
||||
# In the future, this could fetch models dynamically from Anthropic
|
||||
return [
|
||||
{"id": "claude-3-haiku-20240307", "name": "claude-3-haiku"},
|
||||
{"id": "claude-3-opus-20240229", "name": "claude-3-opus"},
|
||||
{"id": "claude-3-sonnet-20240229", "name": "claude-3-sonnet"},
|
||||
# Add other Anthropic models here as they become available
|
||||
]
|
||||
|
||||
async def on_startup(self):
|
||||
print(f"on_startup:{__name__}")
|
||||
pass
|
||||
|
||||
async def on_shutdown(self):
|
||||
print(f"on_shutdown:{__name__}")
|
||||
pass
|
||||
|
||||
async def on_valves_updated(self):
|
||||
# This function is called when the valves are updated.
|
||||
self.client = Anthropic(api_key=self.valves.ANTHROPIC_API_KEY)
|
||||
pass
|
||||
|
||||
# Pipelines are the models that are available in the manifold.
|
||||
# It can be a list or a function that returns a list.
|
||||
def pipelines(self) -> List[dict]:
|
||||
return self.get_anthropic_models()
|
||||
|
||||
def pipe(
|
||||
self, user_message: str, model_id: str, messages: List[dict], body: dict
|
||||
) -> Union[str, Generator, Iterator]:
|
||||
try:
|
||||
if body.get("stream", False):
|
||||
return self.stream_response(model_id, messages, body)
|
||||
else:
|
||||
return self.get_completion(model_id, messages, body)
|
||||
except (RateLimitError, APIStatusError, APIConnectionError) as e:
|
||||
return f"Error: {e}"
|
||||
|
||||
def stream_response(
|
||||
self, model_id: str, messages: List[dict], body: dict
|
||||
) -> Generator:
|
||||
max_tokens = (
|
||||
body.get("max_tokens") if body.get("max_tokens") is not None else 4096
|
||||
)
|
||||
temperature = (
|
||||
body.get("temperature") if body.get("temperature") is not None else 0.8
|
||||
)
|
||||
top_k = body.get("top_k") if body.get("top_k") is not None else 40
|
||||
top_p = body.get("top_p") if body.get("top_p") is not None else 0.9
|
||||
stop_sequences = body.get("stop") if body.get("stop") is not None else []
|
||||
|
||||
stream = self.client.messages.create(
|
||||
model=model_id,
|
||||
messages=messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
stop_sequences=stop_sequences,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
for chunk in stream:
|
||||
if chunk.type == "content_block_start":
|
||||
yield chunk.content_block.text
|
||||
elif chunk.type == "content_block_delta":
|
||||
yield chunk.delta.text
|
||||
|
||||
def get_completion(self, model_id: str, messages: List[dict], body: dict) -> str:
|
||||
max_tokens = (
|
||||
body.get("max_tokens") if body.get("max_tokens") is not None else 4096
|
||||
)
|
||||
temperature = (
|
||||
body.get("temperature") if body.get("temperature") is not None else 0.8
|
||||
)
|
||||
top_k = body.get("top_k") if body.get("top_k") is not None else 40
|
||||
top_p = body.get("top_p") if body.get("top_p") is not None else 0.9
|
||||
stop_sequences = body.get("stop") if body.get("stop") is not None else []
|
||||
|
||||
response = self.client.messages.create(
|
||||
model=model_id,
|
||||
messages=messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
stop_sequences=stop_sequences,
|
||||
)
|
||||
return response.content[0].text
|
||||
90
examples/applescript_pipeline.py
Normal file
90
examples/applescript_pipeline.py
Normal file
@@ -0,0 +1,90 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
import requests
|
||||
|
||||
|
||||
from subprocess import call
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "applescript_pipeline"
|
||||
self.name = "AppleScript Pipeline"
|
||||
pass
|
||||
|
||||
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
|
||||
|
||||
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__}")
|
||||
|
||||
OLLAMA_BASE_URL = "http://localhost:11434"
|
||||
MODEL = "llama3"
|
||||
|
||||
if body.get("title", False):
|
||||
print("Title Generation")
|
||||
return "AppleScript Pipeline"
|
||||
else:
|
||||
if "user" in body:
|
||||
print("######################################")
|
||||
print(f'# User: {body["user"]["name"]} ({body["user"]["id"]})')
|
||||
print(f"# Message: {user_message}")
|
||||
print("######################################")
|
||||
|
||||
commands = user_message.split(" ")
|
||||
|
||||
if commands[0] == "volume":
|
||||
|
||||
try:
|
||||
commands[1] = int(commands[1])
|
||||
if 0 <= commands[1] <= 100:
|
||||
call(
|
||||
[f"osascript -e 'set volume output volume {commands[1]}'"],
|
||||
shell=True,
|
||||
)
|
||||
except:
|
||||
pass
|
||||
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are an agent of the AppleScript Pipeline. You have the power to control the volume of the system.",
|
||||
},
|
||||
{"role": "user", "content": user_message},
|
||||
],
|
||||
"stream": body["stream"],
|
||||
}
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
url=f"{OLLAMA_BASE_URL}/v1/chat/completions",
|
||||
json=payload,
|
||||
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}"
|
||||
59
examples/azure_openai_pipeline.py
Normal file
59
examples/azure_openai_pipeline.py
Normal file
@@ -0,0 +1,59 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
import requests
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "azure_openai_pipeline"
|
||||
self.name = "Azure OpenAI Pipeline"
|
||||
pass
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
AZURE_OPENAI_API_KEY = "your-azure-openai-api-key-here"
|
||||
AZURE_OPENAI_ENDPOINT = "your-azure-openai-endpoint-here"
|
||||
DEPLOYMENT_NAME = "your-deployment-name-here"
|
||||
MODEL = "gpt-3.5-turbo"
|
||||
|
||||
headers = {"api-key": AZURE_OPENAI_API_KEY, "Content-Type": "application/json"}
|
||||
|
||||
url = f"{AZURE_OPENAI_ENDPOINT}/openai/deployments/{DEPLOYMENT_NAME}/chat/completions?api-version=2023-10-01-preview"
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
url=url,
|
||||
json={**body, "model": MODEL},
|
||||
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}"
|
||||
154
examples/cohere_manifold_pipeline.py
Normal file
154
examples/cohere_manifold_pipeline.py
Normal file
@@ -0,0 +1,154 @@
|
||||
"""
|
||||
title: Cohere Manifold Pipeline
|
||||
author: justinh-rahb
|
||||
date: 2024-05-28
|
||||
version: 1.0
|
||||
license: MIT
|
||||
description: A pipeline for generating text using the Anthropic API.
|
||||
requirements: requests
|
||||
environment_variables: COHERE_API_KEY
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
from schemas import OpenAIChatMessage
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from pydantic import BaseModel
|
||||
import requests
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
self.type = "manifold"
|
||||
self.id = "cohere"
|
||||
self.name = "cohere/"
|
||||
|
||||
class Valves(BaseModel):
|
||||
COHERE_API_BASE_URL: str = "https://api.cohere.com/v1"
|
||||
COHERE_API_KEY: str
|
||||
|
||||
self.valves = Valves(**{"COHERE_API_KEY": os.getenv("COHERE_API_KEY")})
|
||||
|
||||
self.pipelines = self.get_cohere_models()
|
||||
|
||||
async def on_startup(self):
|
||||
print(f"on_startup:{__name__}")
|
||||
pass
|
||||
|
||||
async def on_shutdown(self):
|
||||
print(f"on_shutdown:{__name__}")
|
||||
pass
|
||||
|
||||
async def on_valves_updated(self):
|
||||
# This function is called when the valves are updated.
|
||||
|
||||
self.pipelines = self.get_cohere_models()
|
||||
|
||||
pass
|
||||
|
||||
def get_cohere_models(self):
|
||||
if self.valves.COHERE_API_KEY:
|
||||
try:
|
||||
headers = {}
|
||||
headers["Authorization"] = f"Bearer {self.valves.COHERE_API_KEY}"
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
r = requests.get(
|
||||
f"{self.valves.COHERE_API_BASE_URL}/models", headers=headers
|
||||
)
|
||||
|
||||
models = r.json()
|
||||
return [
|
||||
{
|
||||
"id": model["name"],
|
||||
"name": model["name"] if "name" in model else model["name"],
|
||||
}
|
||||
for model in models["models"]
|
||||
]
|
||||
except Exception as e:
|
||||
|
||||
print(f"Error: {e}")
|
||||
return [
|
||||
{
|
||||
"id": self.id,
|
||||
"name": "Could not fetch models from Cohere, please update the API Key in the valves.",
|
||||
},
|
||||
]
|
||||
else:
|
||||
return []
|
||||
|
||||
def pipe(
|
||||
self, user_message: str, model_id: str, messages: List[dict], body: dict
|
||||
) -> Union[str, Generator, Iterator]:
|
||||
try:
|
||||
if body.get("stream", False):
|
||||
return self.stream_response(user_message, model_id, messages, body)
|
||||
else:
|
||||
return self.get_completion(user_message, model_id, messages, body)
|
||||
except Exception as e:
|
||||
return f"Error: {e}"
|
||||
|
||||
def stream_response(
|
||||
self, user_message: str, model_id: str, messages: List[dict], body: dict
|
||||
) -> Generator:
|
||||
|
||||
headers = {}
|
||||
headers["Authorization"] = f"Bearer {self.valves.COHERE_API_KEY}"
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
r = requests.post(
|
||||
url=f"{self.valves.COHERE_API_BASE_URL}/chat",
|
||||
json={
|
||||
"model": model_id,
|
||||
"chat_history": [
|
||||
{
|
||||
"role": "USER" if message["role"] == "user" else "CHATBOT",
|
||||
"message": message["content"],
|
||||
}
|
||||
for message in messages[:-1]
|
||||
],
|
||||
"message": user_message,
|
||||
"stream": True,
|
||||
},
|
||||
headers=headers,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
|
||||
for line in r.iter_lines():
|
||||
if line:
|
||||
try:
|
||||
line = json.loads(line)
|
||||
if line["event_type"] == "text-generation":
|
||||
yield line["text"]
|
||||
except:
|
||||
pass
|
||||
|
||||
def get_completion(
|
||||
self, user_message: str, model_id: str, messages: List[dict], body: dict
|
||||
) -> str:
|
||||
headers = {}
|
||||
headers["Authorization"] = f"Bearer {self.valves.COHERE_API_KEY}"
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
r = requests.post(
|
||||
url=f"{self.valves.COHERE_API_BASE_URL}/chat",
|
||||
json={
|
||||
"model": model_id,
|
||||
"chat_history": [
|
||||
{
|
||||
"role": "USER" if message["role"] == "user" else "CHATBOT",
|
||||
"message": message["content"],
|
||||
}
|
||||
for message in messages[:-1]
|
||||
],
|
||||
"message": user_message,
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
return data["text"] if "text" in data else "No response from Cohere."
|
||||
83
examples/detoxify_filter_pipeline.py
Normal file
83
examples/detoxify_filter_pipeline.py
Normal file
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
title: Detoxify Filter Pipeline
|
||||
author: open-webui
|
||||
date: 2024-05-30
|
||||
version: 1.0
|
||||
license: MIT
|
||||
description: A pipeline for filtering out toxic messages using the Detoxify library.
|
||||
requirements: detoxify
|
||||
"""
|
||||
|
||||
from typing import List, Optional
|
||||
from schemas import OpenAIChatMessage
|
||||
from pydantic import BaseModel
|
||||
from detoxify import Detoxify
|
||||
import os
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Pipeline filters are only compatible with Open WebUI
|
||||
# You can think of filter pipeline as a middleware that can be used to edit the form data before it is sent to the OpenAI API.
|
||||
self.type = "filter"
|
||||
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "detoxify_filter_pipeline"
|
||||
self.name = "Detoxify Filter"
|
||||
|
||||
class Valves(BaseModel):
|
||||
# List target pipeline ids (models) that this filter will be connected to.
|
||||
# If you want to connect this filter to all pipelines, you can set pipelines to ["*"]
|
||||
# e.g. ["llama3:latest", "gpt-3.5-turbo"]
|
||||
pipelines: List[str] = []
|
||||
|
||||
# Assign a priority level to the filter pipeline.
|
||||
# The priority level determines the order in which the filter pipelines are executed.
|
||||
# The lower the number, the higher the priority.
|
||||
priority: int = 0
|
||||
|
||||
# Initialize
|
||||
self.valves = Valves(
|
||||
**{
|
||||
"pipelines": ["*"], # Connect to all pipelines
|
||||
}
|
||||
)
|
||||
|
||||
self.model = None
|
||||
|
||||
pass
|
||||
|
||||
async def on_startup(self):
|
||||
# This function is called when the server is started.
|
||||
print(f"on_startup:{__name__}")
|
||||
|
||||
self.model = Detoxify("original")
|
||||
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.
|
||||
pass
|
||||
|
||||
async def inlet(self, body: dict, user: Optional[dict] = None) -> dict:
|
||||
# This filter is applied to the form data before it is sent to the OpenAI API.
|
||||
print(f"inlet:{__name__}")
|
||||
|
||||
print(body)
|
||||
user_message = body["messages"][-1]["content"]
|
||||
|
||||
# Filter out toxic messages
|
||||
toxicity = self.model.predict(user_message)
|
||||
print(toxicity)
|
||||
|
||||
if toxicity["toxicity"] > 0.5:
|
||||
raise Exception("Toxic message detected")
|
||||
|
||||
return body
|
||||
58
examples/example_pipeline.py
Normal file
58
examples/example_pipeline.py
Normal file
@@ -0,0 +1,58 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "pipeline_example"
|
||||
self.name = "Pipeline Example"
|
||||
|
||||
pass
|
||||
|
||||
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.
|
||||
pass
|
||||
|
||||
async def inlet(self, body: dict, user: dict) -> dict:
|
||||
# This function is called before the OpenAI API request is made. You can modify the form data before it is sent to the OpenAI API.
|
||||
print(f"inlet:{__name__}")
|
||||
|
||||
print(body)
|
||||
print(user)
|
||||
|
||||
return body
|
||||
|
||||
async def outlet(self, body: dict, user: dict) -> dict:
|
||||
# This function is called after the OpenAI API response is completed. You can modify the messages after they are received from the OpenAI API.
|
||||
print(f"outlet:{__name__}")
|
||||
|
||||
print(body)
|
||||
print(user)
|
||||
|
||||
return body
|
||||
|
||||
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)
|
||||
print(body)
|
||||
|
||||
return f"{__name__} response to: {user_message}"
|
||||
63
examples/filter_pipeline.py
Normal file
63
examples/filter_pipeline.py
Normal file
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
title: Filter Pipeline
|
||||
author: open-webui
|
||||
date: 2024-05-30
|
||||
version: 1.1
|
||||
license: MIT
|
||||
description: Example of a filter pipeline that can be used to edit the form data before it is sent to the OpenAI API.
|
||||
requirements: requests
|
||||
"""
|
||||
|
||||
from typing import List, Optional
|
||||
from pydantic import BaseModel
|
||||
from schemas import OpenAIChatMessage
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Pipeline filters are only compatible with Open WebUI
|
||||
# You can think of filter pipeline as a middleware that can be used to edit the form data before it is sent to the OpenAI API.
|
||||
self.type = "filter"
|
||||
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "filter_pipeline"
|
||||
self.name = "Filter"
|
||||
|
||||
class Valves(BaseModel):
|
||||
# List target pipeline ids (models) that this filter will be connected to.
|
||||
# If you want to connect this filter to all pipelines, you can set pipelines to ["*"]
|
||||
pipelines: List[str] = []
|
||||
|
||||
# Assign a priority level to the filter pipeline.
|
||||
# The priority level determines the order in which the filter pipelines are executed.
|
||||
# The lower the number, the higher the priority.
|
||||
priority: int = 0
|
||||
|
||||
# Add your custom parameters here
|
||||
pass
|
||||
|
||||
self.valves = Valves(**{"pipelines": ["llama3:latest"]})
|
||||
|
||||
pass
|
||||
|
||||
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 inlet(self, body: dict, user: Optional[dict] = None) -> dict:
|
||||
# This filter is applied to the form data before it is sent to the OpenAI API.
|
||||
print(f"inlet:{__name__}")
|
||||
|
||||
print(body)
|
||||
print(user)
|
||||
|
||||
return body
|
||||
98
examples/haystack_pipeline.py
Normal file
98
examples/haystack_pipeline.py
Normal file
@@ -0,0 +1,98 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
import os
|
||||
import asyncio
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
self.basic_rag_pipeline = None
|
||||
|
||||
async def on_startup(self):
|
||||
os.environ["OPENAI_API_KEY"] = "your_openai_api_key_here"
|
||||
|
||||
from haystack.components.embedders import SentenceTransformersDocumentEmbedder
|
||||
from haystack.components.embedders import SentenceTransformersTextEmbedder
|
||||
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
|
||||
from haystack.components.builders import PromptBuilder
|
||||
from haystack.components.generators import OpenAIGenerator
|
||||
|
||||
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
||||
|
||||
from datasets import load_dataset
|
||||
from haystack import Document
|
||||
from haystack import Pipeline
|
||||
|
||||
document_store = InMemoryDocumentStore()
|
||||
|
||||
dataset = load_dataset("bilgeyucel/seven-wonders", split="train")
|
||||
docs = [Document(content=doc["content"], meta=doc["meta"]) for doc in dataset]
|
||||
|
||||
doc_embedder = SentenceTransformersDocumentEmbedder(
|
||||
model="sentence-transformers/all-MiniLM-L6-v2"
|
||||
)
|
||||
doc_embedder.warm_up()
|
||||
|
||||
docs_with_embeddings = doc_embedder.run(docs)
|
||||
document_store.write_documents(docs_with_embeddings["documents"])
|
||||
|
||||
text_embedder = SentenceTransformersTextEmbedder(
|
||||
model="sentence-transformers/all-MiniLM-L6-v2"
|
||||
)
|
||||
|
||||
retriever = InMemoryEmbeddingRetriever(document_store)
|
||||
|
||||
template = """
|
||||
Given the following information, answer the question.
|
||||
|
||||
Context:
|
||||
{% for document in documents %}
|
||||
{{ document.content }}
|
||||
{% endfor %}
|
||||
|
||||
Question: {{question}}
|
||||
Answer:
|
||||
"""
|
||||
|
||||
prompt_builder = PromptBuilder(template=template)
|
||||
|
||||
generator = OpenAIGenerator(model="gpt-3.5-turbo")
|
||||
|
||||
self.basic_rag_pipeline = Pipeline()
|
||||
# Add components to your pipeline
|
||||
self.basic_rag_pipeline.add_component("text_embedder", text_embedder)
|
||||
self.basic_rag_pipeline.add_component("retriever", retriever)
|
||||
self.basic_rag_pipeline.add_component("prompt_builder", prompt_builder)
|
||||
self.basic_rag_pipeline.add_component("llm", generator)
|
||||
|
||||
# Now, connect the components to each other
|
||||
self.basic_rag_pipeline.connect(
|
||||
"text_embedder.embedding", "retriever.query_embedding"
|
||||
)
|
||||
self.basic_rag_pipeline.connect("retriever", "prompt_builder.documents")
|
||||
self.basic_rag_pipeline.connect("prompt_builder", "llm")
|
||||
|
||||
pass
|
||||
|
||||
async def on_shutdown(self):
|
||||
# This function is called when the server is stopped.
|
||||
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 RAG pipeline.
|
||||
# Typically, you would retrieve relevant information from your knowledge base and synthesize it to generate a response.
|
||||
|
||||
print(messages)
|
||||
print(user_message)
|
||||
|
||||
question = user_message
|
||||
response = self.basic_rag_pipeline.run(
|
||||
{
|
||||
"text_embedder": {"text": question},
|
||||
"prompt_builder": {"question": question},
|
||||
}
|
||||
)
|
||||
|
||||
return response["llm"]["replies"][0]
|
||||
92
examples/langfuse_filter_pipeline.py
Normal file
92
examples/langfuse_filter_pipeline.py
Normal file
@@ -0,0 +1,92 @@
|
||||
from typing import List, Optional
|
||||
from schemas import OpenAIChatMessage
|
||||
import os
|
||||
|
||||
|
||||
from pydantic import BaseModel
|
||||
from langfuse import Langfuse
|
||||
from langfuse.decorators import langfuse_context, observe
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Pipeline filters are only compatible with Open WebUI
|
||||
# You can think of filter pipeline as a middleware that can be used to edit the form data before it is sent to the OpenAI API.
|
||||
self.type = "filter"
|
||||
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "langfuse_filter_pipeline"
|
||||
self.name = "Langfuse Filter"
|
||||
|
||||
class Valves(BaseModel):
|
||||
# List target pipeline ids (models) that this filter will be connected to.
|
||||
# If you want to connect this filter to all pipelines, you can set pipelines to ["*"]
|
||||
# e.g. ["llama3:latest", "gpt-3.5-turbo"]
|
||||
pipelines: List[str] = []
|
||||
|
||||
# Assign a priority level to the filter pipeline.
|
||||
# The priority level determines the order in which the filter pipelines are executed.
|
||||
# The lower the number, the higher the priority.
|
||||
priority: int = 0
|
||||
|
||||
# Valves
|
||||
secret_key: str
|
||||
public_key: str
|
||||
host: str
|
||||
|
||||
# Initialize
|
||||
self.valves = Valves(
|
||||
**{
|
||||
"pipelines": ["*"], # Connect to all pipelines
|
||||
"secret_key": os.getenv("LANGFUSE_SECRET_KEY"),
|
||||
"public_key": os.getenv("LANGFUSE_PUBLIC_KEY"),
|
||||
"host": os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com"),
|
||||
}
|
||||
)
|
||||
|
||||
self.langfuse = None
|
||||
pass
|
||||
|
||||
async def on_startup(self):
|
||||
# This function is called when the server is started.
|
||||
print(f"on_startup:{__name__}")
|
||||
self.set_langfuse()
|
||||
pass
|
||||
|
||||
async def on_shutdown(self):
|
||||
# This function is called when the server is stopped.
|
||||
print(f"on_shutdown:{__name__}")
|
||||
self.langfuse.flush()
|
||||
pass
|
||||
|
||||
async def on_valves_updated(self):
|
||||
# This function is called when the valves are updated.
|
||||
|
||||
self.set_langfuse()
|
||||
pass
|
||||
|
||||
def set_langfuse(self):
|
||||
self.langfuse = Langfuse(
|
||||
secret_key=self.valves.secret_key,
|
||||
public_key=self.valves.public_key,
|
||||
host=self.valves.host,
|
||||
debug=True,
|
||||
)
|
||||
self.langfuse.auth_check()
|
||||
|
||||
async def inlet(self, body: dict, user: Optional[dict] = None) -> dict:
|
||||
print(f"inlet:{__name__}")
|
||||
|
||||
trace = self.langfuse.trace(
|
||||
name=f"filter:{__name__}",
|
||||
input=body,
|
||||
user_id=user["id"],
|
||||
metadata={"name": user["name"]},
|
||||
)
|
||||
|
||||
print(trace.get_trace_url())
|
||||
|
||||
return body
|
||||
140
examples/libretranlsate_filter_pipeline.py
Normal file
140
examples/libretranlsate_filter_pipeline.py
Normal file
@@ -0,0 +1,140 @@
|
||||
from typing import List, Optional
|
||||
from schemas import OpenAIChatMessage
|
||||
from pydantic import BaseModel
|
||||
import requests
|
||||
import os
|
||||
|
||||
from utils.main import get_last_user_message, get_last_assistant_message
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Pipeline filters are only compatible with Open WebUI
|
||||
# You can think of filter pipeline as a middleware that can be used to edit the form data before it is sent to the OpenAI API.
|
||||
self.type = "filter"
|
||||
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "libretranslate_filter_pipeline"
|
||||
self.name = "LibreTranslate Filter"
|
||||
|
||||
class Valves(BaseModel):
|
||||
# List target pipeline ids (models) that this filter will be connected to.
|
||||
# If you want to connect this filter to all pipelines, you can set pipelines to ["*"]
|
||||
# e.g. ["llama3:latest", "gpt-3.5-turbo"]
|
||||
pipelines: List[str] = []
|
||||
|
||||
# Assign a priority level to the filter pipeline.
|
||||
# The priority level determines the order in which the filter pipelines are executed.
|
||||
# The lower the number, the higher the priority.
|
||||
priority: int = 0
|
||||
|
||||
# Valves
|
||||
libretranslate_url: str
|
||||
|
||||
# Source and target languages
|
||||
# User message will be translated from source_user to target_user
|
||||
source_user: Optional[str] = "auto"
|
||||
target_user: Optional[str] = "en"
|
||||
|
||||
# Assistant languages
|
||||
# Assistant message will be translated from source_assistant to target_assistant
|
||||
source_assistant: Optional[str] = "en"
|
||||
target_assistant: Optional[str] = "es"
|
||||
|
||||
# Initialize
|
||||
self.valves = Valves(
|
||||
**{
|
||||
"pipelines": ["*"], # Connect to all pipelines
|
||||
"libretranslate_url": os.getenv(
|
||||
"LIBRETRANSLATE_API_BASE_URL", "http://localhost:5000"
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
pass
|
||||
|
||||
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.
|
||||
pass
|
||||
|
||||
def translate(self, text: str, source: str, target: str) -> str:
|
||||
payload = {
|
||||
"q": text,
|
||||
"source": source,
|
||||
"target": target,
|
||||
}
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
f"{self.valves.libretranslate_url}/translate", json=payload
|
||||
)
|
||||
r.raise_for_status()
|
||||
|
||||
data = r.json()
|
||||
return data["translatedText"]
|
||||
except Exception as e:
|
||||
print(f"Error translating text: {e}")
|
||||
return text
|
||||
|
||||
async def inlet(self, body: dict, user: Optional[dict] = None) -> dict:
|
||||
print(f"inlet:{__name__}")
|
||||
|
||||
messages = body["messages"]
|
||||
user_message = get_last_user_message(messages)
|
||||
|
||||
print(f"User message: {user_message}")
|
||||
|
||||
# Translate user message
|
||||
translated_user_message = self.translate(
|
||||
user_message,
|
||||
self.valves.source_user,
|
||||
self.valves.target_user,
|
||||
)
|
||||
|
||||
print(f"Translated user message: {translated_user_message}")
|
||||
|
||||
for message in reversed(messages):
|
||||
if message["role"] == "user":
|
||||
message["content"] = translated_user_message
|
||||
break
|
||||
|
||||
body = {**body, "messages": messages}
|
||||
return body
|
||||
|
||||
async def outlet(self, body: dict, user: Optional[dict] = None) -> dict:
|
||||
print(f"outlet:{__name__}")
|
||||
|
||||
messages = body["messages"]
|
||||
assistant_message = get_last_assistant_message(messages)
|
||||
|
||||
print(f"Assistant message: {assistant_message}")
|
||||
|
||||
# Translate assistant message
|
||||
translated_assistant_message = self.translate(
|
||||
assistant_message,
|
||||
self.valves.source_assistant,
|
||||
self.valves.target_assistant,
|
||||
)
|
||||
|
||||
print(f"Translated assistant message: {translated_assistant_message}")
|
||||
|
||||
for message in reversed(messages):
|
||||
if message["role"] == "assistant":
|
||||
message["content"] = translated_assistant_message
|
||||
break
|
||||
|
||||
body = {**body, "messages": messages}
|
||||
return body
|
||||
93
examples/litellm_manifold_pipeline.py
Normal file
93
examples/litellm_manifold_pipeline.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
from pydantic import BaseModel
|
||||
import requests
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# You can also set the pipelines that are available in this pipeline.
|
||||
# Set manifold to True if you want to use this pipeline as a manifold.
|
||||
# Manifold pipelines can have multiple pipelines.
|
||||
self.type = "manifold"
|
||||
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "litellm_manifold"
|
||||
|
||||
# Optionally, you can set the name of the manifold pipeline.
|
||||
self.name = "LiteLLM: "
|
||||
|
||||
class Valves(BaseModel):
|
||||
LITELLM_BASE_URL: str
|
||||
|
||||
# Initialize rate limits
|
||||
self.valves = Valves(**{"LITELLM_BASE_URL": "http://localhost:4001"})
|
||||
self.pipelines = []
|
||||
pass
|
||||
|
||||
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.
|
||||
|
||||
self.pipelines = self.get_litellm_models()
|
||||
pass
|
||||
|
||||
def get_litellm_models(self):
|
||||
if self.valves.LITELLM_BASE_URL:
|
||||
try:
|
||||
r = requests.get(f"{self.valves.LITELLM_BASE_URL}/v1/models")
|
||||
models = r.json()
|
||||
return [
|
||||
{
|
||||
"id": model["id"],
|
||||
"name": model["name"] if "name" in model else model["id"],
|
||||
}
|
||||
for model in models["data"]
|
||||
]
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return [
|
||||
{
|
||||
"id": self.id,
|
||||
"name": "Could not fetch models from LiteLLM, please update the URL in the valves.",
|
||||
},
|
||||
]
|
||||
else:
|
||||
return []
|
||||
|
||||
def pipe(
|
||||
self, user_message: str, model_id: str, messages: List[dict], body: dict
|
||||
) -> Union[str, Generator, Iterator]:
|
||||
if "user" in body:
|
||||
print("######################################")
|
||||
print(f'# User: {body["user"]["name"]} ({body["user"]["id"]})')
|
||||
print(f"# Message: {user_message}")
|
||||
print("######################################")
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
url=f"{self.valves.LITELLM_BASE_URL}/v1/chat/completions",
|
||||
json={**body, "model": model_id, "user_id": body["user"]["id"]},
|
||||
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}"
|
||||
201
examples/litellm_subprocess_manifold_pipeline.py
Normal file
201
examples/litellm_subprocess_manifold_pipeline.py
Normal file
@@ -0,0 +1,201 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
from pydantic import BaseModel
|
||||
import requests
|
||||
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
import subprocess
|
||||
import yaml
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# You can also set the pipelines that are available in this pipeline.
|
||||
# Set manifold to True if you want to use this pipeline as a manifold.
|
||||
# Manifold pipelines can have multiple pipelines.
|
||||
self.type = "manifold"
|
||||
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "litellm_subprocess_manifold"
|
||||
|
||||
# Optionally, you can set the name of the manifold pipeline.
|
||||
self.name = "LiteLLM: "
|
||||
|
||||
class Valves(BaseModel):
|
||||
LITELLM_CONFIG_DIR: str = "./litellm/config.yaml"
|
||||
LITELLM_PROXY_PORT: int = 4001
|
||||
LITELLM_PROXY_HOST: str = "127.0.0.1"
|
||||
litellm_config: dict = {}
|
||||
|
||||
# Initialize Valves
|
||||
self.valves = Valves(**{"LITELLM_CONFIG_DIR": f"./litellm/config.yaml"})
|
||||
self.background_process = None
|
||||
pass
|
||||
|
||||
async def on_startup(self):
|
||||
# This function is called when the server is started.
|
||||
print(f"on_startup:{__name__}")
|
||||
|
||||
# Check if the config file exists
|
||||
if not os.path.exists(self.valves.LITELLM_CONFIG_DIR):
|
||||
with open(self.valves.LITELLM_CONFIG_DIR, "w") as file:
|
||||
yaml.dump(
|
||||
{
|
||||
"general_settings": {},
|
||||
"litellm_settings": {},
|
||||
"model_list": [],
|
||||
"router_settings": {},
|
||||
},
|
||||
file,
|
||||
)
|
||||
|
||||
print(
|
||||
f"Config file not found. Created a default config file at {self.valves.LITELLM_CONFIG_DIR}"
|
||||
)
|
||||
|
||||
with open(self.valves.LITELLM_CONFIG_DIR, "r") as file:
|
||||
litellm_config = yaml.safe_load(file)
|
||||
|
||||
self.valves.litellm_config = litellm_config
|
||||
|
||||
asyncio.create_task(self.start_litellm_background())
|
||||
pass
|
||||
|
||||
async def on_shutdown(self):
|
||||
# This function is called when the server is stopped.
|
||||
print(f"on_shutdown:{__name__}")
|
||||
await self.shutdown_litellm_background()
|
||||
pass
|
||||
|
||||
async def on_valves_updated(self):
|
||||
# This function is called when the valves are updated.
|
||||
|
||||
print(f"on_valves_updated:{__name__}")
|
||||
|
||||
with open(self.valves.LITELLM_CONFIG_DIR, "r") as file:
|
||||
litellm_config = yaml.safe_load(file)
|
||||
|
||||
self.valves.litellm_config = litellm_config
|
||||
|
||||
await self.shutdown_litellm_background()
|
||||
await self.start_litellm_background()
|
||||
pass
|
||||
|
||||
async def run_background_process(self, command):
|
||||
print("run_background_process")
|
||||
|
||||
try:
|
||||
# Log the command to be executed
|
||||
print(f"Executing command: {command}")
|
||||
|
||||
# Execute the command and create a subprocess
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*command,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
self.background_process = process
|
||||
print("Subprocess started successfully.")
|
||||
|
||||
# Capture STDERR for debugging purposes
|
||||
stderr_output = await process.stderr.read()
|
||||
stderr_text = stderr_output.decode().strip()
|
||||
if stderr_text:
|
||||
print(f"Subprocess STDERR: {stderr_text}")
|
||||
|
||||
# log.info output line by line
|
||||
async for line in process.stdout:
|
||||
print(line.decode().strip())
|
||||
|
||||
# Wait for the process to finish
|
||||
returncode = await process.wait()
|
||||
print(f"Subprocess exited with return code {returncode}")
|
||||
except Exception as e:
|
||||
print(f"Failed to start subprocess: {e}")
|
||||
raise # Optionally re-raise the exception if you want it to propagate
|
||||
|
||||
async def start_litellm_background(self):
|
||||
print("start_litellm_background")
|
||||
# Command to run in the background
|
||||
command = [
|
||||
"litellm",
|
||||
"--port",
|
||||
str(self.valves.LITELLM_PROXY_PORT),
|
||||
"--host",
|
||||
self.valves.LITELLM_PROXY_HOST,
|
||||
"--telemetry",
|
||||
"False",
|
||||
"--config",
|
||||
self.valves.LITELLM_CONFIG_DIR,
|
||||
]
|
||||
|
||||
await self.run_background_process(command)
|
||||
|
||||
async def shutdown_litellm_background(self):
|
||||
print("shutdown_litellm_background")
|
||||
|
||||
if self.background_process:
|
||||
self.background_process.terminate()
|
||||
await self.background_process.wait() # Ensure the process has terminated
|
||||
print("Subprocess terminated")
|
||||
self.background_process = None
|
||||
|
||||
def get_litellm_models(self):
|
||||
if self.background_process:
|
||||
try:
|
||||
r = requests.get(
|
||||
f"http://{self.valves.LITELLM_PROXY_HOST}:{self.valves.LITELLM_PROXY_PORT}/v1/models"
|
||||
)
|
||||
models = r.json()
|
||||
return [
|
||||
{
|
||||
"id": model["id"],
|
||||
"name": model["name"] if "name" in model else model["id"],
|
||||
}
|
||||
for model in models["data"]
|
||||
]
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return [
|
||||
{
|
||||
"id": self.id,
|
||||
"name": "Could not fetch models from LiteLLM, please update the URL in the valves.",
|
||||
},
|
||||
]
|
||||
else:
|
||||
return []
|
||||
|
||||
# Pipelines are the models that are available in the manifold.
|
||||
# It can be a list or a function that returns a list.
|
||||
def pipelines(self) -> List[dict]:
|
||||
return self.get_litellm_models()
|
||||
|
||||
def pipe(
|
||||
self, user_message: str, model_id: str, messages: List[dict], body: dict
|
||||
) -> Union[str, Generator, Iterator]:
|
||||
if "user" in body:
|
||||
print("######################################")
|
||||
print(f'# User: {body["user"]["name"]} ({body["user"]["id"]})')
|
||||
print(f"# Message: {user_message}")
|
||||
print("######################################")
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
url=f"http://{self.valves.LITELLM_PROXY_HOST}:{self.valves.LITELLM_PROXY_PORT}/v1/chat/completions",
|
||||
json={**body, "model": model_id, "user_id": body["user"]["id"]},
|
||||
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}"
|
||||
51
examples/llama_cpp_pipeline.py
Normal file
51
examples/llama_cpp_pipeline.py
Normal file
@@ -0,0 +1,51 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "llama_cpp_pipeline"
|
||||
|
||||
self.name = "Llama C++ Pipeline"
|
||||
self.llm = None
|
||||
pass
|
||||
|
||||
async def on_startup(self):
|
||||
# This function is called when the server is started.
|
||||
print(f"on_startup:{__name__}")
|
||||
from llama_cpp import Llama
|
||||
|
||||
self.llm = Llama(
|
||||
model_path="./models/llama3.gguf",
|
||||
# n_gpu_layers=-1, # Uncomment to use GPU acceleration
|
||||
# seed=1337, # Uncomment to set a specific seed
|
||||
# n_ctx=2048, # Uncomment to increase the context window
|
||||
)
|
||||
|
||||
pass
|
||||
|
||||
async def on_shutdown(self):
|
||||
# This function is called when the server is stopped.
|
||||
print(f"on_shutdown:{__name__}")
|
||||
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)
|
||||
print(body)
|
||||
|
||||
response = self.llm.create_chat_completion_openai_v1(
|
||||
messages=messages,
|
||||
stream=body["stream"],
|
||||
)
|
||||
|
||||
return response
|
||||
84
examples/llamaindex_ollama_github_pipeline.py
Normal file
84
examples/llamaindex_ollama_github_pipeline.py
Normal file
@@ -0,0 +1,84 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
import os
|
||||
import asyncio
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
self.documents = None
|
||||
self.index = None
|
||||
|
||||
async def on_startup(self):
|
||||
from llama_index.embeddings.ollama import OllamaEmbedding
|
||||
from llama_index.llms.ollama import Ollama
|
||||
from llama_index.core import VectorStoreIndex, Settings
|
||||
from llama_index.readers.github import GithubRepositoryReader, GithubClient
|
||||
|
||||
Settings.embed_model = OllamaEmbedding(
|
||||
model_name="nomic-embed-text",
|
||||
base_url="http://localhost:11434",
|
||||
)
|
||||
Settings.llm = Ollama(model="llama3")
|
||||
|
||||
global index, documents
|
||||
|
||||
github_token = os.environ.get("GITHUB_TOKEN")
|
||||
owner = "open-webui"
|
||||
repo = "plugin-server"
|
||||
branch = "main"
|
||||
|
||||
github_client = GithubClient(github_token=github_token, verbose=True)
|
||||
|
||||
reader = GithubRepositoryReader(
|
||||
github_client=github_client,
|
||||
owner=owner,
|
||||
repo=repo,
|
||||
use_parser=False,
|
||||
verbose=False,
|
||||
filter_file_extensions=(
|
||||
[
|
||||
".png",
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".gif",
|
||||
".svg",
|
||||
".ico",
|
||||
"json",
|
||||
".ipynb",
|
||||
],
|
||||
GithubRepositoryReader.FilterType.EXCLUDE,
|
||||
),
|
||||
)
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
|
||||
reader._loop = loop
|
||||
|
||||
try:
|
||||
# Load data from the branch
|
||||
self.documents = await asyncio.to_thread(reader.load_data, branch=branch)
|
||||
self.index = VectorStoreIndex.from_documents(self.documents)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
print(self.documents)
|
||||
print(self.index)
|
||||
|
||||
async def on_shutdown(self):
|
||||
# This function is called when the server is stopped.
|
||||
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 RAG pipeline.
|
||||
# Typically, you would retrieve relevant information from your knowledge base and synthesize it to generate a response.
|
||||
|
||||
print(messages)
|
||||
print(user_message)
|
||||
|
||||
query_engine = self.index.as_query_engine(streaming=True)
|
||||
response = query_engine.query(user_message)
|
||||
|
||||
return response.response_gen
|
||||
44
examples/llamaindex_ollama_pipeline.py
Normal file
44
examples/llamaindex_ollama_pipeline.py
Normal file
@@ -0,0 +1,44 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
self.documents = None
|
||||
self.index = None
|
||||
|
||||
async def on_startup(self):
|
||||
from llama_index.embeddings.ollama import OllamaEmbedding
|
||||
from llama_index.llms.ollama import Ollama
|
||||
from llama_index.core import Settings, VectorStoreIndex, SimpleDirectoryReader
|
||||
|
||||
Settings.embed_model = OllamaEmbedding(
|
||||
model_name="nomic-embed-text",
|
||||
base_url="http://localhost:11434",
|
||||
)
|
||||
Settings.llm = Ollama(model="llama3")
|
||||
|
||||
# This function is called when the server is started.
|
||||
global documents, index
|
||||
|
||||
self.documents = SimpleDirectoryReader("./data").load_data()
|
||||
self.index = VectorStoreIndex.from_documents(self.documents)
|
||||
pass
|
||||
|
||||
async def on_shutdown(self):
|
||||
# This function is called when the server is stopped.
|
||||
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 RAG pipeline.
|
||||
# Typically, you would retrieve relevant information from your knowledge base and synthesize it to generate a response.
|
||||
|
||||
print(messages)
|
||||
print(user_message)
|
||||
|
||||
query_engine = self.index.as_query_engine(streaming=True)
|
||||
response = query_engine.query(user_message)
|
||||
|
||||
return response.response_gen
|
||||
39
examples/llamaindex_pipeline.py
Normal file
39
examples/llamaindex_pipeline.py
Normal file
@@ -0,0 +1,39 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
self.documents = None
|
||||
self.index = None
|
||||
|
||||
async def on_startup(self):
|
||||
import os
|
||||
|
||||
# Set the OpenAI API key
|
||||
os.environ["OPENAI_API_KEY"] = "your-api-key-here"
|
||||
|
||||
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
|
||||
|
||||
self.documents = SimpleDirectoryReader("./data").load_data()
|
||||
self.index = VectorStoreIndex.from_documents(self.documents)
|
||||
# This function is called when the server is started.
|
||||
pass
|
||||
|
||||
async def on_shutdown(self):
|
||||
# This function is called when the server is stopped.
|
||||
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 RAG pipeline.
|
||||
# Typically, you would retrieve relevant information from your knowledge base and synthesize it to generate a response.
|
||||
|
||||
print(messages)
|
||||
print(user_message)
|
||||
|
||||
query_engine = self.index.as_query_engine(streaming=True)
|
||||
response = query_engine.query(user_message)
|
||||
|
||||
return response.response_gen
|
||||
52
examples/manifold_pipeline.py
Normal file
52
examples/manifold_pipeline.py
Normal file
@@ -0,0 +1,52 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# You can also set the pipelines that are available in this pipeline.
|
||||
# Set manifold to True if you want to use this pipeline as a manifold.
|
||||
# Manifold pipelines can have multiple pipelines.
|
||||
self.type = "manifold"
|
||||
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "manifold_pipeline"
|
||||
|
||||
# Optionally, you can set the name of the manifold pipeline.
|
||||
self.name = "Manifold: "
|
||||
self.pipelines = [
|
||||
{
|
||||
"id": "pipeline-1", # This will turn into `manifold_pipeline.pipeline-1`
|
||||
"name": "Pipeline 1", # This will turn into `Manifold: Pipeline 1`
|
||||
},
|
||||
{
|
||||
"id": "pipeline-2",
|
||||
"name": "Pipeline 2",
|
||||
},
|
||||
]
|
||||
pass
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
print(body)
|
||||
|
||||
return f"{model_id} response to: {user_message}"
|
||||
109
examples/mlx_pipeline.py
Normal file
109
examples/mlx_pipeline.py
Normal file
@@ -0,0 +1,109 @@
|
||||
"""
|
||||
title: MLX Pipeline
|
||||
author: justinh-rahb
|
||||
date: 2024-05-27
|
||||
version: 1.1
|
||||
license: MIT
|
||||
description: A pipeline for generating text using Apple MLX Framework.
|
||||
requirements: requests, mlx-lm, huggingface-hub
|
||||
environment_variables: MLX_HOST, MLX_PORT, MLX_MODEL, MLX_STOP, MLX_SUBPROCESS, HUGGINGFACE_TOKEN
|
||||
"""
|
||||
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
import requests
|
||||
import os
|
||||
import subprocess
|
||||
import logging
|
||||
from huggingface_hub import login
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "mlx_pipeline"
|
||||
self.name = "MLX Pipeline"
|
||||
self.host = os.getenv("MLX_HOST", "localhost")
|
||||
self.port = os.getenv("MLX_PORT", "8080")
|
||||
self.model = os.getenv("MLX_MODEL", "mistralai/Mistral-7B-Instruct-v0.2")
|
||||
self.stop_sequence = os.getenv("MLX_STOP", "[INST]").split(
|
||||
","
|
||||
) # Default stop sequence is [INST]
|
||||
self.subprocess = os.getenv("MLX_SUBPROCESS", "true").lower() == "true"
|
||||
self.huggingface_token = os.getenv("HUGGINGFACE_TOKEN", None)
|
||||
|
||||
if self.huggingface_token:
|
||||
login(self.huggingface_token)
|
||||
|
||||
if self.subprocess:
|
||||
self.start_mlx_server()
|
||||
|
||||
def start_mlx_server(self):
|
||||
if not os.getenv("MLX_PORT"):
|
||||
self.port = self.find_free_port()
|
||||
command = f"mlx_lm.server --model {self.model} --port {self.port}"
|
||||
self.server_process = subprocess.Popen(command, shell=True)
|
||||
logging.info(f"Started MLX server on port {self.port}")
|
||||
|
||||
def find_free_port(self):
|
||||
import socket
|
||||
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
s.bind(("", 0))
|
||||
port = s.getsockname()[1]
|
||||
s.close()
|
||||
return port
|
||||
|
||||
async def on_startup(self):
|
||||
logging.info(f"on_startup:{__name__}")
|
||||
|
||||
async def on_shutdown(self):
|
||||
if self.subprocess and hasattr(self, "server_process"):
|
||||
self.server_process.terminate()
|
||||
logging.info(f"Terminated MLX server on port {self.port}")
|
||||
|
||||
def pipe(
|
||||
self, user_message: str, model_id: str, messages: List[dict], body: dict
|
||||
) -> Union[str, Generator, Iterator]:
|
||||
logging.info(f"pipe:{__name__}")
|
||||
|
||||
url = f"http://{self.host}:{self.port}/v1/chat/completions"
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
# Extract and validate parameters from the request body
|
||||
max_tokens = body.get("max_tokens", 4096)
|
||||
if not isinstance(max_tokens, int) or max_tokens < 0:
|
||||
max_tokens = 4096 # Default to 4096 if invalid
|
||||
|
||||
temperature = body.get("temperature", 0.8)
|
||||
if not isinstance(temperature, (int, float)) or temperature < 0:
|
||||
temperature = 0.8 # Default to 0.8 if invalid
|
||||
|
||||
repeat_penalty = body.get("repeat_penalty", 1.0)
|
||||
if not isinstance(repeat_penalty, (int, float)) or repeat_penalty < 0:
|
||||
repeat_penalty = 1.0 # Default to 1.0 if invalid
|
||||
|
||||
payload = {
|
||||
"messages": messages,
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
"repetition_penalty": repeat_penalty,
|
||||
"stop": self.stop_sequence,
|
||||
"stream": body.get("stream", False),
|
||||
}
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
url, headers=headers, json=payload, stream=body.get("stream", False)
|
||||
)
|
||||
r.raise_for_status()
|
||||
|
||||
if body.get("stream", False):
|
||||
return r.iter_lines()
|
||||
else:
|
||||
return r.json()
|
||||
except Exception as e:
|
||||
return f"Error: {e}"
|
||||
92
examples/ollama_manifold_pipeline.py
Normal file
92
examples/ollama_manifold_pipeline.py
Normal file
@@ -0,0 +1,92 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
from pydantic import BaseModel
|
||||
import requests
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# You can also set the pipelines that are available in this pipeline.
|
||||
# Set manifold to True if you want to use this pipeline as a manifold.
|
||||
# Manifold pipelines can have multiple pipelines.
|
||||
self.type = "manifold"
|
||||
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "ollama_manifold"
|
||||
|
||||
# Optionally, you can set the name of the manifold pipeline.
|
||||
self.name = "Ollama: "
|
||||
|
||||
class Valves(BaseModel):
|
||||
OLLAMA_BASE_URL: str
|
||||
|
||||
self.valves = Valves(**{"OLLAMA_BASE_URL": "http://localhost:11435"})
|
||||
self.pipelines = []
|
||||
pass
|
||||
|
||||
async def on_startup(self):
|
||||
# This function is called when the server is started.
|
||||
print(f"on_startup:{__name__}")
|
||||
self.pipelines = self.get_ollama_models()
|
||||
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_ollama_models()
|
||||
pass
|
||||
|
||||
def get_ollama_models(self):
|
||||
if self.valves.OLLAMA_BASE_URL:
|
||||
try:
|
||||
r = requests.get(f"{self.valves.OLLAMA_BASE_URL}/api/tags")
|
||||
models = r.json()
|
||||
return [
|
||||
{"id": model["model"], "name": model["name"]}
|
||||
for model in models["models"]
|
||||
]
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return [
|
||||
{
|
||||
"id": self.id,
|
||||
"name": "Could not fetch models from Ollama, please update the URL in the valves.",
|
||||
},
|
||||
]
|
||||
else:
|
||||
return []
|
||||
|
||||
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.'
|
||||
|
||||
if "user" in body:
|
||||
print("######################################")
|
||||
print(f'# User: {body["user"]["name"]} ({body["user"]["id"]})')
|
||||
print(f"# Message: {user_message}")
|
||||
print("######################################")
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
url=f"{self.valves.OLLAMA_BASE_URL}/v1/chat/completions",
|
||||
json={**body, "model": model_id},
|
||||
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}"
|
||||
55
examples/ollama_pipeline.py
Normal file
55
examples/ollama_pipeline.py
Normal file
@@ -0,0 +1,55 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
import requests
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "ollama_pipeline"
|
||||
self.name = "Ollama Pipeline"
|
||||
pass
|
||||
|
||||
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
|
||||
|
||||
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__}")
|
||||
|
||||
OLLAMA_BASE_URL = "http://localhost:11434"
|
||||
MODEL = "llama3"
|
||||
|
||||
if "user" in body:
|
||||
print("######################################")
|
||||
print(f'# User: {body["user"]["name"]} ({body["user"]["id"]})')
|
||||
print(f"# Message: {user_message}")
|
||||
print("######################################")
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
url=f"{OLLAMA_BASE_URL}/v1/chat/completions",
|
||||
json={**body, "model": MODEL},
|
||||
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}"
|
||||
57
examples/openai_pipeline.py
Normal file
57
examples/openai_pipeline.py
Normal file
@@ -0,0 +1,57 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
import requests
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
# Assign a unique identifier to the 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 = "OpenAI Pipeline"
|
||||
pass
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
OPENAI_API_KEY = "your-openai-api-key-here"
|
||||
MODEL = "gpt-3.5-turbo"
|
||||
|
||||
headers = {}
|
||||
headers["Authorization"] = f"Bearer {OPENAI_API_KEY}"
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
url="https://api.openai.com/v1/chat/completions",
|
||||
json={**body, "model": MODEL},
|
||||
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}"
|
||||
47
examples/python_code_pipeline.py
Normal file
47
examples/python_code_pipeline.py
Normal file
@@ -0,0 +1,47 @@
|
||||
from typing import List, Union, Generator, Iterator
|
||||
from schemas import OpenAIChatMessage
|
||||
import subprocess
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Optionally, you can set the id and name of the pipeline.
|
||||
self.id = "python_code_pipeline"
|
||||
self.name = "Python Code Pipeline"
|
||||
pass
|
||||
|
||||
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
|
||||
|
||||
def execute_python_code(self, code):
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["python", "-c", code], capture_output=True, text=True, check=True
|
||||
)
|
||||
stdout = result.stdout.strip()
|
||||
return stdout, result.returncode
|
||||
except subprocess.CalledProcessError as e:
|
||||
return e.output.strip(), e.returncode
|
||||
|
||||
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)
|
||||
|
||||
if body.get("title", False):
|
||||
print("Title Generation")
|
||||
return "Python Code Pipeline"
|
||||
else:
|
||||
stdout, return_code = self.execute_python_code(user_message)
|
||||
return stdout
|
||||
117
examples/rate_limit_filter_pipeline.py
Normal file
117
examples/rate_limit_filter_pipeline.py
Normal file
@@ -0,0 +1,117 @@
|
||||
from typing import List, Optional
|
||||
from pydantic import BaseModel
|
||||
from schemas import OpenAIChatMessage
|
||||
import time
|
||||
|
||||
|
||||
class Pipeline:
|
||||
def __init__(self):
|
||||
# Pipeline filters are only compatible with Open WebUI
|
||||
# You can think of filter pipeline as a middleware that can be used to edit the form data before it is sent to the OpenAI API.
|
||||
self.type = "filter"
|
||||
|
||||
# Assign a unique identifier to the 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 = "rate_limit_filter_pipeline"
|
||||
self.name = "Rate Limit Filter"
|
||||
|
||||
class Valves(BaseModel):
|
||||
# List target pipeline ids (models) that this filter will be connected to.
|
||||
# If you want to connect this filter to all pipelines, you can set pipelines to ["*"]
|
||||
pipelines: List[str] = []
|
||||
|
||||
# Assign a priority level to the filter pipeline.
|
||||
# The priority level determines the order in which the filter pipelines are executed.
|
||||
# The lower the number, the higher the priority.
|
||||
priority: int = 0
|
||||
|
||||
# Valves for rate limiting
|
||||
requests_per_minute: Optional[int] = None
|
||||
requests_per_hour: Optional[int] = None
|
||||
sliding_window_limit: Optional[int] = None
|
||||
sliding_window_minutes: Optional[int] = None
|
||||
|
||||
# Initialize rate limits
|
||||
self.valves = Valves(
|
||||
**{
|
||||
"pipelines": ["*"], # Connect to all pipelines
|
||||
"requests_per_minute": 10,
|
||||
"requests_per_hour": 1000,
|
||||
"sliding_window_limit": 100,
|
||||
"sliding_window_minutes": 15,
|
||||
}
|
||||
)
|
||||
|
||||
# Tracking data - user_id -> (timestamps of requests)
|
||||
self.user_requests = {}
|
||||
|
||||
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
|
||||
|
||||
def prune_requests(self, user_id: str):
|
||||
"""Prune old requests that are outside of the sliding window period."""
|
||||
now = time.time()
|
||||
if user_id in self.user_requests:
|
||||
self.user_requests[user_id] = [
|
||||
req
|
||||
for req in self.user_requests[user_id]
|
||||
if (
|
||||
(self.valves.requests_per_minute is not None and now - req < 60)
|
||||
or (self.valves.requests_per_hour is not None and now - req < 3600)
|
||||
or (
|
||||
self.valves.sliding_window_limit is not None
|
||||
and now - req < self.valves.sliding_window_minutes * 60
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
def log_request(self, user_id: str):
|
||||
"""Log a new request for a user."""
|
||||
now = time.time()
|
||||
if user_id not in self.user_requests:
|
||||
self.user_requests[user_id] = []
|
||||
self.user_requests[user_id].append(now)
|
||||
|
||||
def rate_limited(self, user_id: str) -> bool:
|
||||
"""Check if a user is rate limited."""
|
||||
self.prune_requests(user_id)
|
||||
|
||||
user_reqs = self.user_requests.get(user_id, [])
|
||||
|
||||
if self.valves.requests_per_minute is not None:
|
||||
requests_last_minute = sum(1 for req in user_reqs if time.time() - req < 60)
|
||||
if requests_last_minute >= self.valves.requests_per_minute:
|
||||
return True
|
||||
|
||||
if self.valves.requests_per_hour is not None:
|
||||
requests_last_hour = sum(1 for req in user_reqs if time.time() - req < 3600)
|
||||
if requests_last_hour >= self.valves.requests_per_hour:
|
||||
return True
|
||||
|
||||
if self.valves.sliding_window_limit is not None:
|
||||
requests_in_window = len(user_reqs)
|
||||
if requests_in_window >= self.valves.sliding_window_limit:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
async def inlet(self, body: dict, user: Optional[dict] = None) -> dict:
|
||||
print(f"pipe:{__name__}")
|
||||
print(body)
|
||||
print(user)
|
||||
|
||||
if user.get("role", "admin") == "user":
|
||||
user_id = user["id"] if user and "id" in user else "default_user"
|
||||
if self.rate_limited(user_id):
|
||||
raise Exception("Rate limit exceeded. Please try again later.")
|
||||
|
||||
self.log_request(user_id)
|
||||
return body
|
||||
Reference in New Issue
Block a user