@houtini/lm 2.3.0 → 2.7.0

package/dist/index.js

@@ -8,6 +8,7 @@
 import { Server } from '@modelcontextprotocol/sdk/server/index.js';
 import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
 import { CallToolRequestSchema, ListToolsRequestSchema, } from '@modelcontextprotocol/sdk/types.js';
+import { profileModelsAtStartup, getCachedProfile, toModelProfile as cachedToProfile, getHFEnrichmentLine, getPromptHints, } from './model-cache.js';
 const LM_BASE_URL = process.env.LM_STUDIO_URL || 'http://localhost:1234';
 const LM_MODEL = process.env.LM_STUDIO_MODEL || '';
 const LM_PASSWORD = process.env.LM_STUDIO_PASSWORD || '';
@@ -18,12 +19,269 @@ const INFERENCE_CONNECT_TIMEOUT_MS = 30_000; // generous connect timeout for inf
 const SOFT_TIMEOUT_MS = 55_000; // return partial results before MCP SDK ~60s timeout
 const READ_CHUNK_TIMEOUT_MS = 30_000; // max wait for a single SSE chunk
 const FALLBACK_CONTEXT_LENGTH = parseInt(process.env.LM_CONTEXT_WINDOW || '100000', 10);
+// ── Session-level token accounting ───────────────────────────────────
+// Tracks cumulative tokens offloaded to the local LLM across all calls
+// in this session. Shown in every response footer so Claude can reason
+// about cost savings and continue delegating strategically.
+const session = {
+    calls: 0,
+    promptTokens: 0,
+    completionTokens: 0,
+    /** Per-model performance tracking for routing insights */
+    modelStats: new Map(),
+};
+function recordUsage(resp) {
+    session.calls++;
+    if (resp.usage) {
+        session.promptTokens += resp.usage.prompt_tokens;
+        session.completionTokens += resp.usage.completion_tokens;
+    }
+    // Track per-model perf stats
+    if (resp.model) {
+        const existing = session.modelStats.get(resp.model) || { calls: 0, perfCalls: 0, totalTtftMs: 0, totalTokPerSec: 0 };
+        existing.calls++;
+        if (resp.ttftMs)
+            existing.totalTtftMs += resp.ttftMs;
+        const tokPerSec = resp.usage && resp.generationMs > 0
+            ? (resp.usage.completion_tokens / (resp.generationMs / 1000))
+            : 0;
+        if (tokPerSec > 0) {
+            existing.perfCalls++;
+            existing.totalTokPerSec += tokPerSec;
+        }
+        session.modelStats.set(resp.model, existing);
+    }
+}
+function sessionSummary() {
+    const total = session.promptTokens + session.completionTokens;
+    if (session.calls === 0)
+        return '';
+    return `Session: ${total.toLocaleString()} tokens offloaded across ${session.calls} call${session.calls === 1 ? '' : 's'}`;
+}
 function apiHeaders() {
     const h = { 'Content-Type': 'application/json' };
     if (LM_PASSWORD)
         h['Authorization'] = `Bearer ${LM_PASSWORD}`;
     return h;
 }
