feat: add Google PaLM 2 API support

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,164 @@
+import axios from "axios";
+import { retry } from "@lifeomic/attempt";
+
+export type Palm2Model =
+    | "text-bison-001"
+    | "text-bison-002"
+    | "text-bison"
+    | "text-unicorn-001"
+    | "chat-bison-001"
+    | "chat-bison-002"
+    | "chat-bison";
+
+type Palm2SafetyCategory =
+    | "HARM_CATEGORY_UNSPECIFIED"
+    | "HARM_CATEGORY_DEROGATORY"
+    | "HARM_CATEGORY_TOXICITY"
+    | "HARM_CATEGORY_VIOLENCE"
+    | "HARM_CATEGORY_SEXUAL"
+    | "HARM_CATEGORY_MEDICAL"
+    | "HARM_CATEGORY_DANGEROUS";
+
+type Palm2SafetyProbability = "NEGLIGIBLE" | "LOW" | "MEDIUM" | "HIGH";
+
+interface Palm2SafetySetting {
+    category: Palm2SafetyCategory;
+    threshold: Palm2SafetyProbability;
+}
+
+interface Palm2Prompt {
+    text: string;
+}
+
+interface Palm2SafetyRating {
+    category: Palm2SafetyCategory;
+    probability: Palm2SafetyProbability;
+    blocked?: boolean;
+}
+
+interface Palm2Candidate {
+    output: string;
+    safetyRatings: Palm2SafetyRating[];
+}
+
+interface Palm2Filter {
+    reason: string;
+    message: string;
+}
+
+interface Palm2SafetyFeedback {
+    rating: Palm2SafetyRating;
+    setting: Palm2SafetySetting;
+}
+
+interface Palm2GenerateResponse {
+    candidates: Palm2Candidate[];
+    filters: Palm2Filter[];
+    safetyFeedback: Palm2SafetyFeedback[];
+}
+
+interface Palm2ConstructionOptions {
+    apiKey?: string;
+}
+
+interface Palm2ChatOptions {
+    model?: Palm2Model;
+    prompt: string;
+    temperature?: number;
+    max_output_tokens?: number;
+    top_p?: number;
+    top_k?: number;
+    candidate_count?: number;
+    safety_settings?: Palm2SafetySetting[];
+    max_retry?: number;
+    delay?: number;
+}
+
+interface Palm2ChatReturnOptions {
+    content: string;
+}
+
+export class Palm2AI {
+    apiKey: string;
+
+    constructor(options: Palm2ConstructionOptions) {
+        this.apiKey = options.apiKey || process.env.PALM2_API_KEY || "";
+        this.checkKey();
+    }
+
+    private checkKey(): void {
+        if (!this.apiKey) {
+            console.error(
+                "API key is missing. Please provide a valid Google PaLM2 API key. You can add it in .env file as PALM2_API_KEY"
+            );
+        }
+    }
+
+    private getUrl(model: string): string {
+        return `https://generativelanguage.googleapis.com/v1beta2/models/${model}:generateText`;
+    }
+
+    async generateText(chatOptions: Palm2ChatOptions): Promise<Palm2ChatReturnOptions> {
+        const url = this.getUrl(chatOptions.model || "text-bison-001");
+
+        const requestBody = {
+            prompt: {
+                text: chatOptions.prompt,
+            },
+            temperature: chatOptions.temperature ?? 0.7,
+            max_output_tokens: chatOptions.max_output_tokens || 1024,
+            top_p: chatOptions.top_p ?? 0.95,
+            top_k: chatOptions.top_k ?? 40,
+            candidate_count: chatOptions.candidate_count || 1,
+            safety_settings: chatOptions.safety_settings || [],
+        };
+
+        const response = await retry(
+            async () => {
+                return (await axios.post(url, requestBody, {
+                    headers: {
+                        "Content-Type": "application/json",
+                    },
+                    params: {
+                        key: this.apiKey,
+                    },
+                })).data as Palm2GenerateResponse;
+            },
+            { maxAttempts: chatOptions.max_retry || 3, delay: chatOptions.delay || 200 }
+        );
+
+        if (!response.candidates || response.candidates.length === 0) {
+            throw new Error("No candidates returned from Palm2 API");
+        }
+
+        return { content: response.candidates[0].output };
+    }
+
+    async chat(chatOptions: Palm2ChatOptions): Promise<Palm2ChatReturnOptions> {
+        return this.generateText(chatOptions);
+    }
+
+    async generateEmbeddings({ input, model }: { input: string[]; model: string }): Promise<any> {
+        const url = `https://generativelanguage.googleapis.com/v1beta2/models/${model}:embedText`;
+
+        const response = await retry(
+            async () => {
+                return (await axios.post(
+                    url,
+                    { text: input[0] },
+                    {
+                        headers: {
+                            "Content-Type": "application/json",
+                        },
+                        params: {
+                            key: this.apiKey,
+                        },
+                    }
+                )).data;
+            },
+            { maxAttempts: 3, delay: 200 }
+        );
+
+        return response;
+    }
+}

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

