adapter.test.ts 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177
  1. import { describe, expect } from "bun:test"
  2. import { Effect, Schema, Stream } from "effect"
  3. import { LLM } from "../src"
  4. import { Route, Endpoint, LLMClient, Protocol, type RouteModelInput, type FramingDef } from "../src/route"
  5. import { ModelRef } from "../src/schema"
  6. import { testEffect } from "./lib/effect"
  7. import { dynamicResponse } from "./lib/http"
  8. const updateModel = (model: ModelRef, patch: Partial<ModelRef.Input>) => ModelRef.update(model, patch)
  9. const Json = Schema.fromJsonString(Schema.Unknown)
  10. const encodeJson = Schema.encodeSync(Json)
  11. type FakeBody = {
  12. readonly body: string
  13. }
  14. const FakeEvent = Schema.Union([
  15. Schema.Struct({ type: Schema.Literal("text"), text: Schema.String }),
  16. Schema.Struct({ type: Schema.Literal("finish"), reason: Schema.Literal("stop") }),
  17. ])
  18. type FakeEvent = Schema.Schema.Type<typeof FakeEvent>
  19. const decodeFakeEvents = Schema.decodeUnknownEffect(Schema.fromJsonString(Schema.Array(FakeEvent)))
  20. const fakeFraming: FramingDef<FakeEvent> = {
  21. id: "fake-json-array",
  22. frame: (bytes) =>
  23. Stream.fromEffect(
  24. bytes.pipe(
  25. Stream.decodeText(),
  26. Stream.runFold(
  27. () => "",
  28. (text, event) => text + event,
  29. ),
  30. Effect.flatMap(decodeFakeEvents),
  31. Effect.orDie,
  32. ),
  33. ).pipe(Stream.flatMap(Stream.fromIterable)),
  34. }
  35. const request = LLM.request({
  36. id: "req_1",
  37. model: LLM.model({
  38. id: "fake-model",
  39. provider: "fake-provider",
  40. route: "fake",
  41. baseURL: "https://fake.local",
  42. }),
  43. prompt: "hello",
  44. })
  45. const raiseEvent = (event: FakeEvent): import("../src/schema").LLMEvent =>
  46. event.type === "finish"
  47. ? { type: "request-finish", reason: event.reason }
  48. : { type: "text-delta", id: "text-0", text: event.text }
  49. const fakeProtocol = Protocol.make<FakeBody, FakeEvent, FakeEvent, void>({
  50. id: "fake",
  51. body: {
  52. schema: Schema.Struct({
  53. body: Schema.String,
  54. }),
  55. from: (request) =>
  56. Effect.succeed({
  57. body: [
  58. ...request.messages
  59. .flatMap((message) => message.content)
  60. .filter((part) => part.type === "text")
  61. .map((part) => part.text),
  62. ...request.tools.map((tool) => `tool:${tool.name}:${tool.description}`),
  63. ].join("\n"),
  64. }),
  65. },
  66. stream: {
  67. event: FakeEvent,
  68. initial: () => undefined,
  69. step: (state, event) => Effect.succeed([state, [raiseEvent(event)]] as const),
  70. },
  71. })
  72. const fake = Route.make({
  73. id: "fake",
  74. protocol: fakeProtocol,
  75. endpoint: Endpoint.path("/chat"),
  76. framing: fakeFraming,
  77. })
  78. const gemini = Route.make({
  79. id: "gemini-fake",
  80. protocol: fakeProtocol,
  81. endpoint: Endpoint.path("/chat"),
  82. framing: fakeFraming,
  83. })
  84. const echoLayer = dynamicResponse(({ text, respond }) =>
  85. Effect.succeed(
  86. respond(
  87. encodeJson([
  88. { type: "text", text: `echo:${text}` },
  89. { type: "finish", reason: "stop" },
  90. ]),
  91. ),
  92. ),
  93. )
  94. const it = testEffect(echoLayer)
  95. describe("llm route", () => {
  96. it.effect("stream and generate use the route pipeline", () =>
  97. Effect.gen(function* () {
  98. const llm = yield* LLMClient.Service
  99. const events = Array.from(yield* llm.stream(request).pipe(Stream.runCollect))
  100. const response = yield* llm.generate(request)
  101. expect(events.map((event) => event.type)).toEqual(["text-delta", "request-finish"])
  102. expect(response.events.map((event) => event.type)).toEqual(["text-delta", "request-finish"])
  103. }),
  104. )
  105. it.effect("selects routes by request route", () =>
  106. Effect.gen(function* () {
  107. const llm = yield* LLMClient.Service
  108. const prepared = yield* llm.prepare(
  109. LLM.updateRequest(request, { model: updateModel(request.model, { route: "gemini-fake" }) }),
  110. )
  111. expect(prepared.route).toBe("gemini-fake")
  112. }),
  113. )
  114. it.effect("maps model input before building refs", () =>
  115. Effect.gen(function* () {
  116. const mapped = Route.model<RouteModelInput & { readonly region?: string }>(
  117. fake,
  118. { provider: "fake-provider", baseURL: "https://fake.local" },
  119. {
  120. mapInput: (input) => {
  121. const { region, ...rest } = input
  122. return { ...rest, native: { region } }
  123. },
  124. },
  125. )
  126. expect(mapped({ id: "fake-model", region: "us-east-1" }).native).toEqual({ region: "us-east-1" })
  127. }),
  128. )
  129. it.effect("rejects duplicate route ids", () =>
  130. Effect.gen(function* () {
  131. expect(() =>
  132. Route.make({
  133. id: "fake",
  134. protocol: Protocol.make({
  135. ...fakeProtocol,
  136. body: {
  137. ...fakeProtocol.body,
  138. from: () => Effect.succeed({ body: "late-default" }),
  139. },
  140. }),
  141. endpoint: Endpoint.path("/chat"),
  142. framing: fakeFraming,
  143. }),
  144. ).toThrow('Duplicate LLM route id "fake"')
  145. }),
  146. )
  147. it.effect("rejects missing route", () =>
  148. Effect.gen(function* () {
  149. const llm = yield* LLMClient.Service
  150. const error = yield* llm
  151. .prepare(LLM.updateRequest(request, { model: updateModel(request.model, { route: "missing" }) }))
  152. .pipe(Effect.flip)
  153. expect(error.message).toContain("No LLM route")
  154. }),
  155. )
  156. })