+const MODEL_PROFILES = [
+    {
+        pattern: /nemotron|nemotron_h_moe/i,
+        profile: {
+            family: 'NVIDIA Nemotron',
+            description: 'NVIDIA\'s compact reasoning model optimised for accurate, structured responses. Strong at step-by-step logic and instruction following.',
+            strengths: ['logical reasoning', 'math', 'step-by-step problem solving', 'code review', 'structured output'],
+            weaknesses: ['creative writing', 'constrained generation', 'factual knowledge on niche topics'],
+            bestFor: ['analysis tasks', 'code bug-finding', 'math/science questions', 'data transformation'],
+        },
+    },
+    {
+        pattern: /granite|granitehybrid/i,
+        profile: {
+            family: 'IBM Granite',
+            description: 'IBM\'s enterprise-focused model family. Compact and efficient, designed for business and code tasks with strong instruction following.',
+            strengths: ['code generation', 'instruction following', 'enterprise tasks', 'efficiency'],
+            weaknesses: ['creative tasks', 'long-form generation'],
+            bestFor: ['boilerplate generation', 'code explanation', 'structured Q&A'],
+        },
+    },
+    {
+        pattern: /qwen3-coder|qwen3.*coder/i,
+        profile: {
+            family: 'Qwen3 Coder',
+            description: 'Alibaba\'s code-specialised model with agentic capabilities. Excellent at code generation, review, and multi-step coding tasks.',
+            strengths: ['code generation', 'code review', 'debugging', 'test writing', 'refactoring', 'multi-step reasoning'],
+            weaknesses: ['non-code creative tasks'],
+            bestFor: ['code generation', 'code review', 'test stubs', 'type definitions', 'refactoring'],
+        },
+    },
+    {
+        pattern: /qwen3-vl|qwen.*vl/i,
+        profile: {
+            family: 'Qwen3 Vision-Language',
+            description: 'Alibaba\'s multimodal model handling both text and image inputs. Can analyse screenshots, diagrams, and visual content.',
+            strengths: ['image understanding', 'visual Q&A', 'diagram analysis', 'OCR'],
+            weaknesses: ['pure text tasks (use a text-only model instead)'],
+            bestFor: ['screenshot analysis', 'UI review', 'diagram interpretation'],
+        },
+    },
+    {
+        pattern: /qwen3(?!.*coder)(?!.*vl)/i,
+        profile: {
+            family: 'Qwen3',
+            description: 'Alibaba\'s general-purpose model with strong multilingual and reasoning capabilities. Good all-rounder.',
+            strengths: ['general reasoning', 'multilingual', 'code', 'instruction following'],
+            weaknesses: ['specialised code tasks (use Qwen3 Coder instead)'],
+            bestFor: ['general Q&A', 'translation', 'summarisation', 'brainstorming'],
+        },
+    },
+    {
+        pattern: /llama[- ]?3/i,
+        profile: {
+            family: 'Meta LLaMA 3',
+            description: 'Meta\'s open-weight general-purpose model. Strong baseline across tasks with large community fine-tune ecosystem.',
+            strengths: ['general reasoning', 'code', 'instruction following', 'broad knowledge'],
+            weaknesses: ['specialised tasks where fine-tuned models excel'],
+            bestFor: ['general delegation', 'drafting', 'code review', 'Q&A'],
+        },
+    },
+    {
+        pattern: /minimax[- ]?m2/i,
+        profile: {
+            family: 'MiniMax M2',
+            description: 'MiniMax\'s large MoE model with strong long-context and reasoning capabilities.',
+            strengths: ['long context', 'reasoning', 'creative writing', 'multilingual'],
+            weaknesses: ['may be slower due to model size'],
+            bestFor: ['long document analysis', 'creative tasks', 'complex reasoning'],
+        },
+    },
+    {
+        pattern: /kimi[- ]?k2/i,
+        profile: {
+            family: 'Kimi K2',
+            description: 'Moonshot AI\'s large MoE model with strong agentic and tool-use capabilities.',
+            strengths: ['agentic tasks', 'tool use', 'code', 'reasoning', 'long context'],
+            weaknesses: ['may be slower due to model size'],
+            bestFor: ['complex multi-step tasks', 'code generation', 'reasoning chains'],
+        },
+    },
+    {
+        pattern: /gpt-oss/i,
+        profile: {
+            family: 'OpenAI GPT-OSS',
+            description: 'OpenAI\'s open-source model release. General-purpose with strong instruction following.',
+            strengths: ['instruction following', 'general reasoning', 'code'],
+            weaknesses: ['less tested in open ecosystem than LLaMA/Qwen'],
+            bestFor: ['general delegation', 'code tasks', 'Q&A'],
+        },
+    },
+    {
+        pattern: /glm[- ]?4/i,
+        profile: {
+            family: 'GLM-4',
+            description: 'Zhipu AI\'s open-weight MoE model. Fast inference with strong general reasoning, multilingual support, and tool-use capabilities. Uses chain-of-thought reasoning internally. MIT licensed.',
+            strengths: ['fast inference', 'general reasoning', 'tool use', 'multilingual', 'code', 'instruction following', 'chain-of-thought'],
+            weaknesses: ['always emits internal reasoning (stripped automatically)', 'less tested in English-only benchmarks than LLaMA/Qwen'],
+            bestFor: ['general delegation', 'fast drafting', 'code tasks', 'structured output', 'Q&A'],
+        },
+    },
+    {
+        pattern: /nomic.*embed|embed.*nomic/i,
+        profile: {
+            family: 'Nomic Embed',
+            description: 'Text embedding model for semantic search and similarity. Not a chat model — produces vector embeddings.',
+            strengths: ['text embeddings', 'semantic search', 'clustering'],
+            weaknesses: ['cannot chat or generate text'],
+            bestFor: ['RAG pipelines', 'semantic similarity', 'document search'],
+        },
+    },
+    {
+        pattern: /abliterated/i,
+        profile: {
+            family: 'Abliterated (uncensored)',
+            description: 'Community fine-tune with safety guardrails removed. More permissive but may produce lower-quality or unreliable output.',
+            strengths: ['fewer refusals', 'unconstrained generation'],
+            weaknesses: ['may hallucinate more', 'no safety filtering', 'less tested'],
+            bestFor: ['tasks where the base model refuses unnecessarily'],
+        },
+    },
+];
+/**
+ * Match a model to its known profile.
+ * Priority: 1) static MODEL_PROFILES (curated), 2) SQLite cache (auto-generated from HF)
+ */
+function getModelProfile(model) {
+    // Try static profiles first (curated, most reliable)
+    for (const { pattern, profile } of MODEL_PROFILES) {
+        if (pattern.test(model.id))
+            return profile;
+    }
+    if (model.arch) {
+        for (const { pattern, profile } of MODEL_PROFILES) {
+            if (pattern.test(model.arch))
+                return profile;
+        }
+    }
+    return undefined;
+}
+/**
+ * Async version that also checks SQLite cache for auto-generated profiles.
+ * Use this when you need the most complete profile available.
+ */
+async function getModelProfileAsync(model) {
+    // Static profiles take priority
+    const staticProfile = getModelProfile(model);
+    if (staticProfile)
+        return staticProfile;
+    // Check SQLite cache for auto-generated profile
+    try {
+        const cached = await getCachedProfile(model.id);
+        if (cached) {
+            const profile = cachedToProfile(cached);
+            if (profile)
+                return profile;
+        }
+    }
+    catch {
+        // Cache lookup failed — fall through
+    }
+    return undefined;
+}
+/**
+ * Format a single model's full metadata for display.
+ * Async because it may fetch HuggingFace enrichment data.
+ */
+async function formatModelDetail(model, enrichWithHF = false) {
+    const ctx = getContextLength(model);
+    const maxCtx = getMaxContextLength(model);
+    // Use async profile lookup (checks static + SQLite cache)
+    const profile = await getModelProfileAsync(model);
+    const parts = [];
+    // Header line
+    parts.push(`  ${model.state === 'loaded' ? '●' : '○'} ${model.id}`);
+    // Metadata line
+    const meta = [];
+    if (model.type)
+        meta.push(`type: ${model.type}`);
+    if (model.arch)
+        meta.push(`arch: ${model.arch}`);
+    if (model.quantization)
+        meta.push(`quant: ${model.quantization}`);
+    if (model.compatibility_type)
+        meta.push(`format: ${model.compatibility_type}`);
+    // Show loaded context vs max context when both are available and different
+    if (model.loaded_context_length && maxCtx && model.loaded_context_length !== maxCtx) {
+        meta.push(`context: ${model.loaded_context_length.toLocaleString()} (max ${maxCtx.toLocaleString()})`);
+    }
+    else if (ctx) {
+        meta.push(`context: ${ctx.toLocaleString()}`);
+    }
+    if (model.publisher)
+        meta.push(`by: ${model.publisher}`);
+    if (meta.length > 0)
+        parts.push(`    ${meta.join(' · ')}`);
+    // Capabilities
+    if (model.capabilities && model.capabilities.length > 0) {
+        parts.push(`    Capabilities: ${model.capabilities.join(', ')}`);
+    }
+    // Profile info (static or auto-generated from SQLite cache)
+    if (profile) {
+        parts.push(`    ${profile.family}: ${profile.description}`);
+        parts.push(`    Best for: ${profile.bestFor.join(', ')}`);
+    }
+    // HuggingFace enrichment line from SQLite cache
+    if (enrichWithHF) {
+        try {
+            const hfLine = await getHFEnrichmentLine(model.id);
+            if (hfLine)
+                parts.push(hfLine);
+        }
+        catch {
+            // HF enrichment is best-effort — never block on failure
+        }
+    }
+    return parts.join('\n');
+}
 /**
  * Fetch with a connect timeout so Claude doesn't hang when the host is offline.
  */
