From 39da1c982a0ab3769311ae0d29e3309bff3a03aa Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Tue, 21 May 2024 18:47:46 -0700
Subject: [PATCH] fix

---
 config.py | 4 ----
 main.py   | 9 ++++++---
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/config.py b/config.py
index c74f8a2..29918dd 100644
--- a/config.py
+++ b/config.py
@@ -10,7 +10,3 @@ try:
     load_dotenv(find_dotenv("./.env"))
 except ImportError:
     print("dotenv not installed, skipping...")
-
-
-MODEL_ID = os.environ.get("MODEL_ID", "plugin-id")
-MODEL_NAME = os.environ.get("MODEL_NAME", "Plugin Model")
diff --git a/main.py b/main.py
index eaf291e..dbc34cf 100644
--- a/main.py
+++ b/main.py
@@ -14,7 +14,6 @@ import uuid
 
 from utils import get_last_user_message, stream_message_template
 from schemas import OpenAIChatCompletionForm
-from config import MODEL_ID, MODEL_NAME
 
 import os
 import importlib.util
@@ -61,6 +60,7 @@ async def lifespan(app: FastAPI):
     yield
 
     for pipeline in PIPELINES.values():
+
         if hasattr(pipeline["module"], "on_shutdown"):
             await pipeline["module"].on_shutdown()
 
@@ -124,6 +124,7 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
         )
 
     def job():
+        print(form_data.model)
         get_response = app.state.PIPELINES[form_data.model]["module"].get_response
 
         if form_data.stream:
@@ -131,6 +132,8 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
             def stream_content():
                 res = get_response(user_message, messages=form_data.messages)
 
+                print(res)
+
                 if isinstance(res, str):
                     message = stream_message_template(res)
                     yield f"data: {json.dumps(message)}\n\n"
@@ -145,7 +148,7 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
                     "id": f"{form_data.model}-{str(uuid.uuid4())}",
                     "object": "chat.completion.chunk",
                     "created": int(time.time()),
-                    "model": MODEL_ID,
+                    "model": form_data.model,
                     "choices": [
                         {
                             "index": 0,
@@ -175,7 +178,7 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
                 "id": f"{form_data.model}-{str(uuid.uuid4())}",
                 "object": "chat.completion",
                 "created": int(time.time()),
-                "model": MODEL_ID,
+                "model": form_data.model,
                 "choices": [
                     {
                         "index": 0,