supports pdf/txt file uploading & embedding

Author: c0sogi

.vscode/launch.json

@@ -8,16 +8,5 @@
             "program": "frontend/lib/main.dart",
             "deviceId": "edge"
         },
-        {
-            "name": "(Windows) Launch",
-            "type": "cppvsdbg",
-            "request": "launch",
-            "program": "enter program name, for example ${workspaceFolder}/a.exe",
-            "args": [],
-            "stopAtEntry": false,
-            "cwd": "${fileDirname}",
-            "environment": [],
-            "console": "externalTerminal"
-        }
     ]
 }

app/utils/chatgpt/chatgpt_commands.py

@@ -6,7 +6,6 @@
 from uuid import uuid4
 
 from fastapi import WebSocket
-
 from app.errors.api_exceptions import InternalServerError
 from app.utils.chatgpt.chatgpt_buffer import BufferedUserContext
 from app.utils.chatgpt.chatgpt_cache_manager import ChatGptCacheManager
@@ -534,7 +533,7 @@ async def testchaining(chain_size: int, buffer: BufferedUserContext) -> Tuple[st
         return f"/testchaining {chain_size-1}", ResponseType.REPEAT_COMMAND
 
     @staticmethod
-    @CommandResponse.send_message_and_stop
+    @CommandResponse.handle_gpt
     async def query(query: str, /, buffer: BufferedUserContext) -> None:
         """Query from redis vectorstore\n
         /query <query>"""
@@ -566,6 +565,5 @@ async def embed(text_to_embed: str, /) -> str:
         /embed <text_to_embed>"""
         await VectorStoreManager.create_documents(
             text=text_to_embed,
-            chunk_size=500,
         )
         return "Embedding successful!"

app/utils/chatgpt/chatgpt_fileloader.py

@@ -0,0 +1,37 @@
+import io
+from typing import IO, Any
+
+from langchain.document_loaders.unstructured import UnstructuredBaseLoader
+from langchain.docstore.document import Document
+from unstructured.partition.auto import partition
+
+
+class UnstructuredFileIOLoader(UnstructuredBaseLoader):
+    """Loader that uses unstructured to load file IO objects."""
+
+    def __init__(self, file: IO, filename: str, mode: str = "single", **unstructured_kwargs: Any):
+        """Initialize with file path."""
+        self.file = file
+        self.filename = filename
+        super().__init__(mode=mode, **unstructured_kwargs)
+
+    def _get_elements(self) -> list:
+        return partition(file=self.file, file_filename=self.filename, **self.unstructured_kwargs)
+
+    def _get_metadata(self) -> dict:
+        return {}
+
+
+def read_bytes_to_documents(file: bytes, filename: str) -> list[Document]:
+    return UnstructuredFileIOLoader(file=io.BytesIO(file), filename=filename).load()
+
+
+def read_bytes_to_text(file: bytes, filename: str) -> str:
+    return "\n\n".join([doc.page_content for doc in read_bytes_to_documents(file=file, filename=filename)])
+
+
+if __name__ == "__main__":
+    with open(r"test.pdf", "rb") as f:
+        file = f.read()
+    text = read_bytes_to_text(file, "test.pdf")
+    print(text)

app/utils/chatgpt/chatgpt_stream_manager.py

@@ -1,16 +1,20 @@
 from fastapi import WebSocket, WebSocketDisconnect
-
+from fastapi.concurrency import run_in_threadpool
+from pydantic import ValidationError
 from app.errors.gpt_exceptions import GptOtherException, GptTextGenerationException, GptTooMuchTokenException
+from app.utils.chatgpt.chatgpt_buffer import BufferedUserContext
 from app.utils.chatgpt.chatgpt_cache_manager import ChatGptCacheManager
 from app.utils.chatgpt.chatgpt_commands import (
+    ChatGptCommands,
+    command_handler,
     create_new_chat_room,
     get_contexts_sorted_from_recent_to_past,
-    command_handler,
 )
+from app.utils.chatgpt.chatgpt_fileloader import read_bytes_to_text
 from app.utils.chatgpt.chatgpt_message_manager import MessageManager
+from app.utils.chatgpt.chatgpt_vectorstore_manager import VectorStoreManager
 from app.utils.chatgpt.chatgpt_websocket_manager import HandleMessage, SendToWebsocket
 from app.utils.logger import api_logger
-from app.utils.chatgpt.chatgpt_buffer import BufferedUserContext
 from app.viewmodels.base_models import MessageFromWebsocket, MessageToWebsocket
 from app.viewmodels.gpt_models import GptRoles
 
@@ -33,8 +37,20 @@ async def begin_chat(
 
     while True:  # loop until connection is closed
         try:
-            # receive message from websocket
-            received: MessageFromWebsocket = MessageFromWebsocket.parse_raw(await websocket.receive_text())
+            rcvd: dict = await websocket.receive_json()
+            assert isinstance(rcvd, dict)
+            if "filename" in rcvd:
+                text: str = await run_in_threadpool(
+                    read_bytes_to_text, await websocket.receive_bytes(), rcvd["filename"]
+                )
+                docs: list[str] = await VectorStoreManager.create_documents(text)
+                await SendToWebsocket.message(
+                    websocket=websocket,
+                    msg=f"Successfully embedded documents. You uploaded file begins with...\n\n```{docs[0][:50]}```...",
+                    chat_room_id=buffer.current_chat_room_id,
+                )
+                continue
+            received: MessageFromWebsocket = MessageFromWebsocket(**rcvd)
 
             if received.chat_room_id != buffer.current_chat_room_id:  # change chat room
                 index: int | None = buffer.find_index_of_chatroom(received.chat_room_id)
@@ -76,6 +92,20 @@ async def begin_chat(
 
         except WebSocketDisconnect:
             raise WebSocketDisconnect(code=1000, reason="client disconnected")
+        except (AssertionError, ValidationError):
+            await SendToWebsocket.message(
+                websocket=websocket,
+                msg="Invalid message. Message is not in the correct format, maybe frontend - backend version mismatch?",
+                chat_room_id=buffer.current_chat_room_id,
+            )
+            continue
+        except ValueError:
+            await SendToWebsocket.message(
+                websocket=websocket,
+                msg="Invalid file type.",
+                chat_room_id=buffer.current_chat_room_id,
+            )
+            continue
         except GptTextGenerationException:
             await MessageManager.rpop_message_history_safely(
                 user_gpt_context=buffer.current_user_gpt_context, role=GptRoles.USER

app/utils/chatgpt/chatgpt_vectorstore_manager.py

@@ -1,5 +1,6 @@
 from asyncio import gather
 from langchain.text_splitter import TokenTextSplitter
+from langchain.document_loaders import UnstructuredFileLoader
 from app.utils.langchain.redis_vectorstore import Document
 from app.database.connection import cache
 
@@ -8,17 +9,17 @@ class VectorStoreManager:
     @staticmethod
     async def create_documents(
         text: str,
-        chunk_size: int,
+        chunk_size: int = 500,
         chunk_overlap: int = 0,
         tokenizer_model: str = "gpt-3.5-turbo",
     ) -> list[str]:
-        return await cache.vectorstore.aadd_texts(
-            texts=TokenTextSplitter(
-                chunk_size=chunk_size,
-                chunk_overlap=chunk_overlap,
-                model_name=tokenizer_model,
-            ).split_text(text)
-        )
+        texts = TokenTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            model_name=tokenizer_model,
+        ).split_text(text)
+        await cache.vectorstore.aadd_texts(texts=texts)
+        return texts
 
     @staticmethod
     async def asimilarity_search(queries: list[str], k: int = 1) -> list[list[Document]]:

app/utils/langchain/redis_vectorstore.py

@@ -3,9 +3,9 @@
 
 import json
 import logging
-from uuid import uuid4
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Type, Union
+from uuid import uuid4
 
 import numpy as np
 import pkg_resources
@@ -16,6 +16,10 @@
 from langchain.vectorstores.base import VectorStore
 from pydantic import BaseModel, root_validator
 
+try:
+    from starlette.concurrency import run_in_threadpool
+except ImportError:
+    raise ImportError("Please install starlette to use the Redis vector store. " "pip install starlette")
 try:
     import redis
 
@@ -312,7 +316,7 @@ async def aadd_texts(
         **kwargs: Any,
     ) -> List[str]:
         """Add texts data to an existing index."""
-        ids, pipeline = self._add_texts(texts, metadatas, **kwargs)
+        ids, pipeline = await run_in_threadpool(self._add_texts, texts, metadatas, **kwargs)
         await pipeline.execute()  # type: ignore
         return ids
 
@@ -448,7 +452,7 @@ async def asimilarity_search_with_score(self, query: str, k: int = 4) -> List[Tu
         Returns:
             List of Documents most similar to the query and score for each
         """
-        redis_query, params_dict = self._similarity_search_with_score(query, k=k)
+        redis_query, params_dict = await run_in_threadpool(self._similarity_search_with_score, query, k)
 
         # perform vector search
         results = await self.client.ft(self.index_name).search(redis_query, params_dict)  # type: ignore

app/viewmodels/base_models.py

@@ -118,9 +118,6 @@ class MessageFromWebsocket(BaseModel):
     translate: bool
     chat_room_id: str
 
-    class Config:
-        orm_mode = True
-
 
 class CreateChatRoom(BaseModel):  # stub
     chat_room_type: str

app/web/.last_build_id

@@ -1 +1 @@
-6fe4c07d32d7ee85313bac15705fffd4
+40f065f758a32903d20b775d3e3f717e

app/web/assets/NOTICES

@@ -6949,6 +6949,30 @@ THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 --------------------------------------------------------------------------------
+file_picker
+
+MIT License
+
+Copyright (c) 2018 Miguel Ruivo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NON INFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+--------------------------------------------------------------------------------
 flatbuffers
 
 Apache License
@@ -7225,6 +7249,7 @@ SOFTWARE.
 --------------------------------------------------------------------------------
 flutter_lints
 flutter_markdown
+flutter_plugin_android_lifecycle
 path_provider
 path_provider_android
 path_provider_foundation

app/web/flutter_service_worker.js

@@ -10,7 +10,7 @@ const RESOURCES = {
 "assets/assets/images/user_profile.png": "62ec3ddb01749f1173ca5e2972504c03",
 "assets/FontManifest.json": "dc3d03800ccca4601324923c0b1d6d57",
 "assets/fonts/MaterialIcons-Regular.otf": "e7069dfd19b331be16bed984668fe080",
-"assets/NOTICES": "5f187725cbd4ba575eb9ce64fefe3a4a",
+"assets/NOTICES": "b67701d66ab4415275394db0a602d216",
 "assets/packages/cupertino_icons/assets/CupertinoIcons.ttf": "6d342eb68f170c97609e9da345464e5e",
 "canvaskit/canvaskit.js": "97937cb4c2c2073c968525a3e08c86a3",
 "canvaskit/canvaskit.wasm": "3de12d898ec208a5f31362cc00f09b9e",
@@ -22,9 +22,9 @@ const RESOURCES = {
 "icons/Icon-512.png": "c9242a8fc92d16303b04ac108c9ac581",
 "icons/Icon-maskable-192.png": "710e5d264b437fe72bdb99ba514e2266",
 "icons/Icon-maskable-512.png": "c9242a8fc92d16303b04ac108c9ac581",
-"index.html": "97477ffe63d8e831119b544c3173d940",
-"/": "97477ffe63d8e831119b544c3173d940",
-"main.dart.js": "c5d0085759139c968f0397db84d37860",
+"index.html": "00bbbe3b446dd2ffa5ee43d4e0902e3e",
+"/": "00bbbe3b446dd2ffa5ee43d4e0902e3e",
+"main.dart.js": "0214b9f8354e4111b7edb14011b9364c",
 "manifest.json": "7e9cc88c98034abad925a9f0f78c53d4",
 "site_logo.png": "0fb92091fe8696aa5cfeafb0d1a01cbc",
 "version.json": "9094aacdae789dccd67fa32109ff1a18"

app/web/index.html

@@ -22,7 +22,7 @@
 
   <script>
     // The value below is injected by flutter build, do not touch.
-    var serviceWorkerVersion = '1091727571';
+    var serviceWorkerVersion = '1925346675';
   </script>
   <!-- This script adds the flutter initialization JS code -->
   <script src="flutter.js" defer></script>