@lota-sdk/core 0.1.17 → 0.1.19

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lota-sdk/core",
-  "version": "0.1.17",
+  "version": "0.1.19",
   "type": "module",
   "main": "./src/index.ts",
   "types": "./src/index.ts",
@@ -32,7 +32,7 @@
     "@chat-adapter/slack": "^4.23.0",
     "@chat-adapter/state-ioredis": "^4.23.0",
     "@logtape/logtape": "^2.0.5",
-    "@lota-sdk/shared": "0.1.17",
+    "@lota-sdk/shared": "0.1.19",
     "@mendable/firecrawl-js": "^4.17.0",
     "@surrealdb/node": "^3.0.3",
     "ai": "^6.0.137",

package/src/ai-gateway/ai-gateway.ts

@@ -0,0 +1,431 @@
+import { devToolsMiddleware } from '@ai-sdk/devtools'
+import { createOpenAI } from '@ai-sdk/openai'
+import { wrapLanguageModel } from 'ai'
+import type { LanguageModelMiddleware } from 'ai'
+
+import { getRuntimeConfig } from '../runtime/runtime-config'
+import { isRecord, readString } from '../utils/string'
+import { buildAiGatewayCacheHeaders } from './cache-headers'
+
+type AiGatewayLanguageModel = Parameters<typeof wrapLanguageModel>[0]['model']
+type AiGatewayExtraParams = Record<string, unknown>
+type AiGatewayChatResponse = { body?: unknown }
+type AiGatewayTransformParamsOptions = Parameters<NonNullable<LanguageModelMiddleware['transformParams']>>[0]
+type WrapStreamOptions = Parameters<NonNullable<LanguageModelMiddleware['wrapStream']>>[0]
+type AiGatewayCallOptions = WrapStreamOptions['params']
+type AiGatewayGenerateResult = Awaited<ReturnType<WrapStreamOptions['doGenerate']>>
+type AiGatewayStreamResult = Awaited<ReturnType<WrapStreamOptions['doStream']>>
+type AiGatewayGeneratedContent = AiGatewayGenerateResult['content'][number]
+type AiGatewayStreamPart = AiGatewayStreamResult['stream'] extends ReadableStream<infer T> ? T : never
+type AiGatewayConfig = { apiKey: string; baseURL: string }
+type AiGatewayProviderOptions = NonNullable<AiGatewayCallOptions['providerOptions']>
+
+const EXPECTED_GATEWAY_KEY_PREFIX = 'sk-bf-'
+const AI_GATEWAY_VIRTUAL_KEY_HEADER = 'x-bf-vk'
+const AI_GATEWAY_EXTRA_PARAMS_HEADER = 'x-bf-passthrough-extra-params'
+const DEFAULT_AI_GATEWAY_URL = 'https://ai-gateway.gobrainy.ai' as const
+const OPENROUTER_RESPONSE_HEALING_EXTRA_PARAMS = {
+  plugins: [{ id: 'response-healing' }],
+} as const satisfies AiGatewayExtraParams
+
+function toAiGatewayCacheKeyPart(value: string): string {
+  const normalized = value
+    .trim()
+    .toLowerCase()
+    .replace(/[^a-z0-9:_-]+/g, '-')
+    .replace(/-+/g, '-')
+  return normalized.replace(/^-+|-+$/g, '') || 'request'
+}
+
+function mergeAiGatewayHeaders(
+  existingHeaders: AiGatewayCallOptions['headers'] | undefined,
+  additionalHeaders: Record<string, string>,
+): Record<string, string> {
+  const merged = new Headers(existingHeaders as HeadersInit | undefined)
+  for (const [key, value] of Object.entries(additionalHeaders)) {
+    if (!merged.has(key)) {
+      merged.set(key, value)
+    }
+  }
+  return Object.fromEntries(merged.entries())
+}
+
+function withDefaultAiGatewayCacheHeaders(params: AiGatewayCallOptions, modelId: string): AiGatewayCallOptions {
+  return {
+    ...params,
+    headers: mergeAiGatewayHeaders(
+      params.headers,
+      buildAiGatewayCacheHeaders(`model:${toAiGatewayCacheKeyPart(modelId)}`),
+    ),
+  }
+}
+
+function normalizeAiGatewayUrl(value: string): string {
+  const trimmed = value.trim()
+  if (!trimmed) {
+    throw new Error('[ai-gateway] AI gateway URL is required.')
+  }
+
+  const normalized = trimmed.replace(/\/+$/, '')
+  return normalized.endsWith('/v1') ? normalized : `${normalized}/v1`
+}
+
+function readDirectEnvAiGatewayConfig(): AiGatewayConfig {
+  const apiKey = (process.env.AI_GATEWAY_KEY ?? '').trim()
+  if (!apiKey) {
+    throw new Error(
+      '[ai-gateway] Missing AI gateway key. Set AI_GATEWAY_KEY, or configure createLotaRuntime({ aiGateway: { key } }).',
+    )
+  }
+
+  return { apiKey, baseURL: normalizeAiGatewayUrl(process.env.AI_GATEWAY_URL?.trim() || DEFAULT_AI_GATEWAY_URL) }
+}
+
+function readAiGatewayConfig(): AiGatewayConfig {
+  try {
+    const { aiGateway } = getRuntimeConfig()
+    return { apiKey: aiGateway.key.trim(), baseURL: normalizeAiGatewayUrl(aiGateway.url) }
+  } catch {
+    return readDirectEnvAiGatewayConfig()
+  }
+}
+
+function readReasoningDetailsText(value: unknown): string | null {
+  if (!Array.isArray(value)) return null
+
+  const textParts = value
+    .map((item) => (isRecord(item) ? readString(item.text) : null))
+    .filter((item): item is string => item !== null)
+
+  if (textParts.length === 0) return null
+
+  return textParts.join('\n\n')
+}
+
+function readReasoningDeltaText(value: unknown): string | null {
+  return typeof value === 'string' && value.length > 0 ? value : null
+}
+
+function readAiGatewayChatReasoningText(message: Record<string, unknown>): string | null {
+  return (
+    readString(message.reasoning) ??
+    readString(message.reasoning_content) ??
+    readReasoningDetailsText(message.reasoning_details)
+  )
+}
+
+export function extractAiGatewayChatReasoningText(responseBody: unknown): string | null {
+  if (!isRecord(responseBody) || !Array.isArray(responseBody.choices)) return null
+
+  for (const choice of responseBody.choices) {
+    if (!isRecord(choice) || !isRecord(choice.message)) continue
+
+    const reasoningText = readAiGatewayChatReasoningText(choice.message)
+    if (reasoningText) return reasoningText
+  }
+
+  return null
+}
+
+export function extractAiGatewayChatReasoningDeltaText(rawChunk: unknown): string | null {
+  if (!isRecord(rawChunk) || !Array.isArray(rawChunk.choices)) return null
+
+  for (const choice of rawChunk.choices) {
+    if (!isRecord(choice) || !isRecord(choice.delta)) continue
+
+    const reasoningText =
+      readReasoningDeltaText(choice.delta.reasoning) ??
+      readReasoningDeltaText(choice.delta.reasoning_content) ??
+      readReasoningDetailsText(choice.delta.reasoning_details)
+    if (reasoningText) return reasoningText
+  }
+
+  return null
+}
+
+type AiGatewayResponsesReasoningDelta = { id: string; delta: string; itemId: string }
+
+export function extractAiGatewayResponsesReasoningDelta(rawChunk: unknown): AiGatewayResponsesReasoningDelta | null {
+  if (!isRecord(rawChunk) || rawChunk.type !== 'response.reasoning_summary_text.delta') return null
+  if ('summary_index' in rawChunk) return null
+
+  const itemId = readString(rawChunk.item_id)
+  const delta = readReasoningDeltaText(rawChunk.delta)
+  if (!itemId || !delta) return null
+
+  return { id: `${itemId}:0`, delta, itemId }
+}
+
+export function injectAiGatewayChatReasoningContent(
+  content: readonly AiGatewayGeneratedContent[],
+  response?: AiGatewayChatResponse,
+): AiGatewayGeneratedContent[] {
+  if (content.some((part) => part.type === 'reasoning')) {
+    return [...content]
+  }
+
+  const reasoningText = extractAiGatewayChatReasoningText(response?.body)
+  if (!reasoningText) return [...content]
+
+  return [{ type: 'reasoning', text: reasoningText }, ...content]
+}
+
+function isReasoningEnabled(params: AiGatewayCallOptions): boolean {
+  if (!isRecord(params.providerOptions) || !isRecord(params.providerOptions.openai)) return false
+
+  const openaiOptions = params.providerOptions.openai
+  if (openaiOptions.forceReasoning === true) return true
+  if (typeof openaiOptions.reasoningSummary === 'string' && openaiOptions.reasoningSummary.length > 0) return true
+  return typeof openaiOptions.reasoningEffort === 'string' && openaiOptions.reasoningEffort !== 'none'
+}
+
+function shouldCloseInjectedReasoning(chunk: AiGatewayStreamPart): boolean {
+  return chunk.type !== 'stream-start' && chunk.type !== 'response-metadata' && chunk.type !== 'raw'
+}
+
+export function injectAiGatewayChatReasoningStream(
+  stream: ReadableStream<AiGatewayStreamPart>,
+): ReadableStream<AiGatewayStreamPart> {
+  const reasoningId = 'ai-gateway-reasoning-0'
+  let reasoningOpen = false
+  let reasoningClosed = false
+
+  return stream.pipeThrough(
+    new TransformStream<AiGatewayStreamPart, AiGatewayStreamPart>({
+      transform(chunk, controller) {
+        const closeReasoning = () => {
+          if (!reasoningOpen || reasoningClosed) return
+
+          controller.enqueue({ type: 'reasoning-end', id: reasoningId } satisfies AiGatewayStreamPart)
+          reasoningClosed = true
+        }
+
+        if (chunk.type === 'raw') {
+          const reasoningDelta = reasoningClosed ? null : extractAiGatewayChatReasoningDeltaText(chunk.rawValue)
+          controller.enqueue(chunk)
+
+          if (reasoningDelta) {
+            if (!reasoningOpen) {
+              controller.enqueue({ type: 'reasoning-start', id: reasoningId } satisfies AiGatewayStreamPart)
+              reasoningOpen = true
+            }
+
+            controller.enqueue({
+              type: 'reasoning-delta',
+              id: reasoningId,
+              delta: reasoningDelta,
+            } satisfies AiGatewayStreamPart)
+          }
+          return
+        }
+
+        if (shouldCloseInjectedReasoning(chunk)) {
+          closeReasoning()
+        }
+
+        controller.enqueue(chunk)
+      },
+      flush(controller) {
+        if (!reasoningOpen || reasoningClosed) return
+        controller.enqueue({ type: 'reasoning-end', id: reasoningId } satisfies AiGatewayStreamPart)
+      },
+    }),
+  )
+}
+
+export function injectAiGatewayResponsesReasoningStream(
+  stream: ReadableStream<AiGatewayStreamPart>,
+): ReadableStream<AiGatewayStreamPart> {
+  return stream.pipeThrough(
+    new TransformStream<AiGatewayStreamPart, AiGatewayStreamPart>({
+      transform(chunk, controller) {
+        controller.enqueue(chunk)
+
+        if (chunk.type !== 'raw') return
+
+        const reasoningDelta = extractAiGatewayResponsesReasoningDelta(chunk.rawValue)
+        if (!reasoningDelta) return
+
+        controller.enqueue({
+          type: 'reasoning-delta',
+          id: reasoningDelta.id,
+          delta: reasoningDelta.delta,
+          providerMetadata: { openai: { itemId: reasoningDelta.itemId } },
+        } satisfies AiGatewayStreamPart)
+      },
+    }),
+  )
+}
+
+function addAiGatewayReasoningRawChunks(
+  params: AiGatewayCallOptions,
+  type: AiGatewayTransformParamsOptions['type'],
+): AiGatewayCallOptions {
+  if (type !== 'stream' || !isReasoningEnabled(params) || params.includeRawChunks === true) {
+    return params
+  }
+
+  return { ...params, includeRawChunks: true }
+}
+
+export function normalizeAiGatewayChatProviderOptions(params: AiGatewayCallOptions): AiGatewayCallOptions {
+  const providerOptions = isRecord(params.providerOptions)
+    ? ({ ...params.providerOptions } as AiGatewayProviderOptions)
+    : ({} as AiGatewayProviderOptions)
+  const openaiOptions = isRecord(providerOptions.openai)
+    ? { ...providerOptions.openai }
+    : ({} as Record<string, unknown>)
+
+  if (openaiOptions.systemMessageMode === 'system') {
+    return params
+  }
+
+  return {
+    ...params,
+    providerOptions: {
+      ...providerOptions,
+      openai: {
+        ...openaiOptions,
+        ...(openaiOptions.systemMessageMode === 'remove' ? {} : { systemMessageMode: 'system' }),
+      },
+    },
+  }
+}
+
+export function injectAiGatewayExtraParamsRequestBody(
+  body: BodyInit | null | undefined,
+  extraParams: AiGatewayExtraParams,
+): BodyInit | null | undefined {
+  if (typeof body !== 'string') return body
+
+  let parsed: unknown
+  try {
+    parsed = JSON.parse(body)
+  } catch {
+    return body
+  }
+
+  if (!isRecord(parsed)) return body
+
+  const mergedExtraParams = isRecord(parsed.extra_params)
+    ? { ...parsed.extra_params, ...extraParams }
+    : { ...extraParams }
+
+  return JSON.stringify({ ...parsed, extra_params: mergedExtraParams })
+}
+
+function createAiGatewayFetchWithExtraParams(extraParams: AiGatewayExtraParams): typeof fetch {
+  const fetchWithExtraParams = (input: RequestInfo | URL, init?: RequestInit | BunFetchRequestInit) =>
+    globalThis.fetch(input, { ...init, body: injectAiGatewayExtraParamsRequestBody(init?.body, extraParams) })
+
+  return Object.assign(fetchWithExtraParams, { preconnect: globalThis.fetch.preconnect.bind(globalThis.fetch) })
+}
+
+function createAiGatewayProvider(extraParams?: AiGatewayExtraParams) {
+  const { apiKey, baseURL } = readAiGatewayConfig()
+  if (!apiKey.startsWith(EXPECTED_GATEWAY_KEY_PREFIX)) {
+    throw new Error(`[ai-gateway] Gateway keys must use the ${EXPECTED_GATEWAY_KEY_PREFIX}* format.`)
+  }
+
+  return createOpenAI({
+    baseURL,
+    apiKey,
+    headers: {
+      [AI_GATEWAY_VIRTUAL_KEY_HEADER]: apiKey,
+      ...(extraParams ? { [AI_GATEWAY_EXTRA_PARAMS_HEADER]: 'true' } : {}),
+    },
+    ...(extraParams ? { fetch: createAiGatewayFetchWithExtraParams(extraParams) } : {}),
+  })
+}
+
+function withAiGatewayDevTools<TModel extends AiGatewayLanguageModel>(model: TModel): TModel {
+  return wrapLanguageModel({ model, middleware: devToolsMiddleware() }) as TModel
+}
+
+let provider: ReturnType<typeof createOpenAI> | null = null
+let openRouterResponseHealingProvider: ReturnType<typeof createOpenAI> | null = null
+
+export function getAiGatewayProvider() {
+  if (provider) return provider
+
+  provider = createAiGatewayProvider()
+
+  return provider
+}
+
+export function getAiGatewayOpenRouterResponseHealingProvider() {
+  if (openRouterResponseHealingProvider) return openRouterResponseHealingProvider
+
+  openRouterResponseHealingProvider = createAiGatewayProvider(OPENROUTER_RESPONSE_HEALING_EXTRA_PARAMS)
+
+  return openRouterResponseHealingProvider
+}
+
+export function aiGatewayModel(modelId: string) {
+  return withAiGatewayDevTools(
+    wrapLanguageModel({
+      model: getAiGatewayProvider()(modelId),
+      middleware: {
+        specificationVersion: 'v3',
+        transformParams: async ({ params, type }) =>
+          withDefaultAiGatewayCacheHeaders(addAiGatewayReasoningRawChunks(params, type), modelId),
+        wrapStream: async ({ doStream, params }) => {
+          const result = await doStream()
+          if (!isReasoningEnabled(params)) return result
+
+          return { ...result, stream: injectAiGatewayResponsesReasoningStream(result.stream) }
+        },
+      },
+    }),
+  )
+}
+
+export function aiGatewayOpenRouterResponseHealingModel(modelId: string) {
+  return withAiGatewayDevTools(
+    wrapLanguageModel({
+      model: getAiGatewayOpenRouterResponseHealingProvider()(modelId),
+      middleware: {
+        specificationVersion: 'v3',
+        transformParams: async ({ params }) => withDefaultAiGatewayCacheHeaders(params, modelId),
+      },
+    }),
+  )
+}
+
+export function aiGatewayChatModel(modelId: string) {
+  return withAiGatewayDevTools(
+    wrapLanguageModel({
+      model: getAiGatewayProvider().chat(modelId),
+      middleware: {
+        specificationVersion: 'v3',
+        transformParams: async ({ params, type }) =>
+          normalizeAiGatewayChatProviderOptions(
+            withDefaultAiGatewayCacheHeaders(addAiGatewayReasoningRawChunks(params, type), modelId),
+          ),
+        wrapGenerate: async ({ doGenerate }) => {
+          const result = await doGenerate()
+
+          return {
+            ...result,
+            content: injectAiGatewayChatReasoningContent(
+              result.content,
+              result.response as AiGatewayChatResponse | undefined,
+            ),
+          }
+        },
+        wrapStream: async ({ doStream, params }) => {
+          const result = await doStream()
+          if (!isReasoningEnabled(params)) return result
+
+          return { ...result, stream: injectAiGatewayChatReasoningStream(result.stream) }
+        },
+      },
+    }),
+  )
+}
+
+export function aiGatewayEmbeddingModel(modelId: string) {
+  return getAiGatewayProvider().embeddingModel(modelId)
+}
+
+export { DEFAULT_AI_GATEWAY_URL, normalizeAiGatewayUrl }

