OpenWebUI/pipelines
1
0
Fork 0
mirror of https://github.com/open-webui/pipelines synced 2025-06-26 18:15:58 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy J. Baek 2024-05-26 23:54:26 -07:00
parent c697e265e7
commit 8c3f262ffd
1 changed files with 11 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
main.py
View File

@ -23,6 +23,7 @@ import logging
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
PIPELINES = {} PIPELINES = {}
PIPELINE_MODULES = {}
def on_startup(): def on_startup():
@ -41,9 +42,10 @@ def on_startup():
logging.info("Loaded:", loaded_module.__name__) logging.info("Loaded:", loaded_module.__name__)
pipeline = loaded_module.Pipeline() pipeline = loaded_module.Pipeline()
pipeline_id = pipeline.id if hasattr(pipeline, "id") else loaded_module.__name__ pipeline_id = pipeline.id if hasattr(pipeline, "id") else loaded_module.__name__
PIPELINE_MODULES[pipeline_id] = pipeline
if hasattr(pipeline, "manifold") and pipeline.manifold: if hasattr(pipeline, "manifold") and pipeline.manifold:
for p in pipeline.pipelines: for p in pipeline.pipelines:
manifold_pipeline_id = f'{pipeline_id}.{p["id"]}' manifold_pipeline_id = f'{pipeline_id}.{p["id"]}'
@ -53,14 +55,14 @@ def on_startup():
manifold_pipeline_name = f"{pipeline.name}{manifold_pipeline_name}" manifold_pipeline_name = f"{pipeline.name}{manifold_pipeline_name}"
PIPELINES[manifold_pipeline_id] = { PIPELINES[manifold_pipeline_id] = {
"module": pipeline, "module": pipeline_id,
"id": manifold_pipeline_id, "id": manifold_pipeline_id,
"name": manifold_pipeline_name, "name": manifold_pipeline_name,
"manifold": True, "manifold": True,
} }
else: else:
PIPELINES[loaded_module.__name__] = { PIPELINES[loaded_module.__name__] = {
"module": pipeline, "module": pipeline_id,
"id": pipeline_id, "id": pipeline_id,
"name": ( "name": (
pipeline.name pipeline.name
@ -78,14 +80,14 @@ from contextlib import asynccontextmanager
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
for pipeline in PIPELINES.values(): for module in PIPELINE_MODULES.values():
if hasattr(pipeline["module"], "on_startup"): if hasattr(module, "on_startup"):
await pipeline["module"].on_startup() await module.on_startup()
yield yield
for pipeline in PIPELINES.values(): for module in PIPELINE_MODULES.values():
if hasattr(pipeline["module"], "on_shutdown"): if hasattr(module, "on_shutdown"):
await pipeline["module"].on_shutdown() await module.on_shutdown()
app = FastAPI(docs_url="/docs", redoc_url=None, lifespan=lifespan) app = FastAPI(docs_url="/docs", redoc_url=None, lifespan=lifespan)