-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Add Ollama Model Management Features (#163)
* fix: CLI hardcoded version and solving issues with containers and images * feat: new global var to deal with compatible version * Release v0.10.0-beta.0 [skip ci] * fix: removing default value from .env.example * feat: new update ollama command --------- Co-authored-by: github-actions[bot] <github-actions[bot]@genlayer.com>
- Loading branch information
Showing
11 changed files
with
264 additions
and
115 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,26 @@ | ||
| import { Command } from "commander"; | ||
| import { OllamaAction } from "./ollama"; | ||
|
|
||
| export function initializeUpdateCommands(program: Command) { | ||
| const updateCommand = program | ||
| .command("update") | ||
| .description("Update resources like models or configurations"); | ||
|
|
||
| updateCommand | ||
| .command("ollama") | ||
| .description("Manage Ollama models (update or remove)") | ||
| .option("--model [model-name]", "Specify the model to update or remove") | ||
| .option("--remove", "Remove the specified model instead of updating") | ||
| .action(async (options) => { | ||
| const modelName = options.model || "default-model"; | ||
| const ollamaAction = new OllamaAction(); | ||
|
|
||
| if (options.remove) { | ||
| await ollamaAction.removeModel(modelName); | ||
| } else { | ||
| await ollamaAction.updateModel(modelName); | ||
| } | ||
| }); | ||
|
|
||
| return program; | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| import Docker from "dockerode" | ||
|
|
||
| export class OllamaAction { | ||
| private docker: Docker; | ||
|
|
||
| constructor() { | ||
| this.docker = new Docker(); | ||
| } | ||
|
|
||
| async updateModel(modelName: string) { | ||
| await this.executeModelCommand("pull", modelName, `Model "${modelName}" updated successfully`); | ||
| } | ||
|
|
||
| async removeModel(modelName: string) { | ||
| await this.executeModelCommand("rm", modelName, `Model "${modelName}" removed successfully`); | ||
| } | ||
|
|
||
| private async executeModelCommand(command: string, modelName: string, successMessage: string) { | ||
| try { | ||
| const ollamaContainer = this.docker.getContainer("ollama"); | ||
| const exec = await ollamaContainer.exec({ | ||
| Cmd: ["ollama", command, modelName], | ||
| AttachStdout: true, | ||
| AttachStderr: true, | ||
| }); | ||
| const stream = await exec.start({ Detach: false, Tty: false }); | ||
|
|
||
| stream.on("data", (chunk: any) => { | ||
| console.log(chunk.toString()); | ||
| }); | ||
|
|
||
| await new Promise<void>((resolve, reject) => { | ||
| stream.on("end", resolve); | ||
| stream.on("error", reject); | ||
| }); | ||
|
|
||
| console.log(successMessage); | ||
| } catch (error) { | ||
| console.error(`Error executing command "${command}" on model "${modelName}":`, error); | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,129 @@ | ||
| import {describe, test, vi, beforeEach, afterEach, expect, Mock} from "vitest"; | ||
| import { OllamaAction } from "../../src/commands/update/ollama"; | ||
| import Docker from "dockerode"; | ||
|
|
||
| vi.mock("dockerode"); | ||
|
|
||
| describe("OllamaAction", () => { | ||
| let ollamaAction: OllamaAction; | ||
| let mockGetContainer: Mock; | ||
| let mockExec: Mock; | ||
| let mockStart: Mock; | ||
| let mockStream: any; | ||
|
|
||
| beforeEach(() => { | ||
| vi.clearAllMocks(); | ||
| ollamaAction = new OllamaAction(); | ||
|
|
||
| mockGetContainer = vi.mocked(Docker.prototype.getContainer); | ||
| mockExec = vi.fn(); | ||
| mockStart = vi.fn(); | ||
|
|
||
| mockStream = { | ||
| on: vi.fn(), | ||
| }; | ||
|
|
||
| mockExec.mockResolvedValue({ | ||
| start: mockStart, | ||
| }); | ||
|
|
||
| mockStart.mockResolvedValue(mockStream); | ||
|
|
||
| mockGetContainer.mockReturnValue({ | ||
| exec: mockExec, | ||
| } as unknown as Docker.Container); | ||
|
|
||
| Docker.prototype.getContainer = mockGetContainer; | ||
| }); | ||
|
|
||
| afterEach(() => { | ||
| vi.restoreAllMocks(); | ||
| }); | ||
|
|
||
| test("should update the model using 'pull'", async () => { | ||
| mockStream.on.mockImplementation((event: any, callback:any) => { | ||
| if (event === "data") callback(Buffer.from("Mocked output")); | ||
| if (event === "end") callback(); | ||
| }); | ||
|
|
||
| console.log = vi.fn(); | ||
|
|
||
| await ollamaAction.updateModel("mocked_model"); | ||
|
|
||
| expect(mockGetContainer).toHaveBeenCalledWith("ollama"); | ||
| expect(mockExec).toHaveBeenCalledWith({ | ||
| Cmd: ["ollama", "pull", "mocked_model"], | ||
| AttachStdout: true, | ||
| AttachStderr: true, | ||
| }); | ||
| expect(mockStart).toHaveBeenCalledWith({ Detach: false, Tty: false }); | ||
| expect(mockStream.on).toHaveBeenCalledWith("data", expect.any(Function)); | ||
| expect(mockStream.on).toHaveBeenCalledWith("end", expect.any(Function)); | ||
| expect(console.log).toHaveBeenCalledWith("Mocked output"); | ||
| expect(console.log).toHaveBeenCalledWith('Model "mocked_model" updated successfully'); | ||
| }); | ||
|
|
||
| test("should remove the model using 'rm'", async () => { | ||
| mockStream.on.mockImplementation((event:any, callback:any) => { | ||
| if (event === "data") callback(Buffer.from("Mocked output")); | ||
| if (event === "end") callback(); | ||
| }); | ||
|
|
||
| console.log = vi.fn(); | ||
|
|
||
| await ollamaAction.removeModel("mocked_model"); | ||
|
|
||
| expect(mockGetContainer).toHaveBeenCalledWith("ollama"); | ||
| expect(mockExec).toHaveBeenCalledWith({ | ||
| Cmd: ["ollama", "rm", "mocked_model"], | ||
| AttachStdout: true, | ||
| AttachStderr: true, | ||
| }); | ||
| expect(mockStart).toHaveBeenCalledWith({ Detach: false, Tty: false }); | ||
| expect(mockStream.on).toHaveBeenCalledWith("data", expect.any(Function)); | ||
| expect(mockStream.on).toHaveBeenCalledWith("end", expect.any(Function)); | ||
| expect(console.log).toHaveBeenCalledWith("Mocked output"); | ||
| expect(console.log).toHaveBeenCalledWith('Model "mocked_model" removed successfully'); | ||
| }); | ||
|
|
||
| test("should log an error if an exception occurs during 'pull'", async () => { | ||
| const error = new Error("Mocked error"); | ||
| mockGetContainer.mockReturnValueOnce( | ||
| { | ||
| exec: () => { | ||
| throw new Error("Mocked error"); | ||
| } | ||
| } | ||
| ); | ||
| console.error = vi.fn(); | ||
|
|
||
| await ollamaAction.updateModel("mocked_model"); | ||
|
|
||
| expect(mockGetContainer).toHaveBeenCalledWith("ollama"); | ||
| expect(console.error).toHaveBeenCalledWith( | ||
| 'Error executing command "pull" on model "mocked_model":', | ||
| error | ||
| ); | ||
| }); | ||
|
|
||
| test("should log an error if an exception occurs during 'rm'", async () => { | ||
| const error = new Error("Mocked error"); | ||
| mockGetContainer.mockReturnValueOnce( | ||
| { | ||
| exec: () => { | ||
| throw new Error("Mocked error"); | ||
| } | ||
| } | ||
| ); | ||
|
|
||
| console.error = vi.fn(); | ||
|
|
||
| await ollamaAction.removeModel("mocked_model"); | ||
|
|
||
| expect(mockGetContainer).toHaveBeenCalledWith("ollama"); | ||
| expect(console.error).toHaveBeenCalledWith( | ||
| 'Error executing command "rm" on model "mocked_model":', | ||
| error | ||
| ); | ||
| }); | ||
| }); |
Oops, something went wrong.