feat: add websocket

main.py

@@ -40,7 +40,7 @@ EMBEDDINGS = HuggingFaceInstructEmbeddings(model_name=EMBEDDING_MODEL_NAME, mode
 DB = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=EMBEDDINGS, client_settings=CHROMA_SETTINGS)
 RETRIEVER = DB.as_retriever()
 
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[])
+LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True)
 
 prompt, memory = get_prompt_template(promptTemplate_type="llama", history=True)
 
@@ -211,24 +211,7 @@ async def create_upload_file(file: UploadFile):
 
 @api_app.websocket("/ws/{user_id}")
 async def websocket_endpoint_student(websocket: WebSocket, user_id: str):
-    DB_USER = Chroma(
-        persist_directory=PERSIST_DIRECTORY,
-        embedding_function=EMBEDDINGS,
-        client_settings=CHROMA_SETTINGS,
-    )
-
-    RETRIEVER = DB_USER.as_retriever()
-
-    QA = RetrievalQA.from_chain_type(
-        llm=LLM,
-        chain_type="stuff",
-        retriever=RETRIEVER,
-        return_source_documents=SHOW_SOURCES,
-        chain_type_kwargs={
-            "prompt": prompt,
-            "memory": memory
-        },
-    )
+    global QA
 
     message = {
         "message": f"Student {user_id} connected"
@@ -258,24 +241,7 @@ async def websocket_endpoint_student(websocket: WebSocket, user_id: str):
 
 @api_app.websocket("/ws/{room_id}/{user_id}")
 async def websocket_endpoint_room(websocket: WebSocket, room_id: str, user_id: str):
-    DB_ROOM = Chroma(
-        persist_directory=PERSIST_DIRECTORY,
-        embedding_function=EMBEDDINGS,
-        client_settings=CHROMA_SETTINGS,
-    )
-
-    RETRIEVER = DB_ROOM.as_retriever()
-
-    QA = RetrievalQA.from_chain_type(
-        llm=LLM,
-        chain_type="stuff",
-        retriever=RETRIEVER,
-        return_source_documents=SHOW_SOURCES,
-        chain_type_kwargs={
-            "prompt": prompt,
-            "memory": memory,
-        },
-    )
+    global QA
 
     message = {
         "message": f"Student {user_id} connected to the classroom"