Spaces:
Running
Running
File size: 3,832 Bytes
c837b7c |
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 |
import axios from "axios";
import type { PreviewServer, ViteDevServer } from "vite";
import { fetchSearXNG } from "./fetchSearXNG";
import { rankSearchResults } from "./rankSearchResults";
import {
incrementGraphicalSearchesSinceLastRestart,
incrementTextualSearchesSinceLastRestart,
} from "./searchesSinceLastRestart";
import { verifyTokenAndRateLimit } from "./verifyTokenAndRateLimit";
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.end(JSON.stringify({ error: "Missing query parameter" }));
return;
}
const { isAuthorized, statusCode, error } =
await verifyTokenAndRateLimit(token);
if (!isAuthorized && statusCode && error) {
response.statusCode = statusCode;
response.end(JSON.stringify({ error }));
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[];
const rankedResults = await rankSearchResults(query, results);
incrementTextualSearchesSinceLastRestart();
response.setHeader("Content-Type", "application/json");
response.end(JSON.stringify(rankedResults));
} else {
const results = searxngResults as ImageResult[];
const rankedResults = await rankSearchResults(
query,
results.map(
([title, url, , sourceUrl]) =>
[
title.slice(0, 100),
sourceUrl.slice(0, 100),
url.slice(0, 100),
] as TextResult,
),
);
const processedResults = (
await Promise.all(
results
.filter((_, index) =>
rankedResults.some(([title]) => title === results[index][0]),
)
.map(async ([title, url, thumbnailSource, sourceUrl]) => {
try {
const axiosResponse = await axios.get(thumbnailSource, {
responseType: "arraybuffer",
});
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.end(JSON.stringify({ error: "Internal server error" }));
}
});
}
|