From 8e06f43e881d6397a29eb01709bee98d2867534b Mon Sep 17 00:00:00 2001 From: Neel Desai Date: Mon, 11 May 2026 09:16:15 -0400 Subject: [PATCH] feat: finish palm2 and gemini support --- JS/edgechains/arakoodev/src/ai/src/index.ts | 13 ++ .../arakoodev/src/ai/src/lib/gemini/gemini.ts | 108 +++------- .../arakoodev/src/ai/src/lib/palm2/palm2.ts | 190 ++++++++++++++++++ .../src/ai/src/tests/palm2/palm2.test.ts | 171 ++++++++++++++++ JS/edgechains/examples/palm2-chat/README.md | 29 +++ .../examples/palm2-chat/jsonnet/main.jsonnet | 10 + .../palm2-chat/jsonnet/secrets.jsonnet | 5 + .../examples/palm2-chat/package.json | 18 ++ .../examples/palm2-chat/src/index.ts | 23 +++ .../examples/palm2-chat/tsconfig.json | 14 ++ 10 files changed, 505 insertions(+), 76 deletions(-) create mode 100644 JS/edgechains/arakoodev/src/ai/src/lib/palm2/palm2.ts create mode 100644 JS/edgechains/arakoodev/src/ai/src/tests/palm2/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/package.json create mode 100644 JS/edgechains/examples/palm2-chat/src/index.ts create mode 100644 JS/edgechains/examples/palm2-chat/tsconfig.json diff --git a/JS/edgechains/arakoodev/src/ai/src/index.ts b/JS/edgechains/arakoodev/src/ai/src/index.ts index 2c98f37dc..485cfd32c 100644 --- a/JS/edgechains/arakoodev/src/ai/src/index.ts +++ b/JS/edgechains/arakoodev/src/ai/src/index.ts @@ -1,5 +1,18 @@ export { OpenAI } from "./lib/openai/openai.js"; export { GeminiAI } from "./lib/gemini/gemini.js"; +export { Palm2AI } from "./lib/palm2/palm2.js"; +export type { + Palm2Candidate, + Palm2ChatOptions, + Palm2ConstructionOptions, + Palm2Content, + Palm2ContentPart, + Palm2EmbeddingOptions, + Palm2EmbeddingResponse, + Palm2GenerateResponse, + ResponseMimeType, + Palm2TextOptions, +} 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/gemini/gemini.ts b/JS/edgechains/arakoodev/src/ai/src/lib/gemini/gemini.ts index 13f1bfcd1..a3e1d0969 100644 --- a/JS/edgechains/arakoodev/src/ai/src/lib/gemini/gemini.ts +++ b/JS/edgechains/arakoodev/src/ai/src/lib/gemini/gemini.ts @@ -1,97 +1,53 @@ -import axios from "axios"; -import { retry } from "@lifeomic/attempt"; -const url = "https://generativelanguage.googleapis.com/v1/models/gemini-pro:generateContent"; +import { Palm2AI, type Palm2GenerateResponse, type ResponseMimeType } from "../palm2/palm2.js"; interface GeminiAIConstructionOptions { apiKey?: string; + baseUrl?: string; } -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; -}; - -type Content = { - parts: ContentPart[]; - role: string; -}; - -type Candidate = { - content: Content; - finishReason: string; - index: number; - safetyRatings: SafetyRating[]; -}; - -type UsageMetadata = { - promptTokenCount: number; - candidatesTokenCount: number; - totalTokenCount: number; -}; - -type Response = { - candidates: Candidate[]; - usageMetadata: UsageMetadata; -}; - -type responseMimeType = "text/plain" | "application/json"; - interface GeminiAIChatOptions { model?: string; max_output_tokens?: number; temperature?: number; + topP?: number; + topK?: number; + candidateCount?: number; + stopSequences?: string[]; prompt: string; max_retry?: number; - responseType?: responseMimeType; + responseType?: ResponseMimeType; delay?: number; } export class GeminiAI { apiKey: string; - constructor(options: GeminiAIConstructionOptions) { - this.apiKey = options.apiKey || process.env.GEMINI_API_KEY || ""; + baseUrl?: string; + + constructor(options: GeminiAIConstructionOptions = {}) { + this.apiKey = + options.apiKey ?? + process.env.GEMINI_API_KEY ?? + process.env.GOOGLE_API_KEY ?? + process.env.PALM2_API_KEY ?? + ""; + this.baseUrl = options.baseUrl; } - async chat(chatOptions: GeminiAIChatOptions): Promise { - let data = JSON.stringify({ - contents: [ - { - role: "user", - parts: [ - { - text: chatOptions.prompt, - }, - ], - }, - ], + async chat(chatOptions: GeminiAIChatOptions): Promise { + const client = new Palm2AI({ apiKey: this.apiKey, baseUrl: this.baseUrl }); + + return client.chat({ + model: chatOptions.model ?? "gemini-pro", + prompt: chatOptions.prompt, + temperature: chatOptions.temperature, + topP: chatOptions.topP, + topK: chatOptions.topK, + candidateCount: chatOptions.candidateCount, + stopSequences: chatOptions.stopSequences, + responseMimeType: chatOptions.responseType ?? "text/plain", + maxOutputTokens: chatOptions.max_output_tokens ?? 1024, + maxRetries: chatOptions.max_retry ?? 3, + delay: chatOptions.delay ?? 200, }); - - let config = { - method: "post", - maxBodyLength: Infinity, - url, - 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, - }; - return await retry( - async () => { - return (await axios.request(config)).data; - }, - { maxAttempts: chatOptions.max_retry || 3, delay: chatOptions.delay || 200 } - ); } } 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..2777419e6 --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/lib/palm2/palm2.ts @@ -0,0 +1,190 @@ +import axios from "axios"; +import { retry } from "@lifeomic/attempt"; + +const defaultBaseUrl = "https://generativelanguage.googleapis.com/v1beta"; + +export type ResponseMimeType = "text/plain" | "application/json"; + +export interface Palm2ConstructionOptions { + apiKey?: string; + baseUrl?: string; +} + +export interface Palm2ContentPart { + text: string; +} + +export interface Palm2Content { + role?: "user" | "model"; + parts: Palm2ContentPart[]; +} + +interface Palm2GenerationConfig { + temperature?: number; + topP?: number; + topK?: number; + maxOutputTokens?: number; + candidateCount?: number; + stopSequences?: string[]; + responseMimeType?: ResponseMimeType; +} + +export interface Palm2ChatOptions extends Palm2GenerationConfig { + model?: string; + prompt?: string; + contents?: Palm2Content[]; + maxRetries?: number; + delay?: number; +} + +export interface Palm2TextOptions { + model?: string; + prompt: string; + temperature?: number; + topP?: number; + topK?: number; + maxOutputTokens?: number; + candidateCount?: number; + stopSequences?: string[]; + maxRetries?: number; + delay?: number; +} + +export interface Palm2EmbeddingOptions { + model?: string; + text: string; + maxRetries?: number; + delay?: number; +} + +export interface Palm2Candidate { + content?: Palm2Content; + output?: string; + finishReason?: string; + index?: number; +} + +export interface Palm2GenerateResponse { + candidates: Palm2Candidate[]; + usageMetadata?: { + promptTokenCount?: number; + candidatesTokenCount?: number; + totalTokenCount?: number; + }; +} + +export interface Palm2EmbeddingResponse { + embedding: { + value?: number[]; + values?: number[]; + }; +} + +const stripUndefined = >(value: T): Partial => + Object.fromEntries(Object.entries(value).filter(([, entry]) => entry !== undefined)) as Partial; + +export class Palm2AI { + private readonly apiKey: string; + private readonly baseUrl: string; + + constructor(options: Palm2ConstructionOptions = {}) { + this.apiKey = + options.apiKey ?? + process.env.GOOGLE_API_KEY ?? + process.env.GEMINI_API_KEY ?? + process.env.PALM2_API_KEY ?? + ""; + this.baseUrl = (options.baseUrl ?? defaultBaseUrl).replace(/\/$/, ""); + } + + async chat(options: Palm2ChatOptions): Promise { + const model = options.model ?? "gemini-pro"; + const contents = options.contents ?? [ + { + role: "user" as const, + parts: [{ text: options.prompt ?? "" }], + }, + ]; + + return this.request({ + endpoint: `models/${model}:generateContent`, + data: { + contents, + generationConfig: stripUndefined({ + temperature: options.temperature, + topP: options.topP, + topK: options.topK, + maxOutputTokens: options.maxOutputTokens, + candidateCount: options.candidateCount, + stopSequences: options.stopSequences, + responseMimeType: options.responseMimeType, + }), + }, + maxRetries: options.maxRetries, + delay: options.delay, + }); + } + + async generateText(options: Palm2TextOptions): Promise { + const model = options.model ?? "text-bison-001"; + + return this.request({ + endpoint: `models/${model}:generateText`, + data: stripUndefined({ + prompt: { + text: options.prompt, + }, + temperature: options.temperature, + topP: options.topP, + topK: options.topK, + maxOutputTokens: options.maxOutputTokens, + candidateCount: options.candidateCount, + stopSequences: options.stopSequences, + }), + maxRetries: options.maxRetries, + delay: options.delay, + }); + } + + async generateEmbedding(options: Palm2EmbeddingOptions): Promise { + const model = options.model ?? "embedding-gecko-001"; + + return this.request({ + endpoint: `models/${model}:embedText`, + data: { + text: options.text, + }, + maxRetries: options.maxRetries, + delay: options.delay, + }); + } + + private async request({ + endpoint, + data, + maxRetries, + delay, + }: { + endpoint: string; + data: object; + maxRetries?: number; + delay?: number; + }): Promise { + if (!this.apiKey) { + throw new Error("Google Generative Language API key is required"); + } + + return retry( + async () => { + const response = await axios.post(`${this.baseUrl}/${endpoint}`, data, { + headers: { + "Content-Type": "application/json", + "x-goog-api-key": this.apiKey, + }, + }); + return response.data; + }, + { maxAttempts: maxRetries ?? 3, delay: delay ?? 200 } + ); + } +} diff --git a/JS/edgechains/arakoodev/src/ai/src/tests/palm2/palm2.test.ts b/JS/edgechains/arakoodev/src/ai/src/tests/palm2/palm2.test.ts new file mode 100644 index 000000000..f10b71fd6 --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/tests/palm2/palm2.test.ts @@ -0,0 +1,171 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { createServer, type IncomingMessage, type ServerResponse } from "node:http"; +import type { AddressInfo } from "node:net"; +import { GeminiAI } from "../../lib/gemini/gemini.js"; +import { Palm2AI } from "../../lib/palm2/palm2.js"; + +type CapturedRequest = { + method?: string; + url?: string; + headers: IncomingMessage["headers"]; + body: unknown; +}; + +async function createJsonServer( + responseBody: unknown +): Promise<{ baseUrl: string; requests: CapturedRequest[]; close: () => Promise }> { + const requests: CapturedRequest[] = []; + + const server = createServer((req: IncomingMessage, res: ServerResponse) => { + const chunks: Buffer[] = []; + + req.on("data", (chunk) => chunks.push(Buffer.from(chunk))); + req.on("end", () => { + const rawBody = Buffer.concat(chunks).toString("utf8"); + requests.push({ + method: req.method, + url: req.url, + headers: req.headers, + body: rawBody ? JSON.parse(rawBody) : undefined, + }); + + res.statusCode = 200; + res.setHeader("Content-Type", "application/json"); + res.end(JSON.stringify(responseBody)); + }); + }); + + await new Promise((resolve) => server.listen(0, "127.0.0.1", () => resolve())); + + const address = server.address() as AddressInfo; + + return { + baseUrl: `http://127.0.0.1:${address.port}`, + requests, + close: () => + new Promise((resolve, reject) => + server.close((error) => (error ? reject(error) : resolve())) + ), + }; +} + +describe("Palm2AI", () => { + afterEach(() => { + delete process.env.GOOGLE_API_KEY; + delete process.env.GEMINI_API_KEY; + delete process.env.PALM2_API_KEY; + }); + + it("sends generateContent requests with generationConfig", async () => { + const server = await createJsonServer({ + candidates: [{ content: { parts: [{ text: "Hello from Gemini" }] } }], + }); + + try { + const client = new Palm2AI({ apiKey: "test-key", baseUrl: server.baseUrl }); + const response = await client.chat({ + prompt: "Say hello", + temperature: 0.2, + maxOutputTokens: 64, + responseMimeType: "application/json", + }); + + expect(response.candidates[0].content?.parts[0].text).toBe("Hello from Gemini"); + expect(server.requests).toHaveLength(1); + expect(server.requests[0].url).toBe("/models/gemini-pro:generateContent"); + expect(server.requests[0].headers["x-goog-api-key"]).toBe("test-key"); + expect(server.requests[0].body).toEqual({ + contents: [{ role: "user", parts: [{ text: "Say hello" }] }], + generationConfig: { + temperature: 0.2, + maxOutputTokens: 64, + responseMimeType: "application/json", + }, + }); + } finally { + await server.close(); + } + }); + + it("sends legacy generateText requests for palm/text-bison models", async () => { + const server = await createJsonServer({ + candidates: [{ output: "Legacy response" }], + }); + + try { + const client = new Palm2AI({ apiKey: "test-key", baseUrl: server.baseUrl }); + const response = await client.generateText({ + prompt: "Explain EdgeChains", + temperature: 0.5, + topK: 10, + }); + + expect(response.candidates[0].output).toBe("Legacy response"); + expect(server.requests).toHaveLength(1); + expect(server.requests[0].url).toBe("/models/text-bison-001:generateText"); + expect(server.requests[0].body).toEqual({ + prompt: { text: "Explain EdgeChains" }, + temperature: 0.5, + topK: 10, + }); + } finally { + await server.close(); + } + }); + + it("sends embedText requests for embedding models", async () => { + const server = await createJsonServer({ + embedding: { + values: [0.1, 0.2, 0.3], + }, + }); + + try { + const client = new Palm2AI({ apiKey: "test-key", baseUrl: server.baseUrl }); + const response = await client.generateEmbedding({ + text: "Embed EdgeChains", + }); + + expect(response.embedding.values).toEqual([0.1, 0.2, 0.3]); + expect(server.requests).toHaveLength(1); + expect(server.requests[0].url).toBe("/models/embedding-gecko-001:embedText"); + expect(server.requests[0].body).toEqual({ + text: "Embed EdgeChains", + }); + } finally { + await server.close(); + } + }); + + it("maps GeminiAI snake_case options onto the Google request body", async () => { + const server = await createJsonServer({ + candidates: [{ content: { parts: [{ text: "Structured output" }] } }], + }); + + try { + const gemini = new GeminiAI({ apiKey: "test-key", baseUrl: server.baseUrl }); + const response = await gemini.chat({ + model: "gemini-1.5-pro", + prompt: "Return JSON", + max_output_tokens: 32, + responseType: "application/json", + temperature: 0.7, + max_retry: 1, + }); + + expect(response.candidates[0].content?.parts[0].text).toBe("Structured output"); + expect(server.requests).toHaveLength(1); + expect(server.requests[0].url).toBe("/models/gemini-1.5-pro:generateContent"); + expect(server.requests[0].body).toEqual({ + contents: [{ role: "user", parts: [{ text: "Return JSON" }] }], + generationConfig: { + temperature: 0.7, + maxOutputTokens: 32, + responseMimeType: "application/json", + }, + }); + } finally { + await server.close(); + } + }); +}); diff --git a/JS/edgechains/examples/palm2-chat/README.md b/JS/edgechains/examples/palm2-chat/README.md new file mode 100644 index 000000000..e94c928dc --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/README.md @@ -0,0 +1,29 @@ +# Palm2 / Gemini Chat Example + +This example uses a Jsonnet prompt file and the `Palm2AI` client from `@arakoodev/edgechains.js/ai`. + +## Setup + +1. Install dependencies: + + ```bash + npm install + ``` + +2. Add your Google API key in `jsonnet/secrets.jsonnet`: + + ```jsonnet + local GOOGLE_API_KEY = "AIza..."; + ``` + +3. The prompt stays in `jsonnet/main.jsonnet`, so you can change the question without editing TypeScript. + +## Usage + +Run the example: + +```bash +npm run start -- "Explain EdgeChains in one paragraph" +``` + +If you do not pass a question, the example uses a default prompt. 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..faf52895c --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/jsonnet/main.jsonnet @@ -0,0 +1,10 @@ +local promptTemplate = ||| + You are a helpful assistant. Answer the following question clearly: + {question} +|||; + +local question = std.extVar("question"); + +{ + prompt: std.strReplace(promptTemplate, "{question}", question + "\n"), +} 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..2d2e09833 --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/jsonnet/secrets.jsonnet @@ -0,0 +1,5 @@ +local GOOGLE_API_KEY = "AIza***"; + +{ + "google_api_key": GOOGLE_API_KEY, +} diff --git a/JS/edgechains/examples/palm2-chat/package.json b/JS/edgechains/examples/palm2-chat/package.json new file mode 100644 index 000000000..97ee85f12 --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/package.json @@ -0,0 +1,18 @@ +{ + "name": "palm2-chat-example", + "version": "1.0.0", + "description": "Example Google Generative Language chat client for EdgeChains", + "main": "index.js", + "type": "module", + "scripts": { + "start": "tsc && node ./dist/index.js" + }, + "dependencies": { + "@arakoodev/edgechains.js": "file:../../arakoodev", + "@arakoodev/jsonnet": "^0.24.0" + }, + "devDependencies": { + "@types/node": "^22.7.4", + "typescript": "^5.6.3" + } +} diff --git a/JS/edgechains/examples/palm2-chat/src/index.ts b/JS/edgechains/examples/palm2-chat/src/index.ts new file mode 100644 index 000000000..0ec665a77 --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/src/index.ts @@ -0,0 +1,23 @@ +import Jsonnet from "@arakoodev/jsonnet"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; +import { Palm2AI } from "@arakoodev/edgechains.js/ai"; + +const jsonnet = new Jsonnet(); +const __dirname = path.dirname(fileURLToPath(import.meta.url)); +const question = process.argv.slice(2).join(" ") || "Explain EdgeChains in one paragraph."; + +jsonnet.extString("question", question); + +const prompt = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../jsonnet/main.jsonnet")) +).prompt; + +const secrets = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../jsonnet/secrets.jsonnet")) +); + +const client = new Palm2AI({ apiKey: secrets.google_api_key }); +const response = await client.chat({ prompt }); + +console.log(JSON.stringify(response, null, 2)); diff --git a/JS/edgechains/examples/palm2-chat/tsconfig.json b/JS/edgechains/examples/palm2-chat/tsconfig.json new file mode 100644 index 000000000..5721c810b --- /dev/null +++ b/JS/edgechains/examples/palm2-chat/tsconfig.json @@ -0,0 +1,14 @@ +{ + "compilerOptions": { + "target": "ES2022", + "moduleResolution": "NodeNext", + "module": "NodeNext", + "rootDir": "./src", + "outDir": "./dist", + "esModuleInterop": true, + "forceConsistentCasingInFileNames": true, + "strict": true, + "skipLibCheck": true + }, + "exclude": ["./**/*.test.ts", "vitest.config.ts"] +}