Feature/fix lmstudio (#34)

* Add gemma tokencounter, more conservative estimator

* LM Studio fixes

* Set correct default endpoint for LM Studio

package-lock.json

@@ -11,6 +11,7 @@
 			"dependencies": {
 				"@electric-sql/pglite": "^0.2.17",
 				"@lenml/char-card-reader": "^1.0.6",
+				"@lenml/tokenizer-gemma": "^3.4.2",
 				"@lmstudio/sdk": "^1.2.1",
 				"@lucide/svelte": "^0.511.0",
 				"@proj-airi/duckdb-wasm": "^0.4.27",
@@ -1126,6 +1127,21 @@
 			"resolved": "https://registry.npmjs.org/@lenml/char-card-reader/-/char-card-reader-1.0.6.tgz",
 			"integrity": "sha512-e6rg/HVeqKx9pGmVthz3mvAoz8HNeaAmZXIkbTKD8+swy/YMoSwgQd3dMxBy3IhaqaJpfFChPSV+Mc3LFIArgw=="
 		},
+		"node_modules/@lenml/tokenizer-gemma": {
+			"version": "3.4.2",
+			"resolved": "https://registry.npmjs.org/@lenml/tokenizer-gemma/-/tokenizer-gemma-3.4.2.tgz",
+			"integrity": "sha512-tyRO/VZ/gmDlJ3rG4qdR52aAWgy7pwzrp2B3ovdUphVP9Gy4wFC/IMvZ51/30/xpc7q2VYF5iw7sy4d788OfHw==",
+			"license": "Apache-2.0",
+			"dependencies": {
+				"@lenml/tokenizers": "^3.4.0"
+			}
+		},
+		"node_modules/@lenml/tokenizers": {
+			"version": "3.4.0",
+			"resolved": "https://registry.npmjs.org/@lenml/tokenizers/-/tokenizers-3.4.0.tgz",
+			"integrity": "sha512-BRDGqddggcsd4YwubhXzHuUdpkPq5GxcYLp+40SxfyV6/kkSQU7rNaQ8iS+HMD+XfqzXaYsHuHNPlaLA6gnM5A==",
+			"license": "Apache-2.0"
+		},
 		"node_modules/@lmstudio/lms-isomorphic": {
 			"version": "0.4.5",
 			"resolved": "https://registry.npmjs.org/@lmstudio/lms-isomorphic/-/lms-isomorphic-0.4.5.tgz",

package.json

@@ -75,6 +75,7 @@
 	"dependencies": {
 		"@electric-sql/pglite": "^0.2.17",
 		"@lenml/char-card-reader": "^1.0.6",
+		"@lenml/tokenizer-gemma": "^3.4.2",
 		"@lmstudio/sdk": "^1.2.1",
 		"@lucide/svelte": "^0.511.0",
 		"@proj-airi/duckdb-wasm": "^0.4.27",

src/lib/client/components/sidebars/ConnectionsSidebar.svelte

@@ -545,7 +545,7 @@
 					bind:value={newConnectionType}
 				>
 					{#each CONNECTION_TYPES as t}
-						<option value={t.value} disabled={t.value===CONNECTION_TYPE.LM_STUDIO}>{t.label}</option>
+						<option value={t.value}>{t.label}</option>
 					{/each}
 				</select>
 			</div>

src/lib/client/connectionForms/LMStudioForm.svelte

@@ -15,7 +15,6 @@
 	let testResult: { ok: boolean; error?: string; models?: any[] } | null =
 		$state(null)
 	let extraFields = $state(extraJsonToExtraFields(connection.extraJson || {}))
-	let availableModels = $derived.by(() => availableLMStudioModels)
 
 	function handleRefreshModels() {
 		socket.emit("refreshModels", {
@@ -34,7 +33,7 @@
 		return {
 			stream: extraJson?.stream ?? true,
 			think: extraJson?.think ?? false,
-			keepAlive: extraJson?.keepAlive ?? "300ms",
+			ttl: extraJson?.ttl ?? 60,
 			raw: extraJson?.raw ?? true,
 			useChat: extraJson?.useChat ?? true
 		}
@@ -45,7 +44,7 @@
 			...connection.extraJson,
 			stream: extraFields.stream,
 			think: extraFields.think,
-			keepAlive: extraFields.keepAlive,
+			ttl: extraFields.ttl,
 			raw: extraFields.raw,
 			useChat: extraFields.useChat
 		}
@@ -53,6 +52,9 @@
 
 	socket.on("refreshModels", (msg: Sockets.RefreshModels.Response) => {
 		if (msg.models) availableLMStudioModels = msg.models
+		if (!connection.model && msg.models.length > 0) {
+			connection.model = msg.models[0].model
+		}
 	})
 
 	socket.on("testConnection", (msg: Sockets.TestConnection.Response) => {
@@ -178,7 +180,7 @@
 					id="baseUrl"
 					type="text"
 					bind:value={connection.baseUrl}
-					placeholder="http://localhost:11434/"
+					placeholder="ws://localhost:1234"
 					required
 					class="input"
 				/>
@@ -206,25 +208,19 @@
 					}}
 				/>
 			</div>
+			<div class="mt-2 flex flex-col gap-1">
+				<label class="font-semibold" for="ttl">Keep Alive (seconds)</label>
+				<input
+					id="ttl"
+					type="number"
+					bind:value={extraFields.ttl}
+					class="input"
+					placeholder="60"
+					min="1"
+				/>
+			</div>
 		</div>
 	</details>
-	<!-- <div>
-		<label class="font-semibold" for="keepAlive">Keep Alive</label>
-		<input
-			id="keepAlive"
-			type="text"
-			bind:value={extraFields.keepAlive}
-			class="input"
-			placeholder="e.g. 300ms"
-			onchange={() => {
-				connection.extraJson = {
-					...connection.extraJson,
-					keepAlive: extraFields.keepAlive
-				}
-				handleChange()
-			}}
-		/>
-	</div> -->
 	{#if testResult?.error}
 		<div class="text-error mt-1 text-xs">{testResult.error}</div>
 	{/if}

src/lib/server/connectionAdapters/BaseConnectionAdapter.ts

@@ -75,8 +75,9 @@ export abstract class BaseConnectionAdapter {
 		})
 	}
 
-	compilePrompt(args: {}): Promise<CompiledPrompt> {
-		return this.promptBuilder.compilePrompt(args)
+	async compilePrompt(args: {}): Promise<CompiledPrompt> {
+		this.promptBuilder.tokenLimit = await this.getContextTokenLimit()
+		return await this.promptBuilder.compilePrompt(args)
 	}
 
 	abstract generate(): Promise<
@@ -90,6 +91,10 @@ export abstract class BaseConnectionAdapter {
 	abort() {
 		this.isAborting = true
 	}
+
+	async getContextTokenLimit(): Promise<number> {
+		return this.sampling.contextTokensEnabled ? this.sampling.contextTokens || 4096 : 4096
+	}
 }
 
 export interface AdapterExports {

src/lib/server/connectionAdapters/LMStudioAdapter.ts

@@ -10,11 +10,9 @@ import {
 } from "./BaseConnectionAdapter"
 import {
 	type BaseLoadModelOpts,
-	type ChatLike,
 	type LLM,
 	type LLMLoadModelConfig,
 	type LLMPredictionOpts,
-	type LLMRespondOpts,
 	LMStudioClient,
 	type OngoingPrediction
 } from "@lmstudio/sdk"
@@ -56,10 +54,7 @@ class LMStudioAdapter extends BaseConnectionAdapter {
 			tokenCounter: new TokenCounters(
 				connection.tokenCounter || TokenCounterOptions.ESTIMATE
 			),
-			tokenLimit:
-				typeof sampling.contextTokens === "number"
-					? sampling.contextTokens
-					: 2048,
+			tokenLimit: 0, // This is set dynamically based on the LM Studio API
 			contextThresholdPercent: 0.9
 		})
 	}
@@ -97,18 +92,50 @@ class LMStudioAdapter extends BaseConnectionAdapter {
 			const name = modelName || this.connection.model
 			if (!name || typeof name !== "string")
 				throw new Error("Model name required for getModelClient")
-			// TODO, keep alive?
+
+			// Check available models first
+			try {
+				const availableModels =
+					await client.system.listDownloadedModels()
+				const modelExists = availableModels.some(
+					(model) => model.modelKey === name
+				)
+				if (!modelExists) {
+					throw new Error(
+						`Model "${name}" is not downloaded in LM Studio. Available models: ${availableModels.map((m) => m.modelKey).join(", ")}`
+					)
+				}
+			} catch (error) {
+				console.warn("Could not check available models:", error)
+			}
+
 			const opts: BaseLoadModelOpts<LLMLoadModelConfig> = {
 				config: {
-					contextLength: this.sampling.contextTokensEnabled
-						? this.sampling.contextTokens || 2048
-						: 2048,
+					contextLength: this.promptBuilder.tokenLimit,
 					keepModelInMemory: false // TODO: make configurable?
 				},
-				ttl: 1 // TODO: TTL is off for now to force reloading models
+				ttl: this.connection.extraJson.ttl || 60 // Increased TTL to avoid frequent reloading
 			}
 			console.log("LM Studio getModelClient opts", opts)
-			this._modelClient = await client.llm.model(name, opts)
+
+			try {
+				this._modelClient = await client.llm.model(name, opts)
+				const modelInstCtxLength =
+					await this._modelClient.getContextLength()
+				console.log(
+					"Model loaded successfully with context length:",
+					modelInstCtxLength
+				)
+			} catch (error) {
+				const errorMsg =
+					error instanceof Error ? error.message : String(error)
+				if (errorMsg.includes("Error loading model")) {
+					throw new Error(
+						`Failed to load model "${name}" in LM Studio. This may be due to insufficient VRAM/RAM or context length mismatch. Requested context: ${this.promptBuilder.tokenLimit} tokens. Try using a smaller model, reducing context length, or check LM Studio settings. Original error: ${errorMsg}`
+					)
+				}
+				throw error
+			}
 		}
 		return this._modelClient
 	}
@@ -140,15 +167,13 @@ class LMStudioAdapter extends BaseConnectionAdapter {
 			throw new Error("LMStudioAdapter: model must be a string")
 
 		// Prepare stop strings for LM Studio
-		const promptContext = {
-			format: this.connection.promptFormat || "chatml",
-			characters: this.chat.chatCharacters?.map((c) => c.character) || [],
-			personas: this.chat.chatPersonas?.map((p) => p.persona) || []
-		}
+		const promptFormat = this.connection.promptFormat || "chatml"
 		const stopStrings = StopStrings.get({
 			format: promptFormat,
-			characters: this.chat.chatCharacters?.map((cc) => cc.character),
-			personas: this.chat.chatPersonas?.map((cp) => cp.persona),
+			characters: this.chat.chatCharacters?.map(
+				(cc: any) => cc.character
+			),
+			personas: this.chat.chatPersonas?.map((cp: any) => cp.persona),
 			currentCharacterId: this.currentCharacterId
 		})
 		const characterName =
@@ -167,7 +192,7 @@ class LMStudioAdapter extends BaseConnectionAdapter {
 		// Use PromptBuilder for prompt construction
 		const compiledPrompt: CompiledPrompt = await this.compilePrompt({})
 
-		let useChat = this.connection.extraJson?.useChat ?? true
+		const useChat = this.connection.extraJson?.useChat ?? true
 		let prompt: string = ""
 		let messages: any[] | undefined = undefined
 
@@ -177,11 +202,12 @@ class LMStudioAdapter extends BaseConnectionAdapter {
 			prompt = compiledPrompt.prompt!
 		}
 
-		let options: LLMPredictionOpts<unknown> = {
+		const options: LLMPredictionOpts<unknown> = {
 			stopStrings: stop,
 			maxTokens: this.sampling.responseTokensEnabled
-				? this.sampling.responseTokens || 2048
-				: 2048,
+				? this.sampling.responseTokens || 250
+				: 250,
+			contextOverflowPolicy: "truncateMiddle",
 			...this.mapSamplingConfig()
 		}
 
@@ -241,7 +267,7 @@ class LMStudioAdapter extends BaseConnectionAdapter {
 			const content = await (async () => {
 				try {
 					if (useChat && messages) {
-						this.prediction = modelClient.chat(messages, options)
+						this.prediction = modelClient.respond(messages, options)
 						const result = await this.prediction
 						if (
 							result &&
@@ -283,18 +309,35 @@ class LMStudioAdapter extends BaseConnectionAdapter {
 			this.prediction.cancel()
 		}
 	}
+
+	async getContextTokenLimit(): Promise<number> {
+		const limit = await super.getContextTokenLimit()
+
+		const models = await this.getClient().system.listDownloadedModels()
+		const modelName = this.connection.model
+		const modelInfo = models.find((m) => m.modelKey === modelName)
+		if (!modelInfo) {
+			console.warn(
+				`LM Studio getContextTokenLimit: Model "${modelName}" not found in downloaded models`
+			)
+		} else if (modelInfo.maxContextLength < limit) {
+			console.warn(
+				`LM Studio getContextTokenLimit: The configured context limit ${limit} exceeds the model's maximum context length of ${modelInfo.maxContextLength}. This may cause the model to crash.`
+			)
+		}
+
+		return limit
+	}
 }
 
 const connectionDefaults = {
 	baseUrl: "ws://localhost:1234",
 	promptFormat: PromptFormats.VICUNA,
-	tokenCounter: TokenCounterOptions.ESTIMATE,
+	tokenCounter: TokenCounterOptions.ESTIMATE, // Use Gemma tokenizer for better accuracy with Gemma models
 	extraJson: {
 		useChat: true, // Use chat (response api)
 		stream: true,
-		// think: false,
-		keepAlive: "300ms"
-		// raw: true
+		ttl: 60
 	}
 }
 
@@ -314,15 +357,38 @@ const samplingKeyMap: Record<string, string> = {
 async function testConnection(
 	connection: SelectConnection
 ): Promise<{ ok: boolean; error?: string }> {
-	const client = new LMStudioClient({ baseUrl: connection.baseUrl || "" })
-	const res = await client.system.getLMStudioVersion()
-	if (res && typeof res === "object" && "version" in res) {
-		return {
-			ok: true
+	try {
+		const client = new LMStudioClient({ baseUrl: connection.baseUrl || "" })
+		const res = await client.system.getLMStudioVersion()
+		if (res && typeof res === "object" && "version" in res) {
+			// Also check if any models are available
+			try {
+				const models = await client.system.listDownloadedModels()
+				if (!models || models.length === 0) {
+					return {
+						ok: true,
+						error: "LM Studio is running but no models are downloaded. Please download a model in LM Studio first."
+					}
+				}
+			} catch (modelError) {
+				console.warn(
+					"Could not check models during connection test:",
+					modelError
+				)
+			}
+			return {
+				ok: true
+			}
+		} else {
+			return {
+				ok: false,
+				error: "Could not get LM Studio version. Make sure LM Studio server is running on the specified URL."
+			}
 		}
-	} else {
+	} catch (error) {
 		return {
-			ok: false
+			ok: false,
+			error: `Connection failed: ${error instanceof Error ? error.message : String(error)}`
 		}
 	}
 }
@@ -330,27 +396,35 @@ async function testConnection(
 async function listModels(
 	connection: SelectConnection
 ): Promise<{ models: any[]; error?: string }> {
-	const client = new LMStudioClient({ baseUrl: connection.baseUrl || "" })
-	const res = await client.system.listDownloadedModels()
-	if (res && Array.isArray(res)) {
-		const models = res.map((model) => {
+	try {
+		const client = new LMStudioClient({ baseUrl: connection.baseUrl || "" })
+		const res = await client.system.listDownloadedModels()
+		if (res && Array.isArray(res)) {
+			const models = res.map((model) => {
+				return {
+					model: model.modelKey,
+					name: model.displayName
+				}
+			})
 			return {
-				model: model.modelKey,
-				name: model.displayName
+				models: models,
+				error: undefined
+			}
+		} else {
+			console.error(
+				"LM Studio listModels error: Unexpected response format",
+				res
+			)
+			return {
+				models: [],
+				error: "Unexpected response format from LM Studio API"
 			}
-		})
-		return {
-			models: models,
-			error: undefined
 		}
-	} else {
-		console.error(
-			"LM Studio listModels error: Unexpected response format",
-			res
-		)
+	} catch (error) {
+		console.error("LM Studio listModels error:", error)
 		return {
 			models: [],
-			error: "Unexpected response format from LM Studio API"
+			error: "Failed to list models from LM Studio API, is the server running?"
 		}
 	}
 }

src/lib/server/sockets/connections.ts

@@ -53,6 +53,11 @@ export async function createConnection(
 	if ("id" in data) delete data.id
 	try {
 		const modelsRes = await Adapter.listModels(data as any)
+		if (modelsRes.error) {
+			const res = { error: modelsRes.error }
+			emitToUser("error", res)
+			return
+		}
 		if (!modelsRes.models || modelsRes.models.length === 0) {
 			const res = { error: "No models found for this connection." }
 			emitToUser("error", res)
@@ -161,6 +166,12 @@ export async function testConnection(
 		let error: string | null = null
 		if (result.ok) {
 			const modelsRes = await listModels(message.connection)
+			if (modelsRes.error) {
+				emitToUser("error", {
+					error: modelsRes.error,
+				})
+				return
+			}
 			models = modelsRes.models || []
 			error = modelsRes.error || null
 		} else {
@@ -192,7 +203,12 @@ export async function refreshModels(
 
 	try {
 		const result = await listModels(message.connection)
-		if (!result.models) {
+		if (result.error) {
+			const res = {
+				error: result.error,
+			}
+			emitToUser("error", res)
+		} else if (!result.models) {
 			const res: Sockets.RefreshModels.Response = {
 				error: "Failed to refresh models.",
 				models: []

src/lib/server/utils/TokenCounterManager.ts

@@ -6,14 +6,19 @@ import llamaTokenizer from 'llama-tokenizer-js'
 import llama3Tokenizer from 'llama3-tokenizer-js'
 import mistralTokenizer from 'mistral-tokenizer-js'
 import { TokenCounterOptions } from "$lib/shared/constants/TokenCounters"
+import { fromPreTrained as getGemmaTokenizer } from "@lenml/tokenizer-gemma";
 
 export interface TokenCounter {
     countTokens(text: string): Promise<number> | number
 }
 
+/**
+ * Estimate should always be the most conservative option
+ * It should never overestimate tokens, but can underestimate
+ */
 export class EstimateTokenCounter implements TokenCounter {
     countTokens(text: string): number {
-        return Math.ceil(text.length / 4)
+        return Math.ceil(text.length / 3.4);
     }
 }
 
@@ -67,7 +72,9 @@ export class AnthropicClaudeTokenCounter implements TokenCounter {
 
 export class CohereTokenCounter implements TokenCounter {
     countTokens(text: string): number {
-        return Math.ceil(text.length / 5)
+        const tokenizer = getGemmaTokenizer()
+        return tokenizer.encode(text, {add_special_tokens: false}).length
+        //return Math.ceil(text.length / 5)
     }
 }
 
@@ -77,6 +84,14 @@ export class GeminiTokenCounter implements TokenCounter {
     }
 }
 
+export class GemmaTokenCounter implements TokenCounter {
+    countTokens(text: string): number {
+        // Gemma models tend to have slightly different tokenization
+        // Use a more conservative estimate that's closer to actual Gemma tokenization
+        return Math.ceil(text.length / 3.5) // More conservative than /4
+    }
+}
+
 export type TokenCounterDescriptor = {
     label: string
     counter: TokenCounter
@@ -128,6 +143,14 @@ export class TokenCounters {
             label: TokenCounterOptions.options.find(o => o.value === TokenCounterOptions.GEMINI)!.label,
             counter: new GeminiTokenCounter()
         }],
+        [TokenCounterOptions.GEMMA, {
+            label: TokenCounterOptions.options.find(o => o.value === TokenCounterOptions.GEMMA)!.label,
+            counter: new GemmaTokenCounter()
+        }],
+        [TokenCounterOptions.GEMMA, {
+            label: TokenCounterOptions.options.find(o => o.value === TokenCounterOptions.GEMMA)!.label,
+            counter: new GemmaTokenCounter()
+        }],
     ])
 
     private active: string

src/lib/shared/constants/ConnectionTypes.ts

@@ -19,7 +19,7 @@ const lmStudioDesc = `
 const lmStudioDiff = "Beginner (GUI) - Minimal setup required"
 
 const ollamaDesc = `
-<p>Serene Pub supports Ollama through their <a class="text-primary-500 hover:underline" href="https://github.com/ollama/ollama/blob/main/docs/api.md#generate-a-completion" target="_blank">"Generate a completion" API.</a></p>
+<p>Serene Pub supports Ollama through its <a class="text-primary-500 hover:underline" href="https://github.com/ollama/ollama/blob/main/docs/api.md" target="_blank">native API.</a></p>
 <p>It provides a simple API for generating completions and supports various model formats.</p>
 <p>To download Ollama, visit their <a class="text-primary-500 hover:underline" href="https://ollama.com/" target="_blank">official website</a>.</p>
 <p>Ollama is simple to setup and run, manages your models automatically, but requires minimal command line usage.</p>

src/lib/shared/constants/TokenCounters.ts

@@ -10,6 +10,7 @@ export class TokenCounterOptions {
     static readonly ANTHROPIC_CLAUDE = "anthropic-claude";
     static readonly COHERE = "cohere";
     static readonly GEMINI = "gemini";
+    static readonly GEMMA = "gemma";
 
     static readonly keys = [
         TokenCounterOptions.ESTIMATE,
@@ -22,7 +23,8 @@ export class TokenCounterOptions {
         TokenCounterOptions.MISTRAL,
         TokenCounterOptions.ANTHROPIC_CLAUDE,
         TokenCounterOptions.COHERE,
-        TokenCounterOptions.GEMINI
+        TokenCounterOptions.GEMINI,
+        TokenCounterOptions.GEMMA
     ];
 
     static readonly options = [
@@ -36,6 +38,7 @@ export class TokenCounterOptions {
         { value: TokenCounterOptions.MISTRAL, label: "Mistral/Mixtral" },
         { value: TokenCounterOptions.ANTHROPIC_CLAUDE, label: "Anthropic Claude" },
         { value: TokenCounterOptions.COHERE, label: "Cohere" },
-        { value: TokenCounterOptions.GEMINI, label: "Google Gemini/PaLM" }
+        { value: TokenCounterOptions.GEMINI, label: "Google Gemini/PaLM" },
+        { value: TokenCounterOptions.GEMMA, label: "Google Gemma" },
     ];
 }