Add ml model service (#469)

Author: DTCurrie

src/main.ts

@@ -199,6 +199,8 @@ export { GantryClient, type Gantry } from './components/gantry';
  */
 export * as gantryApi from './gen/component/gantry/v1/gantry_pb';
 
+export { MLModelClient, type MLModel } from './services/ml-model';
+
 export { MotorClient, type Motor } from './components/motor';
 /**
  * Raw Protobuf interfaces for a Motor component.

src/robot/client.ts

@@ -39,6 +39,7 @@ import { clientHeaders } from '../utils';
 import GRPCConnectionManager from './grpc-connection-manager';
 import type { Robot } from './robot';
 import SessionManager from './session-manager';
+import { MLModelService } from '../gen/service/mlmodel/v1/mlmodel_connect';
 
 interface WebRTCOptions {
   enabled: boolean;
@@ -116,6 +117,10 @@ export class RobotClient extends EventDispatcher implements Robot {
     | PromiseClient<typeof GripperService>
     | undefined;
 
+  private mlModelServiceClient:
+    | PromiseClient<typeof MLModelService>
+    | undefined;
+
   private movementSensorServiceClient:
     | PromiseClient<typeof MovementSensorService>
     | undefined;
@@ -318,6 +323,13 @@ export class RobotClient extends EventDispatcher implements Robot {
     return this.gripperServiceClient;
   }
 
+  get mlModelService() {
+    if (!this.mlModelServiceClient) {
+      throw new Error(RobotClient.notConnectedYetStr);
+    }
+    return this.mlModelServiceClient;
+  }
+
   get movementSensorService() {
     if (!this.movementSensorServiceClient) {
       throw new Error(RobotClient.notConnectedYetStr);
@@ -582,6 +594,10 @@ export class RobotClient extends EventDispatcher implements Robot {
         GripperService,
         clientTransport
       );
+      this.mlModelServiceClient = createPromiseClient(
+        MLModelService,
+        clientTransport
+      );
       this.movementSensorServiceClient = createPromiseClient(
         MovementSensorService,
         clientTransport

src/services/ml-model/client.ts

@@ -0,0 +1,49 @@
+import type { CallOptions, PromiseClient } from '@connectrpc/connect';
+import type { FlatTensors, MLModel } from './ml-model';
+import { Struct, type Options } from '../../types';
+import type { RobotClient } from '../../robot';
+import {
+  InferRequest,
+  MetadataRequest,
+} from '../../gen/service/mlmodel/v1/mlmodel_pb';
+import { MLModelService } from '../../gen/service/mlmodel/v1/mlmodel_connect';
+
+export class MLModelClient implements MLModel {
+  private client: PromiseClient<typeof MLModelService>;
+  private readonly name: string;
+  private readonly options: Options;
+  public callOptions: CallOptions = { headers: {} as Record<string, string> };
+
+  constructor(client: RobotClient, name: string, options: Options = {}) {
+    this.client = client.createServiceClient(MLModelService);
+    this.name = name;
+    this.options = options;
+  }
+
+  async metadata(extra = {}, callOptions = this.callOptions) {
+    const request = new MetadataRequest({
+      name: this.name,
+      extra: Struct.fromJson(extra),
+    });
+
+    this.options.requestLogger?.(request);
+
+    return this.client.metadata(request, callOptions);
+  }
+
+  async infer(
+    inputTensors: FlatTensors,
+    extra = {},
+    callOptions = this.callOptions
+  ) {
+    const request = new InferRequest({
+      name: this.name,
+      inputTensors,
+      extra: Struct.fromJson(extra),
+    });
+
+    this.options.requestLogger?.(request);
+
+    return this.client.infer(request, callOptions);
+  }
+}

src/services/ml-model/index.ts

@@ -0,0 +1,2 @@
+export type { MLModel, FlatTensors } from './ml-model';
+export { MLModelClient } from './client';

src/services/ml-model/ml-model.ts

@@ -0,0 +1,13 @@
+import type { Struct } from '@bufbuild/protobuf';
+import type * as mlModelAPI from '../../gen/service/mlmodel/v1/mlmodel_pb';
+
+export type FlatTensors = mlModelAPI.FlatTensors;
+
+export interface MLModel {
+  metadata: (extra?: Struct) => Promise<mlModelAPI.MetadataResponse>;
+
+  infer: (
+    inputTensors: FlatTensors,
+    extra?: Struct
+  ) => Promise<mlModelAPI.InferResponse>;
+}