From 3fea11675d27d081dbcffa79bc95055ae832c683 Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Mon, 30 Dec 2024 17:30:41 +0800
Subject: [PATCH] websearch integrated into lite, changed to POST

---
 kcpp_docs.embd |  49 ++++++++++++++
 klite.embd     | 176 ++++++++++++++++++++++++++++++++++++++++++++-----
 koboldcpp.py   |  49 ++++++--------
 3 files changed, 230 insertions(+), 44 deletions(-)

diff --git a/kcpp_docs.embd b/kcpp_docs.embd
index 56107ca04..7a1ff1f8f 100644
--- a/kcpp_docs.embd
+++ b/kcpp_docs.embd
@@ -1380,6 +1380,55 @@
                       ]
                    }
                 },
+                "/api/extra/websearch": {
+                   "post": {
+                      "description": "Searches the web using DuckDuckGo and returns the top 3 results.",
+                      "requestBody": {
+                         "content": {
+                            "application/json": {
+                               "example": {
+                                  "q": "What is KoboldCpp"
+                               },
+                               "schema": {
+                                  "properties": {
+                                     "q": {
+                                        "type": "string",
+                                        "description": "The search query string"
+                                     }
+                                  },
+                                  "type": "object"
+                               }
+                            }
+                         },
+                         "required": true
+                      },
+                      "responses": {
+                         "200": {
+                            "content": {
+                               "application/json": {
+                                  "example": [
+                                       {
+                                          "title": "KoboldCpp Wiki",
+                                          "url": "https://github.com/LostRuins/koboldcpp/wiki",
+                                          "desc": "KoboldCpp is a program to run LLMs",
+                                          "content": "KoboldCpp is a program to run LLMs using GGUF files"
+                                       }
+                                    ],
+                                  "schema": {
+                                    "properties": {},
+                                    "type": "object"
+                                 }
+                               }
+                            },
+                            "description": "Successful request"
+                         }
+                      },
+                      "summary": "Searches the web using DuckDuckGo and returns the top 3 results.",
+                      "tags": [
+                         "api/extra"
+                      ]
+                   }
+                },
                 "/.well-known/serviceinfo": {
                    "get": {
                       "responses": {
diff --git a/klite.embd b/klite.embd
index 082667bb5..cdcab17f1 100644
--- a/klite.embd
+++ b/klite.embd
@@ -12,7 +12,7 @@ Current version indicated by LITEVER below.
 -->
 
 <script>
-	const LITEVER = 197;
+	const LITEVER = 198;
 	const urlParams = new URLSearchParams(window.location.search);
 	var localflag = true;
 	const STORAGE_PREFIX = (localflag?"e_":"")+"kaihordewebui_";
@@ -3203,6 +3203,22 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 		}
 	}
 
+	function removeAndReplacePlaceholders(text) {
+		// Replace {{user}} and other placeholders
+		text = replace_placeholders(text);
+
+		// Remove any instruct tags as needed to ensure a more accurate search
+		text = replaceAll(text, get_instruct_starttag(false), "");
+		text = replaceAll(text, get_instruct_endtag(false), "");
+		text = replaceAll(text, get_instruct_systag(false), "");
+		text = replaceAll(text, get_instruct_starttag(false).trim(), "");
+		text = replaceAll(text, get_instruct_endtag(false).trim(), "");
+		text = replaceAll(text, get_instruct_systag(false).trim(), "");
+		text = text.replace(/\{\{\[INPUT\]\}\}/g, "").replace(/\{\{\[OUTPUT\]\}\}/g, "");
+
+		return text;
+	}
+
 	function readTavernPngFromBlob(blob, onDone)
 	{
 		var fileReader = new FileReader();
@@ -4177,6 +4193,7 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 	const koboldcpp_transcribe_endpoint = "/api/extra/transcribe";
 	const koboldcpp_tokenize_endpoint = "/api/extra/tokencount";
 	const koboldcpp_perf_endpoint = "/api/extra/perf";
+	const koboldcpp_websearch_endpoint = "/api/extra/websearch"
 
 	const oai_models_endpoint = "/models";
 	const oai_submit_endpoint = "/completions";
@@ -4278,6 +4295,7 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 	var documentdb_searchrange = 300;
 	var documentdb_chunksize = 800;
 	var documentdb_data = "";
+	var websearch_enabled = false;
 	var generateimagesinterval = 750; //if generated images is enabled, it will trigger after every 700 new characters in context.
 	var nextgeneratedimagemilestone = generateimagesinterval; //used to keep track of when to generate the next image
 	var image_db = {}; //stores a dictionary of pending images
@@ -4319,6 +4337,9 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 	var koboldcpp_version_obj = {};
 	var koboldcpp_has_vision = false;
 	var koboldcpp_has_multiplayer = false;
+	var koboldcpp_has_websearch = false;
+	var lastSearchQuery = "";
+	var lastSearchResults = [];
 	var multiplayer_active = false;
 	var multiplayer_pinged = false;
 	var multiplayer_override_name = "";
@@ -5927,6 +5948,10 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 	{
 		return (custom_kobold_endpoint!="" && koboldcpp_version && koboldcpp_version!="" && compare_version_str(koboldcpp_version, "1.78") >= 0 && koboldcpp_has_multiplayer);
 	}
+	function is_using_kcpp_with_websearch()
+	{
+		return (custom_kobold_endpoint!="" && koboldcpp_version && koboldcpp_version!="" && compare_version_str(koboldcpp_version, "1.80") >= 0 && koboldcpp_has_websearch);
+	}
 
 
 	//0 is none, 1 is poll-streaming, 2 is sse-streaming
@@ -6416,6 +6441,7 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 		new_save_storyobj.documentdb_searchrange = documentdb_searchrange;
 		new_save_storyobj.documentdb_chunksize = documentdb_chunksize;
 		new_save_storyobj.documentdb_data = documentdb_data;
+		new_save_storyobj.websearch_enabled = websearch_enabled;
 
 		if (export_settings) {
 			new_save_storyobj.savedsettings = JSON.parse(JSON.stringify(localsettings));
@@ -6723,6 +6749,10 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 				{
 					documentdb_data = storyobj.documentdb_data;
 				}
+				if(storyobj.websearch_enabled)
+				{
+					websearch_enabled = storyobj.websearch_enabled;
+				}
 			} else {
 				//v2 load
 				if(storyobj.prompt != "")
@@ -9492,6 +9522,7 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 									koboldcpp_has_vision = (data.vision?true:false);
 									koboldcpp_has_whisper = (data.transcribe?true:false);
 									koboldcpp_has_multiplayer = (data.multiplayer?true:false);
+									koboldcpp_has_websearch = (data.websearch?true:false);
 									let has_password = (data.protected?true:false);
 									let has_txt2img = (data.txt2img?true:false);
 									let no_txt_model = (mdlname=="inactive");
@@ -11381,6 +11412,7 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 		documentdb_searchrange = cleannum(documentdb_searchrange,0,1024);
 		documentdb_chunksize = parseInt(documentdb_chunksize);
 		documentdb_chunksize = cleannum(documentdb_chunksize,32,2048);
+		websearch_enabled = document.getElementById("websearch_enabled").checked?true:false;
 	}
 
 	function set_personal_notes()
@@ -11638,6 +11670,7 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 			documentdb_searchrange = 300;
 			documentdb_chunksize = 800;
 			documentdb_data = "";
+			websearch_enabled = false;
 		}
 		warn_on_quit = false;
 		show_corpo_leftpanel(false);
@@ -12674,10 +12707,19 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 		return Math.max(1, Math.floor(((maxctxlen - maxgenamt)) * chars_per_token) - 12);
 	}
 
