packages/llm/test/llm.test.ts

import { describe, expect, test } from "bun:test"
import { CacheHint, LLM, LLMResponse } from "../src"
import * as OpenAIChat from "../src/protocols/openai-chat"
import * as OpenAIResponses from "../src/protocols/openai-responses"
import { LLMRequest, Message, Model, ToolCallPart, ToolChoice, ToolDefinition, ToolResultPart } from "../src/schema"

const chatRoute = OpenAIChat.route
const responsesRoute = OpenAIResponses.route

describe("llm constructors", () => {
  test("builds canonical schema classes from ergonomic input", () => {
    const request = LLM.request({
      id: "req_1",
      model: Model.make({ id: "fake-model", provider: "fake", route: chatRoute }),
      system: "You are concise.",
      prompt: "Say hello.",
    })

    expect(request).toBeInstanceOf(LLMRequest)
    expect(request.model).toBeInstanceOf(Model)
    expect(request.messages[0]).toBeInstanceOf(Message)
    expect(request.system).toEqual([{ type: "text", text: "You are concise." }])
    expect(request.messages[0]?.content).toEqual([{ type: "text", text: "Say hello." }])
    expect(request.generation).toBeUndefined()
    expect(request.tools).toEqual([])
  })

  test("updates requests without spreading schema class instances", () => {
    const base = LLM.request({
      id: "req_1",
      model: Model.make({ id: "fake-model", provider: "fake", route: chatRoute }),
      prompt: "Say hello.",
    })
    const updated = LLM.updateRequest(base, {
      generation: { maxTokens: 20 },
      messages: [...base.messages, Message.assistant("Hi.")],
    })

    expect(updated).toBeInstanceOf(LLMRequest)
    expect(updated.id).toBe("req_1")
    expect(updated.model).toEqual(base.model)
    expect(updated.generation).toEqual({ maxTokens: 20 })
    expect(updated.messages.map((message) => message.role)).toEqual(["user", "assistant"])
  })

  test("keeps request options separate from route defaults", () => {
    const request = LLM.request({
      model: Model.make({
        id: "fake-model",
        provider: "fake",
        route: chatRoute.with({
          generation: { maxTokens: 100, temperature: 1 },
          providerOptions: { openai: { store: false, metadata: { model: true } } },
          http: { body: { metadata: { model: true } }, headers: { "x-shared": "model" }, query: { model: "1" } },
        }),
      }),
      prompt: "Say hello.",
      generation: { temperature: 0 },
      providerOptions: { openai: { store: true, metadata: { request: true } } },
      http: { body: { metadata: { request: true } }, headers: { "x-shared": "request" }, query: { request: "1" } },
    })

    expect(request.generation).toEqual({ temperature: 0 })
    expect(request.providerOptions).toEqual({ openai: { store: true, metadata: { request: true } } })
    expect(request.http).toEqual({
      body: { metadata: { request: true } },
      headers: { "x-shared": "request" },
      query: { request: "1" },
    })
  })

  test("updates canonical requests from the request datatype", () => {
    const base = LLM.request({
      id: "req_1",
      model: Model.make({ id: "fake-model", provider: "fake", route: chatRoute }),
      prompt: "Say hello.",
    })
    const updated = LLMRequest.update(base, { messages: [...base.messages, Message.assistant("Hi.")] })

    expect(updated).toBeInstanceOf(LLMRequest)
    expect(updated.id).toBe("req_1")
    expect(LLMRequest.input(updated).id).toBe("req_1")
    expect(updated.messages.map((message) => message.role)).toEqual(["user", "assistant"])
    expect(LLMRequest.update(updated, {})).toBe(updated)
  })

  test("updates canonical models from the model datatype", () => {
    const base = Model.make({
      id: "fake-model",
      provider: "fake",
      route: chatRoute,
    })
    const updated = Model.update(base, {
      route: responsesRoute,
      defaults: { generation: { maxTokens: 20 } },
      compatibility: { toolSchema: "gemini" },
    })
    const updatedInput = Model.input(updated)

    expect(updated).toBeInstanceOf(Model)
    expect(String(updated.id)).toBe("fake-model")
    expect(updated.route).toBe(responsesRoute)
    expect(updated.defaults?.generation).toEqual({ maxTokens: 20 })
    expect(updated.compatibility).toEqual({ toolSchema: "gemini" })
    expect(updatedInput.defaults).toBe(updated.defaults)
    expect(updatedInput.compatibility).toBe(updated.compatibility)
    expect(String(updatedInput.provider)).toBe("fake")
    expect(Model.update(updated, {})).toBe(updated)
  })

  test("carries model defaults and compatibility through route model selection", () => {
    const model = chatRoute.model({
      id: "kimi-k2",
      defaults: {
        limits: { context: 128_000, output: 8_192 },
        generation: { maxTokens: 1_024, stop: ["END"] },
        providerOptions: { openai: { parallelToolCalls: false } },
        http: { body: { extra_body: true } },
      },
      compatibility: { toolSchema: "moonshot" },
    })
    const request = LLM.request({ model, prompt: "Say hello." })

    expect(request.model.defaults?.limits).toEqual({ context: 128_000, output: 8_192 })
    expect(request.model.defaults?.generation).toEqual({ maxTokens: 1_024, stop: ["END"] })
    expect(request.model.defaults?.providerOptions).toEqual({ openai: { parallelToolCalls: false } })
    expect(request.model.defaults?.http).toEqual({ body: { extra_body: true } })
    expect(request.model.compatibility).toEqual({ toolSchema: "moonshot" })
    expect(request.generation).toBeUndefined()
    expect(request.providerOptions).toBeUndefined()
    expect(request.http).toBeUndefined()
  })

  test("builds tool choices from names and tools", () => {
    const tool = ToolDefinition.make({ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } })

    expect(tool).toBeInstanceOf(ToolDefinition)
    expect(ToolChoice.make("lookup")).toEqual(new ToolChoice({ type: "tool", name: "lookup" }))
    expect(ToolChoice.named("required")).toEqual(new ToolChoice({ type: "tool", name: "required" }))
    expect(ToolChoice.make(tool)).toEqual(new ToolChoice({ type: "tool", name: "lookup" }))
  })

  test("builds tool choice modes from reserved strings", () => {
    expect(ToolChoice.make("auto")).toEqual(new ToolChoice({ type: "auto" }))
    expect(ToolChoice.make("none")).toEqual(new ToolChoice({ type: "none" }))
    expect(ToolChoice.make("required")).toEqual(new ToolChoice({ type: "required" }))
    expect(
      LLM.request({
        model: Model.make({
          id: "fake-model",
          provider: "fake",
          route: chatRoute,
        }),
        prompt: "Use tools if needed.",
        toolChoice: "required",
      }).toolChoice,
    ).toEqual(new ToolChoice({ type: "required" }))
  })

  test("builds assistant tool calls and tool result messages", () => {
    const call = ToolCallPart.make({ id: "call_1", name: "lookup", input: { query: "weather" } })
    const result = ToolResultPart.make({ id: "call_1", name: "lookup", result: { temperature: 72 } })

    expect(Message.assistant([call]).content).toEqual([call])
    expect(Message.tool(result).content).toEqual([
      { type: "tool-result", id: "call_1", name: "lookup", result: { type: "json", value: { temperature: 72 } } },
    ])
  })

  test("builds chronological text-only system updates separately from the initial system prompt", () => {
    const update = Message.system([
      { type: "text", text: "Use parameterized SQL.", cache: new CacheHint({ type: "ephemeral" }) },
    ])
    const request = LLM.request({
      model: Model.make({ id: "fake-model", provider: "fake", route: chatRoute }),
      system: "Initial operator prompt.",
      messages: [Message.user("Review this."), update],
    })

    expect(update).toBeInstanceOf(Message)
    expect(update).toEqual({
      role: "system",
      content: [{ type: "text", text: "Use parameterized SQL.", cache: { type: "ephemeral" } }],
    })
    expect(request.system).toEqual([{ type: "text", text: "Initial operator prompt." }])
    expect(request.messages.map((message) => message.role)).toEqual(["user", "system"])
  })

  test("extracts output text from response events", () => {
    expect(
      LLMResponse.text({
        events: [
          { type: "text-delta", id: "text-0", text: "hi" },
          { type: "finish", reason: "stop" },
        ],
      }),
    ).toBe("hi")
  })
})