import { beforeAll, describe, expect, it } from "vitest";
import { initializeLogger } from "../log.js";
import { DEFAULT_API_CONNECT_OPTIONS } from "../types.js";
import { delay } from "../utils.js";
import { ChatContext, FunctionCall } from "./chat_context.js";
import { LLM, LLMStream } from "./llm.js";
class MockLLMStream extends LLMStream {
  constructor(llm, opts, chunks) {
    super(llm, opts);
    this.chunks = chunks;
  }
  async run() {
    for (const chunk of this.chunks) {
      this.queue.put(chunk);
      await delay(1);
    }
  }
}
class MockLLM extends LLM {
  constructor(chunks) {
    super();
    this.chunks = chunks;
  }
  label() {
    return "mock-llm";
  }
  chat(opts) {
    return new MockLLMStream(
      this,
      {
        chatCtx: opts.chatCtx,
        toolCtx: opts.toolCtx,
        connOptions: opts.connOptions ?? DEFAULT_API_CONNECT_OPTIONS
      },
      this.chunks
    );
  }
}
describe("LLMStream.collect", () => {
  beforeAll(() => {
    initializeLogger({ pretty: false });
    process.on("unhandledRejection", () => {
    });
  });
  it("joins content parts and trims surrounding whitespace", async () => {
    const llm = new MockLLM([
      { id: "1", delta: { role: "assistant", content: "  Hello" } },
      { id: "1", delta: { role: "assistant", content: ", " } },
      { id: "1", delta: { role: "assistant", content: "world!  " } }
    ]);
    const response = await llm.chat({ chatCtx: new ChatContext() }).collect();
    expect(response.text).toBe("Hello, world!");
    expect(response.toolCalls).toHaveLength(0);
    expect(response.usage).toBeUndefined();
    expect(response.extra).toEqual({});
  });
  it("accumulates tool calls across chunks", async () => {
    const callA = new FunctionCall({
      callId: "call_a",
      name: "get_weather",
      args: '{"city":"SF"}'
    });
    const callB = new FunctionCall({
      callId: "call_b",
      name: "play_song",
      args: '{"name":"x"}'
    });
    const llm = new MockLLM([
      { id: "1", delta: { role: "assistant", toolCalls: [callA] } },
      { id: "1", delta: { role: "assistant", toolCalls: [callB] } }
    ]);
    const response = await llm.chat({ chatCtx: new ChatContext() }).collect();
    expect(response.text).toBe("");
    expect(response.toolCalls).toHaveLength(2);
    expect(response.toolCalls[0].callId).toBe("call_a");
    expect(response.toolCalls[1].callId).toBe("call_b");
  });
  it("captures the latest usage and merges extra data", async () => {
    var _a, _b;
    const llm = new MockLLM([
      { id: "1", delta: { role: "assistant", content: "hi", extra: { a: 1 } } },
      {
        id: "1",
        delta: { role: "assistant", content: " there", extra: { b: 2 } },
        usage: {
          completionTokens: 2,
          promptTokens: 5,
          promptCachedTokens: 0,
          totalTokens: 7
        }
      },
      {
        id: "1",
        usage: {
          completionTokens: 3,
          promptTokens: 5,
          promptCachedTokens: 0,
          totalTokens: 8
        }
      }
    ]);
    const response = await llm.chat({ chatCtx: new ChatContext() }).collect();
    expect(response.text).toBe("hi there");
    expect((_a = response.usage) == null ? void 0 : _a.completionTokens).toBe(3);
    expect((_b = response.usage) == null ? void 0 : _b.totalTokens).toBe(8);
    expect(response.extra).toEqual({ a: 1, b: 2 });
  });
  it("returns empty response for an empty stream", async () => {
    const llm = new MockLLM([]);
    const response = await llm.chat({ chatCtx: new ChatContext() }).collect();
    expect(response.text).toBe("");
    expect(response.toolCalls).toHaveLength(0);
    expect(response.usage).toBeUndefined();
    expect(response.extra).toEqual({});
  });
});
//# sourceMappingURL=llm.test.js.map