@@ -0,0 +1,157 @@
+import axios from "axios";
+import { Palm2AI } from "../../../../dist/palm2/src/lib/endpoints/Palm2Endpoint.js";
+
+jest.mock("axios");
+
+describe("Palm2AI", () => {
+    describe("generateText", () => {
+        test("should generate text from Palm2 API", async () => {
+            const mockResponse = {
+                candidates: [
+                    {
+                        output: "Test Palm2 response",
+                        safetyRatings: [
+                            {
+                                category: "HARM_CATEGORY_UNSPECIFIED",
+                                probability: "NEGLIGIBLE",
+                            },
+                        ],
+                    },
+                ],
+                filters: [],
+                safetyFeedback: [],
+            };
+
+            axios.post = jest.fn().mockResolvedValueOnce({ data: mockResponse });
+            const palm2 = new Palm2AI({ apiKey: "test_api_key" });
+            const response = await palm2.generateText({ prompt: "test prompt" });
+            expect(response.content).toEqual("Test Palm2 response");
+        });
+
+        test("should use default model when not specified", async () => {
+            const mockResponse = {
+                candidates: [
+                    {
+                        output: "Default model response",
+                        safetyRatings: [],
+                    },
+                ],
+                filters: [],
+                safetyFeedback: [],
+            };
+
+            axios.post = jest.fn().mockResolvedValueOnce({ data: mockResponse });
+            const palm2 = new Palm2AI({ apiKey: "test_api_key" });
+            await palm2.generateText({ prompt: "test prompt" });
+
+            expect(axios.post).toHaveBeenCalledWith(
+                expect.stringContaining("text-bison-001"),
+                expect.any(Object),
+                expect.any(Object)
+            );
+        });
+
+        test("should pass all generation parameters", async () => {
+            const mockResponse = {
+                candidates: [{ output: "Param response", safetyRatings: [] }],
+                filters: [],
+                safetyFeedback: [],
+            };
+
+            axios.post = jest.fn().mockResolvedValueOnce({ data: mockResponse });
+            const palm2 = new Palm2AI({ apiKey: "test_api_key" });
+            await palm2.generateText({
+                prompt: "test prompt",
+                model: "text-bison-002",
+                temperature: 0.5,
+                max_output_tokens: 512,
+                top_p: 0.9,
+                top_k: 20,
+                candidate_count: 2,
+            });
+
+            expect(axios.post).toHaveBeenCalledWith(
+                expect.any(String),
+                expect.objectContaining({
+                    prompt: { text: "test prompt" },
+                    temperature: 0.5,
+                    max_output_tokens: 512,
+                    top_p: 0.9,
+                    top_k: 20,
+                    candidate_count: 2,
+                }),
+                expect.any(Object)
+            );
+        });
+
+        test("should throw error when no candidates returned", async () => {
+            axios.post = jest.fn().mockResolvedValueOnce({
+                data: { candidates: [], filters: [], safetyFeedback: [] },
+            });
+            const palm2 = new Palm2AI({ apiKey: "test_api_key" });
+
+            await expect(palm2.generateText({ prompt: "test prompt" })).rejects.toThrow(
+                "No candidates returned from Palm2 API"
+            );
+        });
+
+        test("should retry on failure", async () => {
+            axios.post = jest
+                .fn()
+                .mockRejectedValueOnce(new Error("Network error"))
+                .mockResolvedValueOnce({
+                    data: {
+                        candidates: [
+                            {
+                                output: "Retry success",
+                                safetyRatings: [],
+                            },
+                        ],
+                        filters: [],
+                        safetyFeedback: [],
+                    },
+                });
+
+            const palm2 = new Palm2AI({ apiKey: "test_api_key" });
+            const response = await palm2.generateText({
+                prompt: "test prompt",
+                max_retry: 2,
+            });
+            expect(response.content).toEqual("Retry success");
+            expect(axios.post).toHaveBeenCalledTimes(2);
+        });
+    });
+
+    describe("chat", () => {
+        test("should alias generateText", async () => {
+            const mockResponse = {
+                candidates: [{ output: "Chat response", safetyRatings: [] }],
+                filters: [],
+                safetyFeedback: [],
+            };
+
+            axios.post = jest.fn().mockResolvedValueOnce({ data: mockResponse });
+            const palm2 = new Palm2AI({ apiKey: "test_api_key" });
+            const response = await palm2.chat({ prompt: "hello" });
+            expect(response.content).toEqual("Chat response");
+        });
+    });
+
+    describe("generateEmbeddings", () => {
+        test("should generate embeddings from Palm2 API", async () => {
+            const mockResponse = {
+                embedding: {
+                    value: [0.1, 0.2, 0.3],
+                },
+            };
+
+            axios.post = jest.fn().mockResolvedValueOnce({ data: mockResponse });
+            const palm2 = new Palm2AI({ apiKey: "test_api_key" });
+            const res = await palm2.generateEmbeddings({
+                input: ["test text"],
+                model: "embedding-gecko-001",
+            });
+            expect(res.embedding.value).toEqual([0.1, 0.2, 0.3]);
+        });
+    });
+});

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

@@ -21,4 +21,13 @@ export type ChatModel =
     | "gpt-3.5-turbo-0125"
     | "gpt-3.5-turbo-16k-0613";
 
+export type Palm2Model =
+    | "text-bison-001"
+    | "text-bison-002"
+    | "text-bison"
+    | "text-unicorn-001"
+    | "chat-bison-001"
+    | "chat-bison-002"
+    | "chat-bison";
+
 export type role = "user" | "assistant" | "system";