Spaces:
Running
Running
File size: 5,092 Bytes
16b7924 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 |
import axios from "axios";
import type { PreviewServer, ViteDevServer } from "vite";
import { handleTokenVerification } from "./handleTokenVerification";
import { rankSearchResults } from "./rankSearchResults";
import { getRerankerStatus } from "./rerankerService";
import {
incrementGraphicalSearchesSinceLastRestart,
incrementTextualSearchesSinceLastRestart,
} from "./searchesSinceLastRestart";
import { fetchSearXNG } from "./webSearchService";
type TextResult = [title: string, content: string, url: string];
type ImageResult = [
title: string,
url: string,
thumbnailSource: string,
sourceUrl: string,
];
export function searchEndpointServerHook<
T extends ViteDevServer | PreviewServer,
>(server: T) {
server.middlewares.use(async (request, response, next) => {
if (!request.url?.startsWith("/search/")) return next();
const url = new URL(request.url, `http://${request.headers.host}`);
const query = url.searchParams.get("q");
const token = url.searchParams.get("token");
const limit = Number(url.searchParams.get("limit")) || 30;
if (!query) {
response.statusCode = 400;
response.setHeader("Content-Type", "application/json");
response.end(JSON.stringify({ error: "Missing query parameter" }));
return;
}
const { shouldContinue } = await handleTokenVerification(token, response);
if (!shouldContinue) return;
try {
const isTextSearch = request.url?.startsWith("/search/text");
const searchType = isTextSearch ? "text" : "images";
const searxngResults = await fetchSearXNG(query, searchType, limit);
if (isTextSearch) {
const results = searxngResults as TextResult[];
let rankedResults: [title: string, content: string, url: string][];
const isRerankerHealthy = await getRerankerStatus();
if (!isRerankerHealthy) {
console.warn(
"Reranker service is not healthy, using unranked results",
);
}
try {
if (isRerankerHealthy) {
rankedResults = await rankSearchResults(query, results, true);
} else {
rankedResults = results;
}
} catch (error) {
console.error(
"Error ranking search results:",
error instanceof Error ? error.message : error,
);
rankedResults = results;
}
incrementTextualSearchesSinceLastRestart();
response.setHeader("Content-Type", "application/json");
response.end(JSON.stringify(rankedResults));
} else {
const results = searxngResults as ImageResult[];
const resultsText = results.map(
([title, url]) => [title.slice(0, 100), "", url] as TextResult,
);
let rankedResults: [title: string, content: string, url: string][];
const isRerankerHealthy = await getRerankerStatus();
if (!isRerankerHealthy) {
console.warn(
"Reranker service is not healthy, using unranked results",
);
}
try {
if (isRerankerHealthy) {
rankedResults = await rankSearchResults(query, resultsText);
} else {
rankedResults = resultsText;
}
} catch (error) {
console.error(
"Error ranking search results:",
error instanceof Error ? error.message : error,
);
rankedResults = resultsText;
}
const processedResults = (
await Promise.all(
rankedResults.map(async ([title, , rankedResultUrl]) => {
const result = results.find(
([, resultUrl]) => resultUrl === rankedResultUrl,
);
if (!result) return null;
const [_, url, thumbnailSource, sourceUrl] = result;
try {
const axiosResponse = await axios.get(thumbnailSource, {
responseType: "arraybuffer",
timeout: 1000,
});
const contentType = axiosResponse.headers["content-type"];
const base64 = Buffer.from(axiosResponse.data).toString(
"base64",
);
return [
title,
url,
`data:${contentType};base64,${base64}`,
sourceUrl,
] as ImageResult;
} catch {
return null;
}
}),
)
).filter((result): result is ImageResult => result !== null);
incrementGraphicalSearchesSinceLastRestart();
response.setHeader("Content-Type", "application/json");
response.end(JSON.stringify(processedResults));
}
} catch (error) {
console.error(
"Error processing search:",
error instanceof Error ? error.message : error,
);
response.statusCode = 500;
response.setHeader("Content-Type", "application/json");
response.end(JSON.stringify({ error: "Internal server error" }));
}
});
}
|