zeitlich 0.2.49 → 0.2.50

package/dist/adapters/thread/google-genai/index.js

@@ -1,3 +1,4 @@
+import { randomBytes } from 'crypto';
 import { Context } from '@temporalio/activity';
 
 // src/adapters/thread/google-genai/adapter-id.ts
@@ -53,11 +54,11 @@ function createThreadManager(config) {
   return {
     async initialize() {
       await redis.del(redisKey);
-      await redis.set(metaKey, "1", "EX", ttlSeconds);
+      await redis.set(metaKey, "1", { EX: ttlSeconds });
     },
     async load() {
       await assertThreadExists();
-      const data = await redis.lrange(redisKey, 0, -1);
+      const data = await redis.lRange(redisKey, 0, -1);
       return data.map(deserialize);
     },
     async append(messages) {
@@ -65,22 +66,18 @@ function createThreadManager(config) {
       await assertThreadExists();
       if (idOf) {
         const dedupId = messages.map(idOf).join(":");
-        await redis.eval(
-          APPEND_IDEMPOTENT_SCRIPT,
-          2,
-          dedupKey(dedupId),
-          redisKey,
-          String(ttlSeconds),
-          ...messages.map(serialize)
-        );
+        await redis.eval(APPEND_IDEMPOTENT_SCRIPT, {
+          keys: [dedupKey(dedupId), redisKey],
+          arguments: [String(ttlSeconds), ...messages.map(serialize)]
+        });
       } else {
-        await redis.rpush(redisKey, ...messages.map(serialize));
+        await redis.rPush(redisKey, messages.map(serialize));
         await redis.expire(redisKey, ttlSeconds);
       }
     },
     async fork(newThreadId) {
       await assertThreadExists();
-      const data = await redis.lrange(redisKey, 0, -1);
+      const data = await redis.lRange(redisKey, 0, -1);
       const stateRaw = await redis.get(stateKey);
       const forked = createThreadManager({
         ...config,
@@ -89,12 +86,12 @@ function createThreadManager(config) {
       await forked.initialize();
       if (data.length > 0) {
         const newKey = getThreadListKey(key, newThreadId);
-        await redis.rpush(newKey, ...data);
+        await redis.rPush(newKey, data);
         await redis.expire(newKey, ttlSeconds);
       }
       if (stateRaw != null) {
         const newStateKey = getThreadStateKey(key, newThreadId);
-        await redis.set(newStateKey, stateRaw, "EX", ttlSeconds);
+        await redis.set(newStateKey, stateRaw, { EX: ttlSeconds });
       }
       return forked;
     },
@@ -105,20 +102,20 @@ function createThreadManager(config) {
           "replaceAll requires the thread manager to be configured with `idOf`"
         );
       }
-      const existing = await redis.lrange(redisKey, 0, -1);
+      const existing = await redis.lRange(redisKey, 0, -1);
       const existingIds = existing.map((raw) => idOf(deserialize(raw))).filter((id) => typeof id === "string");
       await redis.del(redisKey);
       if (existingIds.length > 0) {
-        await redis.del(...existingIds.map(dedupKey));
+        await redis.del(existingIds.map(dedupKey));
       }
       if (messages.length > 0) {
-        await redis.rpush(redisKey, ...messages.map(serialize));
+        await redis.rPush(redisKey, messages.map(serialize));
         await redis.expire(redisKey, ttlSeconds);
       }
       await redis.expire(metaKey, ttlSeconds);
     },
     async delete() {
-      await redis.del(redisKey, metaKey, stateKey);
+      await redis.del([redisKey, metaKey, stateKey]);
     },
     async loadState() {
       const raw = await redis.get(stateKey);
@@ -127,14 +124,14 @@ function createThreadManager(config) {
     },
     async saveState(state) {
       await assertThreadExists();
-      await redis.set(stateKey, JSON.stringify(state), "EX", ttlSeconds);
+      await redis.set(stateKey, JSON.stringify(state), { EX: ttlSeconds });
     },
     async deleteState() {
       await redis.del(stateKey);
     },
     async length() {
       await assertThreadExists();
-      return redis.llen(redisKey);
+      return redis.lLen(redisKey);
     },
     async truncateFromId(messageId) {
       await assertThreadExists();
@@ -143,7 +140,7 @@ function createThreadManager(config) {
           "truncateFromId requires the thread manager to be configured with `idOf`"
         );
       }
-      const data = await redis.lrange(redisKey, 0, -1);
+      const data = await redis.lRange(redisKey, 0, -1);
       let idx = -1;
       const removedIds = [];
       for (let i = 0; i < data.length; i++) {
@@ -158,11 +155,11 @@ function createThreadManager(config) {
         await redis.del(redisKey);
         await redis.expire(metaKey, ttlSeconds);
       } else {
-        await redis.ltrim(redisKey, 0, idx - 1);
+        await redis.lTrim(redisKey, 0, idx - 1);
         await redis.expire(redisKey, ttlSeconds);
       }
       if (removedIds.length > 0) {
-        await redis.del(...removedIds.map(dedupKey));
+        await redis.del(removedIds.map(dedupKey));
       }
     }
   };
@@ -177,7 +174,7 @@ async function encodeSnapshot(config) {
   }
   const listKey = getThreadListKey(threadKey, threadId);
   const stateKey = getThreadStateKey(threadKey, threadId);
-  const messages = await redis.lrange(listKey, 0, -1);
+  const messages = await redis.lRange(listKey, 0, -1);
   const stateRaw = await redis.get(stateKey);
   const state = stateRaw == null ? null : JSON.parse(stateRaw);
   const dedupIds = idOf ? messages.map(idOf) : [];
@@ -197,31 +194,38 @@ async function applySnapshot(config) {
   }
   const listKey = getThreadListKey(threadKey, threadId);
   const stateKey = getThreadStateKey(threadKey, threadId);
-  await redis.del(listKey, stateKey);
-  const pipeline = redis.pipeline();
+  await redis.del([listKey, stateKey]);
+  const pipeline = redis.multi();
   if (snapshot.messages.length > 0) {
-    pipeline.rpush(listKey, ...snapshot.messages);
+    pipeline.rPush(listKey, snapshot.messages);
     pipeline.expire(listKey, ttlSeconds);
   }
   if (snapshot.state != null) {
-    pipeline.set(stateKey, JSON.stringify(snapshot.state), "EX", ttlSeconds);
+    pipeline.set(stateKey, JSON.stringify(snapshot.state), { EX: ttlSeconds });
   }
   for (const id of snapshot.dedupIds) {
-    pipeline.set(getThreadDedupKey(threadId, id), "1", "EX", ttlSeconds);
+    pipeline.set(getThreadDedupKey(threadId, id), "1", { EX: ttlSeconds });
   }
-  const results = await pipeline.exec();
-  if (results) {
-    const firstErr = results.find(([err]) => err)?.[0] ?? null;
-    if (firstErr) {
-      await redis.del(
-        listKey,
-        stateKey,
-        ...snapshot.dedupIds.map((id) => getThreadDedupKey(threadId, id))
-      ).catch(() => void 0);
-      throw firstErr;
-    }
+  try {
+    await pipeline.execAsPipeline();
+  } catch (err) {
+    await redis.del([
+      listKey,
+      stateKey,
+      ...snapshot.dedupIds.map((id) => getThreadDedupKey(threadId, id))
+    ]).catch(() => void 0);
+    throw firstPipelineError(err);
   }
-  await redis.set(metaKey, "1", "EX", ttlSeconds);
+  await redis.set(metaKey, "1", { EX: ttlSeconds });
+}
+function firstPipelineError(err) {
+  if (err != null && typeof err === "object" && "replies" in err && Array.isArray(err.replies)) {
+    const firstErr = err.replies.find(
+      (r) => r instanceof Error
+    );
+    if (firstErr) return firstErr;
+  }
+  return err;
 }
 async function clearHotTier(config) {
   const { redis, threadKey, threadId, dedupIds = [] } = config;
@@ -231,7 +235,7 @@ async function clearHotTier(config) {
     getThreadStateKey(threadKey, threadId),
     ...dedupIds.map((id) => getThreadDedupKey(threadId, id))
   ];
-  await redis.del(...keys);
+  await redis.del(keys);
 }
 
 // src/lib/thread/tiered.ts
@@ -429,7 +433,9 @@ function createGoogleGenAIModelInvoker({
   redis,
   client,
   model,
-  hooks
+  hooks,
+  config: generationConfig,
+  cache: cacheConfig
 }) {
   return async function invokeGoogleGenAIModel2(config) {
     const { threadId, threadKey, state, assistantMessageId } = config;
@@ -444,12 +450,57 @@ function createGoogleGenAIModelInvoker({
     const { contents, systemInstruction } = await thread.prepareForInvocation();
     const functionDeclarations = toFunctionDeclarations(state.tools);
     const tools = functionDeclarations.length > 0 ? [{ functionDeclarations }] : void 0;
+    const {
+      systemInstruction: _si,
+      tools: _t,
+      abortSignal: _as,
+      cachedContent: callerCachedContent,
+      toolConfig: callerToolConfig,
+      ...callerConfig
+    } = generationConfig ?? {};
+    let liveContents = contents;
+    let cachedContentName;
+    let cachedWriteTokens;
+    if (cacheConfig && cacheConfig.splitIndex > 0 && contents.length > cacheConfig.splitIndex) {
+      liveContents = contents.slice(cacheConfig.splitIndex);
+      const ttl = cacheConfig.ttlSeconds ?? 300;
+      const cacheRedisKey = `${threadKey ?? "messages"}:gemini-cache:${model}:${cacheConfig.splitIndex}:thread:${threadId}`;
+      cachedContentName = await redis.get(cacheRedisKey) ?? void 0;
+      if (!cachedContentName) {
+        const cacheInstance = await client.caches.create({
+          model,
+          config: {
+            contents: contents.slice(0, cacheConfig.splitIndex),
+            ...systemInstruction ? { systemInstruction } : {},
+            ...tools ? { tools } : {},
+            ...callerToolConfig ? { toolConfig: callerToolConfig } : {},
+            ttl: `${ttl}s`,
+            abortSignal: signal
+          }
+        });
+        if (!cacheInstance?.name) {
+          throw new Error("Gemini cache creation did not return a cache name");
+        }
+        cachedContentName = cacheInstance.name;
+        cachedWriteTokens = cacheInstance.usageMetadata?.totalTokenCount ?? void 0;
+        const redisTtl = ttl - 5;
+        if (redisTtl > 0) {
+          await redis.set(cacheRedisKey, cachedContentName, { EX: redisTtl });
+        }
+      }
+    }
     const stream = await client.models.generateContentStream({
       model,
-      contents,
+      contents: liveContents,
       config: {
-        ...systemInstruction ? { systemInstruction } : {},
-        ...tools ? { tools } : {},
+        ...callerConfig,
+        ...cachedContentName ? { cachedContent: cachedContentName } : {
+          ...callerCachedContent ? { cachedContent: callerCachedContent } : {
+            ...systemInstruction ? { systemInstruction } : {},
+            ...tools ? { tools } : {}
+          },
+          ...callerToolConfig ? { toolConfig: callerToolConfig } : {}
+        },
         abortSignal: signal
       }
     });
@@ -463,19 +514,27 @@ function createGoogleGenAIModelInvoker({
     if (!lastChunk) {
       throw new Error("Google GenAI stream ended without producing any chunks");
     }
+    for (const part of allParts) {
+      if (part.functionCall && !part.functionCall.id) {
+        part.functionCall.id = randomBytes(8).toString("hex");
+      }
+    }
     const modelContent = { role: "model", parts: allParts };
-    const functionCalls = lastChunk.functionCalls ?? [];
     return {
       message: modelContent,
-      rawToolCalls: functionCalls.map((fc) => ({
-        id: fc.id,
-        name: fc.name ?? "",
-        args: fc.args ?? {}
+      rawToolCalls: allParts.filter(
+        (p) => !!p.functionCall
+      ).map((p) => ({
+        id: p.functionCall.id,
+        name: p.functionCall.name ?? "",
+        args: p.functionCall.args ?? {}
       })),
       usage: {
         inputTokens: lastChunk.usageMetadata?.promptTokenCount,
         outputTokens: lastChunk.usageMetadata?.candidatesTokenCount,
-        cachedReadTokens: lastChunk.usageMetadata?.cachedContentTokenCount
+        cachedWriteTokens,
+        cachedReadTokens: lastChunk.usageMetadata?.cachedContentTokenCount,
+        reasonTokens: lastChunk.usageMetadata?.thoughtsTokenCount
       }
     };
   };
@@ -485,13 +544,17 @@ async function invokeGoogleGenAIModel({
   client,
   model,
   hooks,
-  config
+  config,
+  generationConfig,
+  cache
 }) {
   const invoker = createGoogleGenAIModelInvoker({
     redis,
     client,
     model,
-    hooks
+    hooks,
+    config: generationConfig,
+    cache
   });
   return invoker(config);
 }