From e235e6818ef012017ee3e945b8fc7a0f5fbb3034 Mon Sep 17 00:00:00 2001 From: EdgeChains Bot Date: Tue, 23 Jun 2026 04:32:02 +0800 Subject: [PATCH] feat: add Google PaLM 2 API support - Add Palm2AI class with generateText, chat, and generateEmbeddings methods - Add Palm2Model types to types/index.ts - Add unit tests in tests/palm2.test.ts - Add jsonnet example in examples/palm2-chat/ - Export Palm2AI from index.ts Closes #279 --- JS/edgechains/arakoodev/src/ai/src/index.ts | 1 + .../arakoodev/src/ai/src/lib/palm2/palm2.ts | 164 ++++++++++++++++++ .../arakoodev/src/ai/src/tests/palm2.test.ts | 157 +++++++++++++++++ .../arakoodev/src/ai/src/types/index.ts | 9 + JS/edgechains/examples/palm2-chat/README.md | 61 +++++++ .../examples/palm2-chat/jsonnet/main.jsonnet | 17 ++ .../palm2-chat/jsonnet/secrets.jsonnet | 6 + .../examples/palm2-chat/src/lib/palm2Call.cts | 13 ++ 8 files changed, 428 insertions(+) create mode 100644 JS/edgechains/arakoodev/src/ai/src/lib/palm2/palm2.ts create mode 100644 JS/edgechains/arakoodev/src/ai/src/tests/palm2.test.ts create mode 100644 JS/edgechains/examples/palm2-chat/README.md create mode 100644 JS/edgechains/examples/palm2-chat/jsonnet/main.jsonnet create mode 100644 JS/edgechains/examples/palm2-chat/jsonnet/secrets.jsonnet create mode 100644 JS/edgechains/examples/palm2-chat/src/lib/palm2Call.cts diff --git a/JS/edgechains/arakoodev/src/ai/src/index.ts b/JS/edgechains/arakoodev/src/ai/src/index.ts index 2c98f37dc..d2ccf0eb0 100644 --- a/JS/edgechains/arakoodev/src/ai/src/index.ts +++ b/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"; diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/palm2/palm2.ts b/JS/edgechains/arakoodev/src/ai/src/lib/palm2/palm2.ts new file mode 100644 index 000000000..a242e817c --- /dev/null +++ b/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 { + 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 { + return this.generateText(chatOptions); + } + + async generateEmbeddings({ input, model }: { input: string[]; model: string }): Promise { + 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; + } +} diff --git a/JS/edgechains/arakoodev/src/ai/src/tests/palm2.test.ts b/JS/edgechains/arakoodev/src/ai/src/tests/palm2.test.ts new file mode 100644 index 000000000..6e29b823b --- /dev/null +++ b/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]); + }); + }); +}); diff --git a/JS/edgechains/arakoodev/src/ai/src/types/index.ts b/JS/edgechains/arakoodev/src/ai/src/types/index.ts index 33485f21a..0cdb6c055 100644 --- a/JS/edgechains/arakoodev/src/ai/src/types/index.ts +++ b/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"; diff --git a/JS/edgechains/examples/palm2-chat/README.md b/JS/edgechains/examples/palm2-chat/README.md new file mode 100644 index 000000000..1d116b219 --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/README.md @@ -0,0 +1,61 @@ +# PaLM 2 Chat Example + +This example demonstrates how to use the Google PaLM 2 API with EdgeChains to build a simple question-answering application. + +## Prerequisites + +1. Obtain a Google AI Studio API key from [https://makersuite.google.com/app/apikey](https://makersuite.google.com/app/apikey) +2. Set your API key in `jsonnet/secrets.jsonnet` + +## Usage + +```typescript +import * as path from "path"; +import { createSyncRPC } from "@arakoodev/edgechains.js"; + +const palm2Call = createSyncRPC(path.join(__dirname, "./lib/palm2Call.cjs")); + +// Register the callback +jsonnet.javascriptCallback("palm2Call", palm2Call); + +// Evaluate the jsonnet file +const result = jsonnet.evaluateFile("jsonnet/main.jsonnet"); +``` + +## Configuration + +### Jsonnet Variables + +| Variable | Description | +|----------|-------------| +| `palm2_api_key` | Your Google PaLM 2 API key | +| `question` | The user question to answer | + +### Model Options + +The `Palm2AI` class supports the following models: +- `text-bison-001` (default) +- `text-bison-002` +- `text-bison` +- `text-unicorn-001` +- `chat-bison-001` +- `chat-bison-002` +- `chat-bison` + +### Generation Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `temperature` | number | 0.7 | Controls randomness | +| `max_output_tokens` | number | 1024 | Maximum tokens to generate | +| `top_p` | number | 0.95 | Nucleus sampling parameter | +| `top_k` | number | 40 | Top-k sampling parameter | +| `candidate_count` | number | 1 | Number of candidates to generate | + +## API Response Format + +```json +{ + "content": "The generated response text" +} +``` diff --git a/JS/edgechains/examples/palm2-chat/jsonnet/main.jsonnet b/JS/edgechains/examples/palm2-chat/jsonnet/main.jsonnet new file mode 100644 index 000000000..a757b4d51 --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/jsonnet/main.jsonnet @@ -0,0 +1,17 @@ +local promptTemplate = ||| + You are a helpful assistant powered by Google PaLM 2. + Answer the following question concisely and accurately. + + 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() diff --git a/JS/edgechains/examples/palm2-chat/jsonnet/secrets.jsonnet b/JS/edgechains/examples/palm2-chat/jsonnet/secrets.jsonnet new file mode 100644 index 000000000..842a3b56e --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/jsonnet/secrets.jsonnet @@ -0,0 +1,6 @@ + +local PALM2_API_KEY = "your-palm2-api-key"; + +{ + "palm2_api_key": PALM2_API_KEY, +} diff --git a/JS/edgechains/examples/palm2-chat/src/lib/palm2Call.cts b/JS/edgechains/examples/palm2-chat/src/lib/palm2Call.cts new file mode 100644 index 000000000..2b38e8335 --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/src/lib/palm2Call.cts @@ -0,0 +1,13 @@ +const { Palm2AI } = require("@arakoodev/edgechains.js/ai"); + +async function palm2Call({ prompt, palm2ApiKey }: any) { + try { + const palm2 = new Palm2AI({ apiKey: palm2ApiKey }); + let res = await palm2.chat({ prompt: prompt }); + return JSON.stringify(res); + } catch (error) { + return error; + } +} + +module.exports = palm2Call;