@lota-sdk/core 0.4.47 → 0.4.48

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lota-sdk/core",
-  "version": "0.4.47",
+  "version": "0.4.48",
   "files": [
     "src",
     "infrastructure/schema"
@@ -32,7 +32,7 @@
     "@ai-sdk/provider": "^3.0.9",
     "@chat-adapter/slack": "^4.26.0",
     "@chat-adapter/state-ioredis": "^4.26.0",
-    "@lota-sdk/shared": "0.4.47",
+    "@lota-sdk/shared": "0.4.48",
     "@mendable/firecrawl-js": "^4.20.0",
     "@surrealdb/node": "^3.0.3",
     "ai": "^6.0.170",

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

@@ -1,3 +1,5 @@
+import { createHash } from 'node:crypto'
+
 import { devToolsMiddleware } from '@ai-sdk/devtools'
 import { createOpenAI } from '@ai-sdk/openai'
 import type { JSONSchema7 } from '@ai-sdk/provider'
@@ -38,7 +40,16 @@ class AiGatewayStreamAttemptTag extends Context.Service<
 const EXPECTED_GATEWAY_KEY_PREFIX = 'sk-bf-'
 const AI_GATEWAY_VIRTUAL_KEY_HEADER = 'x-bf-vk'
 const AI_GATEWAY_PASSTHROUGH_EXTRA_PARAMS_HEADER = 'x-bf-passthrough-extra-params'
+const AI_GATEWAY_SESSION_ID_HEADER = 'x-bf-session-id'
+const AI_GATEWAY_SESSION_TTL_HEADER = 'x-bf-session-ttl'
 const AZURE_OPENAI_PROMPT_CACHE_RETENTION = '24h'
+const AZURE_OPENAI_PROMPT_CACHE_SESSION_TTL = '24h'
+const AZURE_OPENAI_PROMPT_CACHE_KEY_PREFIX = 'azpc'
+const AZURE_OPENAI_PROMPT_CACHE_HASH_LENGTH = 48
+const AZURE_OPENAI_PROMPT_CACHE_MAX_STRING_LENGTH = 120_000
+const AZURE_OPENAI_PROMPT_CACHE_MAX_ARRAY_ITEMS = 80
+const AZURE_OPENAI_PROMPT_CACHE_MAX_OBJECT_KEYS = 80
+const AZURE_OPENAI_PROMPT_CACHE_MAX_DEPTH = 8
 const AI_GATEWAY_TIMEOUT_MS = 360_000
 const AI_GATEWAY_STREAM_IDLE_TIMEOUT_MS = 180_000
 const AI_GATEWAY_MAX_RETRIES = 4
@@ -755,6 +766,77 @@ function isAzureOpenAiPromptCacheModel(modelId: string): boolean {
   )
 }
 
