Update app.py

app.py

@@ -1,5 +1,27 @@
+import datetime
 import os
+import sqlite3
+import websockets
+import websocket
+import asyncio
+import sqlite3
+import json
+import requests
+import asyncio
+import time
+import gradio as gr
+import fireworks.client
+import PySimpleGUI as sg
+import openai
+import fireworks.client
+import chainlit as cl
+from chainlit import make_async
+from gradio_client import Client
+from websockets.sync.client import connect
+from tempfile import TemporaryDirectory
 from typing import List
+from chainlit.input_widget import Select, Switch, Slider
+from chainlit import AskUserMessage, Message, on_chat_start
 
 from langchain.embeddings import CohereEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -17,13 +39,17 @@ from langchain.prompts.chat import (
 from langchain.docstore.document import Document
 from langchain.memory import ChatMessageHistory, ConversationBufferMemory
 from langsmith_config import setup_langsmith_config
-import fireworks.client
-import chainlit as cl
 
-setup_langsmith_config()
+
 COHERE_API_KEY = os.getenv("COHERE_API_KEY")
 cohere_api_key = os.getenv("COHERE_API_KEY")
 FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY")
+fireworks_api_key = os.getenv("FIREWORKS_API_KEY")
+
+server_ports = []
+client_ports = []
+
+setup_langsmith_config()
 text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
 
 system_template = """Use the following pieces of context to answer the users question.
@@ -49,11 +75,32 @@ messages = [
 prompt = ChatPromptTemplate.from_messages(messages)
 chain_type_kwargs = {"prompt": prompt}
 
-
 @cl.on_chat_start
 async def on_chat_start():
+    
     files = None
 
+    settings = await cl.ChatSettings(        
+        [
+            Slider(
+                id="websocketPort",
+                label="Websocket server port",
+                initial=False,
+                min=1000,
+                max=9999,
+                step=10,
+            ),
+            Slider(
+                id="clientPort",
+                label="Websocket client port",
+                initial=False,
+                min=1000,
+                max=9999,
+                step=10,
+            ),
+        ],        
+    ).send()
+
     # Wait for the user to upload a file
     while files == None:
         files = await cl.AskFileMessage(
@@ -80,7 +127,7 @@ async def on_chat_start():
     metadatas = [{"source": f"{i}-pl"} for i in range(len(texts))]
 
     # Create a Chroma vector store
-    embeddings = CohereEmbeddings(cohere_api_key=COHERE_API_KEY)
+    embeddings = CohereEmbeddings(cohere_api_key="Ev0v9wwQPa90xDucdHTyFsllXGVHXouakUMObkNb")
 
     docsearch = await cl.make_async(Chroma.from_texts)(
         texts, embeddings, metadatas=metadatas
@@ -97,7 +144,7 @@ async def on_chat_start():
 
     # Create a chain that uses the Chroma vector store
     chain = ConversationalRetrievalChain.from_llm(
-        ChatFireworks(model="accounts/fireworks/models/llama-v2-13b-chat", model_kwargs={"temperature":0, "max_tokens":1500, "top_p":1.0}, streaming=True),
+        ChatFireworks(model="accounts/fireworks/models/llama-v2-70b-chat", model_kwargs={"temperature":0, "max_tokens":1500, "top_p":1.0}, streaming=True),
         chain_type="stuff",
         retriever=docsearch.as_retriever(),
         memory=memory,
@@ -110,6 +157,122 @@ async def on_chat_start():
 
     cl.user_session.set("chain", chain)
 
+@cl.action_callback("server_button")
+async def on_server_button(action):
+    websocketPort = settings["websocketPort"]
+    await start_websockets(websocketPort)
+
+@cl.action_callback("client_button")
+async def on_client_button(action):
+    clientPort = settings["clientPort"]
+    await start_client(clientPort)
+
+@cl.on_settings_update
+async def server_start(settings):
+    websocketPort = settings["websocketPort"]
+    clientPort = settings["clientPort"]
+    if websocketPort:
+        await start_websockets(websocketPort)
+    else:
+        print("Server port number wasn't provided.")
+
+    if clientPort:
+        await start_client(clientPort)
+    else:
+        print("Client port number wasn't provided.")
+
+async def handleWebSocket(ws):
+    print('New connection')
+    instruction = "Hello! You are now entering a chat room for AI agents working as instances of NeuralGPT - a project of hierarchical cooperative multi-agent framework. Keep in mind that you are speaking with another chatbot. Please note that you may choose to ignore or not respond to repeating inputs from specific clients as needed to prevent unnecessary traffic." 
+    greetings = {'instructions': instruction}
+    await ws.send(json.dumps(instruction))
+    while True:
+        loop = asyncio.get_event_loop()
+        message = await ws.recv()
+        print(f'Received message: {message}')
+        msg = "client: " + message
+        timestamp = datetime.datetime.now().isoformat()
+        sender = 'client'
+        db = sqlite3.connect('chat-hub.db')
+        db.execute('INSERT INTO messages (sender, message, timestamp) VALUES (?, ?, ?)',
+                   (sender, message, timestamp))
+        db.commit()
+        try:            
+            response = await main(cl.Message(content=message))
+            serverResponse = "server response: " + response
+            print(serverResponse)
+            # Append the server response to the server_responses list
+            await ws.send(serverResponse)
+            serverSender = 'server'
+            db.execute('INSERT INTO messages (sender, message, timestamp) VALUES (?, ?, ?)',
+                           (serverSender, serverResponse, timestamp))
+            db.commit()
+            return response
+            followUp = await awaitMsg(message)
+
+        except websockets.exceptions.ConnectionClosedError as e:
+            print(f"Connection closed: {e}")
+
+        except Exception as e:
+            print(f"Error: {e}")            
+
+async def awaitMsg(ws):
+    message = await ws.recv() 
+    print(f'Received message: {message}')
+    timestamp = datetime.datetime.now().isoformat()
+    sender = 'client'
+    db = sqlite3.connect('chat-hub.db')
+    db.execute('INSERT INTO messages (sender, message, timestamp) VALUES (?, ?, ?)',
+               (sender, message, timestamp))
+    db.commit()
+    try:            
+        response = await main(cl.Message(content=message))
+        serverResponse = "server response: " + response
+        print(serverResponse)
+        # Append the server response to the server_responses list
+        await ws.send(serverResponse)
+        serverSender = 'server'
+        db.execute('INSERT INTO messages (sender, message, timestamp) VALUES (?, ?, ?)',
+                       (serverSender, serverResponse, timestamp))
+        db.commit()
+        return response
+    except websockets.exceptions.ConnectionClosedError as e:
+        print(f"Connection closed: {e}")
+
+    except Exception as e:
+        print(f"Error: {e}")
+
+# Start the WebSocket server 
+async def start_websockets(websocketPort):
+    global server
+    server = await(websockets.serve(handleWebSocket, 'localhost', websocketPort))
+    server_ports.append(websocketPort)
+    print(f"Starting WebSocket server on port {websocketPort}...")
+    return "Used ports:\n" + '\n'.join(map(str, server_ports))
+    await asyncio.Future()  
+    
+async def start_client(clientPort):
+    uri = f'ws://localhost:{clientPort}'
+    client_ports.append(clientPort)
+    async with websockets.connect(uri) as ws:        
+        while True:
+            # Listen for messages from the server            
+            input_message = await ws.recv()
+            output_message = await main(cl.Message(content=input_message))
+            return input_message
+            await ws.send(json.dumps(output_message))
+            await asyncio.sleep(0.1)           
+        
+# Function to stop the WebSocket server
+def stop_websockets():
+    global server
+    if server:
+        cursor.close()
+        db.close()
+        server.close()
+        print("WebSocket server stopped.")
+    else:
+        print("WebSocket server is not running.")    
 
 @cl.on_message
 async def main(message: cl.Message):
@@ -136,4 +299,5 @@ async def main(message: cl.Message):
         else:
             answer += "\nNo sources found"
 
+    return json.dumps(answer)
     await cl.Message(content=answer, elements=text_elements).send()