feat: implement list models in REST server (#1564)

* feat: implement list models in REST server

Signed-off-by: Jeff MAURY <[email protected]>

* fix: run generate for typecheck

Signed-off-by: Jeff MAURY <[email protected]>

---------

Signed-off-by: Jeff MAURY <[email protected]>

Author: jeffmaury

.eslintignore

@@ -3,3 +3,4 @@ build
 __mocks__
 coverage
 packages/backend/media/**
+src-generated

api/openapi.yaml

@@ -37,3 +37,65 @@ paths:
                     type: string
                 required:
                 - version
+  /api/tags:
+    get:
+      operationId: getModels
+      tags:
+        - models
+      description: List models that are available locally
+      summary: List models that are available locally
+      responses:
+        '200':
+          description: The models were successfully fetched
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ListResponse'
+
+components:
+  schemas:
+    ListResponse:
+      type: object
+      description: Response from a list request
+      properties:
+        models:
+          type: array
+          items:
+            $ref: '#/components/schemas/ListModelResponse'
+
+    ListModelResponse:
+      type: object
+      description: Response from a list request
+      properties:
+        name:
+          type: string
+        model:
+          type: string
+        modified_at:
+          type: string
+          format: date-time
+        size:
+          type: integer
+        digest:
+          type: string
+        details:
+          $ref: '#/components/schemas/ModelDetails'
+
+    ModelDetails:
+      type: object
+      description: Details about a model
+      properties:
+        parent_model:
+          type: string
+        format:
+          type: string
+        family:
+          type: string
+        families:
+          type: array
+          items:
+            type: string
+        parameter_size:
+          type: string
+        quantization_level:
+          type: string

clean.sh

@@ -0,0 +1 @@
+rm -rf node_modules packages/backend/node_modules packages/frontend/node_modules

package.json

@@ -25,7 +25,7 @@
     "test:e2e": "cd tests/playwright && npm run test:e2e",
     "typecheck:shared": "tsc --noEmit --project packages/shared",
     "typecheck:frontend": "tsc --noEmit --project packages/frontend",
-    "typecheck:backend": "tsc --noEmit --project packages/backend",
+    "typecheck:backend": "yarn --cwd packages/backend typecheck",
     "typecheck": "npm run typecheck:shared && npm run typecheck:frontend && npm run typecheck:backend"
   },
   "resolutions": {

packages/backend/package.json

@@ -52,12 +52,14 @@
     }
   },
   "scripts": {
-    "build": "vite build",
+    "generate": "yarn openapi-typescript ../../api/openapi.yaml -o src-generated/openapi.ts",
+    "build": "yarn generate && vite build",
     "test": "vitest run --coverage",
     "test:watch": "vitest watch --coverage",
     "format:check": "prettier --check \"src/**/*.ts\"",
     "format:fix": "prettier --write \"src/**/*.ts\"",
-    "watch": "vite --mode development build -w"
+    "watch": "yarn generate && vite --mode development build -w",
+    "typecheck": "yarn generate && tsc --noEmit"
   },
   "dependencies": {
     "@huggingface/gguf": "^0.1.9",
@@ -79,6 +81,7 @@
     "@types/node": "^20",
     "@types/postman-collection": "^3.5.10",
     "@types/supertest": "^6.0.2",
+    "openapi-typescript": "^7.3.0",
     "supertest": "^7.0.0",
     "vitest": "^2.0.5"
   }

packages/backend/src/managers/apiServer.spec.ts

@@ -22,6 +22,9 @@ import request from 'supertest';
 import type * as podmanDesktopApi from '@podman-desktop/api';
 import path from 'path';
 import type { Server } from 'http';
