Add Palm2 AI provider

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

@@ -1,5 +1,6 @@
 export { OpenAI } from "./lib/openai/openai.js";
 export { GeminiAI } from "./lib/gemini/gemini.js";
+export { Palm2AI } 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,100 @@
+import axios from "axios";
+import { retry } from "@lifeomic/attempt";
+
+const DEFAULT_MODEL = "text-bison-001";
+const DEFAULT_API_VERSION = "v1beta3";
+const GOOGLE_GENERATIVE_LANGUAGE_BASE_URL = "https://generativelanguage.googleapis.com";
+
+export interface Palm2AIConstructionOptions {
+    apiKey?: string;
+    model?: string;
+    apiVersion?: string;
+    baseUrl?: string;
+}
+
+export interface Palm2AIChatOptions {
+    model?: string;
+    prompt: string;
+    temperature?: number;
+    candidate_count?: number;
+    max_output_tokens?: number;
+    top_k?: number;
+    top_p?: number;
+    max_retry?: number;
+    delay?: number;
+}
+
+export type Palm2SafetyRating = {
+    category: string;
+    probability: string;
+};
+
+export type Palm2Candidate = {
+    output: string;
+    safetyRatings?: Palm2SafetyRating[];
+};
+
+export type Palm2Response = {
+    candidates: Palm2Candidate[];
+    filters?: Array<{
+        reason: string;
+        message?: string;
+    }>;
+};
+
+export class Palm2AI {
+    apiKey: string;
+    model: string;
+    apiVersion: string;
+    baseUrl: string;
+
+    constructor(options: Palm2AIConstructionOptions = {}) {
+        this.apiKey = options.apiKey || process.env.PALM2_API_KEY || process.env.GOOGLE_API_KEY || "";
+        this.model = options.model || DEFAULT_MODEL;
+        this.apiVersion = options.apiVersion || DEFAULT_API_VERSION;
+        this.baseUrl = (options.baseUrl || GOOGLE_GENERATIVE_LANGUAGE_BASE_URL).replace(/\/$/, "");
+        this.checkKeys();
+    }
+
+    private checkKeys(): void {
+        if (!this.apiKey) {
+            console.error(
+                "API key is missing. Please provide a valid Google Generative Language API key. You can add it in .env file as PALM2_API_KEY or GOOGLE_API_KEY"
+            );
+        }
+    }
+
+    async chat(chatOptions: Palm2AIChatOptions): Promise<Palm2Response> {
+        const model = chatOptions.model || this.model;
+        const url = `${this.baseUrl}/${this.apiVersion}/models/${model}:generateText?key=${encodeURIComponent(
+            this.apiKey
+        )}`;
+        const data = {
+            prompt: {
+                text: chatOptions.prompt,
+            },
+            temperature: chatOptions.temperature ?? 0.7,
+            candidate_count: chatOptions.candidate_count ?? 1,
+            max_output_tokens: chatOptions.max_output_tokens ?? 1024,
+            top_k: chatOptions.top_k,
+            top_p: chatOptions.top_p,
+        };
+
+        const config = {
+            method: "post",
+            maxBodyLength: Infinity,
+            url,
+            headers: {
+                "Content-Type": "application/json",
+            },
+            data,
+        };
+
+        return await retry(
+            async () => {
+                return (await axios.request(config)).data;
+            },
+            { maxAttempts: chatOptions.max_retry || 3, delay: chatOptions.delay || 200 }
+        );
+    }
+}

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

@@ -0,0 +1,113 @@
+import { createServer, Server } from "http";
+import { AddressInfo } from "net";
+import { Palm2AI } from "../../lib/palm2/palm2";
+
+type CapturedRequest = {
+    method?: string;
+    url?: string;
+    body: any;
+};
+
+describe("Palm2AI", () => {
+    let server: Server;
+    let baseUrl: string;
+    let capturedRequest: CapturedRequest | undefined;
+    const mockResponse = {
+        candidates: [
+            {
+                output: "Test response",
+            },
+        ],
+    };
+
+    beforeEach(async () => {
+        capturedRequest = undefined;
+        server = createServer((req, res) => {
+            let rawBody = "";
+            req.on("data", (chunk) => {
+                rawBody += chunk;
+            });
+            req.on("end", () => {
+                capturedRequest = {
+                    method: req.method,
+                    url: req.url,
+                    body: JSON.parse(rawBody),
+                };
+                res.writeHead(200, {
+                    "Access-Control-Allow-Origin": "*",
+                    "Content-Type": "application/json",
+                });
+                res.end(JSON.stringify(mockResponse));
+            });
+        });
+
+        await new Promise<void>((resolve) => {
+            server.listen(0, "127.0.0.1", resolve);
+        });
+        const { port } = server.address() as AddressInfo;
+        baseUrl = `http://127.0.0.1:${port}`;
+    });
+
+    afterEach(async () => {
+        await new Promise<void>((resolve, reject) => {
+            server.close((error) => {
+                if (error) {
+                    reject(error);
+                    return;
+                }
+                resolve();
+            });
+        });
+    });
+
+    test("should generate text using the default PaLM2 text model", async () => {
+        const palm2 = new Palm2AI({ apiKey: "test_api_key", baseUrl });
+        const response = await palm2.chat({ prompt: "test prompt" });
+
+        expect(capturedRequest).toEqual({
+            method: "POST",
+            url: "/v1beta3/models/text-bison-001:generateText?key=test_api_key",
+            body: expect.objectContaining({
+                prompt: {
+                    text: "test prompt",
+                },
+                temperature: 0.7,
+                candidate_count: 1,
+                max_output_tokens: 1024,
+            }),
+        });
+        expect(response).toEqual(mockResponse);
+    });
+
+    test("should allow overriding generation settings and model", async () => {
+        const palm2 = new Palm2AI({
+            apiKey: "test_api_key",
+            model: "text-bison-001",
+            apiVersion: "v1beta3",
+            baseUrl,
+        });
+        await palm2.chat({
+            prompt: "test prompt",
+            model: "chat-bison-001",
+            temperature: 0.2,
+            candidate_count: 2,
+            max_output_tokens: 256,
+            top_k: 40,
+            top_p: 0.95,
+            max_retry: 1,
+            delay: 1,
+        });
+
+        expect(capturedRequest).toEqual({
+            method: "POST",
+            url: "/v1beta3/models/chat-bison-001:generateText?key=test_api_key",
+            body: expect.objectContaining({
+                temperature: 0.2,
+                candidate_count: 2,
+                max_output_tokens: 256,
+                top_k: 40,
+                top_p: 0.95,
+            }),
+        });
+    });
+});

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

