feat(ai): add PaLM 2 API support

JS/edgechains/arakoodev/src/ai/src/index.ts

@@ -1,5 +1,17 @@
 export { OpenAI } from "./lib/openai/openai.js";
 export { GeminiAI } from "./lib/gemini/gemini.js";
+export { Palm2AI } from "./lib/palm2/palm2.js";
+export type {
+  Palm2AIConstructionOptions,
+  Palm2Candidate,
+  Palm2ChatOptions,
+  Palm2ChatResponse,
+  Palm2EmbeddingOptions,
+  Palm2EmbeddingResponse,
+  Palm2Message,
+  Palm2TextOptions,
+  Palm2TextResponse,
+} from "./lib/palm2/palm2.js";
 export { LlamaAI } from "./lib/llama/llama.js";
 export { RetellAI } from "./lib/retell-ai/retell.js";
 export { RetellWebClient } from "./lib/retell-ai/retellWebClient.js";

JS/edgechains/arakoodev/src/ai/src/lib/palm2/palm2.ts

@@ -0,0 +1,184 @@
+import axios, { AxiosRequestConfig } from "axios";
+import { retry } from "@lifeomic/attempt";
+
+const DEFAULT_BASE_URL = "https://generativelanguage.googleapis.com/v1beta2";
+
+export interface Palm2AIConstructionOptions {
+  apiKey?: string;
+  baseUrl?: string;
+  model?: string;
+}
+
+export interface Palm2Message {
+  author?: string;
+  content: string;
+}
+
+export interface Palm2Candidate {
+  author: string;
+  content: string;
+  [key: string]: unknown;
+}
+
+export interface Palm2ChatResponse {
+  candidates?: Palm2Candidate[];
+  messages?: Palm2Message[];
+  [key: string]: unknown;
+}
+
+export interface Palm2TextResponse {
+  candidates?: Array<{
+    output: string;
+    [key: string]: unknown;
+  }>;
+  [key: string]: unknown;
+}
+
+export interface Palm2EmbeddingResponse {
+  embedding?: {
+    value?: number[];
+    [key: string]: unknown;
+  };
+  [key: string]: unknown;
+}
+
+export interface Palm2ChatOptions {
+  prompt?: string;
+  messages?: Palm2Message[];
+  model?: string;
+  temperature?: number;
+  candidateCount?: number;
+  topP?: number;
+  topK?: number;
+  maxRetry?: number;
+  delay?: number;
+}
+
+export interface Palm2TextOptions {
+  prompt: string;
+  model?: string;
+  temperature?: number;
+  candidateCount?: number;
+  topP?: number;
+  topK?: number;
+  maxOutputTokens?: number;
+  maxRetry?: number;
+  delay?: number;
+}
+
+export interface Palm2EmbeddingOptions {
+  text: string;
+  model?: string;
+  maxRetry?: number;
+  delay?: number;
+}
+
+export class Palm2AI {
+  readonly apiKey: string;
+  readonly baseUrl: string;
+  readonly model: string;
+
+  constructor(options: Palm2AIConstructionOptions = {}) {
+    this.apiKey = options.apiKey || process.env.PALM_API_KEY || "";
+    this.baseUrl = (options.baseUrl || DEFAULT_BASE_URL).replace(/\/$/, "");
+    this.model = options.model || "chat-bison-001";
+
+    if (!this.apiKey) {
+      throw new Error(
+        "PALM_API_KEY is required. Provide it in the constructor or environment.",
+      );
+    }
+  }
+
+  private async request<T>(
+    model: string,
+    method: "generateMessage" | "generateText" | "embedText",
+    body: unknown,
+    maxRetry = 3,
+    delay = 200,
+  ): Promise<T> {
+    const config: AxiosRequestConfig = {
+      params: { key: this.apiKey },
+      headers: { "Content-Type": "application/json" },
+    };
+
+    return retry(
+      async () => {
+        const response = await axios.post(
+          `${this.baseUrl}/models/${encodeURIComponent(model)}:${method}`,
+          body,
+          config,
+        );
+        return response.data as T;
+      },
+      { maxAttempts: maxRetry, delay },
+    );
+  }
+
+  async chat(options: Palm2ChatOptions): Promise<Palm2ChatResponse> {
+    const messages =
+      options.messages ||
+      (options.prompt ? [{ author: "0", content: options.prompt }] : []);
+
+    if (messages.length === 0) {
+      throw new Error("prompt or messages is required");
+    }
+
+    return this.request<Palm2ChatResponse>(
+      options.model || this.model,
+      "generateMessage",
+      {
+        prompt: { messages },
+        ...(options.temperature === undefined
+          ? {}
+          : { temperature: options.temperature }),
+        ...(options.candidateCount === undefined
+          ? {}
+          : { candidate_count: options.candidateCount }),
+        ...(options.topP === undefined ? {} : { topP: options.topP }),
+        ...(options.topK === undefined ? {} : { topK: options.topK }),
+      },
+      options.maxRetry,
+      options.delay,
+    );
+  }
+
+  async generateText(options: Palm2TextOptions): Promise<Palm2TextResponse> {
+    if (!options.prompt) throw new Error("prompt is required");
+
+    return this.request<Palm2TextResponse>(
+      options.model || "text-bison-001",
+      "generateText",
+      {
+        prompt: { text: options.prompt },
+        ...(options.temperature === undefined
+          ? {}
+          : { temperature: options.temperature }),
+        ...(options.candidateCount === undefined
+          ? {}
+          : { candidateCount: options.candidateCount }),
+        ...(options.topP === undefined ? {} : { topP: options.topP }),
+        ...(options.topK === undefined ? {} : { topK: options.topK }),
+        ...(options.maxOutputTokens === undefined
+          ? {}
+          : { maxOutputTokens: options.maxOutputTokens }),
+      },
+      options.maxRetry,
+      options.delay,
+    );
+  }
+
+  async embedText(
+    options: Palm2EmbeddingOptions,
+  ): Promise<Palm2EmbeddingResponse> {
+    if (!options.text) throw new Error("text is required");
+
+    return this.request<Palm2EmbeddingResponse>(
+      options.model || "embedding-gecko-001",
+      "embedText",
+      { text: options.text },
+      options.maxRetry,
+      options.delay,
+    );
+  }
+}

