OpenWebUI/pipelines
1
0
Fork 0
mirror of https://github.com/open-webui/pipelines synced 2025-05-31 18:48:10 +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
PIPELINES = {}
PIPELINE_MODULES = {}
def on_startup():
@ -41,9 +42,10 @@ def on_startup():
logging.info("Loaded:", loaded_module.__name__)
pipeline = loaded_module.Pipeline()
pipeline_id = pipeline.id if hasattr(pipeline, "id") else loaded_module.__name__
PIPELINE_MODULES[pipeline_id] = pipeline
if hasattr(pipeline, "manifold") and pipeline.manifold:
for p in pipeline.pipelines:
manifold_pipeline_id = f'{pipeline_id}.{p["id"]}'
@ -53,14 +55,14 @@ def on_startup():
manifold_pipeline_name = f"{pipeline.name}{manifold_pipeline_name}"
PIPELINES[manifold_pipeline_id] = {
"module": pipeline,
"module": pipeline_id,
"id": manifold_pipeline_id,
"name": manifold_pipeline_name,
"manifold": True,
}
else:
PIPELINES[loaded_module.__name__] = {
"module": pipeline,
"module": pipeline_id,
"id": pipeline_id,
"name": (
pipeline.name
@ -78,14 +80,14 @@ from contextlib import asynccontextmanager
@asynccontextmanager
async def lifespan(app: FastAPI):
for pipeline in PIPELINES.values():
if hasattr(pipeline["module"], "on_startup"):
await pipeline["module"].on_startup()
for module in PIPELINE_MODULES.values():
if hasattr(module, "on_startup"):
await module.on_startup()
yield
for pipeline in PIPELINES.values():
if hasattr(pipeline["module"], "on_shutdown"):
await pipeline["module"].on_shutdown()
for module in PIPELINE_MODULES.values():
if hasattr(module, "on_shutdown"):
await module.on_shutdown()
app = FastAPI(docs_url="/docs", redoc_url=None, lifespan=lifespan)