| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175 |
- import { describe, expect } from "bun:test"
- import { Effect, Schema, Stream } from "effect"
- import { LLM } from "../src"
- import { Route, Endpoint, LLMClient, Protocol, type RouteModelInput, type FramingDef } from "../src/route"
- import { ModelRef } from "../src/schema"
- import { testEffect } from "./lib/effect"
- import { dynamicResponse } from "./lib/http"
- const updateModel = (model: ModelRef, patch: Partial<ModelRef.Input>) => ModelRef.update(model, patch)
- const Json = Schema.fromJsonString(Schema.Unknown)
- const encodeJson = Schema.encodeSync(Json)
- type FakeBody = {
- readonly body: string
- }
- const FakeEvent = Schema.Union([
- Schema.Struct({ type: Schema.Literal("text"), text: Schema.String }),
- Schema.Struct({ type: Schema.Literal("finish"), reason: Schema.Literal("stop") }),
- ])
- type FakeEvent = Schema.Schema.Type<typeof FakeEvent>
- const decodeFakeEvents = Schema.decodeUnknownEffect(Schema.fromJsonString(Schema.Array(FakeEvent)))
- const fakeFraming: FramingDef<FakeEvent> = {
- id: "fake-json-array",
- frame: (bytes) =>
- Stream.fromEffect(
- bytes.pipe(
- Stream.decodeText(),
- Stream.runFold(
- () => "",
- (text, event) => text + event,
- ),
- Effect.flatMap(decodeFakeEvents),
- Effect.orDie,
- ),
- ).pipe(Stream.flatMap(Stream.fromIterable)),
- }
- const request = LLM.request({
- id: "req_1",
- model: LLM.model({
- id: "fake-model",
- provider: "fake-provider",
- route: "fake",
- baseURL: "https://fake.local",
- }),
- prompt: "hello",
- })
- const raiseEvent = (event: FakeEvent): import("../src/schema").LLMEvent =>
- event.type === "finish" ? { type: "request-finish", reason: event.reason } : { type: "text-delta", text: event.text }
- const fakeProtocol = Protocol.make<FakeBody, FakeEvent, FakeEvent, void>({
- id: "fake",
- body: {
- schema: Schema.Struct({
- body: Schema.String,
- }),
- from: (request) =>
- Effect.succeed({
- body: [
- ...request.messages
- .flatMap((message) => message.content)
- .filter((part) => part.type === "text")
- .map((part) => part.text),
- ...request.tools.map((tool) => `tool:${tool.name}:${tool.description}`),
- ].join("\n"),
- }),
- },
- stream: {
- event: FakeEvent,
- initial: () => undefined,
- step: (state, event) => Effect.succeed([state, [raiseEvent(event)]] as const),
- },
- })
- const fake = Route.make({
- id: "fake",
- protocol: fakeProtocol,
- endpoint: Endpoint.path("/chat"),
- framing: fakeFraming,
- })
- const gemini = Route.make({
- id: "gemini-fake",
- protocol: fakeProtocol,
- endpoint: Endpoint.path("/chat"),
- framing: fakeFraming,
- })
- const echoLayer = dynamicResponse(({ text, respond }) =>
- Effect.succeed(
- respond(
- encodeJson([
- { type: "text", text: `echo:${text}` },
- { type: "finish", reason: "stop" },
- ]),
- ),
- ),
- )
- const it = testEffect(echoLayer)
- describe("llm route", () => {
- it.effect("stream and generate use the route pipeline", () =>
- Effect.gen(function* () {
- const llm = yield* LLMClient.Service
- const events = Array.from(yield* llm.stream(request).pipe(Stream.runCollect))
- const response = yield* llm.generate(request)
- expect(events.map((event) => event.type)).toEqual(["text-delta", "request-finish"])
- expect(response.events.map((event) => event.type)).toEqual(["text-delta", "request-finish"])
- }),
- )
- it.effect("selects routes by request route", () =>
- Effect.gen(function* () {
- const llm = yield* LLMClient.Service
- const prepared = yield* llm.prepare(
- LLM.updateRequest(request, { model: updateModel(request.model, { route: "gemini-fake" }) }),
- )
- expect(prepared.route).toBe("gemini-fake")
- }),
- )
- it.effect("maps model input before building refs", () =>
- Effect.gen(function* () {
- const mapped = Route.model<RouteModelInput & { readonly region?: string }>(
- fake,
- { provider: "fake-provider", baseURL: "https://fake.local" },
- {
- mapInput: (input) => {
- const { region, ...rest } = input
- return { ...rest, native: { region } }
- },
- },
- )
- expect(mapped({ id: "fake-model", region: "us-east-1" }).native).toEqual({ region: "us-east-1" })
- }),
- )
- it.effect("rejects duplicate route ids", () =>
- Effect.gen(function* () {
- expect(() =>
- Route.make({
- id: "fake",
- protocol: Protocol.make({
- ...fakeProtocol,
- body: {
- ...fakeProtocol.body,
- from: () => Effect.succeed({ body: "late-default" }),
- },
- }),
- endpoint: Endpoint.path("/chat"),
- framing: fakeFraming,
- }),
- ).toThrow('Duplicate LLM route id "fake"')
- }),
- )
- it.effect("rejects missing route", () =>
- Effect.gen(function* () {
- const llm = yield* LLMClient.Service
- const error = yield* llm
- .prepare(LLM.updateRequest(request, { model: updateModel(request.model, { route: "missing" }) }))
- .pipe(Effect.flip)
- expect(error.message).toContain("No LLM route")
- }),
- )
- })
|