-	function submit_generation()
+	function submit_generation() //wrap websearch into this
+	{
+		let senttext = document.getElementById("input_text").value;
+		document.getElementById("input_text").value = "";
+		PerformWebsearch(senttext,(res)=>{
+			submit_generation_post_search(senttext);
+		});
+	}
+
+	function submit_generation_post_search(senttext)
 	{
 		warn_on_quit = true;
-		let newgen = document.getElementById("input_text").value;
+		let newgen = senttext;
 		let img_gen_trigger_prompt = "";
 		let doNotGenerate = false;
 		retry_in_progress = false;
@@ -12813,7 +12855,6 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 			retry_prev_text = [];
 			retry_preserve_last = true; //initially set to true
 			redo_prev_text = [];
-			document.getElementById("input_text").value = "";
 			pending_response_id = "-1";
 
 			let mtl = document.getElementById("maintxtloader");
@@ -13158,14 +13199,25 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 				}
 			}
 
-			// TextDocuDB Long term memory minisearch
+			//use newest websearch results
+			if (websearch_enabled && is_using_kcpp_with_websearch() && lastSearchResults && lastSearchResults.length>0)
+			{
+				for(let i=0;i<lastSearchResults.length;++i)
+				{
+					let sresult = lastSearchResults[i];
+					let snippet = `\n[Search Snippet: ${sresult.title}\nSource: ${sresult.url}\nExcerpt: ${((sresult.content && sresult.content!="")?sresult.content:sresult.desc)}]`;
+					wistr = snippet + wistr;
+				}
+			}
+
+			// TextDB Long term memory minisearch
 			if (documentdb_enabled)
 			{
 				// Finds the relevant memory fragments, formats them in a similar way to an authors note and inserts them before WI
 				const ltmSearchQuery = !!rawNewText ? rawNewText : gametext_arr.length > 0 ? gametext_arr.slice(-1)[0] : undefined;
 				if(ltmSearchQuery)
 				{
-					let gameText = concat_gametext().replace(/\xA0/g, ' ').trim();
+					let gameText = concat_gametext(true).replace(/\xA0/g, ' ').trim();
 					let maxAllowedCharacters = getMaxAllowedCharacters(gameText, maxctxlen, maxgenamt);
 					if (documentdb_data!="" || gameText.length > maxAllowedCharacters)
 					{
@@ -17732,10 +17784,12 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 		memory_tab = newtab;
 		document.getElementById("memory_tab").classList.remove("active");
 		document.getElementById("wi_tab").classList.remove("active");
+		document.getElementById("websearch_tab").classList.remove("active");
 		document.getElementById("token_tab").classList.remove("active");
 		document.getElementById("documentdb_tab").classList.remove("active");
 		document.getElementById("memory_tab_container").classList.add("hidden");
 		document.getElementById("wi_tab_container").classList.add("hidden");
+		document.getElementById("websearch_tab_container").classList.add("hidden");
 		document.getElementById("token_tab_container").classList.add("hidden");
 		document.getElementById("documentdb_tab_container").classList.add("hidden");
 
@@ -17753,6 +17807,10 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 				document.getElementById("documentdb_tab_container").classList.remove("hidden");
 				break;
 			case 3:
+				document.getElementById("websearch_tab").classList.add("active");
+				document.getElementById("websearch_tab_container").classList.remove("hidden");
+				break;
+			case 4:
 				document.getElementById("token_tab").classList.add("active");
 				document.getElementById("token_tab_container").classList.remove("hidden");
 				break;
@@ -17795,6 +17853,15 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 		document.getElementById("documentdb_searchrange").value = documentdb_searchrange;
 		document.getElementById("documentdb_chunksize").value = documentdb_chunksize;
 		document.getElementById("documentdb_data").value = documentdb_data;
+		document.getElementById("websearch_enabled").checked = websearch_enabled;
+		if(is_using_kcpp_with_websearch())
+		{
+			document.getElementById("websearchunsupporteddiv").classList.add("hidden");
+		}
+		else
+		{
+			document.getElementById("websearchunsupporteddiv").classList.remove("hidden");
+		}
 
 		populate_placeholder_tags();
 		populate_regex_replacers();
@@ -18966,7 +19033,68 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 		document.getElementById('aesthetic_text_preview').innerHTML = render_aesthetic_ui(preview,true);
 	}
 
-	//LTM TextDocuDB Memsnipper searching
+	function PerformWebsearch(webSearchQuery, onDone)
+	{
+		//websearch
+		if (websearch_enabled && is_using_kcpp_with_websearch())
+		{
+			webSearchQuery = webSearchQuery.trim();
+			if(webSearchQuery=="")
+			{
+				webSearchQuery = (gametext_arr.length > 0 ? gametext_arr.slice(-1)[0] : "");
+			}
+			webSearchQuery = removeAndReplacePlaceholders(webSearchQuery);
+			webSearchQuery = webSearchQuery.trim();
+			webSearchQuery = webSearchQuery.replace(/(?:\r\n|\r|\n)/g, '. ');
+			if(webSearchQuery==lastSearchQuery || webSearchQuery=="")
+			{
+				onDone(); //use cached results
+			}
+			else
+			{
+				if(pending_response_id=="")
+				{
+					pending_response_id = "-1";
+					render_gametext(false);
+				}
+				let murl = `${custom_kobold_endpoint}${koboldcpp_websearch_endpoint}`;
+				murl = apply_proxy_url(murl);
+				fetch(murl, {
+					method: 'POST',
+					headers: get_kobold_header(),
+					body: JSON.stringify({q: webSearchQuery}),
+				})
+				.then(x => x.json())
+				.then(values => {
+					lastSearchQuery = webSearchQuery;
+					lastSearchResults = values;
+					if(pending_response_id=="-1")
+					{
+						pending_response_id = "";
+					}
+					onDone();
+				})
+				.catch(error => {
+					console.log("WebSearch Error: " + error);
+					lastSearchResults = [];
+					lastSearchQuery = "";
+					if(pending_response_id=="-1")
+					{
+						pending_response_id = "";
+					}
+					onDone();
+				});
+			}
+		}
+		else
+		{
+			lastSearchResults = [];
+			lastSearchQuery = "";
+			onDone();
+		}
+	}
+
+	//LTM TextDB Memsnipper searching
 	//searches dbText for searchStr and recentTextStr, returns up to 3 results
 	function DatabaseMinisearch(dbText, searchStr, recentTextStr)
 	{	//predefined minisearch constants
@@ -19048,6 +19176,12 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 		let paragraphs = [];
 		let allText = cleanupSpecialTags(dbText);
 		allText = replaceAll(allText,recentTextStr,"");
+
+		// Ensure placeholders are replaced to allow searching for user / character
+		allText = removeAndReplacePlaceholders(allText)
+		searchStr = removeAndReplacePlaceholders(searchStr)
+		recentTextStr = removeAndReplacePlaceholders(recentTextStr)
+
 		let i = 0, startLoc = 0;
 		while (startLoc < allText.length && i < Number.MAX_SAFE_INTEGER) {
 			let actualChunkStart = Math.max(0, startLoc - chunkSizeOverlap);
@@ -20699,8 +20833,9 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 			<div><ul class="nav nav-tabs settingsnav">
 				<li id="memory_tab" class="active"><a class="" href="#" onclick="display_memory_tab(0)">Memory</a></li>
 				<li id="wi_tab"><a class="" href="#" onclick="display_memory_tab(1)">World Info</a></li>
-				<li id="documentdb_tab"><a class="" href="#" onclick="display_memory_tab(2)">TextDocuDB</a></li>
-				<li id="token_tab"><a class="" href="#" onclick="display_memory_tab(3)">Tokens</a></li>
+				<li id="documentdb_tab"><a class="" href="#" onclick="display_memory_tab(2)">TextDB</a></li>
+				<li id="websearch_tab"><a class="" href="#" onclick="display_memory_tab(3)">WebSearch</a></li>
+				<li id="token_tab"><a class="" href="#" onclick="display_memory_tab(4)">Tokens</a></li>
 			  </ul></div>
 
 			<div class="memtabcontainer" id="memory_tab_container">
@@ -20785,10 +20920,10 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 			</div>
 
 			<div class="memtabcontainer" id="documentdb_tab_container">
-				<div class="settinglabel" style="padding: 4px;">Automatically search and include relevant snippets from uploaded documents or history.</div>
+				<div class="settinglabel" style="padding: 4px;">Automatically search and include relevant snippets from a text document or history.</div>
 				<div class="settinglabel" style="padding: 4px;">
-					<div class="justifyleft settingsmall" title="Enable TextDocuDB">Enable TextDocuDB </div>
-					<input title="Enable TextDocuDB" type="checkbox" id="documentdb_enabled" style="margin:0px 0 0;">
+					<div class="justifyleft settingsmall" title="Enable TextDB">Enable TextDB </div>
+					<input title="Enable TextDB" type="checkbox" id="documentdb_enabled" style="margin:0px 0 0;">
 				</div>
 
 				<div class="settinglabel" style="padding: 4px;">
@@ -20821,11 +20956,20 @@ pre code,td,th{padding:0}pre code,table{background-color:transparent}.table,inpu
 						min="0" step="1" pattern="\d+" placeholder="" value="" id="documentdb_chunksize">
 				</div>
 				<div class="settinglabel">
-					<div class="justifyleft"><br>TextDocuDB Storage<span class="helpicon">?<span
-						class="helptext">Paste as much raw text data here as you like. E.g. background information, reference documents, etc. This text will populate the database that will be chunked and searched by TextDocuDB.</span></span></div>
+					<div class="justifyleft"><br>TextDB Storage<span class="helpicon">?<span
+						class="helptext">Paste as much raw text data here as you like. E.g. background information, reference documents, etc. This text will populate the database that will be chunked and searched by TextDB.</span></span></div>
+				</div>
+				<textarea title="Edit TextDB" class="form-control" id="documentdb_data" style="height: 120px;"
+					placeholder="Paste as much text data here as you like. This text will populate the database that will be searched by TextDB."></textarea>
+			</div>
+
+			<div class="memtabcontainer" id="websearch_tab_container">
+				<div class="settinglabel" style="padding: 4px;">Search the Web for relavant information when using instruct mode<br>(Requires WebSearch enabled KoboldCpp)</div>
+				<div id="websearchunsupporteddiv" class="color_red hidden" style="font-weight:bold;padding:3px;font-size:12px">WebSearch Not Supported</div>
+				<div class="settinglabel" style="padding: 4px;">
+					<div class="justifyleft settingsmall" title="Enable WebSearch">Enable WebSearch </div>
+					<input title="Enable WebSearch" type="checkbox" id="websearch_enabled" style="margin:0px 0 0;">
 				</div>
-				<textarea title="Edit TextDocuDB" class="form-control" id="documentdb_data" style="height: 120px;"
-					placeholder="Paste as much text data here as you like. This text will populate the database that will be searched by TextDocuDB."></textarea>
 			</div>
 
 			<div class="memtabcontainer" id="token_tab_container">
diff --git a/koboldcpp.py b/koboldcpp.py
index 08d3c3d66..41aed5a37 100644
--- a/koboldcpp.py
+++ b/koboldcpp.py
@@ -1281,6 +1281,8 @@ def websearch(query):
     from concurrent.futures import ThreadPoolExecutor
     num_results = 3
     searchresults = []
+    if args.debugmode != -1 and not args.quiet:
+        print("Performing websearch...")
 
     def fetch_searched_webpage(url):
         if args.debugmode:
@@ -1312,7 +1314,7 @@ def websearch(query):
 
     def normalize_page_text(text):
         text = re.sub(r'\s+([.,!?])', r'\1', text)  # Remove spaces before punctuation
-        text = re.sub(r'([.,!?])([^\s])', r'\1 \2', text) # Ensure a single space follows punctuation, if not at the end of a line
+        # text = re.sub(r'([.,!?])([^\s])', r'\1 \2', text) # Ensure a single space follows punctuation, if not at the end of a line
         return text
 
     class VisibleTextParser(HTMLParser):
@@ -2084,7 +2086,8 @@ Enter Prompt:<br>
             has_vision = (mmprojpath!="")
             has_password = (password!="")
             has_whisper = (fullwhispermodelpath!="")
-            response_body = (json.dumps({"result":"KoboldCpp","version":KcppVersion, "protected":has_password ,"txt2img":has_txt2img,"vision":has_vision,"transcribe":has_whisper,"multiplayer":has_multiplayer}).encode())
+            has_search = True if args.websearch else False
+            response_body = (json.dumps({"result":"KoboldCpp","version":KcppVersion, "protected":has_password ,"txt2img":has_txt2img,"vision":has_vision,"transcribe":has_whisper,"multiplayer":has_multiplayer,"websearch":has_search}).encode())
 
         elif self.path.endswith(('/api/extra/perf')):
             global last_req_time, start_time
@@ -2172,32 +2175,6 @@ Enter Prompt:<br>
                 },
             }).encode())
 
