From 49d924be207b3e4716ecde85b00741cee75a5c0d Mon Sep 17 00:00:00 2001
From: Hugo H <hugo@haaxman.co.uk>
Date: Sat, 13 Sep 2025 19:53:41 +0100
Subject: [PATCH] Added basic AI generation endpoint

---
 main.py | 37 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index 0993248..e624294 100644
--- a/main.py
+++ b/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