Add websockets

Dockerfile

@@ -1,34 +1,17 @@
-# 
 FROM python:3.11
 
-# 
 WORKDIR /code
 
-# 
 COPY ./requirements.txt /code/requirements.txt
 
-# 
-RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
-
-# 
-COPY ./start_service.sh /code/start_service.sh
-
-# 
 COPY ./app /code/app
 
-RUN useradd -m docker && echo "docker:docker" | chpasswd && adduser docker sudo
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
-# RUN curl -fsSL https://ollama.com/install.sh | sh
+RUN useradd -m docker && echo "docker:docker" | chpasswd && adduser docker sudo
 
 USER docker
 
-# RUN nohup ollama serve & sleep 5
-
-#
-# RUN chmod +x /code/start_service.sh
-
-#
 EXPOSE 7860
 
-# # Run .sh file
-ENTRYPOINT ["sh", "/code/start_service.sh"]
+CMD ["fastapi", "run", "/code/app/main.py", "--port", "7860"]

app/main.py

@@ -1,10 +1,7 @@
 import os
 import shutil
-import tempfile
-from tempfile import NamedTemporaryFile
-from pathlib import Path
-
-from fastapi import FastAPI, UploadFile
+from typing import List
+from fastapi import FastAPI, UploadFile, WebSocket, WebSocketDisconnect
 from fastapi.middleware import Middleware
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
@@ -24,11 +21,47 @@ app = FastAPI(middleware=middleware)
 files_dir = os.path.expanduser("~/wtp_be_files/")
 session_assistant = ChatPDF()
 
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: List[WebSocket] = []
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections.append(websocket)
+
+    def disconnect(self, websocket: WebSocket):
+        self.active_connections.remove(websocket)
+
+    async def send_personal_message(self, message: str, websocket: WebSocket):
+        await websocket.send_text(message)
+
+    async def broadcast(self, message: str):
+        for connection in self.active_connections:
+            await connection.send_text(message)
+
+manager = ConnectionManager()
+
+@app.websocket("/ws/{client_id}")
+async def websocket_endpoint(websocket: WebSocket, client_id: int):
+    await manager.connect(websocket)
+    now = datetime.now()
+    current_time = now.strftime("%H:%M")
+    try:
+        while True:
+            data = await websocket.receive_text()
+            # await manager.send_personal_message(f"You wrote: {data}", websocket)
+            message = {"time":current_time,"clientId":client_id,"message":data}
+            await manager.broadcast(json.dumps(message))
+            
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        message = {"time":current_time,"clientId":client_id,"message":"Offline"}
+        await manager.broadcast(json.dumps(message))
+
 @app.get("/query")
 async def process_input(text: str):
     if text and len(text.strip()) > 0:
         text = text.strip()
-        print("PRINTING STREAM")
         return StreamingResponse(session_assistant.ask(text), media_type='text/event-stream')
 
 

app/rag.py

@@ -22,8 +22,6 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class ChatPDF:
-    query_engine = None
-
     def __init__(self):
         self.text_parser = SentenceSplitter(chunk_size=512, chunk_overlap=24)
 
@@ -45,7 +43,7 @@ class ChatPDF:
             model_url="https://huggingface.co/Qwen/Qwen1.5-1.8B-Chat-GGUF/resolve/main/qwen1_5-1_8b-chat-q4_k_m.gguf",
             temperature=0.1,
             max_new_tokens=256,
-            context_window=3900,
+            context_window=3900, #32k
             # generate_kwargs={},
             # model_kwargs={"n_gpu_layers": -1},
             # messages_to_prompt=self.messages_to_prompt,
@@ -108,10 +106,11 @@ class ChatPDF:
         logger.info("retrieving the response to the query")
         streaming_response = self.query_engine.query(query)
         print(streaming_response)
-        # return streaming_response.response_gen
-        for text in streaming_response.response_gen:
+        print("PRINTING STREAM")
+        generator = streaming_response.response_gen
+        for text in generator:
             print(text)
             yield text
 
     def clear(self):
-        self.query_engine = None
+        pass

requirements.txt

@@ -5,4 +5,5 @@ qdrant-client
 python-dotenv
 llama-index-llms-llama-cpp
 llama-index-embeddings-fastembed
+websockets
 fastembed==0.2.7

start_service.sh

@@ -1,16 +0,0 @@
-#!/bin/sh
-
-# # Start Ollama in the background
-# ollama serve &
-
-# # Wait for Ollama to start
-# sleep 5
-
-# #
-# ollama pull mxbai-embed-large
-
-# # Pull and run <YOUR_MODEL_NAME>
-# ollama pull qwen:1.8b
-
-#
-fastapi run /code/app/main.py --port 7860