export function endpointGenAI()

in src/lib/server/endpoints/google/endpointGenAI.ts [39:144]
70 lines of code
6 McCabe index (conditional complexity)

export function endpointGenAI(input: z.input<typeof endpointGenAIParametersSchema>): Endpoint {
	const { model, apiKey, safetyThreshold, multimodal } = endpointGenAIParametersSchema.parse(input);

	const genAI = new GoogleGenerativeAI(apiKey);

	const safetySettings = safetyThreshold
		? Object.keys(HarmCategory)
				.filter((cat) => cat !== HarmCategory.HARM_CATEGORY_UNSPECIFIED)
				.reduce((acc, val) => {
					acc.push({
						category: val as HarmCategory,
						threshold: safetyThreshold,
					});
					return acc;
				}, [] as SafetySetting[])
		: undefined;

	return async ({ messages, preprompt, generateSettings }) => {
		const parameters = { ...model.parameters, ...generateSettings };

		const generativeModel = genAI.getGenerativeModel({
			model: model.id ?? model.name,
			safetySettings,
			generationConfig: {
				maxOutputTokens: parameters?.max_new_tokens ?? 4096,
				stopSequences: parameters?.stop,
				temperature: parameters?.temperature ?? 1,
			},
		});

		let systemMessage = preprompt;
		if (messages[0].from === "system") {
			systemMessage = messages[0].content;
			messages.shift();
		}

		const genAIMessages = await Promise.all(
			messages.map(async ({ from, content, files }: Omit<Message, "id">): Promise<Content> => {
				return {
					role: from === "user" ? "user" : "model",
					parts: [
						...(await Promise.all(
							(files ?? []).map((file) => fileToImageBlock(file, multimodal.image))
						)),
						{ text: content },
					],
				};
			})
		);

		const result = await generativeModel.generateContentStream({
			contents: genAIMessages,
			systemInstruction:
				systemMessage && systemMessage.trim() !== ""
					? {
							role: "system",
							parts: [{ text: systemMessage }],
						}
					: undefined,
		});

		let tokenId = 0;
		return (async function* () {
			let generatedText = "";

			for await (const data of result.stream) {
				if (!data?.candidates?.length) break; // Handle case where no candidates are present

				const candidate = data.candidates[0];
				if (!candidate.content?.parts?.length) continue; // Skip if no parts are present

				const firstPart = candidate.content.parts.find((part) => "text" in part) as
					| TextPart
					| undefined;
				if (!firstPart) continue; // Skip if no text part is found

				const content = firstPart.text;
				generatedText += content;

				const output: TextGenerationStreamOutput = {
					token: {
						id: tokenId++,
						text: content,
						logprob: 0,
						special: false,
					},
					generated_text: null,
					details: null,
				};
				yield output;
			}

			const output: TextGenerationStreamOutput = {
				token: {
					id: tokenId++,
					text: "",
					logprob: 0,
					special: true,
				},
				generated_text: generatedText,
				details: null,
			};
			yield output;
		})();
	};
}