ollama-plus/scheduler/helpers/ollamaCalls.ts

import crypto from 'node:crypto';
import { base } from './resolve-user';
import fs from 'fs';

export type schedInp = {
	name: string
	when: {
		cron?: string
		start?: string
	}
	oneShot?: boolean
	template?: { name: string }
	parameters: object
	prompt: string
	userId: string
	model: string,
	tools: string[]
	features?: Record<string, boolean>
	chatId?: string
	files?: { fname: string, fkey?: string, content?: string }[]
}

export type ollamaInp = {
	name: string
	displayName: string
	userId: string
	schedule: string
	startAt?: string
	oneShot: boolean
	template: any //idk
	parameters: object
	prompt: string
	model: string
	tools: string[]
	features: Record<string, boolean>
	cookie: string
	chatId?: string
	files?: { fname: string, fkey: string }[]
}

async function makeRequest(cookie: string, path: string, throwOnErr = false) {
	if (!cookie) return [];

	const r = await fetch(`${base}/${path}`, {
		method: "GET",
		headers: {
			'Authorization': `Bearer ${cookie}`,
			"content-type": "application/json",
			"accept": "application/json",
		}
	});

	if (throwOnErr && !r.ok) throw r;

	return await r.text().then(data => {
		try { return JSON.parse(data); }
		catch (_) { return data; }
	}).catch(async err => {
		console.error(err);
		return [];
	});
}

// NO FOREWARD SLASH NO IDK WHY IG /v1/ DOES IT FUCK ME MAN
export const getModels = (cookie: string) => makeRequest(cookie, 'api/models').then((r) => (r.data || []));
export const getTools = (cookie: string) => makeRequest(cookie, 'api/v1/tools/');
export const authCall = (cookie: string) => makeRequest(cookie, 'api/v1/auths/', true);

const buildPrompt = async (def: ollamaInp) => {
	let content = `Please complete this request given the following information:\n`;
	content += `request: ${def.prompt}\nExtra Context:\n\`\`\`json\n${JSON.stringify(def.parameters || {})}\n\`\`\``;

	if (def.files) {
		const files = await Promise.all(def.files.map(async f => {
			const fpath = `/app/data/files/${f.fkey}`,
				r = await fetch(`https://${base}/api/v1/files`, {
					headers: {
						'Content-Type': 'multipart/form-data',
						'filename': f.fname
					},
					method: 'POST',
					body: await fs.readFileSync(fpath)
				}).catch(err => {
					console.error(err);
					return { data: err, ok: false };
				});

			fs.rm(fpath, { recursive: false }, (err: any) => {
				if (err) console.error(err);
			});

			return {
				ok: r.ok,
				data: r.ok ? await (r as Response).json() : (r as any).data
			};
		}));

		if (files.find(v => !v.ok)) {
			throw new Error(`File upload failed for file arr: ${JSON.stringify(files)}`)
		}

		return {
			content, files: files.map(({ data: f }) => ({
				id: f.id,
				type: 'file',
				url: `/api/v1/files/${f.id}`,
				file: f,
				name: f.filename,
				status: 'uploaded', // trust, trust
				size: f.meta.size,
				error: '',
				itemId: crypto.randomUUID()
			}))
		};
	}

	return { content };
};

const getSystemSettings = async (cookie: string) => {
	const r = await makeRequest(cookie, 'api/v1/users/user/settings/'),
		{ notifications, system } = r?.length || {};
	return { notifications, system };
};

const showErr = async (r: any) => {
	console.error('=========================================');
	console.error(r);
	console.error(await r.text());
	console.error('=========================================');
}

// duplicate? Maybe?
const toFeatureFlags = (f: Record<string, boolean> = {}) => ({
	image_generation: !!f.image_generation,
	code_interpreter: !!f.code_interpreter,
	web_search: !!f.web_search,
	memory: !!f.memory
});

const splitThink = (raw: string, startPos: number = 0) => {
	const open = raw.indexOf('<think>', startPos),
		close = raw.indexOf('</think>', startPos);

	if (close >= 0) {
		const start = open >= 0 ? open + '<think>'.length : 0;
		return {
			thinking: raw.slice(start, close).trim(),
			answer: raw.slice(close + '</think>'.length).trim()
		};
	}
	return { thinking: '', answer: raw.trim(), endPos: close };
};

