fix: add websocket in handlerToken

constants.py

@@ -13,6 +13,8 @@ ROOT_DIRECTORY = os.path.dirname(os.path.realpath(__file__))
 
 PATH_NAME_SOURCE_DIRECTORY = "SOURCE_DOCUMENTS"
 
+SHOW_SOURCES=True
+
 # Define the folder for storing database
 SOURCE_DIRECTORY = f"{ROOT_DIRECTORY}/{PATH_NAME_SOURCE_DIRECTORY}"
 
@@ -43,6 +45,7 @@ N_BATCH = 2048
 # N_BATCH = 512
 
 
+
 # https://python.langchain.com/en/latest/_modules/langchain/document_loaders/excel.html#UnstructuredExcelLoader
 DOCUMENT_MAP = {
     ".txt": TextLoader,

main.py

@@ -1,33 +1,29 @@
+from typing import Any, Dict, List
+
 import os
 import glob
 import shutil
 import subprocess
 import asyncio
-
-from typing import Any, Dict, List
+import redis
+import torch
 
 from fastapi import FastAPI, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
 from fastapi.staticfiles import StaticFiles
 
 from pydantic import BaseModel
 
-# import torch
+# langchain
 from langchain.chains import RetrievalQA
 from langchain.embeddings import HuggingFaceInstructEmbeddings
-from langchain.prompts import PromptTemplate
-from langchain.memory import ConversationBufferMemory
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import LLMResult
+from langchain.vectorstores import Chroma
 
 from prompt_template_utils import get_prompt_template
-
-# from langchain.embeddings import HuggingFaceEmbeddings
 from load_models import load_model
 
-# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-from langchain.vectorstores import Chroma
-
-from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME, PATH_NAME_SOURCE_DIRECTORY
+from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME, PATH_NAME_SOURCE_DIRECTORY, SHOW_SOURCES
 
 class Predict(BaseModel):
     prompt: str
@@ -35,54 +31,36 @@ class Predict(BaseModel):
 class Delete(BaseModel):
     filename: str
 
-# if torch.backends.mps.is_available():
-#     DEVICE_TYPE = "mps"
-# elif torch.cuda.is_available():
-#     DEVICE_TYPE = "cuda"
-# else:
-#     DEVICE_TYPE = "cpu"
-
-DEVICE_TYPE = "cuda"
-SHOW_SOURCES = True
+if torch.backends.mps.is_available():
+    DEVICE_TYPE = "mps"
+elif torch.cuda.is_available():
+    DEVICE_TYPE = "cuda"
+else:
+    DEVICE_TYPE = "cpu"
 
 EMBEDDINGS = HuggingFaceInstructEmbeddings(model_name=EMBEDDING_MODEL_NAME, model_kwargs={"device": DEVICE_TYPE})
-
-# load the vectorstore
-DB = Chroma(
-    persist_directory=PERSIST_DIRECTORY,
-    embedding_function=EMBEDDINGS,
-    client_settings=CHROMA_SETTINGS,
-)
-
+DB = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=EMBEDDINGS, client_settings=CHROMA_SETTINGS)
 RETRIEVER = DB.as_retriever()
 
 class MyCustomSyncHandler(BaseCallbackHandler):
-    def __init__(self):
-        self.end = False
-
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
     ) -> None:
-        self.end = False
+        print(f'on_llm_start self {self}')
+        print(f'on_llm_start kwargs {prompts}')
+        print(f'on_llm_start token {kwargs}')
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        self.end = True
+        print(f'on_llm_end self {self}')
+        print(f'on_llm_end kwargs {response}')
+        print(f'on_llm_end token {kwargs}')
 
     def on_llm_new_token(self, token: str, **kwargs) -> Any:
-        print(self)
-        print(kwargs)
-
-
-# Create State
-handlerToken = MyCustomSyncHandler()
+        print(f'on_llm_new_token self {self}')
+        print(f'on_llm_new_token kwargs {kwargs}')
+        print(f'on_llm_new_token token {token}')
 
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[handlerToken])
-
-# template = """You are a helpful, respectful and honest assistant.
-# Always answer in the most helpful and safe way possible without trying to make up an answer, if you don't know the answer just say "I don't know" and don't share false information or topics that were not provided in your training. Use a maximum of 15 sentences. Your answer should be as concise and clear as possible. Always say "thank you for asking!" at the end of your answer.
-# Context: {context}
-# Question: {question}
-# """
+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)
 
@@ -97,7 +75,9 @@ QA = RetrievalQA.from_chain_type(
     },
 )
 
+redisClient = redis.Redis(host='localhost', port=6379, db=0)
 
+r.set('foo', 'bar')
 
 app = FastAPI(title="homepage-app")
 api_app = FastAPI(title="api app")
@@ -146,7 +126,7 @@ def run_ingest_route():
             retriever=RETRIEVER,
             return_source_documents=SHOW_SOURCES,
             chain_type_kwargs={
-                "prompt": QA_CHAIN_PROMPT,
+                "prompt": prompt,
                 "memory": memory
             },
         )
@@ -250,16 +230,21 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
 
     await websocket.accept()
 
-    oldReceiveText = ''
-
     try:
         while True:
             prompt = await websocket.receive_text()
+            QA(
+                inputs=prompt,
+                return_only_outputs=True,
+                callbacks=[MyCustomSyncHandler()],
+                tags=f'{client_id}',
+                run_name=f'{client_id}',
+                include_run_info=True
+            )
+
+            response = redisClient.get('foo')
 
-            if (oldReceiveText != prompt):
-                handlerToken.callback = websocket.send_text
-                oldReceiveText = prompt
-                await QA(prompt)
+            await websocket.send_text(response)
 
     except WebSocketDisconnect:
         print('disconnect')

requirements.txt

@@ -29,6 +29,7 @@ uvicorn
 fastapi
 websockets
 pydantic
+redis
 
 # Streamlit related
 streamlit