adapter.test.ts 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  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 FramingDef } from "../src/route"
  5. import { Model } from "../src/schema"
  6. import { testEffect } from "./lib/effect"
  7. import { dynamicResponse } from "./lib/http"
  8. const updateModel = (model: Model, patch: Partial<Model.Input>) => Model.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 raiseEvent = (event: FakeEvent): import("../src/schema").LLMEvent =>
  36. event.type === "finish"
  37. ? { type: "finish", reason: event.reason }
  38. : { type: "text-delta", id: "text-0", text: event.text }
  39. const fakeProtocol = Protocol.make<FakeBody, FakeEvent, FakeEvent, void>({
  40. id: "fake",
  41. body: {
  42. schema: Schema.Struct({
  43. body: Schema.String,
  44. }),
  45. from: (request) =>
  46. Effect.succeed({
  47. body: [
  48. ...request.messages
  49. .flatMap((message) => message.content)
  50. .filter((part) => part.type === "text")
  51. .map((part) => part.text),
  52. ...request.tools.map((tool) => `tool:${tool.name}:${tool.description}`),
  53. ].join("\n"),
  54. }),
  55. },
  56. stream: {
  57. event: FakeEvent,
  58. initial: () => undefined,
  59. step: (state, event) => Effect.succeed([state, [raiseEvent(event)]] as const),
  60. },
  61. })
  62. const fake = Route.make({
  63. id: "fake",
  64. protocol: fakeProtocol,
  65. endpoint: Endpoint.path("/chat"),
  66. framing: fakeFraming,
  67. })
  68. const configuredFake = fake.with({ endpoint: { baseURL: "https://fake.local" } })
  69. const gemini = Route.make({
  70. id: "gemini-fake",
  71. protocol: fakeProtocol,
  72. endpoint: Endpoint.path("/chat"),
  73. framing: fakeFraming,
  74. })
  75. const configuredGemini = gemini.with({ endpoint: { baseURL: "https://fake.local" } })
  76. const request = LLM.request({
  77. id: "req_1",
  78. model: Model.make({
  79. id: "fake-model",
  80. provider: "fake-provider",
  81. route: configuredFake,
  82. }),
  83. prompt: "hello",
  84. })
  85. const echoLayer = dynamicResponse(({ text, respond }) =>
  86. Effect.succeed(
  87. respond(
  88. encodeJson([
  89. { type: "text", text: `echo:${text}` },
  90. { type: "finish", reason: "stop" },
  91. ]),
  92. ),
  93. ),
  94. )
  95. const it = testEffect(echoLayer)
  96. describe("llm route", () => {
  97. it.effect("stream and generate use the route pipeline", () =>
  98. Effect.gen(function* () {
  99. const llm = yield* LLMClient.Service
  100. const events = Array.from(yield* llm.stream(request).pipe(Stream.runCollect))
  101. const response = yield* llm.generate(request)
  102. expect(events.map((event) => event.type)).toEqual(["text-delta", "finish"])
  103. expect(response.events.map((event) => event.type)).toEqual(["text-delta", "finish"])
  104. }),
  105. )
  106. it.effect("selects routes by model route value", () =>
  107. Effect.gen(function* () {
  108. const llm = yield* LLMClient.Service
  109. const prepared = yield* llm.prepare(
  110. LLM.updateRequest(request, { model: updateModel(request.model, { route: configuredGemini }) }),
  111. )
  112. expect(prepared.route).toBe("gemini-fake")
  113. }),
  114. )
  115. it.effect("builds models from configured routes", () =>
  116. Effect.gen(function* () {
  117. const configured = fake.with({ provider: "fake-provider", endpoint: { baseURL: "https://fake.local" } })
  118. expect(configured.model({ id: "fake-model" })).toMatchObject({
  119. provider: "fake-provider",
  120. })
  121. }),
  122. )
  123. it.effect("does not register duplicate route ids globally", () =>
  124. Effect.gen(function* () {
  125. const duplicate = Route.make({
  126. id: "fake",
  127. protocol: Protocol.make({
  128. ...fakeProtocol,
  129. body: {
  130. ...fakeProtocol.body,
  131. from: () => Effect.succeed({ body: "late-default" }),
  132. },
  133. }),
  134. endpoint: Endpoint.path("/chat", { baseURL: "https://fake.local" }),
  135. framing: fakeFraming,
  136. })
  137. const prepared = yield* (yield* LLMClient.Service).prepare(
  138. LLM.updateRequest(request, { model: updateModel(request.model, { route: duplicate }) }),
  139. )
  140. expect(prepared.body).toEqual({ body: "late-default" })
  141. }),
  142. )
  143. })