Merge branch 'main' into feat/watsonx_reranker_integration

Author: sestinj

.github/workflows/pr_checks.yaml

@@ -303,7 +303,13 @@ jobs:
           cd extensions/vscode
           npm ci
           npm run e2e:compile
-          FILES=$(ls -1 e2e/_output/tests/*.test.js | jq -R . | jq -s .)
+          if [[ "${{ github.event.pull_request.head.repo.fork }}" == "true" || "${{ github.actor }}" == "dependabot[bot]" ]]; then
+            # Exclude SSH tests for forks
+            FILES=$(ls -1 e2e/_output/tests/*.test.js | grep -v "SSH" | jq -R . | jq -s .)
+          else
+            # Include all tests for non-forks
+            FILES=$(ls -1 e2e/_output/tests/*.test.js | jq -R . | jq -s .)
+          fi
           echo "test_file_matrix<<EOF" >> $GITHUB_OUTPUT
           echo "$FILES" >> $GITHUB_OUTPUT
           echo "EOF" >> $GITHUB_OUTPUT

core/config/load.ts

@@ -272,7 +272,7 @@ async function intermediateToFinalConfig(
                 {
                   ...desc,
                   model: modelName,
-                  title: `${llm.title} - ${modelName}`,
+                  title: modelName,
                 },
                 ide.readFile.bind(ide),
                 uniqueId,

core/config/yaml/models.ts

@@ -127,7 +127,7 @@ async function autodetectModels(
           {
             ...model,
             model: modelName,
-            name: `${llm.title} - ${modelName}`,
+            name: modelName,
           },
           uniqueId,
           ideSettings,

core/index.d.ts

@@ -163,6 +163,10 @@ export interface ILLM extends LLMOptions {
   ): string | ChatMessage[];
 }
 
+export interface ModelInstaller {
+  installModel(modelName: string, signal: AbortSignal, progressReporter?: (task: string, increment: number, total: number) => void): Promise<any>;
+}
+
 export type ContextProviderType = "normal" | "query" | "submenu";
 
 export interface ContextProviderDescription {
@@ -543,7 +547,7 @@ export interface CustomLLMWithOptionals {
     signal: AbortSignal,
     options: CompletionOptions,
     fetch: (input: RequestInfo | URL, init?: RequestInit) => Promise<Response>,
-  ) => AsyncGenerator<string>;
+  ) => AsyncGenerator<ChatMessage | string>;
   listModels?: (
     fetch: (input: RequestInfo | URL, init?: RequestInit) => Promise<Response>,
   ) => Promise<string[]>;

core/llm/index.ts

@@ -19,6 +19,7 @@ import {
   LLMFullCompletionOptions,
   LLMOptions,
   ModelCapability,
+  ModelInstaller,
   PromptLog,
   PromptTemplate,
   RequestOptions,
@@ -58,6 +59,17 @@ import {
   toFimBody,
 } from "./openaiTypeConverters.js";
 
+
+export class LLMError extends Error {
+  constructor(message: string, public llm: ILLM) {
+    super(message);
+  }
+}
+
+export function isModelInstaller(provider: any): provider is ModelInstaller {
+  return provider && typeof provider.installModel === 'function';
+}
+
 export abstract class BaseLLM implements ILLM {
   static providerName: string;
   static defaultOptions: Partial<LLMOptions> | undefined = undefined;
@@ -380,9 +392,11 @@ export abstract class BaseLLM implements ILLM {
         if (!resp.ok) {
           let text = await resp.text();
           if (resp.status === 404 && !resp.url.includes("/v1")) {
-            if (text.includes("try pulling it first")) {
-              const model = JSON.parse(text).error.split(" ")[1].slice(1, -1);
+            const error = JSON.parse(text)?.error?.replace(/"/g, "'");
+            let model = error?.match(/model '(.*)' not found/)?.[1];
+            if (model && resp.url.match("127.0.0.1:11434")) {
               text = `The model "${model}" was not found. To download it, run \`ollama run ${model}\`.`;
+              throw new LLMError(text, this);// No need to add HTTP status details
             } else if (text.includes("/api/chat")) {
               text =
                 "The /api/chat endpoint was not found. This may mean that you are using an older version of Ollama that does not support /api/chat. Upgrading to the latest version will solve the issue.";
@@ -442,6 +456,10 @@ export abstract class BaseLLM implements ILLM {
             throw new Error(message);
           }
         }
+        //if e instance of LLMError, rethrow
+        if (e instanceof LLMError) {
+          throw e;
+        }
         throw new Error(e.message);
       }
     };
@@ -763,6 +781,7 @@ export abstract class BaseLLM implements ILLM {
     }
 
     let completion = "";
+    let citations: null | string[] = null
 
     try {
       if (this.templateMessages) {
@@ -790,6 +809,8 @@ export abstract class BaseLLM implements ILLM {
             completion = renderChatMessage(msg);
           } else {
             // Stream true
+            console.log("Streaming");
+
             const stream = this.openaiAdapter.chatCompletionStream(
               {
                 ...body,
@@ -802,6 +823,9 @@ export abstract class BaseLLM implements ILLM {
               if (result) {
                 yield result;
               }
+              if (!citations && (chunk as any).citations && Array.isArray((chunk as any).citations)) {
+                citations = (chunk as any).citations;
+              }
             }
           }
         } else {
@@ -824,6 +848,10 @@ export abstract class BaseLLM implements ILLM {
 
     if (logEnabled && this.writeLog) {
       await this.writeLog(`Completion:\n${completion}\n\n`);
+
+      if (citations) {
+        await this.writeLog(`Citations:\n${citations.map((c, i) => `${i + 1}: ${c}`).join("\n")}\n\n`);
+      }
     }
 
     return {

core/llm/llms/CustomLLM.ts

@@ -1,4 +1,5 @@
 import { ChatMessage, CompletionOptions, CustomLLM } from "../../index.js";
+import { renderChatMessage } from "../../util/messageContent.js";
 import { BaseLLM } from "../index.js";
 
 class CustomLLMClass extends BaseLLM {
@@ -18,7 +19,7 @@ class CustomLLMClass extends BaseLLM {
     signal: AbortSignal,
     options: CompletionOptions,
     fetch: (input: RequestInfo | URL, init?: RequestInit) => Promise<Response>,
-  ) => AsyncGenerator<string>;
+  ) => AsyncGenerator<ChatMessage | string>;
 
   constructor(custom: CustomLLM) {
     super(custom.options || { model: "custom" });
@@ -38,7 +39,11 @@ class CustomLLMClass extends BaseLLM {
         options,
         (...args) => this.fetch(...args),
       )) {
-        yield { role: "assistant", content };
+        if (typeof content === "string") {
+          yield { role: "assistant", content };
+        } else {
+          yield content;
+        }
       }
     } else {
       for await (const update of super._streamChat(messages, signal, options)) {
@@ -68,7 +73,11 @@ class CustomLLMClass extends BaseLLM {
         options,
         (...args) => this.fetch(...args),
       )) {
-        yield content;
+        if (typeof content === "string") {
+          yield content;
+        } else {
+          yield renderChatMessage(content);
+        }
       }
     } else {
       throw new Error(

core/llm/llms/Ollama.ts

@@ -1,9 +1,10 @@
 import { JSONSchema7, JSONSchema7Object } from "json-schema";
 
-import { ChatMessage, CompletionOptions, LLMOptions } from "../../index.js";
+import { ChatMessage, CompletionOptions, ModelInstaller, LLMOptions } from "../../index.js";
 import { renderChatMessage } from "../../util/messageContent.js";
 import { BaseLLM } from "../index.js";
 import { streamResponse } from "../stream.js";
+import { getRemoteModelInfo } from "../../util/ollamaHelper.js";
 
 type OllamaChatMessage = {
   role: "tool" | "user" | "assistant" | "system";
@@ -123,7 +124,7 @@ interface OllamaTool {
   };
 }
 
-class Ollama extends BaseLLM {
+class Ollama extends BaseLLM implements ModelInstaller{
   static providerName = "ollama";
   static defaultOptions: Partial<LLMOptions> = {
     apiBase: "http://localhost:11434/",
@@ -574,6 +575,38 @@ class Ollama extends BaseLLM {
     }
     return embedding;
   }
+
+  public async installModel(modelName: string, signal: AbortSignal, progressReporter?: (task: string, increment: number, total: number) => void): Promise<any> {
+      const modelInfo = await getRemoteModelInfo(modelName, signal);
+      if (!modelInfo) {
+          throw new Error(`'${modelName}' not found in the Ollama registry!`);
+      }
+      const response = await fetch(this.getEndpoint("api/pull"), {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          Authorization: `Bearer ${this.apiKey}`,
+        },
+        body: JSON.stringify({ name: modelName }),
+        signal
+      });
+
+      const reader = response.body?.getReader();
+      //TODO: generate proper progress based on modelInfo size
+      while (true) {
+        const { done, value } = await reader?.read() || { done: true, value: undefined };
+        if (done) {
+          break;
+        }
+
+        const chunk = new TextDecoder().decode(value);
+        const lines = chunk.split('\n').filter(Boolean);
+        for (const line of lines) {
+          const data = JSON.parse(line);
+          progressReporter?.(data.status, data.completed, data.total);
+        }
+      }
+  }
 }
 
 export default Ollama;

core/llm/toolSupport.ts

@@ -57,9 +57,18 @@ export const PROVIDER_TOOL_SUPPORT: Record<
   },
   // https://ollama.com/search?c=tools
   ollama: (model) => {
+    let modelName = "";
+    // Extract the model name after the last slash to support other registries
+    if(model.includes("/")) {
+      let parts = model.split('/');
+      modelName = parts[parts.length - 1];
+    } else {
+      modelName = model;
+    }
+    
     if (
       ["vision", "math", "guard", "mistrallite", "mistral-openorca"].some(
-        (part) => model.toLowerCase().includes(part),
+        (part) => modelName.toLowerCase().includes(part),
       )
     ) {
       return false;
@@ -79,10 +88,11 @@ export const PROVIDER_TOOL_SUPPORT: Record<
         "nemotron",
         "llama3-groq",
         "granite3",
+        "granite-3",
         "aya-expanse",
         "firefunction-v2",
         "mistral",
-      ].some((part) => model.toLowerCase().startsWith(part))
+      ].some((part) => modelName.toLowerCase().includes(part))
     ) {
       return true;
     }

core/util/ollamaHelper.ts

@@ -1,6 +1,13 @@
+import crypto from "crypto";
+import { exec } from "node:child_process";
 import path from "node:path";
 import { IDE } from "..";
-import { exec } from "node:child_process";
+
+export interface ModelInfo {
+    id: string;
+    size: number;
+    digest: string;
+}
 
 export async function isOllamaInstalled(): Promise<boolean> {
     return new Promise((resolve, _reject) => {
@@ -39,3 +46,54 @@ export async function startLocalOllama(ide: IDE): Promise<any> {
         });
     }
 }
+
+export async function getRemoteModelInfo(
+    modelId: string,
+    signal?: AbortSignal,
+  ): Promise<ModelInfo | undefined> {
+    const start = Date.now();
+    const [modelName, tag = "latest"] = modelId.split(":");
+    const url = `https://registry.ollama.ai/v2/library/${modelName}/manifests/${tag}`;
+    try {
+      const sig = signal ? signal : AbortSignal.timeout(3000);
+      const response = await fetch(url, { signal: sig });
+
+      if (!response.ok) {
+        throw new Error(`Failed to fetch the model page: ${response.statusText}`);
+      }
+
+      // First, read the response body as an ArrayBuffer to compute the digest
+      const buffer = await response.arrayBuffer();
+      const digest = getDigest(buffer);
+
+      // Then, decode the ArrayBuffer into a string and parse it as JSON
+      const text = new TextDecoder().decode(buffer);
+      const manifest = JSON.parse(text) as {
+        config: { size: number };
+        layers: { size: number }[];
+      };
+      const modelSize =
+        manifest.config.size +
+        manifest.layers.reduce((sum, layer) => sum + layer.size, 0);
+
+      const data: ModelInfo = {
+        id: modelId,
+        size: modelSize,
+        digest,
+      };
+      // Cache the successful result
+      return data;
+    } catch (error) {
+      console.error(`Error fetching or parsing model info: ${error}`);
+    } finally {
+      const elapsed = Date.now() - start;
+      console.log(`Fetched remote information for ${modelId} in ${elapsed} ms`);
+    }
+    return undefined;
+  }
+
+  function getDigest(buffer: ArrayBuffer): string {
+    const hash = crypto.createHash("sha256");
+    hash.update(new Uint8Array(buffer));
+    return hash.digest("hex");
+  }

core/util/text.ts

@@ -0,0 +1,6 @@
+export const capitalizeFirstLetter = (val: string) => {
+  if (val.length === 0) {
+    return "";
+  }
+  return val[0].toUpperCase() + val.slice(1);
+};