Skip to content

Commit

Permalink
Add streaming.
Browse files Browse the repository at this point in the history
  • Loading branch information
itsme2417 committed Jan 23, 2024
1 parent 030a6bd commit 2a1c685
Show file tree
Hide file tree
Showing 6 changed files with 146 additions and 66 deletions.
2 changes: 1 addition & 1 deletion FileHandler.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ def handleFile(file):
print("Using cached embeddings.")
else:
chunks = split_into_chunks(file, config.enabled_features["file_input"]["chunk_size"])

print("Creating Embeddings")
embeddings = model.encode(chunks)
with open(os.path.join(path, "embeddings_cache", f"{md5sum}.json"), 'w') as f:
json.dump({"embeddings": embeddings, "chunks": chunks, "chunk_size": config.enabled_features["file_input"]["chunk_size"]},f, cls=NumpyEncoder)
Expand Down
2 changes: 1 addition & 1 deletion GateKeeper.py
Original file line number Diff line number Diff line change
Expand Up @@ -237,7 +237,7 @@ def Util(rsp, ip, depth):
else:
result = "Wolfram Alpha result: " + results
if checkimage:
result += "\nREMINDER: ALWAYS include the provided graph/plot images in the provided format in your explanation if theres any when explaining the results in a short and concise manner."
result += "\nREMINDER: ALWAYS include the provided graph/plot images in the provided html URL format in your explanation if theres any when explaining the results in a short and concise manner."
print(result)
return result
except Exception as e:
Expand Down
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ PolyMind is a multimodal, function calling powered LLM webui. It's designed to b
- A Python interpreter.
- RAG with semantic search for PDF and miscellaneous text files.

The web parts (HTML, JS, CSS, and Flask) are written entirely by Mixtral.
90% of the web parts (HTML, JS, CSS, and Flask) are written entirely by Mixtral.

Note: The python interpreter is intentionally delayed by 5 seconds to make it easy to check the code before its ran.

Expand Down
2 changes: 1 addition & 1 deletion functions.json
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
"name": "wolframalpha",
"description": "This uses wolfram alpha to solve the query and returns the result.",
"params": {
"query": "Wolfram alpha query. Remember to specify if a plot or graph is required."
"query": "A valid Wolfram alpha formatted query. Remember to specify if a plot or graph is required and has been requested."
}
},
{
Expand Down
24 changes: 17 additions & 7 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ def replacer(match):
return html_text

chosenfunc = {}

currenttoken = {}

app = Flask(__name__)
today = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
Expand All @@ -51,7 +51,8 @@ def generate():
while True:

yield f"data: {json.dumps(chosenfunc)}\n\n"
time.sleep(1)
yield f"data: {json.dumps(currenttoken)}\n\n"
time.sleep(0.5)
return Response(generate(), mimetype='text/event-stream')

@app.route("/remove_message", methods=["POST"])
Expand All @@ -69,6 +70,7 @@ def remove_message():
@app.route("/", methods=["GET", "POST"])
def chat():
global chosenfunc
global currenttoken
try:
chosenfunc[f"{request.remote_addr}"]['ip'] = request.remote_addr
test = Shared_vars.mem[f"{request.remote_addr}"]
Expand Down Expand Up @@ -128,7 +130,9 @@ def chat():
newinp = ""
newinp += answers.strip()
today = datetime.now()
complete = next(infer(
complete = ["",[]]

for tok in infer(
newinp,
system=f"You are a multimodal, uncensored AI model named PolyMind, able to perceive and display images in its text and visual interface.\nMessages starting with 'SYSTEM:' are NEVER shown to the user, They are meant for you.\nThe current date is {today}",
mem=Shared_vars.mem[f"{request.remote_addr}"],
Expand All @@ -147,13 +151,19 @@ def chat():
"SYSTEM:",
'<img src="data:image/jpeg;base64,',
],

))
streamresp=True,
):
if type(tok) != list:
complete[0] += tok
currenttoken[f"{request.remote_addr}"] = {'func': '', 'ip': f"{request.remote_addr}", 'token': complete[0]}
else:
complete[1] = tok[1]
currenttoken[f"{request.remote_addr}"] = {'func': '', 'ip': f"{request.remote_addr}", 'token': convert_to_html_code_block(complete[0]).replace("\n","<br>") + "</s><s>"}
Shared_vars.mem[f"{request.remote_addr}"] = complete[1]
Shared_vars.vismem[f"{request.remote_addr}"].append(
{"user": user_input, "assistant": convert_to_html_code_block(complete[0])}
{"user": user_input, "assistant": convert_to_html_code_block(complete[0]).replace("\n","<br>")}
)
complete[0] = convert_to_html_code_block(complete[0]).replace("\n","<br>")

