adapter.test.ts 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  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" ? { type: "request-finish", reason: event.reason } : { type: "text-delta", text: event.text }
  47. const fakeProtocol = Protocol.make<FakeBody, FakeEvent, FakeEvent, void>({
  48. id: "fake",
  49. body: {
  50. schema: Schema.Struct({
  51. body: Schema.String,
  52. }),
  53. from: (request) =>
  54. Effect.succeed({
  55. body: [
  56. ...request.messages
  57. .flatMap((message) => message.content)
  58. .filter((part) => part.type === "text")
  59. .map((part) => part.text),
  60. ...request.tools.map((tool) => `tool:${tool.name}:${tool.description}`),
  61. ].join("\n"),
  62. }),
  63. },
  64. stream: {
  65. event: FakeEvent,
  66. initial: () => undefined,
  67. step: (state, event) => Effect.succeed([state, [raiseEvent(event)]] as const),
  68. },
  69. })
  70. const fake = Route.make({
  71. id: "fake",
  72. protocol: fakeProtocol,
  73. endpoint: Endpoint.path("/chat"),
  74. framing: fakeFraming,
  75. })
  76. const gemini = Route.make({
  77. id: "gemini-fake",
  78. protocol: fakeProtocol,
  79. endpoint: Endpoint.path("/chat"),
  80. framing: fakeFraming,
  81. })
  82. const echoLayer = dynamicResponse(({ text, respond }) =>
  83. Effect.succeed(
  84. respond(
  85. encodeJson([
  86. { type: "text", text: `echo:${text}` },
  87. { type: "finish", reason: "stop" },
  88. ]),
  89. ),
  90. ),
  91. )
  92. const it = testEffect(echoLayer)
  93. describe("llm route", () => {
  94. it.effect("stream and generate use the route pipeline", () =>
  95. Effect.gen(function* () {
  96. const llm = yield* LLMClient.Service
  97. const events = Array.from(yield* llm.stream(request).pipe(Stream.runCollect))
  98. const response = yield* llm.generate(request)
  99. expect(events.map((event) => event.type)).toEqual(["text-delta", "request-finish"])
  100. expect(response.events.map((event) => event.type)).toEqual(["text-delta", "request-finish"])
  101. }),
  102. )
  103. it.effect("selects routes by request route", () =>
  104. Effect.gen(function* () {
  105. const llm = yield* LLMClient.Service
  106. const prepared = yield* llm.prepare(
  107. LLM.updateRequest(request, { model: updateModel(request.model, { route: "gemini-fake" }) }),
  108. )
  109. expect(prepared.route).toBe("gemini-fake")
  110. }),
  111. )
  112. it.effect("maps model input before building refs", () =>
  113. Effect.gen(function* () {
  114. const mapped = Route.model<RouteModelInput & { readonly region?: string }>(
  115. fake,
  116. { provider: "fake-provider", baseURL: "https://fake.local" },
  117. {
  118. mapInput: (input) => {
  119. const { region, ...rest } = input
  120. return { ...rest, native: { region } }
  121. },
  122. },
  123. )
  124. expect(mapped({ id: "fake-model", region: "us-east-1" }).native).toEqual({ region: "us-east-1" })
  125. }),
  126. )
  127. it.effect("rejects duplicate route ids", () =>
  128. Effect.gen(function* () {
  129. expect(() =>
  130. Route.make({
  131. id: "fake",
  132. protocol: Protocol.make({
  133. ...fakeProtocol,
  134. body: {
  135. ...fakeProtocol.body,
  136. from: () => Effect.succeed({ body: "late-default" }),
  137. },
  138. }),
  139. endpoint: Endpoint.path("/chat"),
  140. framing: fakeFraming,
  141. }),
  142. ).toThrow('Duplicate LLM route id "fake"')
  143. }),
  144. )
  145. it.effect("rejects missing route", () =>
  146. Effect.gen(function* () {
  147. const llm = yield* LLMClient.Service
  148. const error = yield* llm
  149. .prepare(LLM.updateRequest(request, { model: updateModel(request.model, { route: "missing" }) }))
  150. .pipe(Effect.flip)
  151. expect(error.message).toContain("No LLM route")
  152. }),
  153. )
  154. })