async function newChat(def: ollamaInp, sysMsg?: string) {
	const userId = crypto.randomUUID(),
		assistantId = crypto.randomUUID(), // will be reused in completions
		now = Date.now();

	const userMsg = {
		id: userId,
		parentId: null,
		childrenIds: [assistantId],
		role: 'user',
		...(await buildPrompt(def)),
		timestamp: now,
		models: [def.model]
	};

	const assistantPlaceholder = {
		id: assistantId,
		parentId: userId,
		childrenIds: [],
		role: 'assistant',
		content: '',
		model: def.model,
		modelName: def.model,
		modelIdx: 0,
		timestamp: now
	};

	const reqObj = {
		user_id: def.userId,
		title: 'New Chat',
		chat: {
			id: '',
			title: 'New Chat',
			models: [def.model],
			system: sysMsg || '',
			params: {}, // model params???
			history: {
				messages: {
					[userMsg.id]: userMsg,
					[assistantPlaceholder.id]: assistantPlaceholder
				},
				currentId: assistantId // needed for openwebUI stuff
			},
			messages: [userMsg, assistantPlaceholder],
			tags: []
		},
		share_id: null,
		archived: false,
		pinned: false,
		meta: {},
		folder_id: null
	};

	const r = await fetch(`${base}/api/v1/chats/new`, {
		method: 'POST',
		headers: {
			Authorization: `Bearer ${def.cookie}`,
			'Content-Type': 'application/json'
		},
		body: JSON.stringify(reqObj)
	});

	if (!r.ok) {
		await showErr(r);
		throw new Error('Failed to create new chat!');
	}

	const created = await r.json();
	return { chatId: created.id as string, assistantId };
}

export async function callNewChat(def: ollamaInp) {
	const { cookie } = def;
	if (!cookie) throw new Error('Cookie not found!');

	const sysSettings = await getSystemSettings(cookie),
		{ name: username } = await makeRequest(cookie, 'api/v1/auths/'),
		models = await getModels(cookie),
		model = models.find((m: any) => m.id === def.model);

	if (!model) throw new Error(`Model ${def.model} not found!`);

	const { chatId, assistantId } = def.chatId
		? { chatId: def.chatId, assistantId: crypto.randomUUID() } // TODO: if passing an existing chat, insert a placeholder there first (same shape as above)
		: await newChat(def, sysSettings.system);

	// mock socket ID
	const sessionId = crypto.randomUUID(),
		completeReq = {
			chat_id: chatId,
			id: assistantId,
			stream: false,
			model: def.model,
			messages: [
				{ role: 'system', content: sysSettings.system || '' },
				{ role: 'user', ...(await buildPrompt(def)) }
			],
			features: toFeatureFlags(def.features),
			variables: {
				'{{USER_NAME}}': username,
				'{{USER_LANGUAGE}}': 'en-US'
			},
			session_id: sessionId,
			background_tasks: { title_generation: true, follow_up_generation: true }
		};

	// run completion
	let r = await fetch(`${base}/api/chat/completions`, {
		method: 'POST',
		headers: {
			Authorization: `Bearer ${cookie}`,
			'Content-Type': 'application/json'
		},
		body: JSON.stringify(completeReq)
	});
	if (!r.ok) return showErr(r);

	// bc stream:false, read the whole json here bc ui still wants a finalizer call
	const txt = await r.text();
	let content = txt;

	try { const j = JSON.parse(txt); content = j?.choices?.[0]?.message?.content ?? txt; } catch { }
	const thinkRes = splitThink(content),
		answerArr = [(thinkRes.thinking ? `<details id="__DETAIL_${0}__"/>\n` : '') + thinkRes.answer];

	let counter = thinkRes.thinking ? 1 : 0,
		{ thinking, endPos } = thinkRes;

	while (thinking && counter < 10) {
		const { thinking: newThink, answer: newAnswer, endPos: newEndPos } = splitThink(content, endPos);

		thinking = newThink;

		if (thinking) {
			answerArr.push(`<details id="__DETAIL_${counter}__"/>\n` + newAnswer);
			endPos = newEndPos;
		}
		else break;

		counter++;
	}

	const answer = answerArr.join('\n')

	// fetch current chat, replace the assistant placeholder content with `answer`
	const chat = await (await fetch(`${base}/api/v1/chats/${chatId}`, {
		headers: { Authorization: `Bearer ${def.cookie}` }
	})).json();

	chat.chat.history.messages[assistantId].content = answer;
	const idx = chat.chat.messages.findIndex((m: any) => m.id === assistantId);
	if (idx >= 0) chat.chat.messages[idx].content = answer;

	console.log(JSON.stringify(chat));

	r = await fetch(`${base}/api/v1/chats/${chatId}`, {
		method: 'POST',
		headers: { Authorization: `Bearer ${def.cookie}`, 'Content-Type': 'application/json' },
		body: JSON.stringify({ chat: chat.chat })
	});

	// now finalize
	r = await fetch(`${base}/api/chat/completed`, {
		method: 'POST',
		headers: { Authorization: `Bearer ${def.cookie}`, 'Content-Type': 'application/json' },
		body: JSON.stringify({ chat_id: chatId, id: assistantId, session_id: sessionId, model: def.model })
	});

	console.log('completed:', await r.text());
}