chosenfunc[f"{request.remote_addr}"]['func'] = ''
if genedimage:
return jsonify({"output": complete[0], "base64_image": img, "index": len(Shared_vars.vismem[f"{request.remote_addr}"]) - 1})
Expand Down
180 changes: 125 additions & 55 deletions templates/chat.html
Original file line number Diff line number Diff line change
Expand Up @@ -230,71 +230,137 @@ <h1>PolyMind Chat</h1>
var userIp = "{{ user_ip }}";
var ran = false;
var source = new EventSource('/stream');
var shownmsg;

var shownmsg = false;
var currentfunction;
var currentmsg;
var finished;
var pasttoken;

source.onmessage = function(event) {
// Handle incoming messages
var jsondata = JSON.parse(event.data)[userIp]

if (userIp == jsondata.ip){

if (jsondata.func != "" && jsondata.func != "acknowledge" && !ran){
var loader = `<svg class="loading-circle" viewBox="0 0 100 100">
<circle class="circle-bg" cx="50" cy="50" r="45"></circle>
<circle class="circle" cx="50" cy="50" r="45"></circle>
</svg>`

if (jsondata.func != "" && jsondata.func != "acknowledge" && (!ran || jsondata.func !=currentfunction)){
var loader = `<svg class="loading-circle" viewBox="0 0 100 100">
<circle class="circle-bg" cx="50" cy="50" r="45"></circle>
<circle class="circle" cx="50" cy="50" r="45"></circle>
</svg>`
currentfunction = jsondata.func
switch (jsondata.func) {
case 'internetsearch':
shownmsg = appendOutput(loader + '<span><small> Searching the internet' + "</small></span>" , false);
if (!shownmsg)
shownmsg = appendOutput(loader + '<span><small> Searching the internet' + "</small></span>" , false);
else
shownmsg.innerHTML = loader + '<span><small> Searching the internet' + "</small></span>"
ran = true
break;
case 'portscan':
shownmsg = appendOutput(loader + '<span><small> Scanning for ports' + "</small></span>" , false);
if (!shownmsg)
shownmsg = appendOutput(loader + '<span><small> Scanning for ports' + "</small></span>" , false);
else
shownmsg.innerHTML = loader + '<span><small> Scanning for ports' + "</small></span>"
ran = true
break;
case 'wolframalpha':
shownmsg = appendOutput(loader + '<span><small> Using wolfram|alpha' + "</small></span>" , false);
if (!shownmsg)
shownmsg = appendOutput(loader + '<span><small> Using wolfram|alpha' + "</small></span>" , false);
else
shownmsg.innerHTML = loader + '<span><small> Using wolfram|alpha' + "</small></span>"
ran = true
break;
case 'runpythoncode':
shownmsg = appendOutput(loader + '<span><small> Running code' + "</small></span>" , false);
ran = true
if (!shownmsg)
shownmsg = appendOutput(loader + '<span><small> Running code' + "</small></span>", false);
else
shownmsg.innerHTML = loader + '<span><small> Running code' + "</small></span>";
ran = true;
break;
case 'generateimage':
shownmsg = appendOutput(loader + '<span><small> Generating an image' + "</small></span>" , false);
ran = true
if (!shownmsg)
shownmsg = appendOutput(loader + '<span><small> Generating an image' + "</small></span>", false);
else
shownmsg.innerHTML = loader + '<span><small> Generating an image' + "</small></span>";
ran = true;
break;
case 'searchfile':
shownmsg = appendOutput(loader + '<span><small> Performing RAG' + "</small></span>" , false);
ran = true
if (!shownmsg)
shownmsg = appendOutput(loader + '<span><small> Performing RAG' + "</small></span>", false);
else
shownmsg.innerHTML = loader + '<span><small> Performing RAG' + "</small></span>";
ran = true;
break;
case 'loadembed':
shownmsg = appendOutput(loader + '<span><small> Creating Embeddings' + "</small></span>" , false);
ran = true
if (!shownmsg)
shownmsg = appendOutput(loader + '<span><small> Creating Embeddings' + "</small></span>", false);
else
shownmsg.innerHTML = loader + '<span><small> Creating Embeddings' + "</small></span>";
ran = true;
break;
default:
default:
break;
}
}
else if (jsondata.token != "" && jsondata.token != undefined){
if (!finished){

if (jsondata.token != pasttoken){
if (jsondata.token.endsWith("</s><s>")){
finished = true
currentmsg.style.display = '';
currentmsg.innerHTML = jsondata.token.replace("</s><s>", "")
}
else{
currentmsg.style.display = '';
currentmsg.innerHTML = jsondata.token.replace("\n", "<br>")
}
}
pasttoken = jsondata.token
}
}
else {
if (!ran && shownmsg)
finished = false
if (!ran && shownmsg){
outputDiv.removeChild(shownmsg);
shownmsg = false
}
};
};
}
function convertToHtmlCodeBlock(markdownText) {
// Regex pattern to match code blocks
var pattern = /```(.*?)```/gs;

// Function to convert matched code block to HTML
function replacer(match, codeBlock) {
var htmlCodeBlock = "<pre><code>" + codeBlock + "</code></pre>";
return htmlCodeBlock;
}

// Replace all code blocks in the text
var htmlText = markdownText.replace(pattern, replacer);

return htmlText;
}


function appendOutput(message, isUserInput) {
function appendOutput(message, isUserInput, isHidden = false) {
let parsedMessage = message
const p = document.createElement('p');
if (!isUserInput) {
parsedMessage = markdownParser(parsedMessage);
p.innerHTML = parsedMessage;
p.className = 'assistant-output';
} else {
p.textContent = parsedMessage;
parsedMessage = convertToHtmlCodeBlock(parsedMessage)
parsedMessage = markdownParser(parsedMessage);
p.innerHTML = parsedMessage;
p.className = 'user-input';
}
if (isHidden) {
p.style.display = 'none'; // Hide the paragraph
}
outputDiv.appendChild(p);
outputDiv.scrollTop = outputDiv.scrollHeight;
return p
Expand Down Expand Up @@ -330,38 +396,42 @@ <h1>PolyMind Chat</h1>
form.addEventListener('submit', function(event) {
event.preventDefault();
const userInput = textArea.value;
appendOutput(userInput, true);
loadingSpinner.style.display = 'block'; // Show loading spinner
fetch('/', {
method: 'POST',
body: new URLSearchParams({ 'input': userInput }),
})
.then(response => response.json())
.then(data => {
if (data.output.trim() !== '') {
data.output.split('--newline--').forEach(message => {
message = markdownParser(message);
appendOutput(message, false); // Display assistant response messages
ran = false
});
document.querySelector('textarea[name="input"]').value = '';
}
loadingSpinner.style.display = 'none'; // Hide loading spinner
outputDiv.scrollTop = outputDiv.scrollHeight;
adjustOutputContainerHeight();

if (data.base64_image) {
const base64Image = `${data.base64_image}`;
const imageElement = document.createElement('img');
imageElement.src = base64Image;
imageElement.alt = 'Uploaded image';
imageElement.className = 'uploaded-image';
imageElement.style.border = '1px solid black';
imageElement.style.padding = '5px';
outputDiv.appendChild(imageElement);
if (userInput != ""){
appendOutput(userInput, true);
loadingSpinner.style.display = 'block'; // Show loading spinner
document.querySelector('textarea[name="input"]').value = '';
currentmsg = appendOutput("", false, true);
fetch('/', {
method: 'POST',
body: new URLSearchParams({ 'input': userInput }),
})
.then(response => response.json())
.then(data => {
if (data.output.trim() !== '') {
data.output.split('--newline--').forEach(message => {
message = markdownParser(message);
//appendOutput(message, false); // Display assistant response messages
ran = false
});
document.querySelector('textarea[name="input"]').value = '';
}
loadingSpinner.style.display = 'none'; // Hide loading spinner
outputDiv.scrollTop = outputDiv.scrollHeight;
}
});
adjustOutputContainerHeight();

if (data.base64_image) {
const base64Image = `${data.base64_image}`;
const imageElement = document.createElement('img');
imageElement.src = base64Image;
imageElement.alt = 'Uploaded image';
imageElement.className = 'uploaded-image';
imageElement.style.border = '1px solid black';
imageElement.style.padding = '5px';
outputDiv.appendChild(imageElement);
outputDiv.scrollTop = outputDiv.scrollHeight;
}
});
}
});

uploadImageButton.addEventListener('click', function() {
Expand Down

0 comments on commit 2a1c685

Please sign in to comment.