From bc24c9334c56d33f2a9901ab0e77d3a1e8a9a5fd Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Sun, 24 Dec 2023 17:08:43 +0800
Subject: [PATCH] prevent prompt leakage during usage of check endpoint when
genkey is provided in multiuser mode
---
koboldcpp.py | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/koboldcpp.py b/koboldcpp.py
index 9615010a4..d65a22285 100755
--- a/koboldcpp.py
+++ b/koboldcpp.py
@@ -716,7 +716,7 @@ Enter Prompt:
self.wfile.write(finalhtml)
def do_GET(self):
- global maxctx, maxhordelen, friendlymodelname, KcppVersion, totalgens, preloaded_story, exitcounter
+ global maxctx, maxhordelen, friendlymodelname, KcppVersion, totalgens, preloaded_story, exitcounter, currentusergenkey
self.path = self.path.rstrip('/')
response_body = None
content_type = 'application/json'
@@ -766,7 +766,7 @@ Enter Prompt:
elif self.path.endswith('/api/extra/generate/check'):
pendtxtStr = ""
- if requestsinqueue==0 and totalgens>0:
+ if requestsinqueue==0 and totalgens>0 and (not args.multiuser or currentusergenkey==""):
pendtxt = handle.get_pending_output()
pendtxtStr = ctypes.string_at(pendtxt).decode("UTF-8","ignore")
response_body = (json.dumps({"results": [{"text": pendtxtStr}]}).encode())
@@ -861,7 +861,7 @@ Enter Prompt:
multiuserkey = ""
if totalgens>0:
- if (multiuserkey=="" and requestsinqueue==0) or (multiuserkey!="" and multiuserkey==currentusergenkey):
+ if (multiuserkey=="" and multiuserkey==currentusergenkey) or (multiuserkey=="" and requestsinqueue==0 and not args.multiuser) or (multiuserkey!="" and multiuserkey==currentusergenkey): #avoid leaking prompts in multiuser
pendtxt = handle.get_pending_output()
pendtxtStr = ctypes.string_at(pendtxt).decode("UTF-8","ignore")
response_body = (json.dumps({"results": [{"text": pendtxtStr}]}).encode())