add job output schema

Author: wangzuo

README.md

@@ -18,7 +18,7 @@ npm install fluent-ai zod
 
 fluent-ai includes support for multiple AI providers and modalities.
 
-| Provider  | chat               | embedding          | image              | models             |
+| provider  | chat completion    | embedding          | image generation   | list models        |
 | --------- | ------------------ | ------------------ | ------------------ | ------------------ |
 | anthropic | :white_check_mark: |                    |                    | :white_check_mark: |
 | fal       |                    |                    | :white_check_mark: |                    |

bunfig.toml

@@ -3,8 +3,10 @@ import { load } from "../src";
 const job = load({
   provider: "openai",
   type: "chat",
-  model: "gpt-4o-mini",
-  messages: [{ role: "user", content: "hi" }],
+  input: {
+    model: "gpt-4o-mini",
+    messages: [{ role: "user", content: "hi" }],
+  },
 });
 const result = await job.run();
 console.log(result);

examples/load.ts

@@ -1,4 +1,4 @@
-import { chunkText, openai, userPrompt } from "../src";
+import { text, openai, userPrompt } from "../src";
 
 const job = openai()
   .chat("gpt-4o-mini")
@@ -7,8 +7,7 @@ const job = openai()
 
 const stream = await job.run();
 for await (const chunk of stream) {
-  // console.log(chunk);
-  process.stdout.write(chunkText(chunk));
+  process.stdout.write(text(chunk));
 }
-const result = await job.done();
-console.log(result);
+const output = await job.done();
+console.log(output);

examples/openai-chat-load.ts

@@ -1,7 +1,9 @@
-import { openai, systemPrompt, userPrompt } from "../src";
+import { openai, systemPrompt, text, userPrompt } from "../src";
 
 const job = openai({})
   .chat("gpt-4o-mini")
   .messages([systemPrompt("you are a helpful assistant"), userPrompt("hi")]);
 const result = await job.run();
-console.log(JSON.stringify(result.raw, null, 2));
+console.log(text(result));
+console.log(job.cost);
+console.log(job.performance);

examples/openai-chat-stream.ts

@@ -1,5 +1,7 @@
 import { openai } from "../src";
 
-const job = openai().embedding("text-embedding-3-small").input("hello");
+const job = openai().embedding("text-embedding-3-small").value("hello");
 const result = await job.run();
-console.log(JSON.stringify(result, null, 2));
+console.log(result.embedding);
+console.log(job.cost);
+console.log(job.performance);

examples/openai-chat.ts

@@ -1,5 +1,12 @@
 import { version } from "../../package.json";
-import type { JobOptions, JobProvider } from "./schema";
+import type { Job } from "./load";
+import type {
+  JobCost,
+  JobOptions,
+  JobPerformance,
+  JobProvider,
+  JobType,
+} from "./schema";
 
 export class HTTPError extends Error {
   status: number;
@@ -12,14 +19,19 @@ export class HTTPError extends Error {
   }
 }
 