+function hashAzureOpenAiPromptCacheValue(value: string): string {
+  return createHash('sha256').update(value).digest('hex').slice(0, AZURE_OPENAI_PROMPT_CACHE_HASH_LENGTH)
+}
+
+function stablePromptCacheStringify(value: unknown): string {
+  return JSON.stringify(normalizePromptCacheValue(value, 0))
+}
+
+function normalizePromptCacheValue(value: unknown, depth: number): unknown {
+  if (value === null || typeof value === 'boolean' || typeof value === 'number') return value
+  if (typeof value === 'string') {
+    return value.length > AZURE_OPENAI_PROMPT_CACHE_MAX_STRING_LENGTH
+      ? value.slice(0, AZURE_OPENAI_PROMPT_CACHE_MAX_STRING_LENGTH)
+      : value
+  }
+  if (typeof value === 'bigint') return value.toString()
+  if (value instanceof Uint8Array) {
+    return {
+      type: 'uint8array',
+      length: value.byteLength,
+      sha256: createHash('sha256').update(value).digest('hex').slice(0, AZURE_OPENAI_PROMPT_CACHE_HASH_LENGTH),
+    }
+  }
+  if (Array.isArray(value)) {
+    return value
+      .slice(0, AZURE_OPENAI_PROMPT_CACHE_MAX_ARRAY_ITEMS)
+      .map((item) => normalizePromptCacheValue(item, depth + 1))
+  }
+  if (!isRecord(value) || depth >= AZURE_OPENAI_PROMPT_CACHE_MAX_DEPTH) return null
+
+  return Object.fromEntries(
+    Object.entries(value)
+      .filter(([, item]) => item !== undefined && typeof item !== 'function' && typeof item !== 'symbol')
+      .sort(([left], [right]) => left.localeCompare(right))
+      .slice(0, AZURE_OPENAI_PROMPT_CACHE_MAX_OBJECT_KEYS)
+      .map(([key, item]) => [key, normalizePromptCacheValue(item, depth + 1)]),
+  )
+}
+
+function readExplicitOpenAiPromptCacheKey(openaiOptions: Record<string, unknown>): string | null {
+  const promptCacheKey = openaiOptions.promptCacheKey
+  return typeof promptCacheKey === 'string' && promptCacheKey.trim().length > 0 ? promptCacheKey.trim() : null
+}
+
+function buildAzureOpenAiPromptCacheKey(params: AiGatewayCallOptions, modelId: string): string {
+  const payload = {
+    version: 1,
+    model: modelId.trim().toLowerCase(),
+    prompt: params.prompt,
+    responseFormat: params.responseFormat,
+    tools: params.tools,
+  }
+  return `${AZURE_OPENAI_PROMPT_CACHE_KEY_PREFIX}_${hashAzureOpenAiPromptCacheValue(stablePromptCacheStringify(payload))}`
+}
+
+function buildAzureOpenAiPromptCacheSessionId(promptCacheKey: string): string {
+  if (promptCacheKey.startsWith(`${AZURE_OPENAI_PROMPT_CACHE_KEY_PREFIX}_`)) return promptCacheKey
+  return `${AZURE_OPENAI_PROMPT_CACHE_KEY_PREFIX}_${hashAzureOpenAiPromptCacheValue(promptCacheKey)}`
+}
+
+function withHeaderIfMissing(
+  headers: Record<string, string | undefined>,
+  name: string,
+  value: string,
+): Record<string, string | undefined> {
+  for (const key of Object.keys(headers)) {
+    if (key.toLowerCase() === name.toLowerCase()) return headers
+  }
+  return { ...headers, [name]: value }
+}
+
 export function addAzureOpenAiPromptCacheRetention(
   params: AiGatewayCallOptions,
   modelId?: string,
@@ -765,12 +847,21 @@ export function addAzureOpenAiPromptCacheRetention(
 
   const providerOptions = isRecord(params.providerOptions) ? { ...params.providerOptions } : {}
   const openaiOptions = isRecord(providerOptions.openai) ? { ...providerOptions.openai } : {}
+  const promptCacheKey =
+    readExplicitOpenAiPromptCacheKey(openaiOptions) ?? buildAzureOpenAiPromptCacheKey(params, modelId)
+  const sessionId = buildAzureOpenAiPromptCacheSessionId(promptCacheKey)
+  const headersWithSession = withHeaderIfMissing(
+    withHeaderIfMissing({ ...params.headers }, AI_GATEWAY_SESSION_ID_HEADER, sessionId),
+    AI_GATEWAY_SESSION_TTL_HEADER,
+    AZURE_OPENAI_PROMPT_CACHE_SESSION_TTL,
+  )
 
   return {
     ...params,
+    headers: headersWithSession,
     providerOptions: {
       ...providerOptions,
-      openai: { ...openaiOptions, promptCacheRetention: AZURE_OPENAI_PROMPT_CACHE_RETENTION },
+      openai: { ...openaiOptions, promptCacheKey, promptCacheRetention: AZURE_OPENAI_PROMPT_CACHE_RETENTION },
     } as AiGatewayCallOptions['providerOptions'],
   }
 }