MagicCode/packages/providers/src/litellm/litellm.test.ts

import { mkdtemp, mkdir, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { describe, expect, it } from "vitest";
import type { StreamEvent } from "@maia/core";
import { buildChatBody, parseSse } from "./openai";
import { joinUrl, loadSettings } from "./settings";
import { parseModelSpec } from "../index";

async function* chunks(...parts: string[]): AsyncGenerator<Uint8Array> {
  for (const p of parts) yield new TextEncoder().encode(p);
}
async function collect(it: AsyncIterable<StreamEvent>): Promise<StreamEvent[]> {
  const out: StreamEvent[] = [];
  for await (const e of it) out.push(e);
  return out;
}
const sse = (o: unknown): string => `data: ${JSON.stringify(o)}\n\n`;

describe("settings", () => {
  async function home(content?: string): Promise<string> {
    const h = await mkdtemp(join(tmpdir(), "maia-"));
    if (content !== undefined) {
      await mkdir(join(h, ".maia"));
      await writeFile(join(h, ".maia", "settings.json"), content);
    }
    return h;
  }
  it("charge les paramètres", async () => {
    const h = await home('{"LiteLLM_URL":" http://x:4000 ","OpenAI_EndPoint":"/v1/chat/completions","LiteLLM_Models":"/v1/models"}');
    expect(await loadSettings(h)).toEqual({
      LiteLLM_URL: "http://x:4000",
      OpenAI_EndPoint: "/v1/chat/completions",
      LiteLLM_Models: "/v1/models",
    });
  });
  it("erreurs : absent, JSON invalide, champ manquant", async () => {
    await expect(loadSettings(await home())).rejects.toThrow(/introuvable/);
    await expect(loadSettings(await home("{"))).rejects.toThrow(/JSON valide/);
    await expect(loadSettings(await home('{"LiteLLM_URL":"http://x"}'))).rejects.toThrow(/OpenAI_EndPoint, LiteLLM_Models/);
  });
  it("joinUrl", () => {
    expect(joinUrl("http://x:4000/", "/v1/models")).toBe("http://x:4000/v1/models");
    expect(joinUrl("http://x:4000", "v1/models")).toBe("http://x:4000/v1/models");
  });
  it("parseModelSpec", () => {
    expect(parseModelSpec("litellm:gpt-4o")).toBe("gpt-4o");
    expect(parseModelSpec("bedrock/claude:3")).toBe("bedrock/claude:3");
    expect(() => parseModelSpec("litellm:")).toThrow();
  });
});

describe("format OpenAI", () => {
  it("corps de requête : system, outils, résultats d'outils avant le texte", () => {
    const body = buildChatBody("m", {
      system: "S",
      tools: [{ name: "Read", description: "d", inputSchema: { type: "object" } }],
      maxTokens: 100,
      messages: [
        { role: "user", content: [{ type: "text", text: "hi" }] },
        { role: "assistant", content: [{ type: "text", text: "ok" }, { type: "tool_use", id: "c1", name: "Read", input: { a: 1 } }] },
        { role: "user", content: [{ type: "text", text: "suite" }, { type: "tool_result", toolUseId: "c1", content: "R" }] },
      ],
    });
    expect(body).toMatchObject({ model: "m", stream: true, max_tokens: 100 });
    const msgs = body.messages as { role: string }[];
    expect(msgs.map((m) => m.role)).toEqual(["system", "user", "assistant", "tool", "user"]);
    expect(msgs[2]).toMatchObject({ tool_calls: [{ id: "c1", function: { name: "Read", arguments: '{"a":1}' } }] });
  });

  it("streaming : texte découpé en morceaux arbitraires", async () => {
    const raw =
      sse({ choices: [{ delta: { content: "Bon" } }] }) +
      sse({ choices: [{ delta: { content: "jour" }, finish_reason: "stop" }] }) +
      "data: [DONE]\n\n";
    const half = Math.floor(raw.length / 2);
    const ev = await collect(parseSse(chunks(raw.slice(0, half), raw.slice(half))));
    expect(ev).toEqual([
      { type: "text_delta", text: "Bon" },
      { type: "text_delta", text: "jour" },
      { type: "done", stopReason: "end_turn" },
    ]);
  });

  it("streaming : appels d'outils assemblés (arguments fragmentés)", async () => {
    const raw =
      sse({ choices: [{ delta: { tool_calls: [{ index: 0, id: "c1", function: { name: "Read", arguments: '{"pa' } }] } }] }) +
      sse({ choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: 'th":"a"}' } }] } }] }) +
      sse({ choices: [{ delta: { tool_calls: [{ index: 1, id: "c2", function: { name: "Glob", arguments: "{}" } }] }, finish_reason: "tool_calls" }] });
    const ev = await collect(parseSse(chunks(raw)));
    expect(ev).toEqual([
      { type: "tool_use", id: "c1", name: "Read", input: { path: "a" } },
      { type: "tool_use", id: "c2", name: "Glob", input: {} },
      { type: "done", stopReason: "tool_use" },
    ]);
  });

  it("erreur renvoyée dans le flux", async () => {
    await expect(collect(parseSse(chunks(sse({ error: { message: "boom" } }))))).rejects.toThrow(/boom/);
  });
});