@iinm/plain-agent 1.0.0

package/src/providers/gemini.mjs

@@ -0,0 +1,752 @@
+/**
+ * @import { ModelInput, Message, AssistantMessage, ModelOutput, PartialMessageContent, ProviderTokenUsage } from "../model";
+ * @import { GeminiCachedContents, GeminiContent, GeminiContentPartFunctionCall, GeminiContentPartText, GeminiCreateCachedContentInput as GeminiCreateCachedContentInput, GeminiFunctionContent, GeminiGenerateContentInput, GeminiGeneratedContent, GeminiModelConfig, GeminiModelContent, GeminiSystemContent, GeminiToolConfig, GeminiToolDefinition, GeminiUserContent } from "./gemini";
+ * @import { ToolDefinition } from "../tool";
+ */
+
+import { styleText } from "node:util";
+import { noThrow } from "../utils/noThrow.mjs";
+import { getGoogleCloudAccessToken } from "./platform/googleCloud.mjs";
+
+/**
+ * @callback GeminiModelCaller
+ * @param {GeminiModelConfig} config
+ * @param {ModelInput} input
+ * @param {number=} retryCount
+ * @returns {Promise<ModelOutput | Error>}
+ */
+
+/**
+ * References:
+ * - https://ai.google.dev/gemini-api/docs/caching
+ * - https://ai.google.dev/api/caching
+ * @param {import("../modelDefinition").PlatformConfig} platformConfig
+ * @param {Pick<GeminiModelConfig, "model">} modelConfig
+ * @returns {GeminiModelCaller}
+ */
+export function createCacheEnabledGeminiModelCaller(
+  platformConfig,
+  modelConfig,
+) {
+  const baseURL =
+    platformConfig.baseURL ||
+    "https://generativelanguage.googleapis.com/v1beta";
+
+  const props = {
+    cacheTTL: 2 * 60, // seconds
+    // https://ai.google.dev/gemini-api/docs/caching#considerations
+    minCacheableTokenCount: 2048,
+  };
+
+  /**
+   * @typedef {Object} CacheState
+   * @property {string} name
+   * @property {number} contentsLength - Length of contents without system
+   * @property {Date} expireTime
+   */
+
+  const state = {
+    /** @type {CacheState=} */
+    cache: undefined,
+  };
+
+  /** @type {GeminiModelCaller} */
+  async function modelCaller(config, input, retryCount = 0) {
+    return await noThrow(async () => {
+      const contents = convertGenericMessageToGeminiFormat(input.messages);
+      const tools = convertGenericToolDefinitionToGeminiFormat(
+        input.tools || [],
+      );
+      /** @type {GeminiToolConfig} */
+      const toolConfig = {
+        functionCallingConfig: {
+          // Workaround to prevent MALFORMED_FUNCTION_CALL issues with gemini-3-flash
+          mode: "VALIDATED",
+        },
+      };
+      const systemInstruction = contents.find((c) => c.role === "system");
+      const contentsWithoutSystem = contents.filter((c) => c.role !== "system");
+
+      // Clear cache if messages are cleared
+      if (contentsWithoutSystem.length <= (state.cache?.contentsLength ?? 0)) {
+        state.cache = undefined;
+      }
+
+      const url = (() => {
+        switch (platformConfig.name) {
+          case "gemini":
+            return `${baseURL}/models/${config.model}:streamGenerateContent?alt=sse`;
+          case "vertex-ai":
+            return `${baseURL}/publishers/google/models/${config.model}:streamGenerateContent?alt=sse`;
+          default:
+            throw new Error(`Unsupported platform: ${platformConfig.name}`);
+        }
+      })();
+
+      /** @type {Record<string,string>} */
+      const headers = await (async () => {
+        switch (platformConfig.name) {
+          case "gemini":
+            return {
+              ...platformConfig.customHeaders,
+              "x-goog-api-key": platformConfig.apiKey,
+            };
+          case "vertex-ai":
+            return {
+              ...platformConfig.customHeaders,
+              Authorization: `Bearer ${await getGoogleCloudAccessToken(platformConfig.account)}`,
+            };
+        }
+      })();
+
+      /** @type {Pick<GeminiGenerateContentInput, "generationConfig" | "safetySettings">} */
+      const baseRequest = {
+        // default
+        generationConfig: {
+          temperature: 0,
+        },
+        safetySettings: [
+          {
+            category: "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+            threshold: "BLOCK_NONE",
+          },
+          { category: "HARM_CATEGORY_HATE_SPEECH", threshold: "BLOCK_NONE" },
+          { category: "HARM_CATEGORY_HARASSMENT", threshold: "BLOCK_NONE" },
+          {
+            category: "HARM_CATEGORY_DANGEROUS_CONTENT",
+            threshold: "BLOCK_NONE",
+          },
+        ],
+        ...config.requestConfig,
+      };
+
+      /** @type {GeminiGenerateContentInput} */
+      const request =
+        state.cache && Date.now() < state.cache.expireTime.getTime()
+          ? {
+              ...baseRequest,
+              cachedContent: state.cache.name,
+              contents: contentsWithoutSystem.slice(state.cache.contentsLength),
+            }
+          : {
+              ...baseRequest,
+              system_instruction: systemInstruction,
+              contents: contentsWithoutSystem,
+              tools: tools.length ? tools : undefined,
+              toolConfig,
+            };
+
+      const response = await fetch(url, {
+        method: "POST",
+        headers: {
+          ...headers,
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify(request),
+        signal: AbortSignal.timeout(120 * 1000),
+      });
+
+      if (response.status === 429 || response.status >= 500) {
+        const interval = Math.min(2 * 2 ** retryCount, 16);
+        console.error(
+          styleText(
+            "yellow",
+            `Gemini rate limit exceeded. Retrying in ${interval} seconds...`,
+          ),
+        );
+        await new Promise((resolve) => setTimeout(resolve, interval * 1000));
+        return modelCaller(config, input, retryCount + 1);
+      }
+
+      if (response.status !== 200) {
+        return new Error(
+          `Failed to call Gemini model: status=${response.status}, body=${await response.text()}`,
+        );
+      }
+
+      if (!response.body) {
+        throw new Error("Response body is empty");
+      }
+
+      const reader = response.body.getReader();
+
+      /** @type {GeminiGeneratedContent[]} */
+      const streamContents = [];
+      /** @type {PartialMessageContent | undefined} */
+      let previousPartialContent;
+
+      for await (const streamContent of readGeminiStreamContents(reader)) {
+        streamContents.push(streamContent);
+
+        const partialContents =
+          convertGeminiStreamContentToAgentPartialContents(
+            streamContent,
+            previousPartialContent,
+          );
+        previousPartialContent = partialContents.at(-1);
+
+        if (input.onPartialMessageContent && partialContents.length) {
+          for (const partialContent of partialContents) {
+            input.onPartialMessageContent(partialContent);
+          }
+        }
+      }
+
+      if (input.onPartialMessageContent && previousPartialContent) {
+        input.onPartialMessageContent({
+          type: previousPartialContent.type,
+          position: "stop",
+        });
+      }
+
+      /** @type {GeminiGeneratedContent} */
+      const content = convertGeminiStreamContentsToContent(streamContents);
+
+      /** @type {ProviderTokenUsage} */
+      const tokenUsage = {
+        input:
+          content.usageMetadata.promptTokenCount -
+          (content.usageMetadata.cachedContentTokenCount ?? 0),
+        cached: content.usageMetadata.cachedContentTokenCount ?? 0,
+        output: content.usageMetadata.candidatesTokenCount ?? 0,
+        thought: content.usageMetadata.thoughtsTokenCount ?? 0,
+        total: content.usageMetadata.totalTokenCount,
+      };
+
+      const message = convertGeminiAssistantMessageToGenericFormat(content);
+      if (
+        message instanceof GeminiNoCandidateError ||
+        message instanceof GeminiMalformedFunctionCallError
+      ) {
+        const interval = Math.min(2 * 2 ** retryCount, 16);
+        console.error(
+          styleText(
+            "yellow",
+            `${message.name}: Retrying in ${interval} seconds...`,
+          ),
+        );
+        await new Promise((resolve) => setTimeout(resolve, interval * 1000));
+        return modelCaller(config, input, retryCount + 1);
+      }
+
+      // Create context cache for next request
+      if (
+        props.minCacheableTokenCount < content.usageMetadata.promptTokenCount
+      ) {
+        await updateCache({
+          contentsWithoutSystem: [
+            ...contentsWithoutSystem,
+            /** @type {GeminiModelContent} */ (
+              content.candidates?.at(0)?.content
+            ),
+          ],
+          systemInstruction,
+          tools,
+          toolConfig,
+          headers,
+        });
+      }
+
+      return {
+        message,
+        providerTokenUsage: tokenUsage,
+      };
+    });
+  }
+
+  /**
+   * @typedef {Object} UpdateCacheParams
+   * @property {(GeminiUserContent|GeminiModelContent|GeminiFunctionContent)[]} contentsWithoutSystem
+   * @property {GeminiSystemContent=} systemInstruction
+   * @property {GeminiToolDefinition[]=} tools
+   * @property {GeminiToolConfig=} toolConfig
+   * @property {Record<string,string>} headers
+   */
+
+  /**
+   * @param {UpdateCacheParams} params
+   */
+  async function updateCache({
+    contentsWithoutSystem,
+    systemInstruction,
+    tools,
+    toolConfig,
+    headers,
+  }) {
+    const modelPrefix =
+      platformConfig.name === "vertex-ai"
+        ? `${baseURL.match(/projects\/[^/]+\/locations\/[^/]+/)?.[0] || ""}/publishers/google/models`
+        : "models";
+
+    const url = `${baseURL}/cachedContents`;
+
+    /** @type {GeminiCreateCachedContentInput} */
+    const request = {
+      model: `${modelPrefix}/${modelConfig.model}`,
+      ttl: `${props.cacheTTL}s`,
+      system_instruction: systemInstruction,
+      contents: contentsWithoutSystem,
+      tools,
+      toolConfig,
+    };
+
+    await fetch(url, {
+      method: "POST",
+      headers: {
+        ...headers,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify(request),
+      signal: AbortSignal.timeout(120 * 1000),
+    })
+      .then(async (response) => {
+        if (response.status !== 200) {
+          console.error(
+            styleText(
+              "yellow",
+              `Failed to create Gemini context cache: status=${response.status}, body=${await response.text()}`,
+            ),
+          );
+        } else {
+          /** @type {GeminiCachedContents} */
+          const cachedContents = await response.json();
+
+          // Delete old cache if previous cache is alive
+          if (state.cache && Date.now() < state.cache.expireTime.getTime()) {
+            fetch(
+              `${url}/${state.cache.name.replace(/.*cachedContents\//, "")}`,
+              {
+                method: "DELETE",
+                headers: {
+                  ...headers,
+                  "Content-Type": "application/json",
+                },
+                signal: AbortSignal.timeout(120 * 1000),
+              },
+            )
+              .then(async (response) => {
+                if (response.status !== 200) {
+                  console.error(
+                    styleText(
+                      "yellow",
+                      `Failed to delete Gemini context cache: status=${response.status}, body=${await response.text()}`,
+                    ),
+                  );
+                }
+              })
+              .catch((error) => {
+                console.error(
+                  styleText(
+                    "yellow",
+                    `Failed to delete Gemini context cache: ${error}`,
+                  ),
+                );
+              });
+          }
+
+          state.cache = {
+            name: cachedContents.name,
+            contentsLength: contentsWithoutSystem.length,
+            expireTime: new Date(cachedContents.expireTime),
+          };
+        }
+      })
+      .catch((error) => {
+        console.error(
+          styleText(
+            "yellow",
+            `Failed to create Gemini context cache: ${error}`,
+          ),
+        );
+      });
+  }
+
+  return modelCaller;
+}
+
+/**
+ * @param {Message[]} messages
+ * @returns {GeminiContent[]}
+ */
+function convertGenericMessageToGeminiFormat(messages) {
+  /** @type {GeminiContent[]} */
+  const geminiContents = [];
+  for (const message of messages) {
+    switch (message.role) {
+      case "system": {
+        geminiContents.push({
+          role: "system",
+          parts: message.content.map((part) => ({
+            text: part.text,
+          })),
+        });
+        break;
+      }
+      case "user": {
+        const toolUseResults = message.content.filter(
+          (part) => part.type === "tool_result",
+        );
+        const userContentParts = message.content.filter(
+          (part) => part.type === "text" || part.type === "image",
+        );
+
+        if (toolUseResults.length) {
+          geminiContents.push({
+            role: "user",
+            parts: toolUseResults.map((toolResult) => ({
+              functionResponse: {
+                name: toolResult.toolName,
+                response: {
+                  name: toolResult.toolName,
+                  content: toolResult.content.map((part) => {
+                    switch (part.type) {
+                      case "text":
+                        return { text: part.text };
+                      case "image":
+                        return {
+                          inline_data: {
+                            mime_type: part.mimeType,
+                            data: part.data,
+                          },
+                        };
+                      default:
+                        throw new Error(
+                          `Unsupported content part: ${JSON.stringify(part)}`,
+                        );
+                    }
+                  }),
+                },
+              },
+            })),
+          });
+        }
+
+        if (userContentParts.length) {
+          geminiContents.push({
+            role: "user",
+            parts: userContentParts.map((part) => {
+              if (part.type === "text") {
+                return { text: part.text };
+              }
+              if (part.type === "image") {
+                return {
+                  inline_data: {
+                    mime_type: part.mimeType,
+                    data: part.data,
+                  },
+                };
+              }
+              throw new Error(
+                `Unsupported content part: ${JSON.stringify(part)}`,
+              );
+            }),
+          });
+        }
+
+        break;
+      }
+      case "assistant": {
+        /** @type {(GeminiContentPartText | GeminiContentPartFunctionCall)[]} */
+        const parts = [];
+        for (const part of message.content) {
+          if (part.type === "thinking") {
+            parts.push({
+              text: part.thinking,
+              thought: true,
+              ...(part.provider?.fields || {}),
+            });
+          } else if (part.type === "text") {
+            parts.push({
+              text: part.text,
+              ...(part.provider?.fields || {}),
+            });
+          } else if (part.type === "tool_use") {
+            parts.push({
+              functionCall: {
+                name: part.toolName,
+                args: part.input,
+              },
+              ...(part.provider?.fields || {}),
+            });
+          }
+        }
+        geminiContents.push({
+          role: "model",
+          parts,
+        });
+        break;
+      }
+    }
+  }
+
+  return geminiContents;
+}
+
+/**
+ * @param {ToolDefinition[]} tools
+ * @returns {GeminiToolDefinition[]}
+ */
+function convertGenericToolDefinitionToGeminiFormat(tools) {
+  /** @type {GeminiToolDefinition["functionDeclarations"]} */
+  const functionDeclarations = [];
+  for (const tool of tools) {
+    functionDeclarations.push({
+      name: tool.name,
+      description: tool.description,
+      parametersJsonSchema: tool.inputSchema,
+    });
+  }
+
+  return [
+    {
+      functionDeclarations,
+    },
+  ];
+}
+
+/**
+ * @param {GeminiGeneratedContent} content
+ * @returns {AssistantMessage | GeminiNoCandidateError | GeminiMalformedFunctionCallError}
+ */
+function convertGeminiAssistantMessageToGenericFormat(content) {
+  const candidate = content.candidates?.at(0);
+  if (!candidate) {
+    return new GeminiNoCandidateError(
+      `No candidates found: content=${JSON.stringify(content)}`,
+    );
+  }
+
+  if (candidate.finishReason === "MALFORMED_FUNCTION_CALL") {
+    return new GeminiMalformedFunctionCallError(
+      `Malformed function call: content=${JSON.stringify(content)}`,
+    );
+  }
+
+  /** @type {AssistantMessage["content"]} */
+  const assistantMessageContent = [];
+  for (const part of candidate.content.parts || []) {
+    if ("text" in part) {
+      if (part.thought) {
+        // thought summary
+        assistantMessageContent.push({
+          type: "thinking",
+          thinking: part.text,
+        });
+      } else {
+        assistantMessageContent.push({
+          type: "text",
+          text: part.text,
+          provider: part.thoughtSignature
+            ? { fields: { thoughtSignature: part.thoughtSignature } }
+            : undefined,
+        });
+      }
+    }
+    if ("functionCall" in part) {
+      assistantMessageContent.push({
+        type: "tool_use",
+        toolUseId: part.functionCall.name,
+        toolName: part.functionCall.name,
+        input: part.functionCall.args,
+        provider: part.thoughtSignature
+          ? { fields: { thoughtSignature: part.thoughtSignature } }
+          : undefined,
+      });
+    }
+  }
+
+  return {
+    role: "assistant",
+    content: assistantMessageContent,
+  };
+}
+
+/**
+ * @param {GeminiGeneratedContent[]} events
+ * @returns {GeminiGeneratedContent}
+ */
+function convertGeminiStreamContentsToContent(events) {
+  const firstContent = events.at(0);
+  if (!firstContent) {
+    throw new Error("No content found");
+  }
+
+  /** @type {GeminiGeneratedContent} */
+  const mergedContent = {
+    ...firstContent,
+    // avoid side effects of mutating the original object
+    candidates: (firstContent.candidates || []).map((candidate) => ({
+      ...candidate,
+      content: {
+        ...candidate.content,
+        parts: [...(candidate.content.parts || [])],
+      },
+    })),
+  };
+
+  for (let i = 1; i < events.length; i++) {
+    const event = events[i];
+    if (event.candidates?.length) {
+      const candidate = event.candidates.at(0);
+      if (candidate?.content.parts?.length) {
+        mergedContent.candidates?.[0].content.parts?.push(
+          ...candidate.content.parts,
+        );
+      }
+      if (candidate?.finishReason && mergedContent.candidates?.[0]) {
+        mergedContent.candidates[0].finishReason = candidate.finishReason;
+      }
+      if (candidate?.finishMessage && mergedContent.candidates?.[0]) {
+        mergedContent.candidates[0].finishMessage = candidate.finishMessage;
+      }
+    }
+
+    if (event.usageMetadata.totalTokenCount) {
+      mergedContent.usageMetadata = event.usageMetadata;
+    }
+  }
+
+  return mergedContent;
+}
+
+/**
+ * @param {GeminiGeneratedContent} event
+ * @param {PartialMessageContent | undefined} previousPartialContent
+ * @returns {PartialMessageContent[]}
+ */
+function convertGeminiStreamContentToAgentPartialContents(
+  event,
+  previousPartialContent,
+) {
+  const candiate = event.candidates?.at(0);
+  /** @type {PartialMessageContent[]} */
+  const partialMessageContents = [];
+  if (candiate?.content.parts?.length) {
+    /** @type {string | undefined} */
+    let previousPartType = previousPartialContent?.type;
+    for (const part of candiate.content.parts) {
+      const partType =
+        "text" in part ? (part.thought ? "thinking" : "text") : "tool_use";
+
+      if (previousPartType && previousPartType !== partType) {
+        partialMessageContents.push({
+          type: previousPartType,
+          position: "stop",
+        });
+      }
+
+      if (previousPartType !== partType) {
+        previousPartType = partType;
+        partialMessageContents.push({
+          type: partType,
+          position: "start",
+        });
+      }
+
+      if ("text" in part) {
+        if (part.thought) {
+          partialMessageContents.push({
+            type: "thinking",
+            content: part.text,
+            position: "delta",
+          });
+        } else {
+          partialMessageContents.push({
+            type: "text",
+            content: part.text,
+            position: "delta",
+          });
+        }
+      }
+
+      if ("functionCall" in part) {
+        if (previousPartialContent?.type === "tool_use") {
+          partialMessageContents.push({
+            type: "tool_use",
+            content: "\n",
+            position: "delta",
+          });
+        }
+        partialMessageContents.push({
+          type: "tool_use",
+          content: part.functionCall.name,
+          position: "delta",
+        });
+      }
+    }
+  }
+
+  return partialMessageContents;
+}
+
+class GeminiNoCandidateError extends Error {
+  /**
+   * @param {string} message
+   */
+  constructor(message) {
+    super(message);
+    this.name = "GeminiNoCandidateError";
+  }
+}
+
+class GeminiMalformedFunctionCallError extends Error {
+  /**
+   * @param {string} message
+   */
+  constructor(message) {
+    super(message);
+    this.name = "GeminiMalformedFunctionCallError";
+  }
+}
+
+/**
+ * @param {ReadableStreamDefaultReader<Uint8Array>} reader
+ * @returns {AsyncGenerator<GeminiGeneratedContent>}
+ */
+async function* readGeminiStreamContents(reader) {
+  let buffer = new Uint8Array();
+
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) {
+      break;
+    }
+
+    const nextBuffer = new Uint8Array(buffer.length + value.length);
+    nextBuffer.set(buffer);
+    nextBuffer.set(value, buffer.length);
+    buffer = nextBuffer;
+
+    const carriageReturn = "\r".charCodeAt(0);
+    const lineFeed = "\n".charCodeAt(0);
+
+    const dataEndIndices = [];
+    for (let i = 0; i < buffer.length - 3; i++) {
+      if (
+        buffer[i] === carriageReturn &&
+        buffer[i + 1] === lineFeed &&
+        buffer[i + 2] === carriageReturn &&
+        buffer[i + 3] === lineFeed
+      ) {
+        dataEndIndices.push(i);
+      }
+    }
+
+    for (let i = 0; i < dataEndIndices.length; i++) {
+      const dataStartIndex = i === 0 ? 0 : dataEndIndices[i - 1] + 4;
+      const dataEndIndex = dataEndIndices[i];
+      const data = buffer.slice(dataStartIndex, dataEndIndex);
+      const decodedData = new TextDecoder().decode(data);
+
+      if (decodedData.startsWith("data: {")) {
+        /** @type {GeminiGeneratedContent} */
+        const parsedData = JSON.parse(decodedData.slice("data: ".length));
+        yield parsedData;
+      }
+    }
+
+    if (dataEndIndices.length) {
+      buffer = buffer.slice(dataEndIndices[dataEndIndices.length - 1] + 4);
+    }
+  }
+}