From 0441a8725f14892d4e3b5e3d4c67f94767bb6bc6 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Wed, 15 Jan 2025 21:58:49 +0200
Subject: [PATCH] WebSearch: Add KoboldCpp search endpoint.

---
 src/endpoints/search.js | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/src/endpoints/search.js b/src/endpoints/search.js
index e29e53008..c999f8264 100644
--- a/src/endpoints/search.js
+++ b/src/endpoints/search.js
@@ -4,6 +4,8 @@ import express from 'express';
 import { decode } from 'html-entities';
 import { readSecret, SECRET_KEYS } from './secrets.js';
 import { jsonParser } from '../express-common.js';
+import { trimV1 } from '../util.js';
+import { setAdditionalHeaders } from '../additional-headers.js';
 
 export const router = express.Router();
 
@@ -257,6 +259,41 @@ router.post('/tavily', jsonParser, async (request, response) => {
     }
 });
 
+router.post('/koboldcpp', jsonParser, async (request, response) => {
+    try {
+        const { query, url } = request.body;
+
+        if (!url) {
+            console.error('No URL provided for KoboldCpp search');
+            return response.sendStatus(400);
+        }
+
+        console.debug('KoboldCpp search query', query);
+
+        const baseUrl = trimV1(url);
+        const args = {
+            method: 'POST',
+            headers: {},
+            body: JSON.stringify({ q: query }),
+        };
+
+        setAdditionalHeaders(request, args, baseUrl);
+        const result = await fetch(`${baseUrl}/api/extra/websearch`, args);
+
+        if (!result.ok) {
+            const text = await result.text();
+            console.error('KoboldCpp request failed', result.statusText, text);
+            return response.status(500).send(text);
+        }
+
+        const data = await result.json();
+        return response.json(data);
+    } catch (error) {
+        console.error(error);
+        return response.sendStatus(500);
+    }
+});
+
 router.post('/visit', jsonParser, async (request, response) => {
     try {
         const url = request.body.url;