feat: implement /api/ps endpoint

Fixes #2605

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

Author: jeffmaury

api/openapi.yaml

@@ -162,6 +162,21 @@ paths:
               schema:
                 $ref: '#/components/schemas/ChatResponse'
 
+  /api/ps:
+    get:
+      operationId: getRunningModels
+      tags:
+        - models
+      description: List running models
+      summary: List running models
+      responses:
+        '200':
+          description: The list of running models was successfully fetched
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ProcessResponse'
+
 components:
   schemas:
     ListResponse:
@@ -191,6 +206,35 @@ components:
         details:
           $ref: '#/components/schemas/ModelDetails'
 
+    ProcessResponse:
+      type: object
+      description: Response with a list of running models
+      properties:
+        models:
+          type: array
+          items:
+            $ref: '#/components/schemas/ProcessModelResponse'
+
+    ProcessModelResponse:
+      type: object
+      description: Running model description
+      properties:
+        name:
+          type: string
+        model:
+          type: string
+        size:
+          type: integer
+        digest:
+          type: string
+        details:
+          $ref: '#/components/schemas/ModelDetails'
+        expires_at:
+          type: string
+          format: date-time
+        size_vram:
+          type: integer
+
     ModelDetails:
       type: object
       description: Details about a model

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

@@ -793,3 +793,75 @@ describe.each([undefined, true, false])('stream is %o', stream => {
     });
   });
 });
+
+describe('/api/ps', () => {
+  test('returns an error if the model is not known', async () => {
+    expect(server.getListener()).toBeDefined();
+    vi.mocked(inferenceManager.getServers).mockImplementation(() => {
+      throw new Error('model unknown');
+    });
+    const res = await request(server.getListener()!).get('/api/ps').expect(500);
+    expect(res.body).toMatchObject({ message: 'unable to ps' });
+  });
+
+  test('returns empty result if no servers', async () => {
+    expect(server.getListener()).toBeDefined();
+    vi.mocked(inferenceManager.getServers).mockReturnValue([]);
+    const res = await request(server.getListener()!).get('/api/ps').expect(200);
+    expect(res.body).toEqual({ models: [] });
+  });
+
+  test('returns empty result if server is stopped', async () => {
+    expect(server.getListener()).toBeDefined();
+    vi.mocked(inferenceManager.getServers).mockReturnValue([
+      {
+        models: [
+          {
+            id: 'modelId1',
+            name: 'model-name',
+            description: 'model 1',
+          },
+        ],
+        container: {
+          engineId: 'engine1',
+          containerId: 'container1',
+        },
+        status: 'stopped',
+      } as unknown as InferenceServer,
+    ]);
+    const res = await request(server.getListener()!).get('/api/ps').expect(200);
+    expect(res.body).toEqual({ models: [] });
+  });
+
+  test('returns result if server is started', async () => {
+    expect(server.getListener()).toBeDefined();
+    vi.mocked(inferenceManager.getServers).mockReturnValue([
+      {
+        models: [
+          {
+            id: 'modelId1',
+            name: 'model-name',
+            description: 'model 1',
+            memory: 1_000_000,
+          },
+        ],
+        container: {
+          engineId: 'engine1',
+          containerId: 'container1',
+        },
+        status: 'running',
+      } as unknown as InferenceServer,
+    ]);
+    const res = await request(server.getListener()!).get('/api/ps').expect(200);
+    expect(res.body).toEqual({
+      models: [
+        {
+          name: 'model-name',
+          model: 'model-name',
+          size: 1_000_000,
+          digest: 'b48fa42fa5b28c4363747ec0797532e274650f73004383a3054697137d9d1f30',
+        },
+      ],
+    });
+  });
+});

packages/backend/src/managers/apiServer.ts

@@ -41,6 +41,7 @@ import OpenAI from 'openai';
 import type { ChatCompletionMessageParam } from 'openai/resources';
 import type { ContainerRegistry } from '../registries/ContainerRegistry';
 import type { Stream } from 'openai/streaming';
+import crypto from 'node:crypto';
 
 const SHOW_API_INFO_COMMAND = 'ai-lab.show-api-info';
 const SHOW_API_ERROR_COMMAND = 'ai-lab.show-api-error';
@@ -49,6 +50,7 @@ export const PREFERENCE_RANDOM_PORT = 0;
 
 type ListModelResponse = components['schemas']['ListModelResponse'];
 type Message = components['schemas']['Message'];
+type ProcessModelResponse = components['schemas']['ProcessModelResponse'];
 
 function asListModelResponse(model: ModelInfo): ListModelResponse {
   return {
@@ -61,6 +63,20 @@ function asListModelResponse(model: ModelInfo): ListModelResponse {
   };
 }
 
+// ollama expect at least 12 characters for the digest
+function toDigest(name: string, sha256?: string): string {
+  return sha256 ?? crypto.createHash('sha256').update(name).digest('hex');
+}
+
+function asProcessModelResponse(model: ModelInfo): ProcessModelResponse {
+  return {
+    name: model.name,
+    model: model.name,
+    size: model.memory,
+    digest: toDigest(model.name, model.sha256),
+  };
+}
+
 const LISTENING_ADDRESS = '127.0.0.1';
 
 interface ChatCompletionOptions {
@@ -124,6 +140,7 @@ export class ApiServer implements Disposable {
     router.post('/show', this.show.bind(this));
     router.post('/generate', this.generate.bind(this));
     router.post('/chat', this.chat.bind(this));
+    router.get('/ps', this.ps.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));
@@ -574,4 +591,17 @@ export class ApiServer implements Disposable {
       })
       .catch((err: unknown) => console.error(`unable to check if the inference server is running: ${err}`));
   }
+
+  ps(_req: Request, res: Response): void {
+    try {
+      const models = this.inferenceManager
+        .getServers()
+        .filter(server => server.status === 'running')
+        .flatMap(server => server.models)
+        .map(model => asProcessModelResponse(model));
+      res.status(200).json({ models });
+    } catch (err: unknown) {
+      this.doErr(res, 'unable to ps', err);
+    }
+  }
 }