feat: add palm2 gemini client 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,97 +1,107 @@
 import axios from "axios";
 import { retry } from "@lifeomic/attempt";
-const url = "https://generativelanguage.googleapis.com/v1/models/gemini-pro:generateContent";
 
-interface GeminiAIConstructionOptions {
-    apiKey?: string;
+const baseUrl = "https://generativelanguage.googleapis.com/v1/models";
+
+export interface GeminiAIConstructionOptions {
+  apiKey?: string;
 }
 
-type SafetyRating = {
-    category:
-        | "HARM_CATEGORY_SEXUALLY_EXPLICIT"
-        | "HARM_CATEGORY_HATE_SPEECH"
-        | "HARM_CATEGORY_HARASSMENT"
-        | "HARM_CATEGORY_DANGEROUS_CONTENT";
-    probability: "NEGLIGIBLE" | "LOW" | "MEDIUM" | "HIGH";
+export type SafetyRating = {
+  category:
+    | "HARM_CATEGORY_SEXUALLY_EXPLICIT"
+    | "HARM_CATEGORY_HATE_SPEECH"
+    | "HARM_CATEGORY_HARASSMENT"
+    | "HARM_CATEGORY_DANGEROUS_CONTENT";
+  probability: "NEGLIGIBLE" | "LOW" | "MEDIUM" | "HIGH";
 };
 
-type ContentPart = {
-    text: string;
+export type ContentPart = {
+  text: string;
 };
 
-type Content = {
-    parts: ContentPart[];
-    role: string;
+export type Content = {
+  parts: ContentPart[];
+  role: string;
 };
 
-type Candidate = {
-    content: Content;
-    finishReason: string;
-    index: number;
-    safetyRatings: SafetyRating[];
+export type Candidate = {
+  content: Content;
+  finishReason: string;
+  index: number;
+  safetyRatings: SafetyRating[];
 };
 
-type UsageMetadata = {
-    promptTokenCount: number;
-    candidatesTokenCount: number;
-    totalTokenCount: number;
+export type UsageMetadata = {
+  promptTokenCount: number;
+  candidatesTokenCount: number;
+  totalTokenCount: number;
 };
 
-type Response = {
-    candidates: Candidate[];
-    usageMetadata: UsageMetadata;
+export type GeminiAIResponse = {
+  candidates: Candidate[];
+  usageMetadata: UsageMetadata;
 };
 
-type responseMimeType = "text/plain" | "application/json";
+export type Response = GeminiAIResponse;
+
+export type ResponseMimeType = "text/plain" | "application/json";
 
-interface GeminiAIChatOptions {
-    model?: string;
-    max_output_tokens?: number;
-    temperature?: number;
-    prompt: string;
-    max_retry?: number;
-    responseType?: responseMimeType;
-    delay?: number;
+export interface GeminiAIChatOptions {
+  model?: string;
+  max_output_tokens?: number;
+  temperature?: number;
+  prompt: string;
+  max_retry?: number;
+  responseType?: ResponseMimeType;
+  delay?: number;
 }
 
 export class GeminiAI {
-    apiKey: string;
-    constructor(options: GeminiAIConstructionOptions) {
-        this.apiKey = options.apiKey || process.env.GEMINI_API_KEY || "";
-    }
-
-    async chat(chatOptions: GeminiAIChatOptions): Promise<Response> {
-        let data = JSON.stringify({
-            contents: [
-                {
-                    role: "user",
-                    parts: [
-                        {
-                            text: chatOptions.prompt,
-                        },
-                    ],
-                },
-            ],
-        });
+  apiKey: string;
+  constructor(options: GeminiAIConstructionOptions) {
+    this.apiKey = options.apiKey || process.env.GEMINI_API_KEY || "";
+  }
 
-        let config = {
-            method: "post",
-            maxBodyLength: Infinity,
-            url,
-            headers: {
-                "Content-Type": "application/json",
-                "x-goog-api-key": this.apiKey,
+  async chat(chatOptions: GeminiAIChatOptions): Promise<GeminiAIResponse> {
+    let data = JSON.stringify({
+      contents: [
+        {
+          role: "user",
+          parts: [
+            {
+              text: chatOptions.prompt,
             },
-            temperature: chatOptions.temperature || "0.7",
-            responseMimeType: chatOptions.responseType || "text/plain",
-            max_output_tokens: chatOptions.max_output_tokens || 1024,
-            data: data,
-        };
-        return await retry(
-            async () => {
-                return (await axios.request(config)).data;
-            },
-            { maxAttempts: chatOptions.max_retry || 3, delay: chatOptions.delay || 200 }
-        );
-    }
+          ],
+        },
+      ],
+      generationConfig: {
+        temperature: chatOptions.temperature ?? 0.7,
+        responseMimeType: chatOptions.responseType || "text/plain",
+        maxOutputTokens: chatOptions.max_output_tokens ?? 1024,
+      },
+    });
+
+    let config = {
+      method: "post",
+      maxBodyLength: Infinity,
+      url: `${baseUrl}/${chatOptions.model || "gemini-pro"}:generateContent`,
+      headers: {
+        "Content-Type": "application/json",
+        "x-goog-api-key": this.apiKey,
+      },
+      data: data,
+    };
+    return await retry(
+      async () => {
+        return (await axios.request(config)).data;
+      },
+      {
+        maxAttempts: chatOptions.max_retry || 3,
+        delay: chatOptions.delay || 200,
+      },
+    );
+  }
 }
+
+export class Palm2AI extends GeminiAI {}

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

@@ -0,0 +1,95 @@
+import axios from "axios";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { GeminiAI, Palm2AI } from "../../lib/gemini/gemini.js";
+
+vi.mock("axios", () => ({
+  default: {
+    request: vi.fn(),
+  },
+}));
+
+const mockedRequest = vi.mocked(axios.request);
+
+describe("Palm2/Gemini client", () => {
+  beforeEach(() => {
+    mockedRequest.mockReset();
+  });
+
+  it("sends prompts and generation settings in the Google request body", async () => {
+    mockedRequest.mockResolvedValue({
+      data: {
+        candidates: [],
+        usageMetadata: {
+          promptTokenCount: 1,
+          candidatesTokenCount: 0,
+          totalTokenCount: 1,
+        },
+      },
+    });
+
+    const gemini = new GeminiAI({ apiKey: "test-key" });
+    await gemini.chat({
+      model: "gemini-1.5-flash",
+      prompt: "Summarize EdgeChains in one sentence.",
+      temperature: 0.2,
+      max_output_tokens: 128,
+      responseType: "application/json",
+      max_retry: 1,
+    });
+
+    const config = mockedRequest.mock.calls[0][0] as any;
+    expect(config.url).toBe(
+      "https://generativelanguage.googleapis.com/v1/models/gemini-1.5-flash:generateContent",
+    );
+    expect(config.headers).toMatchObject({
+      "Content-Type": "application/json",
+      "x-goog-api-key": "test-key",
+    });
+    expect(JSON.parse(config.data)).toEqual({
+      contents: [
+        {
+          role: "user",
+          parts: [
+            {
+              text: "Summarize EdgeChains in one sentence.",
+            },
+          ],
+        },
+      ],
+      generationConfig: {
+        temperature: 0.2,
+        responseMimeType: "application/json",
+        maxOutputTokens: 128,
+      },
+    });
+  });
+
+  it("keeps a Palm2AI export compatible with the Gemini implementation", async () => {
+    mockedRequest.mockResolvedValue({
+      data: {
+        candidates: [
+          {
+            content: { role: "model", parts: [{ text: "Hello" }] },
+            finishReason: "STOP",
+            index: 0,
+            safetyRatings: [],
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 1,
+          candidatesTokenCount: 1,
+          totalTokenCount: 2,
+        },
+      },
+    });
+
+    const palm2 = new Palm2AI({ apiKey: "test-key" });
+    const response = await palm2.chat({
+      prompt: "Say hello",
+      max_retry: 1,
+    });
+
+    expect(response.candidates[0].content.parts[0].text).toBe("Hello");
+    expect(mockedRequest).toHaveBeenCalledOnce();
+  });
+});

JS/edgechains/arakoodev/testcases/palm2/main.jsonnet

@@ -0,0 +1,13 @@
+{
+  prompt: |||
+    You are a helpful assistant.
+    Answer the user's question in one short paragraph.
+
+    Question: {question}
+  |||,
+  question: "What is EdgeChains?",
+  model: "gemini-pro",
+  temperature: 0.2,
+  max_output_tokens: 256,
+  responseType: "text/plain",
+}

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

@@ -0,0 +1,11 @@
+# Palm2 Chat Example
+
+This example loads the prompt and model settings from Jsonnet, then calls the
+EdgeChains `Palm2AI` export, which uses Google's Generative Language API.
+
+```bash
+export GEMINI_API_KEY="your-api-key"
+npm install
+npm run build
+npm start -- "How can I use EdgeChains?"
+```

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

@@ -0,0 +1,16 @@
+local promptTemplate = |||
+  You are a helpful assistant.
+  Answer the user's question in one short paragraph.
+
+  Question: {question}
+|||;
+
+local question = std.extVar('question');
+
+{
+  model: 'gemini-pro',
+  prompt: std.strReplace(promptTemplate, '{question}', question),
+  temperature: 0.2,
+  max_output_tokens: 256,
+  responseType: 'text/plain',
+}

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

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

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

@@ -0,0 +1,20 @@
+import { Palm2AI } from "@arakoodev/edgechains.js/ai";
+import Jsonnet from "@arakoodev/jsonnet";
+import fileURLToPath from "file-uri-to-path";
+import path from "path";
+
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+const jsonnet = new Jsonnet();
+
+jsonnet.extString(
+  "question",
+  process.argv.slice(2).join(" ") || "What is EdgeChains?",
+);
+
+const chatOptions = JSON.parse(
+  jsonnet.evaluateFile(path.join(__dirname, "../jsonnet/main.jsonnet")),
+);
+const palm2 = new Palm2AI({ apiKey: process.env.GEMINI_API_KEY });
+
+const response = await palm2.chat(chatOptions);
+console.log(JSON.stringify(response, null, 2));

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

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