@@ -0,0 +1,11 @@
+{
+  name: 'palm2-basic-generate-text',
+  provider: 'palm2',
+  model: 'text-bison-001',
+  input: {
+    prompt: 'Answer in one sentence: what is EdgeChains?',
+    temperature: 0.2,
+    candidate_count: 1,
+    max_output_tokens: 128,
+  },
+}

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

@@ -0,0 +1,15 @@
+local promptTemplate = |||
+  You are a helpful assistant that can answer questions based on the given question.
+  Answer the following question: {question}
+|||;
+
+local key = std.extVar('palm2_api_key');
+local UserQuestion = std.extVar('question');
+
+local promptWithQuestion = std.strReplace(promptTemplate, '{question}', UserQuestion + '\n');
+
+local main() =
+  local response = arakoo.native('palm2Call')({ prompt: promptWithQuestion, palm2ApiKey: key });
+  response;
+
+main()

JS/edgechains/examples/chat-with-palm2/jsonnet/secrets.jsonnet

@@ -0,0 +1,3 @@
+{
+  palm2_api_key: std.extVar('palm2_api_key'),
+}

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

@@ -0,0 +1,18 @@
+{
+    "name": "chat-with-palm2",
+    "version": "1.0.0",
+    "type": "module",
+    "scripts": {
+        "start": "tsc && node --experimental-wasm-modules ./dist/index.js"
+    },
+    "dependencies": {
+        "@arakoodev/edgechains.js": "file:../../arakoodev",
+        "@arakoodev/jsonnet": "^0.24.0",
+        "file-uri-to-path": "^2.0.0",
+        "path": "^0.12.7",
+        "typescript": "^5.7.0-dev.20241007"
+    },
+    "devDependencies": {
+        "@types/node": "^22.7.4"
+    }
+}

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

@@ -0,0 +1,30 @@
+import { ArakooServer } from "@arakoodev/edgechains.js/arakooserver";
+import Jsonnet from "@arakoodev/jsonnet";
+import { createSyncRPC } from "@arakoodev/edgechains.js/sync-rpc";
+import fileURLToPath from "file-uri-to-path";
+import path from "path";
+
+const server = new ArakooServer();
+const app = server.createApp();
+const jsonnet = new Jsonnet();
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+
+const palm2Call = createSyncRPC(path.join(__dirname, "./lib/generateResponse.cjs"));
+
+app.post("/chat", async (c: any) => {
+    try {
+        const { question } = await c.req.json();
+        const key = JSON.parse(
+            jsonnet.evaluateFile(path.join(__dirname, "../jsonnet/secrets.jsonnet"))
+        ).palm2_api_key;
+        jsonnet.extString("palm2_api_key", key);
+        jsonnet.extString("question", question || "");
+        jsonnet.javascriptCallback("palm2Call", palm2Call);
+        const response = jsonnet.evaluateFile(path.join(__dirname, "../jsonnet/main.jsonnet"));
+        return c.json(JSON.parse(response));
+    } catch (error) {
+        console.log("error occured", error);
+    }
+});
+
+server.listen(3000);

JS/edgechains/examples/chat-with-palm2/src/lib/generateResponse.cts

@@ -0,0 +1,16 @@
+const { Palm2AI } = require("@arakoodev/edgechains.js/ai");
+
+async function palm2Call({ prompt, palm2ApiKey }: any) {
+    try {
+        const palm2 = new Palm2AI({ apiKey: palm2ApiKey });
+        const response = await palm2.chat({ prompt });
+        return JSON.stringify({
+            answer: response.candidates?.[0]?.output ?? "",
+            raw: response,
+        });
+    } catch (error) {
+        return error;
+    }
+}
+
+module.exports = palm2Call;

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

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