+import type { ModelsManager } from './modelsManager';
+import type { EventEmitter } from 'node:events';
+import { once } from 'node:events';
 
 class TestApiServer extends ApiServer {
   public override getListener(): Server | undefined {
@@ -33,16 +36,21 @@ const extensionContext = {} as unknown as podmanDesktopApi.ExtensionContext;
 
 let server: TestApiServer;
 
+const modelsManager = {
+  getModelsInfo: vi.fn(),
+} as unknown as ModelsManager;
+
 beforeEach(async () => {
-  server = new TestApiServer(extensionContext);
+  server = new TestApiServer(extensionContext, modelsManager);
   vi.spyOn(server, 'displayApiInfo').mockReturnValue();
   vi.spyOn(server, 'getSpecFile').mockReturnValue(path.join(__dirname, '../../../../api/openapi.yaml'));
   vi.spyOn(server, 'getPackageFile').mockReturnValue(path.join(__dirname, '../../../../package.json'));
   await server.init();
 });
 
-afterEach(() => {
+afterEach(async () => {
   server.dispose();
+  await once(server.getListener() as EventEmitter, 'close');
 });
 
 test('/spec endpoint', async () => {
@@ -101,3 +109,21 @@ test('/api/wrongEndpoint', async () => {
   expect(server.getListener()).toBeDefined();
   await request(server.getListener()!).get('/api/wrongEndpoint').expect(404);
 });
+
+test('/', async () => {
+  expect(server.getListener()).toBeDefined();
+  await request(server.getListener()!).get('/').expect(200);
+});
+
+test('/api/tags', async () => {
+  expect(server.getListener()).toBeDefined();
+  vi.mocked(modelsManager.getModelsInfo).mockReturnValue([]);
+  await request(server.getListener()!).get('/api/tags').expect(200);
+});
+
+test('/api/tags returns error', async () => {
+  expect(server.getListener()).toBeDefined();
+  vi.mocked(modelsManager.getModelsInfo).mockRejectedValue({});
+  const res = await request(server.getListener()!).get('/api/tags').expect(500);
+  expect(res.body.message).toEqual('unable to get models');
+});

packages/backend/src/managers/apiServer.ts

@@ -26,14 +26,33 @@ import { existsSync } from 'fs';
 import { getFreeRandomPort } from '../utils/ports';
 import * as podmanDesktopApi from '@podman-desktop/api';
 import { readFile } from 'fs/promises';
+import type { ModelsManager } from './modelsManager';
+import type { components } from '../../src-generated/openapi';
+import type { ModelInfo } from '@shared/src/models/IModelInfo';
 
 const DEFAULT_PORT = 10434;
 const SHOW_API_INFO_COMMAND = 'ai-lab.show-api-info';
 
+type ListModelResponse = components['schemas']['ListModelResponse'];
+
+function asListModelResponse(model: ModelInfo): ListModelResponse {
+  return {
+    model: model.id,
+    name: model.name,
+    digest: model.sha256,
+    size: model.file?.size,
+    modified_at: model.file?.creation?.toISOString(),
+    details: {},
+  };
+}
+
 export class ApiServer implements Disposable {
   #listener?: Server;
 
-  constructor(private extensionContext: podmanDesktopApi.ExtensionContext) {}
+  constructor(
+    private extensionContext: podmanDesktopApi.ExtensionContext,
+    private modelsManager: ModelsManager,
+  ) {}
 
   protected getListener(): Server | undefined {
     return this.#listener;
@@ -47,6 +66,8 @@ export class ApiServer implements Disposable {
 
     // declare routes
     router.get('/version', this.getVersion.bind(this));
+    router.get('/tags', this.getModels.bind(this));
+    app.get('/', (_res, res) => res.sendStatus(200)); //required for the ollama client to work against us
     app.use('/api', router);
     app.use('/spec', this.getSpec.bind(this));
 
@@ -150,4 +171,22 @@ export class ApiServer implements Disposable {
       doErr(err);
     }
   }
+
+  getModels(_req: Request, res: Response): void {
+    const doErr = (err: unknown) => {
+      res.status(500).json({
+        message: 'unable to get models',
+        errors: [err],
+      });
+    };
+    try {
+      const models = this.modelsManager
+        .getModelsInfo()
+        .filter(model => this.modelsManager.isModelOnDisk(model.id))
+        .map(model => asListModelResponse(model));
+      res.status(200).json({ models: models });
+    } catch (err: unknown) {
+      doErr(err);
+    }
+  }
 }

packages/backend/src/studio.ts

@@ -328,7 +328,7 @@ export class Studio {
     // Register the instance
     this.#rpcExtension.registerInstance<StudioApiImpl>(StudioApiImpl, this.#studioApi);
 
-    const apiServer = new ApiServer(this.#extensionContext);
+    const apiServer = new ApiServer(this.#extensionContext, this.#modelsManager);
     await apiServer.init();
     this.#extensionContext.subscriptions.push(apiServer);
   }