From 24e02a901732e6a2ed9674585f7ed05466f4bd66 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Tue, 21 May 2024 18:18:02 -0700
Subject: [PATCH] refac

---
 main.py                                                | 10 ++++++----
 .../examples/{pipeline.py => pipeline_example.py}      |  0
 2 files changed, 6 insertions(+), 4 deletions(-)
 rename pipelines/examples/{pipeline.py => pipeline_example.py} (100%)

diff --git a/main.py b/main.py
index 43597e2..fbe30fb 100644
--- a/main.py
+++ b/main.py
@@ -64,6 +64,8 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(docs_url="/docs", redoc_url=None, lifespan=lifespan)
 
+app.state.PIPELINES = PIPELINES
+
 
 origins = ["*"]
 
@@ -112,14 +114,14 @@ async def get_models():
 async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
     user_message = get_last_user_message(form_data.messages)
 
-    if form_data.model not in PIPELINES:
+    if form_data.model not in app.state.PIPELINES:
         return HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail=f"Model {form_data.model} not found",
         )
 
-    async def job():
-        get_response = PIPELINES[form_data.model]["module"].get_response
+    def job():
+        get_response = app.state.PIPELINES[form_data.model]["module"].get_response
 
         if form_data.stream:
 
@@ -184,7 +186,7 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
                 ],
             }
 
-    return await job()
+    return await run_in_threadpool(job)
 
 
 @app.get("/")
diff --git a/pipelines/examples/pipeline.py b/pipelines/examples/pipeline_example.py
similarity index 100%
rename from pipelines/examples/pipeline.py
rename to pipelines/examples/pipeline_example.py