setup-recording-env.ts 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542
  1. #!/usr/bin/env bun
  2. import { NodeFileSystem } from "@effect/platform-node"
  3. import * as path from "node:path"
  4. import * as prompts from "@clack/prompts"
  5. import { AwsV4Signer } from "aws4fetch"
  6. import { Config, ConfigProvider, Effect, FileSystem, PlatformError, Redacted } from "effect"
  7. import { FetchHttpClient, HttpClient, HttpClientRequest, type HttpClientResponse } from "effect/unstable/http"
  8. import * as ProviderShared from "../src/protocols/shared"
  9. import * as Cloudflare from "../src/providers/cloudflare"
  10. type Provider = {
  11. readonly id: string
  12. readonly label: string
  13. readonly tier: "core" | "canary" | "compatible" | "optional"
  14. readonly note: string
  15. readonly vars: ReadonlyArray<{
  16. readonly name: string
  17. readonly label?: string
  18. readonly optional?: boolean
  19. readonly secret?: boolean
  20. }>
  21. readonly validate?: (env: Env) => Effect.Effect<string | undefined, unknown, HttpClient.HttpClient>
  22. }
  23. type Env = Record<string, string>
  24. const PROVIDERS: ReadonlyArray<Provider> = [
  25. {
  26. id: "openai",
  27. label: "OpenAI",
  28. tier: "core",
  29. note: "Native OpenAI Chat / Responses recorded tests",
  30. vars: [{ name: "OPENAI_API_KEY" }],
  31. validate: (env) => validateBearer("https://api.openai.com/v1/models", Redacted.make(env.OPENAI_API_KEY)),
  32. },
  33. {
  34. id: "anthropic",
  35. label: "Anthropic",
  36. tier: "core",
  37. note: "Native Anthropic Messages recorded tests",
  38. vars: [{ name: "ANTHROPIC_API_KEY" }],
  39. validate: (env) =>
  40. HttpClientRequest.get("https://api.anthropic.com/v1/models").pipe(
  41. HttpClientRequest.setHeaders({
  42. "anthropic-version": "2023-06-01",
  43. "x-api-key": Redacted.value(Redacted.make(env.ANTHROPIC_API_KEY)),
  44. }),
  45. executeRequest,
  46. ),
  47. },
  48. {
  49. id: "google",
  50. label: "Google Gemini",
  51. tier: "core",
  52. note: "Native Gemini recorded tests",
  53. vars: [{ name: "GOOGLE_GENERATIVE_AI_API_KEY" }],
  54. validate: (env) =>
  55. HttpClientRequest.get(
  56. `https://generativelanguage.googleapis.com/v1beta/models?key=${encodeURIComponent(env.GOOGLE_GENERATIVE_AI_API_KEY)}`,
  57. ).pipe(executeRequest),
  58. },
  59. {
  60. id: "bedrock",
  61. label: "Amazon Bedrock",
  62. tier: "core",
  63. note: "Native Bedrock Converse recorded tests",
  64. vars: [
  65. { name: "AWS_ACCESS_KEY_ID" },
  66. { name: "AWS_SECRET_ACCESS_KEY" },
  67. { name: "AWS_SESSION_TOKEN", optional: true },
  68. { name: "BEDROCK_RECORDING_REGION", optional: true },
  69. { name: "BEDROCK_MODEL_ID", optional: true },
  70. ],
  71. validate: (env) => validateBedrock(env),
  72. },
  73. {
  74. id: "groq",
  75. label: "Groq",
  76. tier: "canary",
  77. note: "Fast OpenAI-compatible canary for text/tool streaming",
  78. vars: [{ name: "GROQ_API_KEY" }],
  79. validate: (env) => validateBearer("https://api.groq.com/openai/v1/models", Redacted.make(env.GROQ_API_KEY)),
  80. },
  81. {
  82. id: "openrouter",
  83. label: "OpenRouter",
  84. tier: "canary",
  85. note: "Router canary for OpenAI-compatible text/tool streaming",
  86. vars: [{ name: "OPENROUTER_API_KEY" }],
  87. validate: (env) =>
  88. validateChat({
  89. url: "https://openrouter.ai/api/v1/chat/completions",
  90. token: Redacted.make(env.OPENROUTER_API_KEY),
  91. model: "openai/gpt-4o-mini",
  92. }),
  93. },
  94. {
  95. id: "xai",
  96. label: "xAI",
  97. tier: "canary",
  98. note: "OpenAI-compatible xAI chat endpoint",
  99. vars: [{ name: "XAI_API_KEY" }],
  100. validate: (env) => validateBearer("https://api.x.ai/v1/models", Redacted.make(env.XAI_API_KEY)),
  101. },
  102. {
  103. id: "cloudflare-ai-gateway",
  104. label: "Cloudflare AI Gateway",
  105. tier: "canary",
  106. note: "Cloudflare Unified/OpenAI-compatible gateway; supports provider/model ids like workers-ai/@cf/...",
  107. vars: [
  108. { name: "CLOUDFLARE_ACCOUNT_ID", label: "Cloudflare account ID", secret: false },
  109. {
  110. name: "CLOUDFLARE_GATEWAY_ID",
  111. label: "Cloudflare AI Gateway ID (defaults to default)",
  112. optional: true,
  113. secret: false,
  114. },
  115. { name: "CLOUDFLARE_API_TOKEN", label: "Cloudflare AI Gateway token" },
  116. ],
  117. validate: (env) =>
  118. validateChat({
  119. url: `${Cloudflare.aiGatewayBaseURL({
  120. accountId: env.CLOUDFLARE_ACCOUNT_ID,
  121. gatewayId: env.CLOUDFLARE_GATEWAY_ID || undefined,
  122. })}/chat/completions`,
  123. token: Redacted.make(envValue(env, Cloudflare.aiGatewayAuthEnvVars)),
  124. tokenHeader: "cf-aig-authorization",
  125. model: "workers-ai/@cf/meta/llama-3.1-8b-instruct",
  126. }),
  127. },
  128. {
  129. id: "cloudflare-workers-ai",
  130. label: "Cloudflare Workers AI",
  131. tier: "canary",
  132. note: "Direct Workers AI OpenAI-compatible endpoint; supports model ids like @cf/meta/...",
  133. vars: [
  134. { name: "CLOUDFLARE_ACCOUNT_ID", label: "Cloudflare account ID", secret: false },
  135. { name: "CLOUDFLARE_API_KEY", label: "Cloudflare Workers AI API token" },
  136. ],
  137. validate: (env) =>
  138. validateChat({
  139. url: `${Cloudflare.workersAIBaseURL({ accountId: env.CLOUDFLARE_ACCOUNT_ID })}/chat/completions`,
  140. token: Redacted.make(envValue(env, Cloudflare.workersAIAuthEnvVars)),
  141. model: "@cf/meta/llama-3.1-8b-instruct",
  142. }),
  143. },
  144. {
  145. id: "deepseek",
  146. label: "DeepSeek",
  147. tier: "compatible",
  148. note: "Existing OpenAI-compatible recorded tests",
  149. vars: [{ name: "DEEPSEEK_API_KEY" }],
  150. validate: (env) => validateBearer("https://api.deepseek.com/models", Redacted.make(env.DEEPSEEK_API_KEY)),
  151. },
  152. {
  153. id: "togetherai",
  154. label: "TogetherAI",
  155. tier: "compatible",
  156. note: "Existing OpenAI-compatible text/tool recorded tests",
  157. vars: [{ name: "TOGETHER_AI_API_KEY" }],
  158. validate: (env) => validateBearer("https://api.together.xyz/v1/models", Redacted.make(env.TOGETHER_AI_API_KEY)),
  159. },
  160. {
  161. id: "mistral",
  162. label: "Mistral",
  163. tier: "optional",
  164. note: "OpenAI-compatible bridge; native reasoning parity is follow-up work",
  165. vars: [{ name: "MISTRAL_API_KEY" }],
  166. validate: (env) => validateBearer("https://api.mistral.ai/v1/models", Redacted.make(env.MISTRAL_API_KEY)),
  167. },
  168. {
  169. id: "perplexity",
  170. label: "Perplexity",
  171. tier: "optional",
  172. note: "OpenAI-compatible bridge; citations/search metadata are follow-up work",
  173. vars: [{ name: "PERPLEXITY_API_KEY" }],
  174. validate: (env) => validateBearer("https://api.perplexity.ai/models", Redacted.make(env.PERPLEXITY_API_KEY)),
  175. },
  176. {
  177. id: "venice",
  178. label: "Venice",
  179. tier: "optional",
  180. note: "OpenAI-compatible bridge",
  181. vars: [{ name: "VENICE_API_KEY" }],
  182. validate: (env) => validateBearer("https://api.venice.ai/api/v1/models", Redacted.make(env.VENICE_API_KEY)),
  183. },
  184. {
  185. id: "cerebras",
  186. label: "Cerebras",
  187. tier: "optional",
  188. note: "OpenAI-compatible bridge",
  189. vars: [{ name: "CEREBRAS_API_KEY" }],
  190. validate: (env) => validateBearer("https://api.cerebras.ai/v1/models", Redacted.make(env.CEREBRAS_API_KEY)),
  191. },
  192. {
  193. id: "deepinfra",
  194. label: "DeepInfra",
  195. tier: "optional",
  196. note: "OpenAI-compatible bridge",
  197. vars: [{ name: "DEEPINFRA_API_KEY" }],
  198. validate: (env) =>
  199. validateBearer("https://api.deepinfra.com/v1/openai/models", Redacted.make(env.DEEPINFRA_API_KEY)),
  200. },
  201. {
  202. id: "fireworks",
  203. label: "Fireworks",
  204. tier: "optional",
  205. note: "OpenAI-compatible bridge",
  206. vars: [{ name: "FIREWORKS_API_KEY" }],
  207. validate: (env) =>
  208. validateBearer("https://api.fireworks.ai/inference/v1/models", Redacted.make(env.FIREWORKS_API_KEY)),
  209. },
  210. {
  211. id: "baseten",
  212. label: "Baseten",
  213. tier: "optional",
  214. note: "OpenAI-compatible bridge",
  215. vars: [{ name: "BASETEN_API_KEY" }],
  216. },
  217. ]
  218. const args = process.argv.slice(2)
  219. const hasFlag = (name: string) => args.includes(name)
  220. const option = (name: string) => {
  221. const index = args.indexOf(name)
  222. if (index === -1) return undefined
  223. return args[index + 1]
  224. }
  225. const envPath = path.resolve(process.cwd(), option("--env") ?? ".env.local")
  226. const checkOnly = hasFlag("--check")
  227. const providerOption = option("--providers")
  228. const interactive = Boolean(process.stdin.isTTY && process.stdout.isTTY)
  229. const envNames = Array.from(new Set(PROVIDERS.flatMap((provider) => provider.vars.map((item) => item.name))))
  230. const providersForOption = (value: string | undefined) => {
  231. if (!value || value === "recommended")
  232. return PROVIDERS.filter((provider) => provider.tier === "core" || provider.tier === "canary")
  233. if (value === "recorded") return PROVIDERS.filter((provider) => provider.tier !== "optional")
  234. if (value === "all") return PROVIDERS
  235. const ids = new Set(
  236. value
  237. .split(",")
  238. .map((item) => item.trim())
  239. .filter(Boolean),
  240. )
  241. return PROVIDERS.filter((provider) => ids.has(provider.id))
  242. }
  243. const chooseProviders = async () => {
  244. if (providerOption) return providersForOption(providerOption)
  245. return providersForOption("recommended")
  246. }
  247. const catchMissingFile = (error: PlatformError.PlatformError) => {
  248. if (error.reason._tag === "NotFound") return Effect.succeed("")
  249. return Effect.fail(error)
  250. }
  251. const readEnvFile = Effect.fn("RecordingEnv.readFile")(function* () {
  252. const fileSystem = yield* FileSystem.FileSystem
  253. return yield* fileSystem.readFileString(envPath).pipe(Effect.catch(catchMissingFile))
  254. })
  255. const readConfigString = (provider: ConfigProvider.ConfigProvider, name: string) =>
  256. Config.string(name)
  257. .parse(provider)
  258. .pipe(
  259. Effect.match({
  260. onFailure: () => undefined,
  261. onSuccess: (value) => value,
  262. }),
  263. )
  264. const parseEnv = Effect.fn("RecordingEnv.parseEnv")(function* (contents: string) {
  265. const provider = ConfigProvider.fromDotEnvContents(contents)
  266. return Object.fromEntries(
  267. (yield* Effect.forEach(envNames, (name) =>
  268. readConfigString(provider, name).pipe(Effect.map((value) => [name, value] as const)),
  269. )).filter((entry): entry is readonly [string, string] => entry[1] !== undefined),
  270. )
  271. })
  272. const quote = (value: string) => JSON.stringify(value)
  273. const status = (name: string, fileEnv: Env) => {
  274. if (fileEnv[name]) return "file"
  275. if (process.env[name]) return "shell"
  276. return "missing"
  277. }
  278. const statusLine = (provider: Provider, fileEnv: Env) =>
  279. [
  280. `${provider.label} (${provider.tier})`,
  281. provider.note,
  282. ...provider.vars.map((item) => {
  283. const value = status(item.name, fileEnv)
  284. const suffix = item.optional ? " optional" : ""
  285. return ` ${value === "missing" ? "missing" : "set"} ${item.name}${suffix}${value === "shell" ? " (shell only)" : ""}`
  286. }),
  287. ].join("\n")
  288. const printStatus = (providers: ReadonlyArray<Provider>, fileEnv: Env) => {
  289. prompts.note(providers.map((provider) => statusLine(provider, fileEnv)).join("\n\n"), `Recording env: ${envPath}`)
  290. }
  291. const exitIfCancel = <A>(value: A | symbol): A => {
  292. if (!prompts.isCancel(value)) return value as A
  293. prompts.cancel("Cancelled")
  294. process.exit(130)
  295. }
  296. const upsertEnv = (contents: string, values: Env) => {
  297. const names = Object.keys(values)
  298. const seen = new Set<string>()
  299. const lines = contents.split(/\r?\n/).map((line) => {
  300. const match = line.match(/^\s*(?:export\s+)?([A-Za-z_][A-Za-z0-9_]*)\s*=/)
  301. if (!match || !names.includes(match[1])) return line
  302. seen.add(match[1])
  303. return `${match[1]}=${quote(values[match[1]])}`
  304. })
  305. const missing = names.filter((name) => !seen.has(name))
  306. if (missing.length === 0) return lines.join("\n").replace(/\n*$/, "\n")
  307. const prefix = lines.join("\n").trimEnd()
  308. const block = [
  309. "",
  310. "# Added by bun run setup:recording-env",
  311. ...missing.map((name) => `${name}=${quote(values[name])}`),
  312. ].join("\n")
  313. return `${prefix}${block}\n`
  314. }
  315. const providerRequiredStatus = (provider: Provider, fileEnv: Env) => {
  316. const required = requiredVars(provider)
  317. if (required.some((item) => status(item.name, fileEnv) === "missing")) return "missing"
  318. if (required.some((item) => status(item.name, fileEnv) === "shell")) return "set in shell"
  319. return "already added"
  320. }
  321. const requiredVars = (provider: Provider) => provider.vars.filter((item) => !item.optional)
  322. const promptVars = (provider: Provider) => provider.vars.filter((item) => !item.optional || item.secret === false)
  323. const processEnv = (): Env =>
  324. Object.fromEntries(Object.entries(process.env).filter((entry): entry is [string, string] => entry[1] !== undefined))
  325. const envValue = (env: Env, names: ReadonlyArray<string>) => names.map((name) => env[name]).find(Boolean) ?? ""
  326. const envWithValues = (fileEnv: Env, values: Env): Env => ({
  327. ...processEnv(),
  328. ...fileEnv,
  329. ...values,
  330. })
  331. const responseError = Effect.fn("RecordingEnv.responseError")(function* (
  332. response: HttpClientResponse.HttpClientResponse,
  333. ) {
  334. if (response.status >= 200 && response.status < 300) return undefined
  335. const body = yield* response.text.pipe(Effect.catch(() => Effect.succeed("")))
  336. return `${response.status}${body ? `: ${body.slice(0, 180)}` : ""}`
  337. })
  338. const executeRequest = Effect.fn("RecordingEnv.executeRequest")(function* (
  339. request: HttpClientRequest.HttpClientRequest,
  340. ) {
  341. const http = yield* HttpClient.HttpClient
  342. return yield* http.execute(request).pipe(Effect.flatMap(responseError))
  343. })
  344. const validateBearer = (url: string, token: Redacted.Redacted<string>, headers: Record<string, string> = {}) =>
  345. HttpClientRequest.get(url).pipe(
  346. HttpClientRequest.setHeaders({ ...headers, authorization: `Bearer ${Redacted.value(token)}` }),
  347. executeRequest,
  348. )
  349. const validateChat = (input: {
  350. readonly url: string
  351. readonly token: Redacted.Redacted<string>
  352. readonly tokenHeader?: string
  353. readonly model: string
  354. readonly headers?: Record<string, string>
  355. }) =>
  356. ProviderShared.jsonPost({
  357. url: input.url,
  358. headers: { ...input.headers, [input.tokenHeader ?? "authorization"]: `Bearer ${Redacted.value(input.token)}` },
  359. body: ProviderShared.encodeJson({
  360. model: input.model,
  361. messages: [{ role: "user", content: "Reply with exactly: ok" }],
  362. max_tokens: 3,
  363. temperature: 0,
  364. }),
  365. }).pipe(executeRequest)
  366. const validateBedrock = (env: Env) =>
  367. Effect.gen(function* () {
  368. const request = yield* Effect.promise(() =>
  369. new AwsV4Signer({
  370. url: `https://bedrock.${env.BEDROCK_RECORDING_REGION || "us-east-1"}.amazonaws.com/foundation-models`,
  371. method: "GET",
  372. service: "bedrock",
  373. region: env.BEDROCK_RECORDING_REGION || "us-east-1",
  374. accessKeyId: env.AWS_ACCESS_KEY_ID,
  375. secretAccessKey: env.AWS_SECRET_ACCESS_KEY,
  376. sessionToken: env.AWS_SESSION_TOKEN || undefined,
  377. }).sign(),
  378. )
  379. return yield* HttpClientRequest.get(request.url.toString()).pipe(
  380. HttpClientRequest.setHeaders(Object.fromEntries(request.headers.entries())),
  381. executeRequest,
  382. )
  383. })
  384. const validateProvider = Effect.fn("RecordingEnv.validateProvider")(function* (provider: Provider, env: Env) {
  385. return yield* (provider.validate?.(env) ?? Effect.succeed("no lightweight validator")).pipe(
  386. Effect.catch((error) => {
  387. if (error instanceof Error) return Effect.succeed(error.message)
  388. return Effect.succeed(String(error))
  389. }),
  390. )
  391. })
  392. const validateProviders = Effect.fn("RecordingEnv.validateProviders")(function* (
  393. providers: ReadonlyArray<Provider>,
  394. env: Env,
  395. ) {
  396. const spinner = prompts.spinner()
  397. spinner.start("Validating credentials")
  398. const results = yield* Effect.forEach(
  399. providers,
  400. (provider) => validateProvider(provider, env).pipe(Effect.map((error) => ({ provider, error }))),
  401. { concurrency: 4 },
  402. )
  403. spinner.stop("Validation complete")
  404. prompts.note(
  405. results
  406. .map(
  407. (result) =>
  408. `${result.error ? "failed" : "ok"} ${result.provider.label}${result.error ? ` - ${result.error}` : ""}`,
  409. )
  410. .join("\n"),
  411. "Credential validation",
  412. )
  413. })
  414. const writeEnvFile = Effect.fn("RecordingEnv.writeFile")(function* (contents: string) {
  415. const fileSystem = yield* FileSystem.FileSystem
  416. yield* fileSystem.makeDirectory(path.dirname(envPath), { recursive: true })
  417. yield* fileSystem.writeFileString(envPath, contents, { mode: 0o600 })
  418. })
  419. const prompt = <A>(run: () => Promise<A | symbol>) => Effect.promise(run).pipe(Effect.map(exitIfCancel))
  420. const chooseConfigurableProviders = Effect.fn("RecordingEnv.chooseConfigurableProviders")(function* (
  421. providers: ReadonlyArray<Provider>,
  422. fileEnv: Env,
  423. ) {
  424. const configurable = providers.filter((provider) => requiredVars(provider).length > 0)
  425. const selected = yield* prompt<ReadonlyArray<string>>(() =>
  426. prompts.multiselect({
  427. message: "Select provider credentials to add or override",
  428. options: configurable.map((provider) => ({
  429. value: provider.id,
  430. label: provider.label,
  431. hint: `${providerRequiredStatus(provider, fileEnv)} - ${requiredVars(provider)
  432. .map((item) => item.name)
  433. .join(", ")}`,
  434. })),
  435. initialValues: configurable
  436. .filter((provider) => providerRequiredStatus(provider, fileEnv) === "missing")
  437. .map((provider) => provider.id),
  438. }),
  439. )
  440. return configurable.filter((provider) => selected.includes(provider.id))
  441. })
  442. const promptEnvVar = (item: Provider["vars"][number]) =>
  443. prompt<string>(() => {
  444. const input = {
  445. message: item.label ?? item.name,
  446. validate: (input: string | undefined) => {
  447. if (item.optional) return undefined
  448. return !input || input.length === 0 ? "Leave blank by pressing Esc/cancel, or paste a value" : undefined
  449. },
  450. }
  451. return item.secret === false ? prompts.text(input) : prompts.password(input)
  452. })
  453. const promptProviderValues = Effect.fn("RecordingEnv.promptProviderValues")(function* (
  454. providers: ReadonlyArray<Provider>,
  455. ) {
  456. const values: Env = {}
  457. for (const provider of providers) {
  458. prompts.log.info(`${provider.label}: ${provider.note}`)
  459. for (const item of promptVars(provider)) {
  460. if (values[item.name]) continue
  461. const value = yield* promptEnvVar(item)
  462. if (value !== "") values[item.name] = value
  463. }
  464. }
  465. return values
  466. })
  467. const main = Effect.fn("RecordingEnv.main")(function* () {
  468. prompts.intro("LLM recording credentials")
  469. const contents = yield* readEnvFile()
  470. const fileEnv = yield* parseEnv(contents)
  471. const providers = yield* Effect.promise(() => chooseProviders())
  472. printStatus(providers, fileEnv)
  473. if (checkOnly) {
  474. prompts.outro("Check complete")
  475. return
  476. }
  477. if (!interactive) {
  478. prompts.outro("Run this command in a terminal to enter credentials")
  479. return
  480. }
  481. const selectedProviders = yield* chooseConfigurableProviders(providers, fileEnv)
  482. const values = yield* promptProviderValues(selectedProviders)
  483. if (Object.keys(values).length === 0) {
  484. prompts.outro("No changes")
  485. return
  486. }
  487. if (
  488. interactive &&
  489. (yield* prompt(() => prompts.confirm({ message: "Validate credentials before saving?", initialValue: true })))
  490. ) {
  491. yield* validateProviders(selectedProviders, envWithValues(fileEnv, values))
  492. }
  493. yield* writeEnvFile(upsertEnv(contents, values))
  494. prompts.log.success(
  495. `Saved ${Object.keys(values).length} value${Object.keys(values).length === 1 ? "" : "s"} to ${envPath}`,
  496. )
  497. prompts.outro("Keep .env.local local. Store shared team credentials in a password manager or vault.")
  498. })
  499. await Effect.runPromise(main().pipe(Effect.provide(NodeFileSystem.layer), Effect.provide(FetchHttpClient.layer)))