feat: implement models() function for OpenAI, Ollama, and Anthropic providers

Author: wangzuo

examples/anthropic-list-models.ts examples/anthropic-models.tsexamples/anthropic-list-models.ts renamed to examples/anthropic-models.ts

@@ -1,6 +1,6 @@
 import { anthropic } from "../src";
 
-const job = anthropic().listModels();
+const job = anthropic().models();
 const result = await job.run();
 
-console.log(result);
+console.log(result);

examples/ollama-list-models.ts examples/ollama-models.tsexamples/ollama-list-models.ts renamed to examples/ollama-models.ts

@@ -1,6 +1,6 @@
 import { ollama } from "../src";
 
-const job = ollama().listModels();
+const job = ollama().models();
 const result = await job.run();
 
 console.log(result);

examples/openai-list-models.ts examples/openai-models.tsexamples/openai-list-models.ts renamed to examples/openai-models.ts

@@ -1,6 +1,6 @@
 import { openai } from "../src";
 
-const job = openai().listModels();
+const job = openai().models();
 const result = await job.run();
 
-console.log(result);
+console.log(result);

src/jobs/job.ts

@@ -1,4 +1,5 @@
 import { z } from "zod";
+import zodToJsonSchema from "zod-to-json-schema";
 import { version } from "../../package.json";
 
 const providerSchema = z.enum([
@@ -146,10 +147,15 @@ const imageJobSchema = baseJobSchema.extend({
   params: imageJobParamsSchema,
 });
 
+const modelsJobSchema = baseJobSchema.extend({ type: z.literal("models") });
+const modelJobSchema = baseJobSchema.extend({ type: z.literal("model") });
+
 export const jobSchema = z.discriminatedUnion("type", [
   chatJobSchema,
   embeddingJobSchema,
   imageJobSchema,
+  modelsJobSchema,
+  modelJobSchema,
 ]);
 
 export class Job {
@@ -180,6 +186,7 @@ export class Job {
   }
 }
 
+export const jsonSchema = zodToJsonSchema(jobSchema);
 export type AIProviderOptions = z.infer<typeof optionsSchema>;
 export type ChatJobParams = z.infer<typeof chatJobParamsSchema>;
 export type EmbeddingJobParams = z.infer<typeof embeddedJobParamsSchema>;

src/jobs/load.ts

@@ -34,6 +34,9 @@ export function load(obj: AIJob): Job {
   if (obj.type === "image" && "image" in provider) {
     return provider.image(obj.model)._setParams(obj.params);
   }
+  if (obj.type === "models" && "models" in provider) {
+    return provider.models();
+  }
 
   throw new Error("Failed to load job");
 }

src/jobs/models.ts

@@ -2,12 +2,12 @@ import { Job } from "./job";
 
 export class ListModelsJob extends Job {
   constructor() {
-	super();
-	this.params = {};
+    super();
+    this.params = {};
   }
 
   dump() {
     const obj = super.dump();
-    return { ...obj, models: { params: this.params } };
+    return { ...obj, type: "models", params: this.params };
   }
 }

src/providers/anthropic.ts

@@ -10,9 +10,9 @@ export function anthropic(options?: AIProviderOptions) {
     chat(model: string) {
       return new AnthropicChatJob(options, model);
     },
-    listModels() {
+    models() {
       return new AnthropicListModelsJob(options);
-    }
+    },
   };
 }
 
@@ -79,4 +79,4 @@ export class AnthropicListModelsJob extends ListModelsJob {
     const json = await response.json();
     return json;
   };
-}
+}

src/providers/ollama.ts

@@ -10,12 +10,12 @@ export function ollama(options?: AIProviderOptions) {
     chat(model: string) {
       return new OllamaChatJob(options, model);
     },
-    listModels() {
-      return new OllamaListModelsJob(options);
-    },
     embedding(model: string) {
       return new OllamaEmbeddingJob(options, model);
     },
+    models() {
+      return new OllamaListModelsJob(options);
+    },
   };
 }
 

src/providers/openai.ts

@@ -21,15 +21,15 @@ export function openai(options?: AIProviderOptions) {
     chat(model: string) {
       return new OpenAIChatJob(options, model);
     },
-    listModels() {
-      return new OpenAIListModelsJob(options);
-    },
     image(model: string) {
       return new OpenAIImageJob(options, model);
     },
     embedding(model: string) {
       return new OpenAIEmbeddingJob(options, model);
     },
+    models() {
+      return new OpenAIListModelsJob(options);
+    },
   };
 }
 
@@ -225,12 +225,11 @@ export class OpenAIChatJob extends ChatJob {
   };
 }
 
-
 export class OpenAIListModelsJob extends ListModelsJob {
   constructor(options: AIProviderOptions) {
     super();
     this.provider = "openai";
-    this.options = options
+    this.options = options;
   }
 
   makeRequest = () => {
@@ -240,7 +239,7 @@ export class OpenAIListModelsJob extends ListModelsJob {
         Authorization: `Bearer ${this.options.apiKey}`,
         "Content-Type": "application/json",
       },
-      method: "GET"
+      method: "GET",
     });
   };
 
@@ -311,4 +310,4 @@ export class OpenAIEmbeddingJob extends EmbeddingJob {
   handleResponse = async (response: Response) => {
     return await response.json();
   };
-}
+}

…st/__snapshots__/listModels.test.ts.snap test/__snapshots__/models.snaptest/__snapshots__/listModels.test.ts.snap renamed to test/__snapshots__/models.snap test/__snapshots__/listModels.test.ts.snap renamed to test/__snapshots__/models.snap

@@ -1,6 +1,15 @@
 // Bun Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`listModels 1`] = `
+exports[`models 1`] = `
+{
+  "body": "",
+  "headers": {},
+  "method": "GET",
+  "url": "http://localhost:11434/api/tags",
+}
+`;
+
+exports[`models 2`] = `
 {
   "body": "",
   "headers": {
@@ -12,7 +21,7 @@ exports[`listModels 1`] = `
 }
 `;
 
-exports[`listModels 2`] = `
+exports[`models 3`] = `
 {
   "body": "",
   "headers": {
@@ -27,24 +36,34 @@ exports[`listModels 2`] = `
 
 exports[`dump 1`] = `
 {
-  "models": {
-    "params": {},
-  },
+  "options": {},
+  "params": {},
+  "provider": "ollama",
+  "type": "models",
+  "version": "0.2.4",
+}
+`;
+
+exports[`dump 2`] = `
+{
   "options": {
     "apiKey": "<key>",
   },
+  "params": {},
   "provider": "openai",
+  "type": "models",
+  "version": "0.2.4",
 }
 `;
 
-exports[`dump 2`] = `
+exports[`dump 3`] = `
 {
-  "models": {
-    "params": {},
-  },
   "options": {
     "apiKey": "<key>",
   },
+  "params": {},
   "provider": "anthropic",
+  "type": "models",
+  "version": "0.2.4",
 }
 `;