openAgent

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