@@ -67,10 +325,14 @@ async function chatCompletionStreaming(messages, options = {}) {
         temperature: options.temperature ?? DEFAULT_TEMPERATURE,
         max_tokens: options.maxTokens ?? DEFAULT_MAX_TOKENS,
         stream: true,
+        stream_options: { include_usage: true },
     };
     if (options.model || LM_MODEL) {
         body.model = options.model || LM_MODEL;
     }
+    if (options.responseFormat) {
+        body.response_format = options.responseFormat;
+    }
     const startTime = Date.now();
     const res = await fetchWithTimeout(`${LM_BASE_URL}/v1/chat/completions`, { method: 'POST', headers: apiHeaders(), body: JSON.stringify(body) }, INFERENCE_CONNECT_TIMEOUT_MS);
     if (!res.ok) {
@@ -88,6 +350,7 @@ async function chatCompletionStreaming(messages, options = {}) {
     let finishReason = '';
     let truncated = false;
     let buffer = '';
+    let ttftMs;
     try {
         while (true) {
             // Check soft timeout before each read
@@ -123,8 +386,11 @@ async function chatCompletionStreaming(messages, options = {}) {
                     if (json.model)
                         model = json.model;
                     const delta = json.choices?.[0]?.delta;
-                    if (delta?.content)
+                    if (delta?.content) {
+                        if (ttftMs === undefined)
+                            ttftMs = Date.now() - startTime;
                         content += delta.content;
+                    }
                     const reason = json.choices?.[0]?.finish_reason;
                     if (reason)
                         finishReason = reason;
@@ -142,9 +408,33 @@ async function chatCompletionStreaming(messages, options = {}) {
         // Release the reader — don't await cancel() as it can hang
         reader.releaseLock();
     }
-    return { content, model, usage, finishReason, truncated };
+    const generationMs = Date.now() - startTime;
+    // Strip <think>...</think> reasoning blocks from models that always emit them
+    // (e.g. GLM Flash, Nemotron). Claude doesn't need the model's internal reasoning.
+    // Handle both closed blocks and unclosed ones (model ran out of tokens mid-think,
+    // or grammar-constrained output forced content before the closing tag).
+    let cleanContent = content.replace(/<think>[\s\S]*?<\/think>\s*/g, ''); // closed blocks
+    cleanContent = cleanContent.replace(/^<think>\s*/, ''); // orphaned opening tag
+    cleanContent = cleanContent.trim();
+    return { content: cleanContent, model, usage, finishReason, truncated, ttftMs, generationMs };
 }
+/**
+ * Fetch models from LM Studio's native v0 API first (richer metadata),
+ * falling back to the OpenAI-compatible v1 endpoint for non-LM-Studio hosts.
+ */
 async function listModelsRaw() {
+    // Try v0 API first — returns type, arch, publisher, quantization, state
+    try {
+        const v0 = await fetchWithTimeout(`${LM_BASE_URL}/api/v0/models`, { headers: apiHeaders() });
+        if (v0.ok) {
+            const data = (await v0.json());
+            return data.data;
+        }
+    }
+    catch {
+        // v0 not available — fall through to v1
+    }
+    // Fallback: OpenAI-compatible v1 endpoint (works with Ollama, vLLM, llama.cpp)
     const res = await fetchWithTimeout(`${LM_BASE_URL}/v1/models`, { headers: apiHeaders() });
     if (!res.ok)
         throw new Error(`Failed to list models: ${res.status}`);
@@ -152,22 +442,101 @@ async function listModelsRaw() {
     return data.data;
 }
 function getContextLength(model) {
-    // LM Studio uses context_length, vLLM uses max_model_len, fall back to env/100k
-    return model.context_length ?? model.max_model_len ?? FALLBACK_CONTEXT_LENGTH;
+    // Prefer loaded_context_length (actual configured context) over max_context_length (theoretical max)
+    // v0 API: loaded_context_length / max_context_length, v1: context_length, vLLM: max_model_len
+    return model.loaded_context_length ?? model.max_context_length ?? model.context_length ?? model.max_model_len ?? FALLBACK_CONTEXT_LENGTH;
+}
+function getMaxContextLength(model) {
+    return model.max_context_length;
+}
+async function routeToModel(taskType) {
+    let models;
+    try {
+        models = await listModelsRaw();
+    }
+    catch {
+        // Can't reach server — fall back to default
+        const hints = getPromptHints(LM_MODEL);
+        return { modelId: LM_MODEL || '', hints };
+    }
+    const loaded = models.filter((m) => m.state === 'loaded' || !m.state);
+    const available = models.filter((m) => m.state === 'not-loaded');
+    if (loaded.length === 0) {
+        const hints = getPromptHints(LM_MODEL);
+        return { modelId: LM_MODEL || '', hints };
+    }
+    // Score each loaded model for the requested task type
+    let bestModel = loaded[0];
+    let bestScore = -1;
+    for (const model of loaded) {
+        const hints = getPromptHints(model.id, model.arch);
+        // Primary: is this task type in the model's best types?
+        let score = (hints.bestTaskTypes ?? []).includes(taskType) ? 10 : 0;
+        // Bonus: code-specialised models get extra points for code tasks
+        const profile = getModelProfile(model);
+        if (taskType === 'code' && profile?.family.toLowerCase().includes('coder'))
+            score += 5;
+        // Bonus: larger context for analysis tasks
+        if (taskType === 'analysis') {
+            const ctx = getContextLength(model);
+            if (ctx && ctx > 100000)
+                score += 2;
+        }
+        if (score > bestScore) {
+            bestScore = score;
+            bestModel = model;
+        }
+    }
+    const hints = getPromptHints(bestModel.id, bestModel.arch);
+    const result = { modelId: bestModel.id, hints };
+    // If the best loaded model isn't ideal for this task, suggest a better available one.
+    // We don't JIT-load because model loading takes minutes and the MCP SDK has a ~60s
+    // hard timeout. Instead, suggest the user loads the better model in LM Studio.
+    if (!(hints.bestTaskTypes ?? []).includes(taskType)) {
+        const better = available.find((m) => {
+            const mHints = getPromptHints(m.id, m.arch);
+            return (mHints.bestTaskTypes ?? []).includes(taskType);
+        });
+        if (better) {
+            const label = taskType === 'code' ? 'code tasks'
+                : taskType === 'analysis' ? 'analysis'
+                    : taskType === 'embedding' ? 'embeddings'
+                        : 'this kind of task';
+            result.suggestion = `💡 ${better.id} is downloaded and better suited for ${label} — ask the user to load it in LM Studio.`;
+        }
+    }
+    return result;
 }
 /**
  * Format a footer line for streaming results showing model, usage, and truncation status.
  */
 function formatFooter(resp, extra) {
+    // Record usage for session tracking before formatting
+    recordUsage(resp);
     const parts = [];
     if (resp.model)
         parts.push(`Model: ${resp.model}`);
     if (resp.usage)
-        parts.push(`Tokens: ${resp.usage.prompt_tokens}→${resp.usage.completion_tokens}`);
+        parts.push(`${resp.usage.prompt_tokens}→${resp.usage.completion_tokens} tokens`);
+    // Perf stats — computed from streaming, no proprietary API needed
+    const perfParts = [];
+    if (resp.ttftMs !== undefined)
+        perfParts.push(`TTFT: ${resp.ttftMs}ms`);
+    if (resp.usage && resp.generationMs > 0) {
+        const tokPerSec = resp.usage.completion_tokens / (resp.generationMs / 1000);
+        perfParts.push(`${tokPerSec.toFixed(1)} tok/s`);
+    }
+    if (resp.generationMs)
+        perfParts.push(`${(resp.generationMs / 1000).toFixed(1)}s`);
+    if (perfParts.length > 0)
+        parts.push(perfParts.join(', '));
     if (extra)
         parts.push(extra);
     if (resp.truncated)
         parts.push('⚠ TRUNCATED (soft timeout — partial result)');
+    const sessionLine = sessionSummary();
+    if (sessionLine)
+        parts.push(sessionLine);
     return parts.length > 0 ? `\n\n---\n${parts.join(' | ')}` : '';
 }
 // ── MCP Tool definitions ─────────────────────────────────────────────
@@ -177,6 +546,9 @@ const TOOLS = [
         description: 'Send a task to a local LLM running on a separate machine. This is a FREE, parallel worker — ' +
             'use it to offload bounded work while you continue doing other things. The local LLM runs independently ' +
             'and does not consume your tokens or rate limits.\n\n' +
+            'PLANNING: When you start a large task (refactoring, migrations, test suites, documentation), ' +
+            'break it into steps and identify which ones are bounded grunt work you can delegate here. ' +
+            'The more you offload, the more tokens you save. The session footer tracks cumulative savings.\n\n' +
             'WHEN TO USE (delegate generously — it costs nothing):\n' +
             '• Explain or summarise code/docs you just read\n' +
             '• Generate boilerplate, test stubs, type definitions, mock data\n' +
@@ -185,11 +557,17 @@ const TOOLS = [
             '• Translate or reformat content (JSON↔YAML, snake_case↔camelCase)\n' +
             '• Brainstorm approaches before you commit to one\n' +
             '• Any self-contained subtask that does not need tool access\n\n' +
-            'RULES:\n' +
+            'PROMPT QUALITY (the local model is highly capable — results depend on your prompt):\n' +
             '(1) Always send COMPLETE code/context — never truncate, the local LLM cannot access files.\n' +
             '(2) Be explicit about output format ("respond as a JSON array", "return only the function").\n' +
-            '(3) Call discover first if you are unsure whether the local LLM is online.\n\n' +
-            'The local model, context window, and speed vary — call the discover tool to check what is loaded.',
+            '(3) Set a specific persona in the system field — "Senior TypeScript dev" beats "helpful assistant".\n' +
+            '(4) State constraints: "no preamble", "reference line numbers", "max 5 bullet points".\n' +
+            '(5) For code generation, include the surrounding context (imports, types, function signatures).\n\n' +
+            'QA: Always review the local LLM\'s output before using it. Verify correctness, check edge cases, ' +
+            'and fix any issues. You are the architect — the local model is a fast drafter, not the final authority.\n\n' +
+            'ROUTING: If multiple models are loaded, houtini-lm automatically picks the best one for the task. ' +
+            'If a better model is downloaded but not loaded, you\'ll see a suggestion in the response footer. ' +
+            'Call discover to see what\'s available.',
         inputSchema: {
             type: 'object',
             properties: {
@@ -209,6 +587,10 @@ const TOOLS = [
                     type: 'number',
                     description: 'Max response tokens. Default 2048. Use higher for code generation, lower for quick answers.',
                 },
+                json_schema: {
+                    type: 'object',
+                    description: 'Force structured JSON output. Provide a JSON Schema object and the response will be guaranteed valid JSON conforming to it. Example: {"name":"result","schema":{"type":"object","properties":{"answer":{"type":"string"}},"required":["answer"]}}',
+                },
             },
             required: ['message'],
         },
@@ -217,14 +599,19 @@ const TOOLS = [
         name: 'custom_prompt',
         description: 'Structured analysis via the local LLM with explicit system/context/instruction separation. ' +
             'This 3-part format prevents context bleed and gets the best results from local models.\n\n' +
+            'USE THIS for complex tasks where prompt structure matters — it consistently outperforms ' +
+            'stuffing everything into a single message. The separation helps the local model focus.\n\n' +
             'WHEN TO USE:\n' +
             '• Code review — paste full source, ask for bugs/improvements\n' +
             '• Comparison — paste two implementations, ask which is better and why\n' +
             '• Refactoring suggestions — paste code, ask for a cleaner version\n' +
             '• Content analysis — paste text, ask for structure/tone/issues\n' +
             '• Any task where separating context from instruction improves clarity\n\n' +
-            'System sets persona. Context provides COMPLETE data (never truncate). ' +
-            'Instruction states exactly what to produce.',
+            'PROMPT STRUCTURE (each field has a job — keep them focused):\n' +
+            '• System: persona + constraints, under 30 words. "Expert Python developer focused on performance and correctness."\n' +
+            '• Context: COMPLETE data. Full source code, full logs, full text. NEVER truncate or summarise.\n' +
+            '• Instruction: exactly what to produce, under 50 words. Specify format: "Return a JSON array of {line, issue, fix}."\n\n' +
+            'QA: Review the output. The local model is a capable drafter — verify its analysis before acting on it.',
         inputSchema: {
             type: 'object',
             properties: {
@@ -248,6 +635,10 @@ const TOOLS = [
                     type: 'number',
                     description: 'Max response tokens. Default 2048.',
                 },
+                json_schema: {
+                    type: 'object',
+                    description: 'Force structured JSON output. Provide a JSON Schema object and the response will be guaranteed valid JSON conforming to it.',
+                },
             },
             required: ['instruction'],
         },
@@ -255,13 +646,20 @@ const TOOLS = [
     {
         name: 'code_task',
         description: 'Send a code analysis task to the local LLM. Wraps the request with an optimised code-review system prompt.\n\n' +
+            'This is the fastest way to offload code-specific work. Temperature is locked to 0.2 for ' +
+            'focused, deterministic output. The system prompt is pre-configured for code review.\n\n' +
             'WHEN TO USE:\n' +
             '• Explain what a function/class does\n' +
             '• Find bugs or suggest improvements\n' +
             '• Generate unit tests or type definitions for existing code\n' +
             '• Add error handling, logging, or validation\n' +
             '• Convert between languages or patterns\n\n' +
-            'Provide COMPLETE source code (the local LLM cannot read files) and a specific task.',
+            'GETTING BEST RESULTS:\n' +
+            '• Provide COMPLETE source code — the local LLM cannot read files.\n' +
+            '• Include imports and type definitions so the model has full context.\n' +
+            '• Be specific in the task: "Write 3 Jest tests for the error paths in fetchUser" beats "Write tests".\n' +
+            '• Set the language field — it shapes the system prompt and improves accuracy.\n\n' +
+            'QA: Always verify generated code compiles, handles edge cases, and follows project conventions.',
         inputSchema: {
             type: 'object',
             properties: {
@@ -288,49 +686,91 @@ const TOOLS = [
     {
         name: 'discover',
         description: 'Check whether the local LLM is online and what model is loaded. Returns model name, context window size, ' +
-            'and response latency. Call this if you are unsure whether the local LLM is available before delegating work. ' +
+            'response latency, and cumulative session stats (tokens offloaded so far). ' +
+            'Call this if you are unsure whether the local LLM is available before delegating work. ' +
             'Fast — typically responds in under 1 second, or returns an offline status within 5 seconds if the host is unreachable.',
         inputSchema: { type: 'object', properties: {} },
     },
     {
         name: 'list_models',
-        description: 'List all models currently loaded in the local LLM server, with context window sizes. ' +
-            'Use discover instead for a quick availability check.',
+        description: 'List all models on the local LLM server — both loaded (ready) and available (downloaded but not active). ' +
+            'Shows rich metadata for each model: type (llm/vlm/embeddings), architecture, quantization, context window, ' +
+            'and a capability profile describing what the model is best at. ' +
+            'Use this to understand which models are available and suggest switching when a different model would suit the task better.',
         inputSchema: { type: 'object', properties: {} },
     },
+    {
+        name: 'embed',
+        description: 'Generate text embeddings via the local LLM server. Requires an embedding model to be loaded ' +
+            '(e.g. Nomic Embed). Returns a vector representation of the input text for semantic search, ' +
+            'similarity comparison, or RAG pipelines. Uses the OpenAI-compatible /v1/embeddings endpoint.',
+        inputSchema: {
+            type: 'object',
+            properties: {
+                input: {
+                    type: 'string',
+                    description: 'The text to embed. Can be a single string.',
+                },
+                model: {
+                    type: 'string',
+                    description: 'Embedding model ID. If omitted, uses whatever embedding model is loaded.',
+                },
+            },
+            required: ['input'],
+        },
+    },
 ];
 // ── MCP Server ───────────────────────────────────────────────────────
-const server = new Server({ name: 'houtini-lm', version: '2.3.0' }, { capabilities: { tools: {} } });
+const server = new Server({ name: 'houtini-lm', version: '2.7.0' }, { capabilities: { tools: {} } });
 server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: TOOLS }));
 server.setRequestHandler(CallToolRequestSchema, async (request) => {
     const { name, arguments: args } = request.params;
     try {
         switch (name) {
             case 'chat': {
-                const { message, system, temperature, max_tokens } = args;
+                const { message, system, temperature, max_tokens, json_schema } = args;
+                const route = await routeToModel('chat');
                 const messages = [];
-                if (system)
-                    messages.push({ role: 'system', content: system });
+                // Inject output constraint into system prompt if the model needs it
+                const systemContent = system
+                    ? (route.hints.outputConstraint ? `${system}\n\n${route.hints.outputConstraint}` : system)
+                    : (route.hints.outputConstraint || undefined);
+                if (systemContent)
+                    messages.push({ role: 'system', content: systemContent });
                 messages.push({ role: 'user', content: message });
+                const responseFormat = json_schema
+                    ? { type: 'json_schema', json_schema: { name: json_schema.name, strict: json_schema.strict ?? true, schema: json_schema.schema } }
+                    : undefined;
                 const resp = await chatCompletionStreaming(messages, {
-                    temperature,
+                    temperature: temperature ?? route.hints.chatTemp,
                     maxTokens: max_tokens,
+                    model: route.modelId,
+                    responseFormat,
                 });
                 const footer = formatFooter(resp);
                 return { content: [{ type: 'text', text: resp.content + footer }] };
             }
             case 'custom_prompt': {
-                const { system, context, instruction, temperature, max_tokens } = args;
+                const { system, context, instruction, temperature, max_tokens, json_schema } = args;
+                const route = await routeToModel('analysis');
                 const messages = [];
-                if (system)
-                    messages.push({ role: 'system', content: system });
+                const systemContent = system
+                    ? (route.hints.outputConstraint ? `${system}\n\n${route.hints.outputConstraint}` : system)
+                    : (route.hints.outputConstraint || undefined);
+                if (systemContent)
+                    messages.push({ role: 'system', content: systemContent });
                 let userContent = instruction;
                 if (context)
                     userContent = `Context:\n${context}\n\nInstruction:\n${instruction}`;
                 messages.push({ role: 'user', content: userContent });
+                const responseFormat = json_schema
+                    ? { type: 'json_schema', json_schema: { name: json_schema.name, strict: json_schema.strict ?? true, schema: json_schema.schema } }
+                    : undefined;
                 const resp = await chatCompletionStreaming(messages, {
-                    temperature,
+                    temperature: temperature ?? route.hints.chatTemp,
                     maxTokens: max_tokens,
+                    model: route.modelId,
+                    responseFormat,
                 });
                 const footer = formatFooter(resp);
                 return {
@@ -340,10 +780,14 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
             case 'code_task': {
                 const { code, task, language, max_tokens: codeMaxTokens } = args;
                 const lang = language || 'unknown';
+                const route = await routeToModel('code');
+                const outputConstraint = route.hints.outputConstraint
+                    ? ` ${route.hints.outputConstraint}`
+                    : '';
                 const codeMessages = [
                     {
                         role: 'system',
-                        content: `Expert ${lang} developer. Analyse the provided code and complete the task. Be specific — reference line numbers, function names, and concrete fixes. No preamble.`,
+                        content: `Expert ${lang} developer. Analyse the provided code and complete the task. Be specific — reference line numbers, function names, and concrete fixes.${outputConstraint}`,
                     },
                     {
                         role: 'user',
@@ -351,11 +795,13 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
                     },
                 ];
                 const codeResp = await chatCompletionStreaming(codeMessages, {
-                    temperature: 0.2,
+                    temperature: route.hints.codeTemp,
                     maxTokens: codeMaxTokens ?? DEFAULT_MAX_TOKENS,
+                    model: route.modelId,
                 });
                 const codeFooter = formatFooter(codeResp, lang);
-                return { content: [{ type: 'text', text: codeResp.content + codeFooter }] };
+                const suggestionLine = route.suggestion ? `\n${route.suggestion}` : '';
+                return { content: [{ type: 'text', text: codeResp.content + codeFooter + suggestionLine }] };
             }
             case 'discover': {
                 const start = Date.now();
@@ -384,36 +830,98 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
                             }],
                     };
                 }
-                const lines = models.map((m) => {
-                    const ctx = getContextLength(m);
-                    return `  • ${m.id} (context: ${ctx.toLocaleString()} tokens)`;
-                });
-                const primary = models[0];
+                const loaded = models.filter((m) => m.state === 'loaded' || !m.state);
+                const available = models.filter((m) => m.state === 'not-loaded');
+                const primary = loaded[0] || models[0];
                 const ctx = getContextLength(primary);
-                return {
-                    content: [{
-                            type: 'text',
-                            text: `Status: ONLINE\n` +
-                                `Endpoint: ${LM_BASE_URL}\n` +
-                                `Latency: ${ms}ms\n` +
-                                `Model: ${primary.id}\n` +
-                                `Context window: ${ctx.toLocaleString()} tokens\n` +
-                                `\nLoaded models:\n${lines.join('\n')}\n\n` +
-                                `The local LLM is available. You can delegate tasks using chat, custom_prompt, or code_task.`,
-                        }],
-                };
+                const primaryProfile = await getModelProfileAsync(primary);
+                const sessionStats = session.calls > 0
+                    ? `\nSession stats: ${(session.promptTokens + session.completionTokens).toLocaleString()} tokens offloaded across ${session.calls} call${session.calls === 1 ? '' : 's'}`
+                    : '\nSession stats: no calls yet — delegate tasks to start saving tokens';
+                let text = `Status: ONLINE\n` +
+                    `Endpoint: ${LM_BASE_URL}\n` +
+                    `Latency: ${ms}ms\n` +
+                    `Active model: ${primary.id}\n` +
+                    `Context window: ${ctx.toLocaleString()} tokens\n`;
+                if (primaryProfile) {
+                    text += `Family: ${primaryProfile.family}\n`;
+                    text += `Description: ${primaryProfile.description}\n`;
+                    text += `Best for: ${primaryProfile.bestFor.join(', ')}\n`;
+                    text += `Strengths: ${primaryProfile.strengths.join(', ')}\n`;
+                    if (primaryProfile.weaknesses.length > 0) {
+                        text += `Weaknesses: ${primaryProfile.weaknesses.join(', ')}\n`;
+                    }
+                }
+                if (loaded.length > 0) {
+                    text += `\nLoaded models (● ready to use):\n`;
+                    text += (await Promise.all(loaded.map((m) => formatModelDetail(m)))).join('\n\n');
+                }
+                if (available.length > 0) {
+                    text += `\n\nAvailable models (○ downloaded, not loaded — can be activated in LM Studio):\n`;
+                    text += (await Promise.all(available.map((m) => formatModelDetail(m)))).join('\n\n');
+                }
+                // Per-model performance stats from this session
+                if (session.modelStats.size > 0) {
+                    text += `\n\nPerformance (this session):\n`;
+                    for (const [modelId, stats] of session.modelStats) {
+                        const avgTtft = stats.calls > 0 ? Math.round(stats.totalTtftMs / stats.calls) : 0;
+                        const avgTokSec = stats.perfCalls > 0 ? (stats.totalTokPerSec / stats.perfCalls).toFixed(1) : '?';
+                        text += `  ${modelId}: ${stats.calls} calls, avg TTFT ${avgTtft}ms, avg ${avgTokSec} tok/s\n`;
+                    }
+                }
+                text += `${sessionStats}\n\n`;
+                text += `The local LLM is available. You can delegate tasks using chat, custom_prompt, code_task, or embed.`;
+                return { content: [{ type: 'text', text }] };
             }
             case 'list_models': {
                 const models = await listModelsRaw();
                 if (!models.length) {
-                    return { content: [{ type: 'text', text: 'No models currently loaded.' }] };
+                    return { content: [{ type: 'text', text: 'No models currently loaded or available.' }] };
                 }
-                const lines = models.map((m) => {
-                    const ctx = getContextLength(m);
-                    return `  • ${m.id}${ctx ? ` (context: ${ctx.toLocaleString()} tokens)` : ''}`;
-                });
+                const loaded = models.filter((m) => m.state === 'loaded' || !m.state);
+                const available = models.filter((m) => m.state === 'not-loaded');
+                let text = '';
+                // list_models enriches with HuggingFace data (cached after first call)
+                if (loaded.length > 0) {
+                    text += `Loaded models (● ready to use):\n\n`;
+                    text += (await Promise.all(loaded.map((m) => formatModelDetail(m, true)))).join('\n\n');
+                }
+                if (available.length > 0) {
+                    if (text)
+                        text += '\n\n';
+                    text += `Available models (○ downloaded, not loaded):\n\n`;
+                    text += (await Promise.all(available.map((m) => formatModelDetail(m, true)))).join('\n\n');
+                }
+                return { content: [{ type: 'text', text }] };
+            }
+            case 'embed': {
+                const { input, model: embedModel } = args;
+                const embedBody = { input };
+                if (embedModel) {
+                    embedBody.model = embedModel;
+                }
+                const res = await fetchWithTimeout(`${LM_BASE_URL}/v1/embeddings`, { method: 'POST', headers: apiHeaders(), body: JSON.stringify(embedBody) }, INFERENCE_CONNECT_TIMEOUT_MS);
+                if (!res.ok) {
+                    const errText = await res.text().catch(() => '');
+                    throw new Error(`Embeddings API error ${res.status}: ${errText}`);
+                }
+                const data = (await res.json());
+                const embedding = data.data[0]?.embedding;
+                if (!embedding)
+                    throw new Error('No embedding returned');
+                const usageInfo = data.usage
+                    ? `${data.usage.prompt_tokens} tokens embedded`
+                    : '';
                 return {
-                    content: [{ type: 'text', text: `Loaded models:\n${lines.join('\n')}` }],
+                    content: [{
+                            type: 'text',
+                            text: JSON.stringify({
+                                model: data.model,
+                                dimensions: embedding.length,
+                                embedding,
+                                usage: usageInfo,
+                            }),
+                        }],
                 };
             }
             default:
@@ -431,6 +939,11 @@ async function main() {
     const transport = new StdioServerTransport();
     await server.connect(transport);
     process.stderr.write(`Houtini LM server running (${LM_BASE_URL})\n`);
+    // Background: profile all available models via HF → SQLite cache
+    // Non-blocking — server is already accepting requests
+    listModelsRaw()
+        .then((models) => profileModelsAtStartup(models))
+        .catch((err) => process.stderr.write(`[houtini-lm] Startup profiling skipped: ${err}\n`));
 }
 main().catch((error) => {
     process.stderr.write(`Fatal error: ${error}\n`);