@iinm/plain-agent 1.0.0

package/src/providers/openaiCompatible.mjs

@@ -0,0 +1,658 @@
+/**
+ * @import { ModelInput, Message, MessageContentText, AssistantMessage, ModelOutput, PartialMessageContent, MessageContentThinking, MessageContentToolUse } from "../model"
+ * @import { OpenAIAssistantMessage, OpenAIMessage, OpenAIMessageToolCall, OpenAICompatibleModelConfig, OpenAIToolDefinition, OpenAIStreamData, OpenAIChatCompletion, OpenAIMessageContentImage, OpenAIChatCompletionRequest } from "./openaiCompatible"
+ * @import { ToolDefinition } from "../tool"
+ */
+
+import { styleText } from "node:util";
+import { Sha256 } from "@aws-crypto/sha256-js";
+import { fromIni } from "@aws-sdk/credential-providers";
+import { HttpRequest } from "@smithy/protocol-http";
+import { SignatureV4 } from "@smithy/signature-v4";
+import { noThrow } from "../utils/noThrow.mjs";
+import { retryOnError } from "../utils/retryOnError.mjs";
+import { readBedrockStreamEvents } from "./platform/bedrock.mjs";
+import { getGoogleCloudAccessToken } from "./platform/googleCloud.mjs";
+
+/**
+ * @param {import("../modelDefinition").PlatformConfig} platformConfig
+ * @param {OpenAICompatibleModelConfig} modelConfig
+ * @param {ModelInput} input
+ * @param {number} retryCount
+ * @returns {Promise<ModelOutput | Error>}
+ */
+export async function callOpenAICompatibleModel(
+  platformConfig,
+  modelConfig,
+  input,
+  retryCount = 0,
+) {
+  const retryInterval = Math.min(2 * 2 ** retryCount, 16);
+
+  return await noThrow(async () => {
+    const messages = convertGenericMessageToOpenAIFormat(input.messages);
+    const tools = convertGenericeToolDefinitionToOpenAIFormat(
+      input.tools || [],
+    );
+
+    const url = (() => {
+      const baseURL = platformConfig.baseURL;
+
+      switch (platformConfig.name) {
+        case "openai":
+          return `${baseURL}/v1/chat/completions`;
+        case "bedrock":
+          return `${baseURL}/model/${modelConfig.model}/invoke-with-response-stream`;
+        case "vertex-ai":
+          return `${baseURL}/endpoints/openapi/chat/completions`;
+        default:
+          throw new Error(`Unsupported platform: ${platformConfig.name}`);
+      }
+    })();
+
+    /** @type {Record<string,string>} */
+    const headers = await (async () => {
+      switch (platformConfig.name) {
+        case "openai":
+          return {
+            ...platformConfig.customHeaders,
+            Authorization: `Bearer ${platformConfig.apiKey}`,
+          };
+        case "bedrock":
+          return platformConfig.customHeaders ?? {};
+        case "vertex-ai":
+          return {
+            ...platformConfig.customHeaders,
+            Authorization: `Bearer ${await getGoogleCloudAccessToken()}`,
+          };
+      }
+    })();
+
+    const { model: _, ...modelConfigWithoutName } = modelConfig;
+    const platformRequest = (() => {
+      switch (platformConfig.name) {
+        case "openai":
+          return {
+            ...modelConfig,
+            stream: true,
+          };
+        case "bedrock":
+          return {
+            ...modelConfigWithoutName,
+          };
+        case "vertex-ai":
+          return {
+            ...modelConfig,
+            model: modelConfig.model,
+            stream: true,
+          };
+      }
+    })();
+
+    /** @type {OpenAIChatCompletionRequest} */
+    const request = {
+      ...platformRequest,
+      messages,
+      tools: tools.length ? tools : undefined,
+      stream_options: {
+        include_usage: true,
+      },
+    };
+
+    const runFetchDefault = async () =>
+      fetch(url, {
+        method: "POST",
+        headers: {
+          ...headers,
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify(request),
+        signal: AbortSignal.timeout(120 * 1000),
+      });
+
+    // bedrock + sso profile
+    const runFetchForBedrock = async () => {
+      const region =
+        url.match(/bedrock-runtime\.([\w-]+)\.amazonaws\.com/)?.[1] ?? "";
+      const urlParsed = new URL(url);
+      const { hostname, pathname } = urlParsed;
+
+      const signer = new SignatureV4({
+        credentials: fromIni({
+          profile:
+            platformConfig.name === "bedrock" ? platformConfig.awsProfile : "",
+        }),
+        region,
+        service: "bedrock",
+        sha256: Sha256,
+      });
+
+      const req = new HttpRequest({
+        protocol: "https:",
+        method: "POST",
+        hostname,
+        path: pathname,
+        headers: {
+          host: hostname,
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify(request),
+      });
+
+      const signed = await signer.sign(req);
+
+      return fetch(url, {
+        method: signed.method,
+        headers: signed.headers,
+        body: signed.body,
+        signal: AbortSignal.timeout(120 * 1000),
+      });
+    };
+
+    const runFetch =
+      platformConfig.name === "bedrock" ? runFetchForBedrock : runFetchDefault;
+
+    const response = await retryOnError(() => runFetch(), {
+      shouldRetry: (err) => err instanceof Error && err.name === "TimeoutError",
+      beforeRetry: (err, interval) => {
+        console.error(
+          styleText(
+            "yellow",
+            `Failed to call model: ${String(err)}; Retry in ${interval} seconds...`,
+          ),
+        );
+        return Promise.resolve();
+      },
+      initialInterval: 2,
+      maxInterval: 16,
+      multiplier: 2,
+      maxAttempt: 5,
+    });
+
+    if (response.status === 429 || response.status >= 500) {
+      console.error(
+        styleText(
+          "yellow",
+          `Model rate limit exceeded. Retry in ${retryInterval} seconds...`,
+        ),
+      );
+      await new Promise((resolve) => setTimeout(resolve, retryInterval * 1000));
+      return callOpenAICompatibleModel(
+        platformConfig,
+        modelConfig,
+        input,
+        retryCount + 1,
+      );
+    }
+
+    if (response.status !== 200) {
+      throw new Error(
+        `Failed to call OpenAI compatible model: status=${response.status}, body=${await response.text()}`,
+      );
+    }
+
+    if (!response.body) {
+      throw new Error("Response body is empty");
+    }
+
+    const reader = response.body.getReader();
+    const eventStreamReader =
+      platformConfig.name === "bedrock"
+        ? /** @type {typeof readOpenAIStreamData} */ (readBedrockStreamEvents)
+        : readOpenAIStreamData;
+
+    /** @type {OpenAIStreamData[]} */
+    const dataList = [];
+    /** @type {PartialMessageContent | undefined} */
+    let previousPartialContent;
+    for await (const data of eventStreamReader(reader)) {
+      dataList.push(data);
+
+      const partialContents = convertOpenAIStreamDataToAgentPartialContent(
+        data,
+        previousPartialContent,
+      );
+
+      if (partialContents.length) {
+        previousPartialContent = partialContents.at(-1);
+      }
+
+      if (input.onPartialMessageContent) {
+        for (const partialContent of partialContents) {
+          input.onPartialMessageContent(partialContent);
+        }
+      }
+    }
+
+    const chatCompletion = convertOpenAIStreamDataToChatCompletion(dataList);
+    if (chatCompletion instanceof Error) {
+      console.error(
+        styleText(
+          "yellow",
+          `Failed to process stream: ${chatCompletion.message}; Retry in ${retryInterval} seconds...`,
+        ),
+      );
+      await new Promise((resolve) => setTimeout(resolve, retryInterval * 1000));
+      return callOpenAICompatibleModel(
+        platformConfig,
+        modelConfig,
+        input,
+        retryCount + 1,
+      );
+    }
+
+    const openAIAssistantMessage = chatCompletion.choices[0].message;
+
+    return {
+      message: convertOpenAIAssistantMessageToGenericFormat(
+        openAIAssistantMessage,
+      ),
+      providerTokenUsage: chatCompletion.usage,
+    };
+  });
+}
+
+/**
+ * @param {Message[]} genericMessages
+ * @returns {OpenAIMessage[]}
+ */
+function convertGenericMessageToOpenAIFormat(genericMessages) {
+  /** @type {OpenAIMessage[]} */
+  const openAIMessages = [];
+  for (const genericMessage of genericMessages) {
+    switch (genericMessage.role) {
+      case "system": {
+        openAIMessages.push({
+          role: "system",
+          content: genericMessage.content.map((part) => ({
+            type: "text",
+            text: part.text,
+          })),
+        });
+        break;
+      }
+      case "user": {
+        const toolResults = genericMessage.content.filter(
+          (part) => part.type === "tool_result",
+        );
+        const userContentParts = genericMessage.content.filter(
+          (part) => part.type === "text" || part.type === "image",
+        );
+
+        // Tool Results
+        let imageIndex = 0;
+        for (const result of toolResults) {
+          const toolResultContentString = result.content
+            .map((part) => {
+              switch (part.type) {
+                case "text":
+                  return part.text;
+                case "image":
+                  imageIndex += 1;
+                  return `(Image [${imageIndex}] omitted. See next message from user.)`;
+                default:
+                  throw new Error(
+                    `Unsupported content part: ${JSON.stringify(part)}`,
+                  );
+              }
+            })
+            .join("\n\n");
+          openAIMessages.push({
+            role: "tool",
+            tool_call_id: result.toolUseId,
+            content: toolResultContentString,
+          });
+        }
+
+        /** @type {OpenAIMessageContentImage[]} */
+        const imageParts = [];
+        for (const result of toolResults) {
+          for (const part of result.content) {
+            if (part.type === "image") {
+              imageParts.push({
+                type: "image_url",
+                image_url: {
+                  url: `data:${part.mimeType};base64,${part.data}`,
+                },
+              });
+            }
+          }
+        }
+
+        if (imageParts.length) {
+          openAIMessages.push({
+            role: "user",
+            content: imageParts,
+          });
+        }
+
+        // User Input Parts
+        if (userContentParts.length) {
+          openAIMessages.push({
+            role: "user",
+            content: userContentParts.map((part) => {
+              if (part.type === "text") {
+                return { type: "text", text: part.text };
+              }
+              if (part.type === "image") {
+                return {
+                  type: "image_url",
+                  image_url: {
+                    url: `data:${part.mimeType};base64,${part.data}`,
+                  },
+                };
+              }
+              throw new Error(
+                `Unsupported content part: ${JSON.stringify(part)}`,
+              );
+            }),
+          });
+        }
+
+        break;
+      }
+      case "assistant": {
+        /** @type {MessageContentThinking[]} */
+        const thinkingParts = genericMessage.content.filter(
+          (part) => part.type === "thinking",
+        );
+        if (thinkingParts.length > 1) {
+          console.error(
+            `OpenAI Unsupported message format: ${JSON.stringify(genericMessage)}`,
+          );
+        }
+        const thinking = thinkingParts.map((part) => part.thinking).join("\n");
+
+        /** @type {MessageContentText[]} */
+        const textParts = genericMessage.content.filter(
+          (part) => part.type === "text",
+        );
+        if (textParts.length > 1) {
+          console.error(
+            `OpenAI Unsupported message format: ${JSON.stringify(genericMessage)}`,
+          );
+        }
+        const text = textParts.map((part) => part.text).join("\n");
+
+        /** @type {MessageContentToolUse[]} */
+        const toolUseParts = genericMessage.content.filter(
+          (part) => part.type === "tool_use",
+        );
+
+        /** @type {OpenAIMessageToolCall[]} */
+        const toolCalls = toolUseParts.map((part) => ({
+          id: part.toolUseId,
+          type: "function",
+          function: {
+            name: part.toolName,
+            arguments: JSON.stringify(part.input),
+          },
+        }));
+
+        openAIMessages.push({
+          role: "assistant",
+          reasoning_content: thinking ? thinking : undefined,
+          content: text ? text : undefined,
+          tool_calls: toolCalls.length ? toolCalls : undefined,
+        });
+      }
+    }
+  }
+
+  return openAIMessages;
+}
+
+/**
+ * @param {ToolDefinition[]} genericToolDefs
+ * @returns {OpenAIToolDefinition[]}
+ */
+function convertGenericeToolDefinitionToOpenAIFormat(genericToolDefs) {
+  /** @type {OpenAIToolDefinition[]} */
+  const openAIToolDefs = [];
+  for (const toolDef of genericToolDefs) {
+    openAIToolDefs.push({
+      type: "function",
+      function: {
+        name: toolDef.name,
+        description: toolDef.description,
+        parameters: toolDef.inputSchema,
+      },
+    });
+  }
+
+  return openAIToolDefs;
+}
+
+/**
+ * @param {OpenAIAssistantMessage} openAIAsistantMessage
+ * @returns {AssistantMessage}
+ */
+function convertOpenAIAssistantMessageToGenericFormat(openAIAsistantMessage) {
+  /** @type {AssistantMessage["content"]} */
+  const content = [];
+  if (openAIAsistantMessage.reasoning_content) {
+    content.push({
+      type: "thinking",
+      thinking: openAIAsistantMessage.reasoning_content,
+    });
+  }
+
+  if (openAIAsistantMessage.content) {
+    content.push({ type: "text", text: openAIAsistantMessage.content });
+  }
+
+  if (openAIAsistantMessage.tool_calls) {
+    for (const toolCall of openAIAsistantMessage.tool_calls) {
+      if (toolCall.type === "function") {
+        /** @type {Record<string, unknown>} */
+        let args;
+        try {
+          args = JSON.parse(toolCall.function.arguments);
+        } catch (err) {
+          args = { err: String(err) };
+        }
+
+        content.push({
+          type: "tool_use",
+          toolUseId: toolCall.id,
+          toolName: toolCall.function.name,
+          input: args,
+        });
+      } else {
+        throw new Error(
+          `Unsupported tool call type: ${JSON.stringify(toolCall)}`,
+        );
+      }
+    }
+  }
+
+  return {
+    role: "assistant",
+    content,
+  };
+}
+
+/**
+ * @param {OpenAIStreamData[]} dataList
+ * @returns {OpenAIChatCompletion | Error}
+ */
+function convertOpenAIStreamDataToChatCompletion(dataList) {
+  const firstData = dataList.at(0);
+  if (!firstData) {
+    return new Error("No data found in the stream");
+  }
+
+  const fistChoice = firstData.choices.at(0);
+  if (!fistChoice) {
+    return new Error("No choice found in the first data");
+  }
+
+  const message = /** @type {OpenAIAssistantMessage} */ (fistChoice.delta);
+
+  /** @type {Partial<OpenAIChatCompletion>} */
+  const chatCompletion = {
+    ...firstData,
+    choices: [
+      {
+        index: fistChoice.index,
+        message,
+        finish_reason: /** @type {string} */ (fistChoice.finish_reason),
+      },
+    ],
+  };
+
+  for (let i = 1; i < dataList.length; i++) {
+    const data = dataList[i];
+    const firstChoice = data?.choices.at(0);
+    if (firstChoice) {
+      const delta = firstChoice.delta;
+
+      if (delta.reasoning_content) {
+        message.reasoning_content =
+          (message.reasoning_content ?? "") + delta.reasoning_content;
+      }
+
+      if (delta.content) {
+        message.content = (message.content ?? "") + delta.content;
+      }
+
+      if (delta.tool_calls) {
+        for (const toolCallDelta of delta.tool_calls) {
+          const toolCall = message.tool_calls?.at(toolCallDelta.index);
+          if (!toolCall) {
+            if (!message.tool_calls) {
+              message.tool_calls = [];
+            }
+            /** @type {OpenAIMessageToolCall[]} */ (message.tool_calls).push(
+              /** @type {OpenAIMessageToolCall} */
+              (toolCallDelta),
+            );
+          }
+          if (toolCall && toolCallDelta.function) {
+            toolCall.function.arguments =
+              (toolCall.function.arguments ?? "") +
+              toolCallDelta.function.arguments;
+          }
+        }
+      }
+
+      if (firstChoice.finish_reason && chatCompletion.choices) {
+        chatCompletion.choices[0].finish_reason = firstChoice.finish_reason;
+      }
+    }
+
+    if (data.usage) {
+      chatCompletion.usage = data.usage;
+    }
+  }
+
+  return /** @type {OpenAIChatCompletion} */ (chatCompletion);
+}
+
+/**
+ * @param {OpenAIStreamData} data
+ * @param {PartialMessageContent | undefined} previousPartialContent
+ * @returns {PartialMessageContent[]}
+ */
+function convertOpenAIStreamDataToAgentPartialContent(
+  data,
+  previousPartialContent,
+) {
+  /** @type {PartialMessageContent[]} */
+  const partialContents = [];
+  const firstChoice = data.choices.at(0);
+
+  if (firstChoice?.delta.reasoning_content) {
+    partialContents.push({
+      type: "thinking",
+      content: firstChoice?.delta.reasoning_content,
+      position: previousPartialContent?.type === "thinking" ? "delta" : "start",
+    });
+  }
+
+  if (firstChoice?.delta.content) {
+    partialContents.push({
+      type: "text",
+      content: firstChoice.delta.content,
+      position: previousPartialContent?.type === "text" ? "delta" : "start",
+    });
+  }
+
+  if (firstChoice?.delta.tool_calls) {
+    partialContents.push({
+      type: "tool_use",
+      content: [
+        firstChoice.delta.tool_calls.at(0)?.function?.name,
+        firstChoice.delta.tool_calls.at(0)?.function?.arguments,
+      ].join(" "),
+      position: previousPartialContent?.type === "tool_use" ? "delta" : "start",
+    });
+  }
+
+  if (firstChoice?.finish_reason) {
+    partialContents.push({
+      type: previousPartialContent?.type || "unknown",
+      position: "stop",
+    });
+  }
+
+  if (
+    partialContents.length &&
+    previousPartialContent &&
+    partialContents[0].position !== "stop" &&
+    partialContents[0].type !== previousPartialContent.type
+  ) {
+    partialContents.unshift({
+      type: previousPartialContent.type,
+      position: "stop",
+    });
+  }
+
+  return partialContents;
+}
+
+/**
+ * @param {ReadableStreamDefaultReader<Uint8Array>} reader
+ */
+async function* readOpenAIStreamData(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 lineFeed = "\n".charCodeAt(0);
+    const dataEndIndices = [];
+    for (let i = 0; i < buffer.length - 1; i++) {
+      if (buffer[i] === lineFeed && buffer[i + 1] === lineFeed) {
+        dataEndIndices.push(i);
+      }
+    }
+
+    for (let i = 0; i < dataEndIndices.length; i++) {
+      const dataStartIndex = i === 0 ? 0 : dataEndIndices[i - 1] + 2;
+      const dataEndIndex = dataEndIndices[i];
+      const data = buffer.slice(dataStartIndex, dataEndIndex);
+      const decodedData = new TextDecoder().decode(data);
+      if (decodedData === "data: [DONE]") {
+        break;
+      }
+      if (decodedData.startsWith("data: ")) {
+        /** @type {OpenAIStreamData} */
+        const parsedData = JSON.parse(decodedData.slice("data: ".length));
+        yield parsedData;
+      }
+    }
+
+    if (dataEndIndices.length) {
+      buffer = buffer.slice(dataEndIndices[dataEndIndices.length - 1] + 2);
+    }
+  }
+}

package/src/providers/platform/azure.mjs

@@ -0,0 +1,42 @@
+import { execFile } from "node:child_process";
+
+/**
+ * @param {{azureConfigDir: string}=} config
+ * @returns {Promise<string>}
+ */
+export async function getAzureAccessToken(config) {
+  /** @type {string} */
+  const stdout = await new Promise((resolve, reject) => {
+    execFile(
+      "az",
+      [
+        "account",
+        "get-access-token",
+        "--resource",
+        "https://cognitiveservices.azure.com",
+        "--query",
+        "accessToken",
+        "--output",
+        "tsv",
+      ],
+      {
+        shell: false,
+        timeout: 10 * 1000,
+        env: config
+          ? {
+              AZURE_CONFIG_DIR: config.azureConfigDir,
+            }
+          : undefined,
+      },
+      (error, stdout, _stderr) => {
+        if (error) {
+          reject(error);
+          return;
+        }
+        resolve(stdout.trim());
+      },
+    );
+  });
+
+  return stdout;
+}