Add PaLM2-compatible Google AI support

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

@@ -1,5 +1,16 @@
 export { OpenAI } from "./lib/openai/openai.js";
-export { GeminiAI } from "./lib/gemini/gemini.js";
+export { GeminiAI, Palm2AI } from "./lib/gemini/gemini.js";
+export type {
+    Candidate,
+    Content,
+    ContentPart,
+    GeminiAIChatOptions,
+    GeminiAIConstructionOptions,
+    GeminiAIResponse,
+    ResponseMimeType,
+    SafetyRating,
+    UsageMetadata,
+} from "./lib/gemini/gemini.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/gemini/gemini.ts

@@ -1,12 +1,12 @@
 import axios from "axios";
 import { retry } from "@lifeomic/attempt";
-const url = "https://generativelanguage.googleapis.com/v1/models/gemini-pro:generateContent";
 
-interface GeminiAIConstructionOptions {
+export interface GeminiAIConstructionOptions {
     apiKey?: string;
+    baseUrl?: string;
 }
 
-type SafetyRating = {
+export type SafetyRating = {
     category:
         | "HARM_CATEGORY_SEXUALLY_EXPLICIT"
         | "HARM_CATEGORY_HATE_SPEECH"
@@ -15,53 +15,59 @@ type SafetyRating = {
     probability: "NEGLIGIBLE" | "LOW" | "MEDIUM" | "HIGH";
 };
 
-type ContentPart = {
+export type ContentPart = {
     text: string;
 };
 
-type Content = {
+export type Content = {
     parts: ContentPart[];
     role: string;
 };
 
-type Candidate = {
+export type Candidate = {
     content: Content;
     finishReason: string;
     index: number;
     safetyRatings: SafetyRating[];
 };
 
-type UsageMetadata = {
+export type UsageMetadata = {
     promptTokenCount: number;
     candidatesTokenCount: number;
     totalTokenCount: number;
 };
 
-type Response = {
+export type GeminiAIResponse = {
     candidates: Candidate[];
     usageMetadata: UsageMetadata;
 };
 
-type responseMimeType = "text/plain" | "application/json";
+export type ResponseMimeType = "text/plain" | "application/json";
 
-interface GeminiAIChatOptions {
+export interface GeminiAIChatOptions {
     model?: string;
     max_output_tokens?: number;
     temperature?: number;
     prompt: string;
     max_retry?: number;
-    responseType?: responseMimeType;
+    responseType?: ResponseMimeType;
     delay?: number;
 }
 
 export class GeminiAI {
     apiKey: string;
+    baseUrl: string;
     constructor(options: GeminiAIConstructionOptions) {
         this.apiKey = options.apiKey || process.env.GEMINI_API_KEY || "";
+        this.baseUrl =
+            options.baseUrl ||
+            process.env.GEMINI_API_BASE_URL ||
+            "https://generativelanguage.googleapis.com/v1beta/models";
     }
 
-    async chat(chatOptions: GeminiAIChatOptions): Promise<Response> {
-        let data = JSON.stringify({
+    async chat(chatOptions: GeminiAIChatOptions): Promise<GeminiAIResponse> {
+        const model = chatOptions.model || "gemini-3.5-flash-lite";
+        const data = {
             contents: [
                 {
                     role: "user",
@@ -72,26 +78,34 @@ export class GeminiAI {
                     ],
                 },
             ],
-        });
+            generationConfig: {
+                temperature: chatOptions.temperature ?? 0.7,
+                responseMimeType: chatOptions.responseType || "text/plain",
+                maxOutputTokens: chatOptions.max_output_tokens ?? 1024,
+            },
+        };
 
-        let config = {
+        const config = {
             method: "post",
             maxBodyLength: Infinity,
-            url,
+            url: `${this.baseUrl}/${encodeURIComponent(model)}:generateContent`,
             headers: {
                 "Content-Type": "application/json",
                 "x-goog-api-key": this.apiKey,
             },
-            temperature: chatOptions.temperature || "0.7",
-            responseMimeType: chatOptions.responseType || "text/plain",
-            max_output_tokens: chatOptions.max_output_tokens || 1024,
-            data: data,
+            data,
         };
-        return await retry(
+        return retry(
             async () => {
                 return (await axios.request(config)).data;
             },
-            { maxAttempts: chatOptions.max_retry || 3, delay: chatOptions.delay || 200 }
+            {
+                maxAttempts: chatOptions.max_retry ?? 3,
+                delay: chatOptions.delay ?? 200,
+            }
         );
     }
 }
+
+/** PaLM2-compatible name backed by Google's maintained Generative Language API. */
+export class Palm2AI extends GeminiAI {}

JS/edgechains/arakoodev/src/ai/src/testcases/palm2/prompt.jsonnet

@@ -0,0 +1,3 @@
+{
+    "prompt": "Explain vector search"
+}

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

@@ -0,0 +1,48 @@
+import axios from "axios";
+import { readFileSync } from "node:fs";
+import { describe, expect, it, vi } from "vitest";
+import { Palm2AI } from "../../lib/gemini/gemini.js";
+
+vi.mock("axios");
+
+describe("Palm2AI", () => {
+    it("sends typed generation settings in the Google request body", async () => {
+        const { prompt } = JSON.parse(
+            readFileSync(new URL("../../testcases/palm2/prompt.jsonnet", import.meta.url), "utf8")
+        ) as { prompt: string };
+        const response = { candidates: [], usageMetadata: {} };
+        vi.mocked(axios.request).mockResolvedValue({ data: response });
+
+        const palm2 = new Palm2AI({ apiKey: "test-key" });
+        await expect(
+            palm2.chat({
+                model: "gemini-3.5-flash-lite",
+                prompt,
+                temperature: 0,
+                max_output_tokens: 128,
+                responseType: "application/json",
+                max_retry: 1,
+            })
+        ).resolves.toBe(response);
+
+        expect(axios.request).toHaveBeenCalledWith(
+            expect.objectContaining({
+                method: "post",
+                url: expect.stringContaining("/gemini-3.5-flash-lite:generateContent"),
+                data: {
+                    contents: [
+                        {
+                            role: "user",
+                            parts: [{ text: prompt }],
+                        },
+                    ],
+                    generationConfig: {
+                        temperature: 0,
+                        responseMimeType: "application/json",
+                        maxOutputTokens: 128,
+                    },
+                },
+            })
+        );
+    });
+});

JS/edgechains/examples/palm2-chat/README.md

@@ -0,0 +1,11 @@
+# PaLM2-compatible Google AI example
+
+Google retired the original PaLM models, so `Palm2AI` keeps the requested class name while using
+the maintained Generative Language `generateContent` API. The prompt lives in
+`jsonnet/main.jsonnet`, not in TypeScript.
+
+```sh
+npm install
+npm run build
+GEMINI_API_KEY=your-key npm start
+```

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

@@ -0,0 +1,5 @@
+{
+  prompt: |||
+    Explain in two sentences how vector search differs from keyword search.
+  |||,
+}

JS/edgechains/examples/palm2-chat/package.json

@@ -0,0 +1,18 @@
+{
+    "name": "palm2-chat-example",
+    "version": "1.0.0",
+    "private": true,
+    "type": "module",
+    "scripts": {
+        "build": "tsc",
+        "start": "node dist/index.js"
+    },
+    "dependencies": {
+        "@arakoodev/edgechains.js": "file:../../arakoodev",
+        "@arakoodev/jsonnet": "^0.25.0"
+    },
+    "devDependencies": {
+        "@types/node": "^20.17.2",
+        "typescript": "^5.6.3"
+    }
+}

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

@@ -0,0 +1,14 @@
+import { Palm2AI } from "@arakoodev/edgechains.js/ai";
+import Jsonnet from "@arakoodev/jsonnet";
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+
+const jsonnet = new Jsonnet();
+const directory = path.dirname(fileURLToPath(import.meta.url));
+const { prompt } = JSON.parse(
+    jsonnet.evaluateFile(path.join(directory, "../jsonnet/main.jsonnet"))
+) as { prompt: string };
+
+const palm2 = new Palm2AI({ apiKey: process.env.GEMINI_API_KEY });
+const response = await palm2.chat({ prompt });
+console.log(response.candidates[0]?.content.parts[0]?.text ?? "No response");

JS/edgechains/examples/palm2-chat/tsconfig.json

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