gemini.test.ts 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365
  1. import { describe, expect } from "bun:test"
  2. import { Effect } from "effect"
  3. import { LLM, LLMError, Usage } from "../../src"
  4. import { LLMClient } from "../../src/route"
  5. import * as Gemini from "../../src/protocols/gemini"
  6. import { it } from "../lib/effect"
  7. import { fixedResponse } from "../lib/http"
  8. import { sseEvents, sseRaw } from "../lib/sse"
  9. const model = Gemini.model({
  10. id: "gemini-2.5-flash",
  11. baseURL: "https://generativelanguage.test/v1beta/",
  12. headers: { "x-goog-api-key": "test" },
  13. })
  14. const request = LLM.request({
  15. id: "req_1",
  16. model,
  17. system: "You are concise.",
  18. prompt: "Say hello.",
  19. generation: { maxTokens: 20, temperature: 0 },
  20. })
  21. describe("Gemini route", () => {
  22. it.effect("prepares Gemini target", () =>
  23. Effect.gen(function* () {
  24. const prepared = yield* LLMClient.prepare(request)
  25. expect(prepared.body).toEqual({
  26. contents: [{ role: "user", parts: [{ text: "Say hello." }] }],
  27. systemInstruction: { parts: [{ text: "You are concise." }] },
  28. generationConfig: { maxOutputTokens: 20, temperature: 0 },
  29. })
  30. }),
  31. )
  32. it.effect("prepares multimodal user input and tool history", () =>
  33. Effect.gen(function* () {
  34. const prepared = yield* LLMClient.prepare(
  35. LLM.request({
  36. id: "req_tool_result",
  37. model,
  38. tools: [
  39. {
  40. name: "lookup",
  41. description: "Lookup data",
  42. inputSchema: { type: "object", properties: { query: { type: "string" } } },
  43. },
  44. ],
  45. toolChoice: { type: "tool", name: "lookup" },
  46. messages: [
  47. LLM.user([
  48. { type: "text", text: "What is in this image?" },
  49. { type: "media", mediaType: "image/png", data: "AAECAw==" },
  50. ]),
  51. LLM.assistant([LLM.toolCall({ id: "call_1", name: "lookup", input: { query: "weather" } })]),
  52. LLM.toolMessage({ id: "call_1", name: "lookup", result: { forecast: "sunny" } }),
  53. ],
  54. }),
  55. )
  56. expect(prepared.body).toEqual({
  57. contents: [
  58. {
  59. role: "user",
  60. parts: [{ text: "What is in this image?" }, { inlineData: { mimeType: "image/png", data: "AAECAw==" } }],
  61. },
  62. {
  63. role: "model",
  64. parts: [{ functionCall: { name: "lookup", args: { query: "weather" } } }],
  65. },
  66. {
  67. role: "user",
  68. parts: [
  69. { functionResponse: { name: "lookup", response: { name: "lookup", content: '{"forecast":"sunny"}' } } },
  70. ],
  71. },
  72. ],
  73. tools: [
  74. {
  75. functionDeclarations: [
  76. {
  77. name: "lookup",
  78. description: "Lookup data",
  79. parameters: { type: "object", properties: { query: { type: "string" } } },
  80. },
  81. ],
  82. },
  83. ],
  84. toolConfig: { functionCallingConfig: { mode: "ANY", allowedFunctionNames: ["lookup"] } },
  85. })
  86. }),
  87. )
  88. it.effect("omits tools when tool choice is none", () =>
  89. Effect.gen(function* () {
  90. const prepared = yield* LLMClient.prepare(
  91. LLM.request({
  92. id: "req_no_tools",
  93. model,
  94. prompt: "Say hello.",
  95. tools: [{ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } }],
  96. toolChoice: { type: "none" },
  97. }),
  98. )
  99. expect(prepared.body).toEqual({
  100. contents: [{ role: "user", parts: [{ text: "Say hello." }] }],
  101. })
  102. }),
  103. )
  104. it.effect("sanitizes integer enums, dangling required, untyped arrays, and scalar object keys", () =>
  105. Effect.gen(function* () {
  106. const prepared = yield* LLMClient.prepare(
  107. LLM.request({
  108. id: "req_schema_patch",
  109. model,
  110. prompt: "Use the tool.",
  111. tools: [
  112. {
  113. name: "lookup",
  114. description: "Lookup data",
  115. inputSchema: {
  116. type: "object",
  117. required: ["status", "missing"],
  118. properties: {
  119. status: { type: "integer", enum: [1, 2] },
  120. tags: { type: "array" },
  121. name: { type: "string", properties: { ignored: { type: "string" } }, required: ["ignored"] },
  122. },
  123. },
  124. },
  125. ],
  126. }),
  127. )
  128. expect(prepared.body).toMatchObject({
  129. tools: [
  130. {
  131. functionDeclarations: [
  132. {
  133. parameters: {
  134. type: "object",
  135. required: ["status"],
  136. properties: {
  137. status: { type: "string", enum: ["1", "2"] },
  138. tags: { type: "array", items: { type: "string" } },
  139. name: { type: "string" },
  140. },
  141. },
  142. },
  143. ],
  144. },
  145. ],
  146. })
  147. }),
  148. )
  149. it.effect("parses text, reasoning, and usage stream fixtures", () =>
  150. Effect.gen(function* () {
  151. const body = sseEvents(
  152. {
  153. candidates: [
  154. {
  155. content: { role: "model", parts: [{ text: "thinking", thought: true }] },
  156. },
  157. ],
  158. },
  159. {
  160. candidates: [
  161. {
  162. content: { role: "model", parts: [{ text: "Hello" }] },
  163. },
  164. ],
  165. },
  166. {
  167. candidates: [
  168. {
  169. content: { role: "model", parts: [{ text: "!" }] },
  170. finishReason: "STOP",
  171. },
  172. ],
  173. },
  174. {
  175. usageMetadata: {
  176. promptTokenCount: 5,
  177. candidatesTokenCount: 2,
  178. totalTokenCount: 7,
  179. thoughtsTokenCount: 1,
  180. cachedContentTokenCount: 1,
  181. },
  182. },
  183. )
  184. const response = yield* LLMClient.generate(request).pipe(Effect.provide(fixedResponse(body)))
  185. expect(response.text).toBe("Hello!")
  186. expect(response.reasoning).toBe("thinking")
  187. expect(response.usage).toMatchObject({
  188. inputTokens: 5,
  189. outputTokens: 3,
  190. nonCachedInputTokens: 4,
  191. cacheReadInputTokens: 1,
  192. reasoningTokens: 1,
  193. totalTokens: 7,
  194. })
  195. expect(response.events).toEqual([
  196. { type: "reasoning-delta", id: "reasoning-0", text: "thinking" },
  197. { type: "text-delta", id: "text-0", text: "Hello" },
  198. { type: "text-delta", id: "text-0", text: "!" },
  199. {
  200. type: "request-finish",
  201. reason: "stop",
  202. usage: new Usage({
  203. inputTokens: 5,
  204. outputTokens: 3,
  205. nonCachedInputTokens: 4,
  206. cacheReadInputTokens: 1,
  207. reasoningTokens: 1,
  208. totalTokens: 7,
  209. providerMetadata: {
  210. google: {
  211. promptTokenCount: 5,
  212. candidatesTokenCount: 2,
  213. totalTokenCount: 7,
  214. thoughtsTokenCount: 1,
  215. cachedContentTokenCount: 1,
  216. },
  217. },
  218. }),
  219. },
  220. ])
  221. }),
  222. )
  223. it.effect("emits streamed tool calls and maps finish reason", () =>
  224. Effect.gen(function* () {
  225. const body = sseEvents({
  226. candidates: [
  227. {
  228. content: {
  229. role: "model",
  230. parts: [{ functionCall: { name: "lookup", args: { query: "weather" } } }],
  231. },
  232. finishReason: "STOP",
  233. },
  234. ],
  235. usageMetadata: { promptTokenCount: 5, candidatesTokenCount: 1 },
  236. })
  237. const response = yield* LLMClient.generate(
  238. LLM.updateRequest(request, {
  239. tools: [{ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } }],
  240. }),
  241. ).pipe(Effect.provide(fixedResponse(body)))
  242. expect(response.toolCalls).toEqual([
  243. { type: "tool-call", id: "tool_0", name: "lookup", input: { query: "weather" } },
  244. ])
  245. expect(response.events).toEqual([
  246. { type: "tool-call", id: "tool_0", name: "lookup", input: { query: "weather" } },
  247. {
  248. type: "request-finish",
  249. reason: "tool-calls",
  250. usage: new Usage({
  251. inputTokens: 5,
  252. outputTokens: 1,
  253. nonCachedInputTokens: 5,
  254. totalTokens: 6,
  255. providerMetadata: { google: { promptTokenCount: 5, candidatesTokenCount: 1 } },
  256. }),
  257. },
  258. ])
  259. }),
  260. )
  261. it.effect("assigns unique ids to multiple streamed tool calls", () =>
  262. Effect.gen(function* () {
  263. const body = sseEvents({
  264. candidates: [
  265. {
  266. content: {
  267. role: "model",
  268. parts: [
  269. { functionCall: { name: "lookup", args: { query: "weather" } } },
  270. { functionCall: { name: "lookup", args: { query: "news" } } },
  271. ],
  272. },
  273. finishReason: "STOP",
  274. },
  275. ],
  276. })
  277. const response = yield* LLMClient.generate(
  278. LLM.updateRequest(request, {
  279. tools: [{ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } }],
  280. }),
  281. ).pipe(Effect.provide(fixedResponse(body)))
  282. expect(response.toolCalls).toEqual([
  283. { type: "tool-call", id: "tool_0", name: "lookup", input: { query: "weather" } },
  284. { type: "tool-call", id: "tool_1", name: "lookup", input: { query: "news" } },
  285. ])
  286. expect(response.events.at(-1)).toMatchObject({ type: "request-finish", reason: "tool-calls" })
  287. }),
  288. )
  289. it.effect("maps length and content-filter finish reasons", () =>
  290. Effect.gen(function* () {
  291. const length = yield* LLMClient.generate(request).pipe(
  292. Effect.provide(
  293. fixedResponse(
  294. sseEvents({ candidates: [{ content: { role: "model", parts: [] }, finishReason: "MAX_TOKENS" }] }),
  295. ),
  296. ),
  297. )
  298. const filtered = yield* LLMClient.generate(request).pipe(
  299. Effect.provide(
  300. fixedResponse(sseEvents({ candidates: [{ content: { role: "model", parts: [] }, finishReason: "SAFETY" }] })),
  301. ),
  302. )
  303. expect(length.events).toEqual([{ type: "request-finish", reason: "length" }])
  304. expect(filtered.events).toEqual([{ type: "request-finish", reason: "content-filter" }])
  305. }),
  306. )
  307. it.effect("leaves total usage undefined when component counts are missing", () =>
  308. Effect.gen(function* () {
  309. const response = yield* LLMClient.generate(request).pipe(
  310. Effect.provide(fixedResponse(sseEvents({ usageMetadata: { thoughtsTokenCount: 1 } }))),
  311. )
  312. expect(response.usage).toMatchObject({ reasoningTokens: 1 })
  313. expect(response.usage?.totalTokens).toBeUndefined()
  314. }),
  315. )
  316. it.effect("fails invalid stream events", () =>
  317. Effect.gen(function* () {
  318. const error = yield* LLMClient.generate(request).pipe(
  319. Effect.provide(fixedResponse(sseRaw("data: {not json}"))),
  320. Effect.flip,
  321. )
  322. expect(error).toBeInstanceOf(LLMError)
  323. expect(error.reason).toMatchObject({ _tag: "InvalidProviderOutput" })
  324. expect(error.message).toContain("Invalid google/gemini stream event")
  325. }),
  326. )
  327. it.effect("rejects unsupported assistant media content", () =>
  328. Effect.gen(function* () {
  329. const error = yield* LLMClient.prepare(
  330. LLM.request({
  331. id: "req_media",
  332. model,
  333. messages: [LLM.assistant({ type: "media", mediaType: "image/png", data: "AAECAw==" })],
  334. }),
  335. ).pipe(Effect.flip)
  336. expect(error.message).toContain(
  337. "Gemini assistant messages only support text, reasoning, and tool-call content for now",
  338. )
  339. }),
  340. )
  341. })