refac: filter pipelines

This commit is contained in:
Timothy J. Baek 2024-05-28 14:59:42 -07:00
parent e25670e6ec
commit 2e899f0e19
4 changed files with 55 additions and 37 deletions

30
main.py
View File

@ -74,9 +74,9 @@ def on_startup():
PIPELINES[manifold_pipeline_id] = { PIPELINES[manifold_pipeline_id] = {
"module": pipeline_id, "module": pipeline_id,
"type": pipeline.type if hasattr(pipeline, "type") else "pipe",
"id": manifold_pipeline_id, "id": manifold_pipeline_id,
"name": manifold_pipeline_name, "name": manifold_pipeline_name,
"manifold": True,
"valves": ( "valves": (
pipeline.valves if hasattr(pipeline, "valves") else None pipeline.valves if hasattr(pipeline, "valves") else None
), ),
@ -84,22 +84,29 @@ def on_startup():
if pipeline.type == "filter": if pipeline.type == "filter":
PIPELINES[pipeline_id] = { PIPELINES[pipeline_id] = {
"module": pipeline_id, "module": pipeline_id,
"type": (pipeline.type if hasattr(pipeline, "type") else "pipe"),
"id": pipeline_id, "id": pipeline_id,
"name": ( "name": (
pipeline.name if hasattr(pipeline, "name") else pipeline_id pipeline.name if hasattr(pipeline, "name") else pipeline_id
), ),
"filter": True,
"pipelines": ( "pipelines": (
pipeline.pipelines if hasattr(pipeline, "pipelines") else [] pipeline.valves.pipelines
if hasattr(pipeline, "valves")
and hasattr(pipeline.valves, "pipelines")
else []
), ),
"priority": ( "priority": (
pipeline.priority if hasattr(pipeline, "priority") else 0 pipeline.valves.priority
if hasattr(pipeline, "valves")
and hasattr(pipeline.valves, "priority")
else 0
), ),
"valves": pipeline.valves if hasattr(pipeline, "valves") else None, "valves": pipeline.valves if hasattr(pipeline, "valves") else None,
} }
else: else:
PIPELINES[pipeline_id] = { PIPELINES[pipeline_id] = {
"module": pipeline_id, "module": pipeline_id,
"type": (pipeline.type if hasattr(pipeline, "type") else "pipe"),
"id": pipeline_id, "id": pipeline_id,
"name": (pipeline.name if hasattr(pipeline, "name") else pipeline_id), "name": (pipeline.name if hasattr(pipeline, "name") else pipeline_id),
"valves": pipeline.valves if hasattr(pipeline, "valves") else None, "valves": pipeline.valves if hasattr(pipeline, "valves") else None,
@ -166,15 +173,13 @@ async def get_models():
"created": int(time.time()), "created": int(time.time()),
"owned_by": "openai", "owned_by": "openai",
"pipeline": { "pipeline": {
"type": pipeline["type"],
**( **(
{ {
"type": (
"pipeline" if not pipeline.get("filter") else "filter"
),
"pipelines": pipeline.get("pipelines", []), "pipelines": pipeline.get("pipelines", []),
"priority": pipeline.get("priority", 0), "priority": pipeline.get("priority", 0),
} }
if pipeline.get("filter", False) if pipeline.get("type", "pipe") == "filter"
else {} else {}
), ),
"valves": pipeline["valves"] != None, "valves": pipeline["valves"] != None,
@ -241,12 +246,13 @@ async def update_valves(pipeline_id: str, form_data: dict):
@app.post("/{pipeline_id}/filter") @app.post("/{pipeline_id}/filter")
async def filter(pipeline_id: str, form_data: FilterForm): async def filter(pipeline_id: str, form_data: FilterForm):
if pipeline_id not in app.state.PIPELINES or not app.state.PIPELINES[ if (
pipeline_id pipeline_id not in app.state.PIPELINES
].get("filter", False): or app.state.PIPELINES[pipeline_id].get("type", "pipe") != "filter"
):
raise HTTPException( raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND, status_code=status.HTTP_404_NOT_FOUND,
detail=f"filter {pipeline_id} not found", detail=f"Filter {pipeline_id} not found",
) )
pipeline = PIPELINE_MODULES[pipeline_id] pipeline = PIPELINE_MODULES[pipeline_id]

View File

@ -1,4 +1,5 @@
from typing import List, Optional from typing import List, Optional
from pydantic import BaseModel
from schemas import OpenAIChatMessage from schemas import OpenAIChatMessage
@ -15,16 +16,21 @@ class Pipeline:
self.id = "filter_pipeline" self.id = "filter_pipeline"
self.name = "Filter" 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. # Assign a priority level to the filter pipeline.
# The priority level determines the order in which the filter pipelines are executed. # The priority level determines the order in which the filter pipelines are executed.
# The lower the number, the higher the priority. # The lower the number, the higher the priority.
self.priority = 0 priority: int = 0
# Add your custom parameters here
pass
self.valves = Valves(**{"pipelines": ["llama3:latest"]})
# List target pipelines (models) that this filter will be connected to.
# If you want to connect this filter to all pipelines, you can set pipelines to ["*"]
self.pipelines = [
{"id": "llama3:latest"},
]
pass pass
async def on_startup(self): async def on_startup(self):

View File

@ -16,15 +16,17 @@ class Pipeline:
self.id = "rate_limit_filter_pipeline" self.id = "rate_limit_filter_pipeline"
self.name = "Rate Limit Filter" 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. # Assign a priority level to the filter pipeline.
# The priority level determines the order in which the filter pipelines are executed. # The priority level determines the order in which the filter pipelines are executed.
# The lower the number, the higher the priority. # The lower the number, the higher the priority.
self.priority = 0 priority: int = 0
# List target pipelines (models) that this filter will be connected to. # Valves for rate limiting
self.pipelines = ["*"]
class Valves(BaseModel):
requests_per_minute: Optional[int] = None requests_per_minute: Optional[int] = None
requests_per_hour: Optional[int] = None requests_per_hour: Optional[int] = None
sliding_window_limit: Optional[int] = None sliding_window_limit: Optional[int] = None
@ -33,6 +35,7 @@ class Pipeline:
# Initialize rate limits # Initialize rate limits
self.valves = Valves( self.valves = Valves(
**{ **{
"pipelines": ["*"], # Connect to all pipelines
"requests_per_minute": 10, "requests_per_minute": 10,
"requests_per_hour": 1000, "requests_per_hour": 1000,
"sliding_window_limit": 100, "sliding_window_limit": 100,

View File

@ -16,15 +16,17 @@ class Pipeline:
self.id = "rate_limit_filter_pipeline" self.id = "rate_limit_filter_pipeline"
self.name = "Rate Limit Filter" 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. # Assign a priority level to the filter pipeline.
# The priority level determines the order in which the filter pipelines are executed. # The priority level determines the order in which the filter pipelines are executed.
# The lower the number, the higher the priority. # The lower the number, the higher the priority.
self.priority = 0 priority: int = 0
# List target pipelines (models) that this filter will be connected to. # Valves for rate limiting
self.pipelines = ["*"]
class Valves(BaseModel):
requests_per_minute: Optional[int] = None requests_per_minute: Optional[int] = None
requests_per_hour: Optional[int] = None requests_per_hour: Optional[int] = None
sliding_window_limit: Optional[int] = None sliding_window_limit: Optional[int] = None
@ -33,6 +35,7 @@ class Pipeline:
# Initialize rate limits # Initialize rate limits
self.valves = Valves( self.valves = Valves(
**{ **{
"pipelines": ["*"], # Connect to all pipelines
"requests_per_minute": 10, "requests_per_minute": 10,
"requests_per_hour": 1000, "requests_per_hour": 1000,
"sliding_window_limit": 100, "sliding_window_limit": 100,