JS/edgechains/arakoodev/src/ai/src/tests/palm2/palm2.test.ts

@@ -0,0 +1,69 @@
+import axios from "axios";
+import { describe, expect, it, vi } from "vitest";
+
+import { Palm2AI } from "../../lib/palm2/palm2.js";
+
+describe("Palm2AI", () => {
+  it("sends a chat prompt using the PaLM REST shape", async () => {
+    const post = vi.spyOn(axios, "post").mockResolvedValueOnce({
+      data: { candidates: [{ author: "1", content: "Hello" }] },
+    } as never);
+    const client = new Palm2AI({
+      apiKey: "test-key",
+      baseUrl: "https://example.test",
+    });
+
+    const response = await client.chat({
+      prompt: "Say hello",
+      temperature: 0.2,
+      candidateCount: 1,
+      topP: 0.8,
+      topK: 10,
+    });
+
+    expect(response.candidates?.[0].content).toBe("Hello");
+    expect(post).toHaveBeenCalledWith(
+      "https://example.test/models/chat-bison-001:generateMessage",
+      {
+        prompt: { messages: [{ author: "0", content: "Say hello" }] },
+        temperature: 0.2,
+        candidate_count: 1,
+        topP: 0.8,
+        topK: 10,
+      },
+      expect.objectContaining({
+        params: { key: "test-key" },
+      }),
+    );
+  });
+
+  it("supports text generation and embeddings", async () => {
+    const post = vi
+      .spyOn(axios, "post")
+      .mockResolvedValueOnce({
+        data: { candidates: [{ output: "Generated" }] },
+      } as never)
+      .mockResolvedValueOnce({
+        data: { embedding: { value: [0.1, 0.2] } },
+      } as never);
+    const client = new Palm2AI({ apiKey: "test-key" });
+
+    const text = await client.generateText({ prompt: "Write one word" });
+    const embedding = await client.embedText({ text: "Embed this" });
+
+    expect(text.candidates?.[0].output).toBe("Generated");
+    expect(embedding.embedding?.value).toEqual([0.1, 0.2]);
+    expect(post.mock.calls.map(([url]) => url)).toEqual([
+      "https://generativelanguage.googleapis.com/v1beta2/models/text-bison-001:generateText",
+      "https://generativelanguage.googleapis.com/v1beta2/models/embedding-gecko-001:embedText",
+    ]);
+  });
+
+  it("requires a prompt or messages", async () => {
+    const client = new Palm2AI({ apiKey: "test-key" });
+
+    await expect(client.chat({})).rejects.toThrow(
+      "prompt or messages is required",
+    );
+  });
+});

JS/edgechains/examples/palm2/README.md

@@ -0,0 +1,12 @@
+# PaLM 2 example
+
+This example keeps the prompt template in `jsonnet/main.jsonnet` and supplies
+the question at runtime.
+
+```bash
+PALM_API_KEY=your-key PALM2_QUESTION="What is a vector database?" npm start
+```
+
+The PaLM 2 REST API is a legacy API. New applications should use the Gemini
+API, but this example remains available for projects that still need the
+PaLM-compatible interface.

JS/edgechains/examples/palm2/jsonnet/main.jsonnet

@@ -0,0 +1,5 @@
+local question = std.extVar("question");
+
+{
+  prompt: "Answer the following question clearly and concisely: " + question,
+}

JS/edgechains/examples/palm2/package.json

@@ -0,0 +1,16 @@
+{
+  "name": "edgechains-palm2-example",
+  "private": true,
+  "type": "module",
+  "scripts": {
+    "start": "tsc && node dist/index.js"
+  },
+  "dependencies": {
+    "@arakoodev/edgechains.js": "file:../../arakoodev",
+    "@arakoodev/jsonnet": "^0.24.0"
+  },
+  "devDependencies": {
+    "@types/node": "^20.14.2",
+    "typescript": "^5.6.3"
+  }
+}

JS/edgechains/examples/palm2/src/index.ts

@@ -0,0 +1,22 @@
+import Jsonnet from "@arakoodev/jsonnet";
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+
+import { Palm2AI } from "@arakoodev/edgechains.js/ai";
+
+const question = process.env.PALM2_QUESTION;
+if (!question) {
+  throw new Error("Set PALM2_QUESTION to run the example.");
+}
+
+const jsonnet = new Jsonnet();
+jsonnet.extString("question", question);
+
+const directory = path.dirname(fileURLToPath(import.meta.url));
+const request = JSON.parse(
+  jsonnet.evaluateFile(path.join(directory, "../jsonnet/main.jsonnet")),
+);
+
+const client = new Palm2AI({ apiKey: process.env.PALM_API_KEY });
+const response = await client.chat(request);
+console.log(response.candidates?.[0]?.content || "No candidate returned");

JS/edgechains/examples/palm2/tsconfig.json

@@ -0,0 +1,13 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "module": "NodeNext",
+    "moduleResolution": "NodeNext",
+    "strict": true,
+    "esModuleInterop": true,
+    "skipLibCheck": true,
+    "outDir": "./dist",
+    "rootDir": "./src"
+  },
+  "include": ["./src/**/*.ts"]
+}