uncloseai.com/src/models.js

// Model registry and selection functionality
import { VLLM_ENDPOINTS } from "./config.js";

// This registry maps a model's ID to the endpoint where it resides.
export const modelRegistry = {};

// Fetch models from each endpoint with caching.
// Cache is busted if the endpoint array changes or TTL.
export async function fetchModelsFromEndpoints() {
	const cacheKey = "modelRegistryCache";
	const endpointsKey = "vllmEndpointsHash";
	const endpointsString = JSON.stringify(VLLM_ENDPOINTS);
	const cachedEndpoints = localStorage.getItem(endpointsKey);
	const cacheItem = localStorage.getItem(cacheKey);
	const now = Date.now();
	const TTL = 300000; // 5 minutes in milliseconds

	if (cacheItem && cachedEndpoints === endpointsString) {
		try {
			const cachedData = JSON.parse(cacheItem);
			if (now - cachedData.timestamp < TTL) {
				// Restore cached modelRegistry
				Object.assign(modelRegistry, cachedData.modelRegistry);
				return cachedData.models;
			}
		} catch (e) {
			console.error("Error reading model registry from cache", e);
		}
	}

	// If no valid cache, fetch models from all endpoints
	const fetchPromises = VLLM_ENDPOINTS.map(async (endpoint) => {
		try {
			const res = await fetch(`${endpoint.url}/models`);
			if (!res.ok)
				throw new Error(
					`HTTP error! status: ${res.status} from ${endpoint.url}`,
				);
			const jsonResponse = await res.json();
			// Expected JSON structure: { data: [ { id, ... }, ... ], object: "list" }
			const models = jsonResponse.data || [];
			// Map each model to include its endpoint ID, unique ID, and model name
			return models.map((model) => ({
				...model,
				modelName: model.id, // Explicitly store model name
				endpointId: endpoint.id,
				uniqueId: `${endpoint.id}-${model.id}`, // Unique ID with endpoint ID first
			}));
		} catch (error) {
			console.error(`Error fetching models from ${endpoint.url}:`, error);
			return [];
		}
	});
	const allModelsArrays = await Promise.all(fetchPromises);
	const models = allModelsArrays.flat();

	// Update modelRegistry with unique model instances
	models.forEach((model) => {
		modelRegistry[model.uniqueId] = {
			url: VLLM_ENDPOINTS.find((e) => e.id === model.endpointId).url,
			endpointId: model.endpointId,
		};
	});

	// Cache the results
	const cacheData = {
		timestamp: now,
		modelRegistry: modelRegistry,
		models: models,
	};
	localStorage.setItem(cacheKey, JSON.stringify(cacheData));
	localStorage.setItem(endpointsKey, endpointsString);

	return models;
}

// Create a dynamic drop-down for model selection
// This function creates a <select> element and fills it with unique model instances.
export async function createModelSelectionDropdown() {
	// Only create dropdown if there's a chat interface or user input
	const userInput = document.getElementById("user-input");
	const chatBox = document.getElementById("chat-box");
	const uncloseaiElements = document.querySelectorAll(".uncloseai");

	if (!userInput && !chatBox && uncloseaiElements.length === 0) {
		console.log("No chat interface found, skipping model dropdown creation");
		return;
	}

	const models = await fetchModelsFromEndpoints();
	if (!models.length) {
		console.warn("No models returned from any endpoint.");
		return;
	}

	const dropdown = document.createElement("select");
	dropdown.id = "model-selection";
	dropdown.style.margin = "10px";

	models.forEach((model) => {
		const option = document.createElement("option");
		option.value = model.uniqueId; // Use unique ID for selection
		option.textContent = `${model.endpointId} | ${model.modelName}`;
		dropdown.appendChild(option);
	});

	// Insert the dropdown directly above the chat input box.
	if (userInput?.parentNode) {
		userInput.parentNode.parentNode.insertBefore(
			dropdown,
			userInput.parentNode,
		);
	} else {
		// Only insert at top of body if there are uncloseai elements
		if (uncloseaiElements.length > 0) {
			document.body.insertBefore(dropdown, document.body.firstChild);
		}
	}
}

