chore: mv examples out of pipelines

This commit is contained in:
Timothy J. Baek
2024-05-30 22:24:29 -07:00
parent 086acb5c29
commit fcc7f8abaa
24 changed files with 3 additions and 15 deletions

View 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

View 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}"

View 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}"

View 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."

View 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

View 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}"

View 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

View 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]

View 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

View 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

View 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}"

View 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}"

View 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

View 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

View 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

View 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

View 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
View 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}"

View 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}"

View 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}"

View 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}"

View 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

View 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