infer

onCheckpoint: async ({ step, infer }) => {
  const res = await infer({
    messages: [
      { role: "user", content: "I can't log in." },
    ],
  });
  console.log(`step=${step} sample=`, await res.text());
}

type Infer = (args: InferArgs) => Promise<Response>;

interface InferArgs {
  messages: ChatMessage[];
  temperature?: number;
  topP?: number;
  maxTokens?: number;
  /** Default: true. Set false to get a single JSON body instead of SSE. */
  stream?: boolean;
  /** OpenAI-compatible function-calling tool definitions. */
  tools?: ToolDefinition[];
  /** "auto" | "none" | "required" | { type: "function", function: { name } } */
  toolChoice?: ToolChoice;
  /** OpenAI-compatible response_format (text / json_object / json_schema). */
  responseFormat?: ResponseFormat;
  /** vLLM structured outputs (regex / choice / grammar) for cases response_format can't express. */
  structuredOutputs?: StructuredOutputs;
  signal?: AbortSignal;
}

onCheckpoint: async ({ infer }) => {
  const res = await infer({
    messages: [
      { role: "user", content: "What's the weather in Tokyo?" },
    ],
    tools: [
      {
        type: "function",
        function: {
          name: "get_weather",
          parameters: {
            type: "object",
            properties: { city: { type: "string" } },
            required: ["city"],
          },
        },
      },
    ],
    toolChoice: "auto",
    stream: false,
  });
  const data = (await res.json()) as { choices: Array<{ message: ChatMessage }> };
  // data.choices[0].message may be { role: "assistant", tool_calls: [...] }
};

const res = await infer({
  messages: [{ role: "user", content: "Extract the user's email." }],
  responseFormat: {
    type: "json_schema",
    json_schema: {
      name: "user",
      schema: {
        type: "object",
        properties: { email: { type: "string", format: "email" } },
        required: ["email"],
        additionalProperties: false,
      },
      strict: true,
    },
  },
});

const res = await infer({
  messages: [{ role: "user", content: "Classify urgency: I can't log in." }],
  structuredOutputs: { choice: ["low", "medium", "high"] },
  stream: false,
});

const res = await infer({
  messages: [{ role: "user", content: "Generate a ticket id." }],
  structuredOutputs: { regex: "^[A-Z]{3}-\\d{4}$" },
  stream: false,
});

const res = await infer({
  messages: [{ role: "user", content: "Spell out the digit." }],
  structuredOutputs: {
    grammar: 'root ::= "zero" | "one" | "two" | "three"',
  },
  stream: false,
});

import type { ChatMessage } from "arkor";

const tools = [
  {
    type: "function" as const,
    function: {
      name: "get_weather",
      parameters: {
        type: "object",
        properties: { city: { type: "string" } },
        required: ["city"],
      },
    },
  },
];

const messages: ChatMessage[] = [
  { role: "user", content: "What's the weather in Tokyo?" },
];

const first = await infer({ messages, tools, toolChoice: "auto", stream: false });
const firstData = (await first.json()) as { choices: Array<{ message: ChatMessage }> };
const reply = firstData.choices[0].message;

if (reply.role === "assistant" && reply.tool_calls?.length) {
  const call = reply.tool_calls[0];
  const args = JSON.parse(call.function.arguments) as { city: string };
  const result = await getWeather(args.city);

  const second = await infer({
    messages: [
      ...messages,
      reply,                                        // assistant turn carrying tool_calls
      { role: "tool", tool_call_id: call.id, content: JSON.stringify(result) },
    ],
    tools,
    toolChoice: "auto",
    stream: false,
  });
  const finalReply = ((await second.json()) as { choices: Array<{ message: ChatMessage }> })
    .choices[0].message;
  // finalReply.content is the user-facing answer.
}

export type ChatMessage =
  | { role: "system"; content: string }
  | { role: "user"; content: string }
  | { role: "assistant"; content: string; tool_calls?: ToolCall[] }
  | {
      role: "assistant";
      /** May be omitted or `null` when the turn is purely a tool call. */
      content?: string | null;
      tool_calls: [ToolCall, ...ToolCall[]];
    }
  | { role: "tool"; content: string; tool_call_id: string };

export interface ToolCall {
  id: string;
  type: "function";
  function: {
    name: string;
    /** JSON-encoded arguments string; partial deltas may be streamed. */
    arguments: string;
  };
}

export interface ToolDefinition {
  type: "function";
  function: {
    name: string;
    description?: string;
    /** JSON Schema describing the tool arguments. */
    parameters?: Record<string, unknown>;
    strict?: boolean;
  };
}

export type ToolChoice =
  | "auto"
  | "none"
  | "required"
  | { type: "function"; function: { name: string } };

export type ResponseFormat =
  | { type: "text" }
  | { type: "json_object" }
  | {
      type: "json_schema";
      json_schema: {
        name: string;
        description?: string;
        schema: Record<string, unknown>;
        strict?: boolean;
      };
    };

// Helper: a record where exactly one key is required and every
// sibling is forbidden (typed `never`). Encodes vLLM's
// "must specify exactly one constraint" invariant at the type level.
// Internal to the SDK; shown here so the StructuredOutputs definition
// below is self-contained.
type ExactlyOne<T> = {
  [K in keyof T]: { [P in K]: T[K] } & {
    [P in Exclude<keyof T, K>]?: never;
  };
}[keyof T];

// Exactly one of `json` / `regex` / `choice` / `grammar` /
// `json_object` must be set. Field names are snake_case to match
// vLLM's wire format verbatim.
export type StructuredOutputs = ExactlyOne<{
  json: Record<string, unknown>;
  regex: string;
  choice: string[];
  grammar: string;
  json_object: true;
}> & {
  disable_any_whitespace?: boolean;
  disable_additional_properties?: boolean;
  whitespace_pattern?: string;
};

// Streaming (default)
const res = await infer({ messages });
for await (const chunk of res.body!) {
  // chunk: Uint8Array of one or more SSE frames
}

// Or read the whole stream at once
const text = await res.text();

// Or, if you set stream: false, parse the JSON body
const res = await infer({ messages, stream: false });
const data = await res.json();

{
  choices: [
    {
      index: 0,
      finish_reason: "stop" | "length" | "tool_calls" | string,
      message: ChatMessage,
    },
  ],
  // plus the OpenAI-standard `id` / `model` / `usage` fields the backend chooses to surface.
}

import { CloudApiError } from "arkor";

try {
  const res = await infer({ messages, ... });
  // Use the Response.
} catch (err) {
  if (err instanceof CloudApiError) {
    // err.status is the HTTP status; err.message carries the upstream message.
  }
}