package/src/ai-gateway/cache-headers.ts

@@ -0,0 +1,33 @@
+const AI_GATEWAY_CACHE_KEY_HEADER = 'x-bf-cache-key'
+const AI_GATEWAY_CACHE_TTL_HEADER = 'x-bf-cache-ttl'
+const AI_GATEWAY_CACHE_THRESHOLD_HEADER = 'x-bf-cache-threshold'
+const AI_GATEWAY_CACHE_TYPE_HEADER = 'x-bf-cache-type'
+
+export const AI_GATEWAY_STRICT_SEMANTIC_CACHE_THRESHOLD = 0.975
+
+export type AiGatewayCacheType = 'direct' | 'semantic'
+
+export function buildAiGatewayCacheHeaders(
+  cacheKey: string,
+  ttl?: string,
+  threshold?: number,
+  cacheType?: AiGatewayCacheType,
+): Record<string, string> {
+  const headers: Record<string, string> = { [AI_GATEWAY_CACHE_KEY_HEADER]: cacheKey }
+  if (ttl) headers[AI_GATEWAY_CACHE_TTL_HEADER] = ttl
+  if (typeof threshold === 'number') headers[AI_GATEWAY_CACHE_THRESHOLD_HEADER] = String(threshold)
+  if (cacheType) headers[AI_GATEWAY_CACHE_TYPE_HEADER] = cacheType
+  return headers
+}
+
+export function buildAiGatewayDirectCacheHeaders(cacheKey: string, ttl?: string): Record<string, string> {
+  return buildAiGatewayCacheHeaders(cacheKey, ttl, undefined, 'direct')
+}
+
+export function buildAiGatewayStrictSemanticCacheHeaders(
+  cacheKey: string,
+  ttl?: string,
+  threshold = AI_GATEWAY_STRICT_SEMANTIC_CACHE_THRESHOLD,
+): Record<string, string> {
+  return buildAiGatewayCacheHeaders(cacheKey, ttl, threshold, 'semantic')
+}