// Function to add the Refresh Models button
export function addRefreshModelsButton() {
	// Only create button if there's a chat interface
	const userInput = document.getElementById("user-input");
	const chatBox = document.getElementById("chat-box");
	const uncloseaiElements = document.querySelectorAll(".uncloseai");

	if (!userInput && !chatBox && uncloseaiElements.length === 0) {
		console.log("No chat interface found, skipping refresh button creation");
		return;
	}

	// Create a container for the button
	const container = document.createElement("div");
	container.id = "refresh-models-container";

	// Position it near the model selection dropdown
	if (userInput?.parentNode) {
		userInput.parentNode.parentNode.insertBefore(
			container,
			userInput.parentNode,
		);
	}

	// Create a button element
	const refreshButton = document.createElement("button");
	refreshButton.textContent = "Refresh Models";
	refreshButton.style.margin = "10px";
	refreshButton.onclick = async () => {
		refreshButton.textContent = "Refreshing...";
		refreshButton.disabled = true;

		try {
			// Clear the model registry cache
			localStorage.removeItem("modelRegistryCache");
			localStorage.removeItem("vllmEndpointsHash");

			// Fetch models again
			const models = await fetchModelsFromEndpoints();

			// Update the dropdown
			const dropdown = document.getElementById("model-selection");
			if (dropdown) {
				// Clear existing options
				dropdown.innerHTML = "";

				// Add options using models array
				models.forEach((model) => {
					const option = document.createElement("option");
					option.value = model.uniqueId;
					option.textContent = `${model.endpointId} | ${model.modelName}`;
					dropdown.appendChild(option);
				});
			}

			alert("Models refreshed successfully!");
		} catch (error) {
			console.error("Error refreshing models:", error);
			alert("Failed to refresh models. Please try again.");
		} finally {
			refreshButton.textContent = "Refresh Models";
			refreshButton.disabled = false;
		}
	};
	container.appendChild(refreshButton);
}

// Helper function to get the selected model from the dropdown.
export function getSelectedModel() {
	// Check modal dropdown first
	const modalDropdown = document.getElementById("hermes-model-selection");
	if (modalDropdown?.value) {
		// Extract just the model name, not the unique ID
		const selectedId = modalDropdown.value;
		if (modelRegistry[selectedId]) {
			return selectedId.split("-").slice(1).join("-"); // Remove endpoint prefix
		}
		return modalDropdown.value;
	}

	// Fallback to main dropdown
	const dropdown = document.getElementById("model-selection");
	if (dropdown?.value) {
		// Extract just the model name, not the unique ID
		const selectedId = dropdown.value;
		if (modelRegistry[selectedId]) {
			return selectedId.split("-").slice(1).join("-"); // Remove endpoint prefix
		}
		return dropdown.value;
	}

	// Return default model if nothing selected
	const defaultModel = Object.keys(modelRegistry)[0];
	if (defaultModel) {
		return defaultModel.split("-").slice(1).join("-"); // Remove endpoint prefix
	}

	// Fallback to hardcoded model
	return "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic";
}

// Helper function to get the API endpoint for the selected model.
export function getSelectedModelEndpoint() {
	// Check modal dropdown first
	const modalDropdown = document.getElementById("modal-model-selection");
	if (modalDropdown?.value && modelRegistry[modalDropdown.value]) {
		return modelRegistry[modalDropdown.value].url;
	}

	// Fallback to main dropdown
	const dropdown = document.getElementById("model-selection");
	if (dropdown?.value && modelRegistry[dropdown.value]) {
		return modelRegistry[dropdown.value].url;
	}

	// Find the first available model in the registry
	const availableModels = Object.keys(modelRegistry);
	if (availableModels.length > 0) {
		return modelRegistry[availableModels[0]].url;
	}

	// Fallback to the first endpoint if model registry is empty
	if (VLLM_ENDPOINTS.length > 0) {
		return VLLM_ENDPOINTS[0].url;
	}

	// Final fallback
	return "https://hermes.ai.unturf.com/v1";
}