feat: add websocket

main.py

@@ -21,7 +21,7 @@ from langchain.vectorstores import Chroma
 from prompt_template_utils import get_prompt_template
 from load_models import load_model
 
-from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME, PATH_NAME_SOURCE_DIRECTORY, SHOW_SOURCES
+from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME, PATH_NAME_SOURCE_DIRECTORY, SHOW_SOURCES, CONTEXT_WINDOW_SIZE, MAX_NEW_TOKENS
 
 class Predict(BaseModel):
     prompt: str
@@ -273,7 +273,11 @@ async def websocket_endpoint_room(websocket: WebSocket, room_id: str, user_id: s
         return_source_documents=SHOW_SOURCES,
         chain_type_kwargs={
             "prompt": prompt,
-            "memory": memory
+            "memory": memory,
+            "n_ctx": CONTEXT_WINDOW_SIZE,
+            "max_tokens": MAX_NEW_TOKENS,
+            "n_batch": MAX_NEW_TOKENS,
+
         },
     )