Merge pull request #49 from vlm-run/sh/add-web-endpoint

Added web predictions

package.json

@@ -1,6 +1,6 @@
 {
   "name": "vlmrun",
-  "version": "0.1.12",
+  "version": "0.1.13",
   "description": "The official TypeScript library for the VlmRun API",
   "author": "VlmRun <[email protected]>",
   "main": "dist/index.js",

src/client/predictions.ts

@@ -5,6 +5,7 @@ import {
   ListParams,
   ImagePredictionParams,
   FilePredictionParams,
+  WebPredictionParams,
 } from "./types";
 import { processImage } from "../utils/image";
 
@@ -138,6 +139,36 @@ export class FilePredictions extends Predictions {
   }
 }
 
+export class WebPredictions extends Predictions {
+  async generate(params: WebPredictionParams): Promise<PredictionResponse> {
+    const { url, model, domain, mode, metadata, callbackUrl, config } = params;
+    const [response] = await this.requestor.request<PredictionResponse>(
+      "POST",
+      `/web/generate`,
+      undefined,
+      {
+        url,
+        model,
+        domain,
+        mode,
+        config: {
+          detail: config?.detail ?? "auto",
+          json_schema: config?.jsonSchema ?? null,
+          confidence: config?.confidence ?? false,
+          grounding: config?.grounding ?? false,
+        },
+        metadata: {
+          environment: metadata?.environment ?? "dev",
+          session_id: metadata?.sessionId,
+          allow_training: metadata?.allowTraining ?? true,
+        },
+        callback_url: callbackUrl,
+      }
+    );
+    return response;
+  }
+}
+
 // Create specialized instances for different file types
 export const DocumentPredictions = (client: Client) =>
   new FilePredictions(client, "document");

src/client/types.ts

@@ -63,7 +63,6 @@ export interface FeedbackSubmitParams {
 export interface PredictionGenerateParams {
   model?: string;
   domain: string;
-  batch?: boolean;
   config?: GenerationConfigParams;
   metadata?: RequestMetadataParams;
   callbackUrl?: string;
@@ -157,14 +156,21 @@ export class GenerationConfig {
 export type GenerationConfigInput = GenerationConfig | GenerationConfigParams;
 
 export interface ImagePredictionParams extends PredictionGenerateParams {
+  batch?: boolean;
   images: string[];
 }
 
 export interface FilePredictionParams extends PredictionGenerateParams {
+  batch?: boolean;
   fileId?: string;
   url?: string;
 }
 
+export interface WebPredictionParams extends PredictionGenerateParams {
+  url: string;
+  mode: "fast" | "accurate";
+}
+
 export class APIError extends Error {
   constructor(
     message: string,

src/index.ts

@@ -7,6 +7,7 @@ import {
   DocumentPredictions,
   AudioPredictions,
   VideoPredictions,
+  WebPredictions,
 } from "./client/predictions";
 import { Feedback } from "./client/feedback";
 
@@ -33,6 +34,7 @@ export class VlmRun {
   readonly document: ReturnType<typeof DocumentPredictions>;
   readonly audio: ReturnType<typeof AudioPredictions>;
   readonly video: ReturnType<typeof VideoPredictions>;
+  readonly web: WebPredictions;
   readonly feedback: Feedback;
 
   constructor(config: VlmRunConfig) {
@@ -49,5 +51,6 @@ export class VlmRun {
     this.audio = AudioPredictions(this.client);
     this.video = VideoPredictions(this.client);
     this.feedback = new Feedback(this.client);
+    this.web = new WebPredictions(this.client);
   }
 }

tests/unit/client/predictions.test.ts

@@ -4,6 +4,7 @@ import {
   DocumentPredictions,
   AudioPredictions,
   VideoPredictions,
+  WebPredictions,
 } from "../../../src/client/predictions";
 import * as imageUtils from "../../../src/utils/image";
 
@@ -261,4 +262,104 @@ describe("Predictions", () => {
       });
     });
   });
+
+  describe("WebPredictions", () => {
+    let webPredictions: WebPredictions;
+
+    beforeEach(() => {
+      webPredictions = new WebPredictions(client);
+    });
+
+    describe("generate", () => {
+      it("should generate web predictions with default options", async () => {
+        const mockResponse = { id: "pred_123", status: "completed" };
+        jest
+          .spyOn(webPredictions["requestor"], "request")
+          .mockResolvedValue([mockResponse, 200, {}]);
+
+        const result = await webPredictions.generate({
+          url: "https://example.com",
+          model: "model1",
+          domain: "domain1",
+          mode: "accurate",
+        });
+
+        expect(result).toEqual(mockResponse);
+        expect(webPredictions["requestor"].request).toHaveBeenCalledWith(
+          "POST",
+          "/web/generate",
+          undefined,
+          {
+            url: "https://example.com",
+            model: "model1",
+            domain: "domain1",
+            mode: "accurate",
+            config: {
+              detail: "auto",
+              json_schema: null,
+              confidence: false,
+              grounding: false,
+            },
+            metadata: {
+              environment: "dev",
+              session_id: undefined,
+              allow_training: true,
+            },
+            callback_url: undefined,
+          }
+        );
+      });
+
+      it("should generate web predictions with custom options", async () => {
+        const mockResponse = { id: "pred_123", status: "completed" };
+        jest
+          .spyOn(webPredictions["requestor"], "request")
+          .mockResolvedValue([mockResponse, 200, {}]);
+
+        const result = await webPredictions.generate({
+          url: "https://example.com",
+          model: "model1",
+          domain: "domain1",
+          mode: "fast",
+          config: {
+            detail: "hi",
+            jsonSchema: { type: "object" },
+            confidence: true,
+            grounding: true,
+          },
+          metadata: {
+            environment: "prod",
+            sessionId: "session123",
+            allowTraining: false,
+          },
+          callbackUrl: "https://callback.example.com",
+        });
+
+        expect(result).toEqual(mockResponse);
+        expect(webPredictions["requestor"].request).toHaveBeenCalledWith(
+          "POST",
+          "/web/generate",
+          undefined,
+          {
+            url: "https://example.com",
+            model: "model1",
+            domain: "domain1",
+            mode: "fast",
+            config: {
+              detail: "hi",
+              json_schema: { type: "object" },
+              confidence: true,
+              grounding: true,
+            },
+            metadata: {
+              environment: "prod",
+              session_id: "session123",
+              allow_training: false,
+            },
+            callback_url: "https://callback.example.com",
+          }
+        );
+      });
+    });
+  });
 });