Compare commits
3 Commits
dea4b79014
...
1abb1b5106
| Author | SHA1 | Date | |
|---|---|---|---|
| 1abb1b5106 | |||
| 49d924be20 | |||
| 0be88678df |
43
main.py
43
main.py
@@ -1,4 +1,5 @@
|
||||
from flask import Flask, jsonify, request, render_template
|
||||
from flask import Flask, jsonify, request, render_template, Response
|
||||
from flask_cors import CORS, cross_origin
|
||||
from pymongo import MongoClient
|
||||
from bson.objectid import ObjectId
|
||||
from bson.json_util import dumps, loads
|
||||
@@ -9,6 +10,7 @@ import string
|
||||
import json
|
||||
import requests
|
||||
from urllib.parse import parse_qs
|
||||
from ollama import chat
|
||||
|
||||
with open("config/settings.json", "r") as f:
|
||||
settings = json.load(f)
|
||||
@@ -44,6 +46,7 @@ except Exception as e:
|
||||
print("Error connecting to MongoDB:", e)
|
||||
|
||||
app = Flask(__name__)
|
||||
CORS(app)
|
||||
|
||||
def checkUserPermission(token, permission):
|
||||
# Find the correct user token in user db
|
||||
@@ -82,6 +85,38 @@ def checkChatPermission(token, chatId, permission):
|
||||
return False, "Incorrect permissions"
|
||||
else:
|
||||
return False, "Invalid Token"
|
||||
|
||||
# Message Generation Endpoint
|
||||
# Generate a new message on a specific chat
|
||||
# Arguments: token (required), chatId (required), message
|
||||
@app.route('/api/chat/<_id>/generate', methods =['POST'])
|
||||
@cross_origin()
|
||||
def generateMessage(_id):
|
||||
# Get user auth token
|
||||
token = request.headers['token']
|
||||
a, userId = checkChatPermission(token, _id, "view")
|
||||
if (a == True):
|
||||
returnedChat = chatCollection.find_one({'_id': ObjectId(_id)})
|
||||
message = request.json['message']
|
||||
messages = returnedChat['messages']
|
||||
messages.append({'role':'user', 'content':message})
|
||||
stream = chat(
|
||||
model=returnedChat['model'],
|
||||
messages=messages,
|
||||
stream=True
|
||||
)
|
||||
def generateStream():
|
||||
response = ""
|
||||
for chunk in stream:
|
||||
print(chunk['message']['content'], end='', flush=True)
|
||||
content = chunk['message']['content']
|
||||
response += content
|
||||
json_data = {"response": response}
|
||||
yield f"data: {json.dumps(json_data)}\n\n"
|
||||
|
||||
return Response(generateStream(), mimetype='text/event-stream')
|
||||
else:
|
||||
return userId
|
||||
|
||||
# Chat List Endpoint:
|
||||
# Get all the chats associated with a user
|
||||
@@ -89,7 +124,7 @@ def checkChatPermission(token, chatId, permission):
|
||||
@app.route('/api/user/chats', methods = ['GET'])
|
||||
def getUserChats():
|
||||
# Get user auth token
|
||||
token = request.json['token']
|
||||
token = request.headers['token']
|
||||
a, userId = checkUserPermission(token, True)
|
||||
if (a == True):
|
||||
returnedChats = list(chatCollection.find({'permissions.' + userId : "view"}))
|
||||
@@ -108,7 +143,7 @@ def getUserChats():
|
||||
@app.route('/api/chat/<_id>/details/<details>', methods = ['GET', 'POST'])
|
||||
def getChatHistory(_id, details):
|
||||
# Get user auth token
|
||||
token = request.json['token']
|
||||
token = request.headers['token']
|
||||
a, userId = checkChatPermission(token, _id, True)
|
||||
if (a == True):
|
||||
# If the user is trying to GET data
|
||||
@@ -154,7 +189,7 @@ def getChatHistory(_id, details):
|
||||
@app.route('/api/chat/create', methods = ['POST'])
|
||||
def createChat():
|
||||
# Get user auth token
|
||||
token = request.json['token']
|
||||
token = request.headers['token']
|
||||
a, userId = checkUserPermission(token, "createChat")
|
||||
if (a == True):
|
||||
name = request.json['name']
|
||||
|
||||
72
templates/chattest.html
Normal file
72
templates/chattest.html
Normal file
@@ -0,0 +1,72 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Streaming Response</title>
|
||||
|
||||
</head>
|
||||
<body>
|
||||
|
||||
<div class="container">
|
||||
<h1>AiThingy</h1>
|
||||
<p>Enter a message, chatId, and auth token below:</p>
|
||||
<input placeholder="Message" id="message"></input>
|
||||
<input placeholder="ChatId" id="chatid"></input>
|
||||
<input placeholder="Token" id="token"></input>
|
||||
<button onclick="startStreaming()">Send Message</button>
|
||||
|
||||
<br>
|
||||
|
||||
<textarea id="response-output" readonly placeholder="Waiting for streaming data..."></textarea>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
async function startStreaming() {
|
||||
const outputArea = document.getElementById('response-output');
|
||||
outputArea.value = 'Connecting to endpoint...\n';
|
||||
const chatId = document.getElementById("chatid").value;
|
||||
const endpointUrl = 'http://127.0.0.1:5000/api/chat/' + chatId + '/generate';
|
||||
const postData = {
|
||||
"message": document.getElementById("message").value
|
||||
};
|
||||
|
||||
try {
|
||||
const response = await fetch(endpointUrl, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'token': document.getElementById("token").value
|
||||
},
|
||||
body: JSON.stringify(postData)
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP error! status: ${response.status}`);
|
||||
}
|
||||
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder('utf-8');
|
||||
let receivedChunks = '';
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
|
||||
if (done) {
|
||||
break;
|
||||
}
|
||||
|
||||
const chunk = decoder.decode(value, { stream: true });
|
||||
const jsonString = chunk.substring(6);
|
||||
const data = JSON.parse(jsonString);
|
||||
outputArea.value = data["response"];
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('An error has occurred:', error);
|
||||
outputArea.value += `\n\n--- An error has occurred: ${error.message} ---`;
|
||||
}
|
||||
}
|
||||
</script>
|
||||
|
||||
</body>
|
||||
</html>
|
||||
Reference in New Issue
Block a user