diff --git a/main.py b/main.py index 3689c9b..09bb741 100644 --- a/main.py +++ b/main.py @@ -74,9 +74,9 @@ def on_startup(): PIPELINES[manifold_pipeline_id] = { "module": pipeline_id, + "type": pipeline.type if hasattr(pipeline, "type") else "pipe", "id": manifold_pipeline_id, "name": manifold_pipeline_name, - "manifold": True, "valves": ( pipeline.valves if hasattr(pipeline, "valves") else None ), @@ -84,22 +84,29 @@ def on_startup(): if pipeline.type == "filter": PIPELINES[pipeline_id] = { "module": pipeline_id, + "type": (pipeline.type if hasattr(pipeline, "type") else "pipe"), "id": pipeline_id, "name": ( pipeline.name if hasattr(pipeline, "name") else pipeline_id ), - "filter": True, "pipelines": ( - pipeline.pipelines if hasattr(pipeline, "pipelines") else [] + pipeline.valves.pipelines + if hasattr(pipeline, "valves") + and hasattr(pipeline.valves, "pipelines") + else [] ), "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, } else: PIPELINES[pipeline_id] = { "module": pipeline_id, + "type": (pipeline.type if hasattr(pipeline, "type") else "pipe"), "id": pipeline_id, "name": (pipeline.name if hasattr(pipeline, "name") else pipeline_id), "valves": pipeline.valves if hasattr(pipeline, "valves") else None, @@ -166,15 +173,13 @@ async def get_models(): "created": int(time.time()), "owned_by": "openai", "pipeline": { + "type": pipeline["type"], **( { - "type": ( - "pipeline" if not pipeline.get("filter") else "filter" - ), "pipelines": pipeline.get("pipelines", []), "priority": pipeline.get("priority", 0), } - if pipeline.get("filter", False) + if pipeline.get("type", "pipe") == "filter" else {} ), "valves": pipeline["valves"] != None, @@ -241,12 +246,13 @@ async def update_valves(pipeline_id: str, form_data: dict): @app.post("/{pipeline_id}/filter") async def filter(pipeline_id: str, form_data: FilterForm): - if pipeline_id not in app.state.PIPELINES or not app.state.PIPELINES[ - pipeline_id - ].get("filter", False): + if ( + pipeline_id not in app.state.PIPELINES + or app.state.PIPELINES[pipeline_id].get("type", "pipe") != "filter" + ): raise HTTPException( 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] diff --git a/pipelines/examples/filter_pipeline.py b/pipelines/examples/filter_pipeline.py index 6cb42c6..d376418 100644 --- a/pipelines/examples/filter_pipeline.py +++ b/pipelines/examples/filter_pipeline.py @@ -1,4 +1,5 @@ from typing import List, Optional +from pydantic import BaseModel from schemas import OpenAIChatMessage @@ -15,16 +16,21 @@ class Pipeline: self.id = "filter_pipeline" self.name = "Filter" - # 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. - self.priority = 0 + 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"]}) - # 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 async def on_startup(self): diff --git a/pipelines/examples/rate_limit_filter_pipeline.py b/pipelines/examples/rate_limit_filter_pipeline.py index 94996bb..d66bbc5 100644 --- a/pipelines/examples/rate_limit_filter_pipeline.py +++ b/pipelines/examples/rate_limit_filter_pipeline.py @@ -16,15 +16,17 @@ class Pipeline: self.id = "rate_limit_filter_pipeline" self.name = "Rate Limit Filter" - # 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. - self.priority = 0 - - # List target pipelines (models) that this filter will be connected to. - self.pipelines = ["*"] - 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 @@ -33,6 +35,7 @@ class Pipeline: # Initialize rate limits self.valves = Valves( **{ + "pipelines": ["*"], # Connect to all pipelines "requests_per_minute": 10, "requests_per_hour": 1000, "sliding_window_limit": 100, diff --git a/pipelines/rate_limit_filter_pipeline.py b/pipelines/rate_limit_filter_pipeline.py index 94996bb..d66bbc5 100644 --- a/pipelines/rate_limit_filter_pipeline.py +++ b/pipelines/rate_limit_filter_pipeline.py @@ -16,15 +16,17 @@ class Pipeline: self.id = "rate_limit_filter_pipeline" self.name = "Rate Limit Filter" - # 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. - self.priority = 0 - - # List target pipelines (models) that this filter will be connected to. - self.pipelines = ["*"] - 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 @@ -33,6 +35,7 @@ class Pipeline: # Initialize rate limits self.valves = Valves( **{ + "pipelines": ["*"], # Connect to all pipelines "requests_per_minute": 10, "requests_per_hour": 1000, "sliding_window_limit": 100,