update openai response and stream handling

bruce-shi · bruce-shi · commit 3f0dc744d0a1 · 2025-04-27T21:05:48.000-07:00
diff --git a/examples/openai-chat-stream.ts b/examples/openai-chat-stream.ts
@@ -4,7 +4,6 @@ const job = openai()
   .chat("gpt-4o-mini")
   .prompt("generate a 50 words text")
   .stream();
-const stream = await job.run();
-for await (const chunk of stream) {
-  process.stdout.write(text(chunk));
+for await (const chunk of job) {
+  console.log(chunk?.message);
 }
diff --git a/examples/openai-chat.ts b/examples/openai-chat.ts
@@ -4,4 +4,4 @@ const job = openai({})
   .chat("gpt-4o-mini")
   .messages([system("you are a helpful assistant"), user("hi")]);
 const result = await job.run();
-console.log(text(result));
+console.log(result?.message);
diff --git a/examples/openai-embedding.ts b/examples/openai-embedding.ts
@@ -2,4 +2,4 @@ import { openai } from "../src";
 
 const job = openai().embedding("text-embedding-3-small").value("hello");
 const result = await job.run();
-console.log(result.embedding);
+console.log(result!.embedding);
diff --git a/examples/openai-models.ts b/examples/openai-models.ts
@@ -2,6 +2,6 @@ import { openai } from "../src";
 
 const job = openai().models();
 const result = await job.run();
-for (const model of result.raw.data) {
-  console.log(model.id);
+for (const model of result!) {
+  console.log(model);
 }
diff --git a/src/jobs/builder.ts b/src/jobs/builder.ts
@@ -22,11 +22,13 @@ export abstract class JobBuilder<Job extends BaseJob> {
   performance?: Job["performance"]; // TODO: track job performance
 
   abstract makeRequest(): Request;
-  abstract handleResponse(
-    response: Response,
-  ): Promise<Job["output"] | AsyncGenerator<Job["output"]>>;
 
-  async run(): Promise<Job["output"] | AsyncGenerator<Job["output"]>> {
+
+  async handleResponse(response: Response): Promise<Job["output"]> {
+    throw new Error("Not implemented");
+  }
+  
+  async run(): Promise<Job["output"]> {
     const request = this.makeRequest!();
     const response = await fetch(request);
     if (!response.ok) {
@@ -43,6 +45,7 @@ export abstract class JobBuilder<Job extends BaseJob> {
     }
     return await this.handleResponse!(response);
   }
+
   dump() {
     return {
       version: version,
diff --git a/src/jobs/chat/builder.ts b/src/jobs/chat/builder.ts
@@ -1,8 +1,9 @@
 import { z } from "zod";
-import { JobBuilder } from "~/jobs/builder";
+import { HTTPError, JobBuilder } from "~/jobs/builder";
 import type {
   ChatJob,
   ChatStreamOptions,
+  ChatToolChoiceSchema,
   Message,
   ResponseFormat,
 } from "./schema";
@@ -22,6 +23,27 @@ export abstract class ChatJobBuilder<
     };
   }
 
+  async *handleStream(response: Response): AsyncGenerator<Job["output"]> {
+    throw new Error("Not implemented");
+  }
+
+  async *stream(options?: ChatStreamOptions): AsyncGenerator<Job["output"]> {
+    this.input.stream = true;
+    this.input.streamOptions = options;
+    if (!this.handleStream) {
+      throw new Error("Stream not supported");
+    }
+    const request = this.makeRequest!();
+    const response = await fetch(request);
+    if (!response.ok) {
+      throw new HTTPError(
+        `Fetch error: ${response.statusText}`,
+        response.status,
+      );
+    }
+    yield* this.handleStream(response);
+  }
+
   system(system: string) {
     this.input.system = system;
     return this;
@@ -73,7 +95,7 @@ export abstract class ChatJobBuilder<
     return this;
   }
 
-  toolChoice(toolChoice: string) {
+  toolChoice(toolChoice: z.infer<typeof ChatToolChoiceSchema>) {
     this.input.toolChoice = toolChoice;
     return this;
   }
@@ -92,12 +114,4 @@ export abstract class ChatJobBuilder<
 
     return this;
   }
-
-  stream(streamOptions?: ChatStreamOptions) {
-    this.input.stream = true;
-    if (streamOptions) {
-      this.input.streamOptions = streamOptions;
-    }
-    return this;
-  }
 }
diff --git a/src/jobs/chat/schema.ts b/src/jobs/chat/schema.ts
@@ -97,6 +97,11 @@ export const JsonSchemaDefSchema = z.object({
 
 export const ChunkSchema = z.object({});
 
+export const ChatToolChoiceSchema = z.object({
+  mode: z.enum(["auto", "none", "any"]),
+  allowed_tools: z.array(z.string()).optional(),
+});
+
 export const ChatInputSchema = z.object({
   model: z.string(),
   temperature: z.number().optional(),
@@ -105,7 +110,8 @@ export const ChatInputSchema = z.object({
   maxTokens: z.number().optional(),
   messages: z.array(MessageSchema),
   tools: z.array(ChatToolSchema).optional(),
-  toolChoice: z.string().optional(),
+  //TODO: support gemini and anthropic tool choice, might need refactor
+  toolChoice: ChatToolChoiceSchema.optional(),
   responseFormat: ResponseFormatSchema.optional(),
   topP: z.number().optional(),
   topK: z.number().optional(),
diff --git a/src/jobs/chat/utils.ts b/src/jobs/chat/utils.ts
@@ -3,7 +3,10 @@ import { parse } from "partial-json";
 import type { ChatToolSchema, Message } from "./schema";
 import { ChatTool } from "./tool";
 
+// TODO: move to providers, different providers have different tool formats
+// @deprecated
 export function convertTools(tools: z.infer<typeof ChatToolSchema>[]) {
+  console.warn("convertTools is deprecated, move to providers");
   return tools.map((tool) => ({
     type: "function",
     function: {
diff --git a/src/jobs/image/builder.ts b/src/jobs/image/builder.ts
@@ -93,9 +93,4 @@ export abstract class ImageJobBuilder<
     this.input.background = background;
     return this;
   }
-
-  stream() {
-    this.input.stream = true;
-    return this;
-  }
 }
diff --git a/src/jobs/models/schema.ts b/src/jobs/models/schema.ts
@@ -3,7 +3,13 @@ import { BaseJobSchema } from "~/jobs/schema";
 
 export const ModelsInputSchema = z.object({});
 
-export const ModelsOuputSchema = z.object({});
+export const ModelsOuputSchema = z.array(
+  z.object({
+    id: z.string(),
+    created: z.number(),
+    owned_by: z.string(),
+  }),
+);
 
 export const ModelsJobSchema = BaseJobSchema.extend({
   type: z.literal("models"),
diff --git a/src/providers/openai/chat.ts b/src/providers/openai/chat.ts
@@ -1,8 +1,21 @@
 import { z } from "zod";
-import { ChatJobBuilder, convertTools, type ChatOutput } from "~/jobs/chat";
+import { ChatJobBuilder, type ChatOutput } from "~/jobs/chat";
 import type { JobOptions } from "~/jobs/schema";
-import { jobStream } from "~/jobs/stream";
 import { OPENAI_BASE_URL, type OpenAIChatJob } from "./schema";
+import { EventSourceParserStream } from "eventsource-parser/stream";
+import type {
+  OpenAIResponse,
+  OpenAIStreamResponse,
+  OpenAIToolCall,
+  OpenAIDelta,
+  FormattedToolCall,
+  OpenAIChatCompletionRequest,
+  OpenAIMessageRequest,
+  OpenAIResponseFormat,
+  OpenAIJSONSchemaResponseFormat,
+  OpenAIToolDefinition,
+} from "./types";
+import { convertToolChoice, convertTools } from "./utils";
 
 export class OpenAIChatJobBuilder extends ChatJobBuilder<OpenAIChatJob> {
   constructor(options: JobOptions, model: string) {
@@ -13,37 +26,56 @@ export class OpenAIChatJobBuilder extends ChatJobBuilder<OpenAIChatJob> {
 
   makeRequest = () => {
     const baseURL = this.options!.baseURL || OPENAI_BASE_URL;
-    const messages = this.input.messages;
+    const messages = this.input.messages as unknown as OpenAIMessageRequest[];
 
     if (this.input.system) {
       messages.unshift({
         role: "system",
         content: this.input.system,
       });
     }
-    const requestBody = {
-      messages: messages,
+
+    const requestBody: OpenAIChatCompletionRequest = {
+      messages,
       model: this.input.model,
-      temperature: this.input.temperature,
       stream: this.input.stream,
-      response_format: this.input.responseFormat,
-    } as any;
+    };
+
+    if (this.input.temperature !== undefined) {
+      requestBody.temperature = this.input.temperature;
+    }
+
+    if (this.input.maxTokens !== undefined) {
+      requestBody.max_tokens = this.input.maxTokens;
+    }
+
+    if (this.input.topP !== undefined) {
+      requestBody.top_p = this.input.topP;
+    }
 
     if (this.input.tools && this.input.tools.length) {
       requestBody.tools = convertTools(this.input.tools);
-      requestBody.tool_choice = this.input.toolChoice;
+
+      if (this.input.toolChoice) {
+        requestBody.tool_choice = convertToolChoice(this.input.toolChoice);
+      }
+    }
+
+    if (this.input.responseFormat) {
+      requestBody.response_format = this.input
+        .responseFormat as OpenAIResponseFormat;
     }
 
     if (this.input.jsonSchema) {
       const schema = z.toJSONSchema(this.input.jsonSchema.schema);
       requestBody.response_format = {
         type: "json_schema",
-        json_schema: {
+        schema: {
           name: this.input.jsonSchema.name,
           description: this.input.jsonSchema.description,
-          schema: schema,
+          ...schema,
         },
-      };
+      } as OpenAIJSONSchemaResponseFormat;
     }
 
     return new Request(`${baseURL}/chat/completions`, {
@@ -57,16 +89,117 @@ export class OpenAIChatJobBuilder extends ChatJobBuilder<OpenAIChatJob> {
   };
 
   async handleResponse(response: Response) {
-    if (this.input.stream) {
-      return jobStream<ChatOutput>(response);
+    // Handle non-streaming response
+    const raw = (await response.json()) as OpenAIResponse;
+
+    // Record cost/usage information
+    if (raw.usage) {
+      this.cost = {
+        promptTokens: raw.usage.prompt_tokens,
+        completionTokens: raw.usage.completion_tokens,
+        totalTokens: raw.usage.total_tokens,
+      };
     }
 
-    const raw = await response.json();
-    this.cost = {
-      promptTokens: raw.usage.prompt_tokens,
-      completionTokens: raw.usage.completion_tokens,
-      totalTokens: raw.usage.total_tokens,
+    // Format the response to match ChatOutput schema
+    const choice = raw.choices?.[0];
+    const message = choice?.message;
+
+    const output: ChatOutput = {
+      raw,
     };
-    return { raw };
+
+    if (message) {
+      output.message = {
+        role: "assistant",
+        content: message.content ?? "",
+      };
+
+      // Handle tool calls if present
+      if (message.tool_calls && message.tool_calls.length > 0) {
+        output.message.tool_calls = message.tool_calls.map(
+          (call: OpenAIToolCall): FormattedToolCall => {
+            const result: FormattedToolCall = {
+              id: call.id,
+              type: call.type,
+              name: call.function?.name || "",
+              arguments: {} as Record<string, unknown>,
+            };
+
+            if (call.function) {
+              // Safely handle arguments
+              if (call.function.arguments) {
+                try {
+                  result.arguments = JSON.parse(call.function.arguments);
+                } catch (e) {
+                  // If parsing fails, return as string
+                  result.arguments = { _raw: call.function.arguments };
+                }
+              }
+            }
+
+            return result;
+          },
+        );
+      }
+    }
+
+    return output;
+  }
+
+  async *handleStream(response: Response) {
+    const eventStream = response
+      .body!.pipeThrough(new TextDecoderStream())
+      .pipeThrough(new EventSourceParserStream());
+    const reader = eventStream.getReader();
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done || value.data === "[DONE]") {
+        break;
+      }
+
+      const chunk = JSON.parse(value.data) as OpenAIStreamResponse;
+      const delta = chunk.choices?.[0]?.delta;
+
+      // Create a properly formatted ChatOutput object
+      const output: ChatOutput = { raw: chunk };
+
+      if (delta) {
+        output.message = {
+          role: "assistant",
+          content: delta.content || "",
+        };
+
+        // Handle tool calls in stream
+        if (delta.tool_calls && delta.tool_calls.length > 0) {
+          output.message.tool_calls = delta.tool_calls.map(
+            (call): FormattedToolCall => {
+              const result: FormattedToolCall = {
+                id: call.id || "",
+                type: call.type || "",
+                name: call.function?.name || "",
+                arguments: {} as Record<string, unknown>,
+              };
+
+              if (call.function) {
+                // Safely handle arguments which might be partial JSON in streaming
+                if (call.function.arguments) {
+                  try {
+                    result.arguments = JSON.parse(call.function.arguments);
+                  } catch (e) {
+                    // If parsing fails (for partial JSON), return as string
+                    result.arguments = { _raw: call.function.arguments };
+                  }
+                }
+              }
+
+              return result;
+            },
+          );
+        }
+      }
+
+      yield output;
+    }
   }
 }
diff --git a/src/providers/openai/models.ts b/src/providers/openai/models.ts
diff --git a/src/providers/openai/types.ts b/src/providers/openai/types.ts
diff --git a/src/providers/openai/utils.ts b/src/providers/openai/utils.ts

Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,6 @@ import { openai } from "../src";`
`2`	`2`
`3`	`3`	`const job = openai().models();`
`4`	`4`	`const result = await job.run();`
`5`		`-for (const model of result.raw.data) {`
`6`		`- console.log(model.id);`
	`5`	`+for (const model of result!) {`
	`6`	`+ console.log(model);`
`7`	`7`	`}`
Original file line number	Diff line number	Diff line change
`@@ -93,9 +93,4 @@ export abstract class ImageJobBuilder<`
`93`	`93`	`this.input.background = background;`
`94`	`94`	`return this;`
`95`	`95`	`}`
`96`		`-`
`97`		`- stream() {`
`98`		`- this.input.stream = true;`
`99`		`- return this;`
`100`		`- }`
`101`	`96`	`}`