-export class JobBuilder {
+export class JobBuilder<Input, Output> {
   provider!: JobProvider;
   options!: JobOptions;
+  type!: JobType;
+  input?: Input;
+  output?: Output;
+  cost?: JobCost;
+  performance?: JobPerformance; // TODO: track job performance
 
   makeRequest?: () => Request;
   handleResponse?: (response: Response) => any;
 
-  async run() {
+  async run(): Promise<Output> {
     const request = this.makeRequest!();
     const response = await fetch(request);
     if (!response.ok) {
@@ -45,6 +57,11 @@ export class JobBuilder {
       version: version,
       provider: this.provider,
       options: this.options,
-    };
+      type: this.type,
+      input: this.input!,
+      output: this.output as any,
+      cost: this.cost,
+      performance: this.performance,
+    } as Job;
   }
 }

examples/openai-embedding.ts

@@ -1,82 +1,81 @@
 import { z } from "zod";
 import { JobBuilder } from "~/jobs/builder";
 import type {
-  ChatJobSchema,
+  ChatInput,
+  ChatOutput,
   ChatStreamOptions,
   Message,
   ResponseFormat,
 } from "./schema";
 import type { ChatTool } from "./tool";
 
-type ChatJob = z.infer<typeof ChatJobSchema>;
-
-export class ChatJobBuilder extends JobBuilder {
-  job: ChatJob;
+export class ChatJobBuilder extends JobBuilder<ChatInput, ChatOutput> {
+  input: ChatInput;
 
   constructor(model: string) {
     super();
-    this.job = {
-      type: "chat",
+    this.type = "chat";
+    this.input = {
       model: model,
       messages: [],
     };
   }
 
   systemPrompt(systemPrompt: string) {
-    this.job.systemPrompt = systemPrompt;
+    this.input.systemPrompt = systemPrompt;
     return this;
   }
 
   messages(messages: Message[]) {
-    this.job.messages = messages;
+    this.input.messages = messages;
     return this;
   }
 
   temperature(temperature: number) {
-    this.job.temperature = temperature;
+    this.input.temperature = temperature;
     return this;
   }
 
   maxTokens(maxTokens: number) {
-    this.job.maxTokens = maxTokens;
+    this.input.maxTokens = maxTokens;
     return this;
   }
 
   topP(topP: number) {
-    this.job.topP = topP;
+    this.input.topP = topP;
     return this;
   }
 
   topK(topK: number) {
-    this.job.topK = topK;
+    this.input.topK = topK;
     return this;
   }
 
   tools(tools: ChatTool[]) {
-    this.job.tools = tools.map((tool) => tool.params);
+    this.input.tools = tools.map((tool) => tool.params);
     return this;
   }
 
   tool(tool: ChatTool) {
-    if (!this.job.tools) {
-      this.job.tools = [];
+    if (!this.input.tools) {
+      this.input.tools = [];
     }
-    this.job.tools.push(tool.params);
+    this.input.tools.push(tool.params);
     return this;
   }
 
   toolChoice(toolChoice: string) {
-    this.job.toolChoice = toolChoice;
+    this.input.toolChoice = toolChoice;
     return this;
   }
 
   responseFormat(responseFormat: ResponseFormat) {
-    this.job.responseFormat = responseFormat;
+    this.input.responseFormat = responseFormat;
     return this;
   }
 
   jsonSchema(schema: z.ZodType, name: string, description?: string) {
-    this.job.jsonSchema = {
+    this.input.jsonSchema = {
       name,
       description,
       schema,
@@ -86,17 +85,10 @@ export class ChatJobBuilder extends JobBuilder {
   }
 
   stream(streamOptions?: ChatStreamOptions) {
-    this.job.stream = true;
+    this.input.stream = true;
     if (streamOptions) {
-      this.job.streamOptions = streamOptions;
+      this.input.streamOptions = streamOptions;
     }
     return this;
   }
-
-  dump() {
-    return {
-      ...super.dump(),
-      ...this.job,
-    };
-  }
 }

src/jobs/builder.ts

@@ -85,8 +85,7 @@ export const ChatResultSchema = z.object({
     .optional(),
 });
 
-export const ChatJobSchema = BaseJobSchema.extend({
-  type: z.literal("chat"),
+export const ChatInputSchema = z.object({
   model: z.string(),
   temperature: z.number().optional(),
   stream: z.boolean().optional(),
@@ -100,6 +99,18 @@ export const ChatJobSchema = BaseJobSchema.extend({
   topK: z.number().optional(),
   systemPrompt: z.string().optional(),
   jsonSchema: JsonSchemaDefSchema.optional(),
+});
 
-  result: ChatResultSchema.optional(),
+export const ChatOutputSchema = z.object({
+  raw: z.any(),
 });
+
+export const ChatJobSchema = BaseJobSchema.extend({
+  type: z.literal("chat"),
+  input: ChatInputSchema,
+  output: ChatOutputSchema.optional(),
+});
+
+export type ChatInput = z.infer<typeof ChatInputSchema>;
+
+export type ChatOutput = z.infer<typeof ChatOutputSchema>;