package/src/{bifrost → ai-gateway}/index.ts

@@ -1,2 +1,2 @@
-export * from './bifrost'
+export * from './ai-gateway'
 export * from './cache-headers'

package/src/config/model-constants.ts

@@ -1,5 +1,5 @@
 export {
-  BIFROST_REASONING_SUMMARY_LEVEL,
+  AI_GATEWAY_REASONING_SUMMARY_LEVEL,
   OPENAI_HIGH_REASONING_PROVIDER_OPTIONS,
   OPENAI_REASONING_MODEL_ID,
   OPENROUTER_DELEGATED_REASONING_MODEL_ID,

package/src/embeddings/provider.ts

@@ -1,7 +1,7 @@
 import { embed, embedMany } from 'ai'
 
+import { aiGatewayEmbeddingModel } from '../ai-gateway/ai-gateway'
 import { getEmbeddingCache } from '../ai/embedding-cache'
-import { bifrostEmbeddingModel } from '../bifrost/bifrost'
 import { getRuntimeConfig } from '../runtime/runtime-config'
 
 const SUPPORTED_EMBEDDING_PREFIXES = ['openai/', 'openrouter/'] as const
@@ -30,7 +30,7 @@ function resolveEmbeddingModel(modelId: string) {
     )
   }
 
-  return bifrostEmbeddingModel(normalized)
+  return aiGatewayEmbeddingModel(normalized)
 }
 
 function normalizeEmbedding(embedding: readonly number[]): number[] {

package/src/index.ts

@@ -1,6 +1,6 @@
 export * from './create-runtime'
 export * from './ai'
-export * from './bifrost'
+export * from './ai-gateway'
 export * from './config'
 export * from './db'
 export * from './document'

package/src/runtime/chat-run-registry.ts

@@ -1,6 +1,10 @@
 export class ChatRunRegistry {
   private controllers = new Map<string, AbortController>()
 
+  has(runId: string): boolean {
+    return this.controllers.has(runId)
+  }
+
   register(runId: string, controller: AbortController): void {
     this.controllers.set(runId, controller)
   }