-        elif self.path.startswith(("/websearch")):
-            if args.websearch:
-                # ensure authorized
-                auth_ok = True
-                if password and password !="":
-                    auth_header = None
-                    auth_ok = False
-                    if 'Authorization' in self.headers:
-                        auth_header = self.headers['Authorization']
-                    elif 'authorization' in self.headers:
-                        auth_header = self.headers['authorization']
-                    if auth_header is not None and auth_header.startswith('Bearer '):
-                        token = auth_header[len('Bearer '):].strip()
-                        if token==password:
-                            auth_ok = True
-                if auth_ok:
-                    parsed_url = urlparse.urlparse(self.path)
-                    parsed_dict = urlparse.parse_qs(parsed_url.query)
-                    searchstr = (parsed_dict['q'][0]) if 'q' in parsed_dict else ""
-                    searchres = websearch(searchstr)
-                    response_body = (json.dumps(searchres).encode())
-                else:
-                    response_body = (json.dumps([]).encode())
-            else:
-                response_body = (json.dumps([]).encode())
-
         elif self.path=="/api" or self.path=="/docs" or self.path.startswith(('/api/?json=','/api?json=','/docs/?json=','/docs?json=')):
             content_type = 'text/html'
             if embedded_kcpp_docs is None:
@@ -2446,6 +2423,22 @@ Enter Prompt:<br>
                     response_code = 400
                     response_body = (json.dumps({"success": False, "error":"Submitted story invalid!"}).encode())
 
+        elif self.path.startswith(("/api/extra/websearch")):
+            if not self.secure_endpoint():
+                return
+            if args.websearch:
+                try:
+                    tempbody = json.loads(body)
+                    searchstr = tempbody.get('q', "")
+                    searchres = websearch(searchstr)
+                    response_body = (json.dumps(searchres).encode())
+                except Exception as e:
+                    utfprint("WebSearch Parse Error: " + str(e))
+                    response_code = 400
+                    response_body = (json.dumps([]).encode())
+            else:
+                response_body = (json.dumps([]).encode())
+
         if response_body is not None:
             self.send_response(response_code)
             self.send_header('content-length', str(len(response_body)))