diff --git a/JS/edgechains/arakoodev/src/ai/src/index.ts b/JS/edgechains/arakoodev/src/ai/src/index.ts index 2c98f37dc..55ef7dda1 100644 --- a/JS/edgechains/arakoodev/src/ai/src/index.ts +++ b/JS/edgechains/arakoodev/src/ai/src/index.ts @@ -1,5 +1,43 @@ -export { OpenAI } from "./lib/openai/openai.js"; export { GeminiAI } 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"; +export { + Router, + NoMatchingDeploymentError, + NoDeploymentsAvailableError, + AxiosHttpClient, + HttpError, + sentryCallback, + posthogCallback, + openAIAdapter, + palmAdapter, + cohereAdapter, + getAdapter, + estimateTokens, +} from "./lib/router/index.js"; +export type { + Deployment, + RouterProvider, + RouterRole, + RouterMessage, + RoutingStrategy, + TokenUsage, + CompletionRequest, + CompletionResult, + FunctionCall, + EmbeddingRequest, + EmbeddingResult, + DeploymentUsage, + RouterEvent, + RouterCallback, + RouterOptions, + HttpClient, + HttpRequestConfig, + HttpResponse, + ProviderAdapter, + NormalizedResponse, + NormalizedEmbedding, + SentryLike, + PostHogLike, +} from "./lib/router/index.js"; diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/openai/openai.ts b/JS/edgechains/arakoodev/src/ai/src/lib/openai/openai.ts deleted file mode 100644 index f60ba522f..000000000 --- a/JS/edgechains/arakoodev/src/ai/src/lib/openai/openai.ts +++ /dev/null @@ -1,306 +0,0 @@ -import axios from "axios"; -import { zodToJsonSchema } from "zod-to-json-schema"; -import { z } from "zod"; -import { ChatModel, role } from "../../types/index"; -const openAI_url = "https://api.openai.com/v1/chat/completions"; - -interface OpenAIConstructionOptions { - apiKey?: string; - orgId?: string; -} - -interface messageOption { - role: role; - content: string; - name?: string; -} - -interface OpenAIChatOptions { - model?: ChatModel; - role?: role; - max_tokens?: number; - temperature?: number; - prompt?: string; - messages?: messageOption[]; - frequency_penalty?: number; -} - -interface chatWithFunctionOptions { - model?: ChatModel; - role?: role; - max_tokens?: number; - temperature?: number; - prompt?: string; - functions?: object | Array; - messages?: messageOption[]; - function_call?: string; -} - -interface ZodSchemaResponseOptions { - model?: ChatModel; - role?: role; - max_tokens?: number; - temperature?: number; - prompt: string; - schema: S; -} - -interface chatWithFunctionReturnOptions { - content: string; - function_call: { - name: string; - arguments: string; - }; -} - -interface OpenAIChatReturnOptions { - content: string; -} - -export class OpenAI { - apiKey: string; - orgId: string; - constructor(options: OpenAIConstructionOptions) { - this.apiKey = options.apiKey || process.env.OPENAI_API_KEY || ""; - this.orgId = options.orgId || process.env.OPENAI_ORG_ID || ""; - this.checkKeys(); - } - - private checkKeys(): void { - if (!this.apiKey) { - console.error( - "API key is missing. Please provide a valid OpenAI API key. You can add it in .env file as OPENAI_API_KEY" - ); - } - if (!this.orgId) { - console.warn( - "Organization ID is missing. Please provide a valid OpenAI Organization ID. You can add it in .env file as OPENAI_ORG_ID" - ); - } - } - - async chat(chatOptions: OpenAIChatOptions): Promise { - const response = await axios - .post( - openAI_url, - { - model: chatOptions.model || "gpt-3.5-turbo", - messages: chatOptions.prompt - ? [ - { - role: chatOptions.role || "user", - content: chatOptions.prompt, - }, - ] - : chatOptions.messages, - max_tokens: chatOptions.max_tokens || 256, - temperature: chatOptions.temperature || 0.7, - frequency_penalty: 1, - }, - { - headers: { - Authorization: "Bearer " + this.apiKey, - "content-type": "application/json", - "OpenAI-Organization": this.orgId, - }, - } - ) - .then((response) => { - return response.data.choices; - }) - .catch((error) => { - if (error.response) { - console.log("Server responded with status code:", error.response.status); - console.log("Response data:", error.response.data); - } else if (error.request) { - console.log("No response received:", error); - } else { - console.log("Error creating request:", error.message); - } - }); - return response[0].message; - } - - async streamedChat(chatOptions: OpenAIChatOptions): Promise { - const response = await axios - .post( - openAI_url, - { - model: chatOptions.model || "gpt-3.5-turbo", - messages: chatOptions.prompt - ? [ - { - role: chatOptions.role || "user", - content: chatOptions.prompt, - }, - ] - : chatOptions.messages, - max_tokens: chatOptions.max_tokens || 256, - temperature: chatOptions.temperature || 0.7, - frequency_penalty: chatOptions.frequency_penalty || 1, - stream: true, - }, - { - headers: { - Authorization: "Bearer " + this.apiKey, - "content-type": "application/json", - "OpenAI-Organization": this.orgId, - }, - } - ) - .then((response) => { - return response.data.choices; - }) - .catch((error) => { - if (error.response) { - console.log("Server responded with status code:", error.response.status); - console.log("Response data:", error.response.data); - } else if (error.request) { - console.log("No response received:", error); - } else { - console.log("Error creating request:", error.message); - } - }); - return response[0].message; - } - - async chatWithFunction( - chatOptions: chatWithFunctionOptions - ): Promise { - const response = await axios - .post( - openAI_url, - { - model: chatOptions.model || "gpt-3.5-turbo", - messages: chatOptions.prompt - ? [ - { - role: chatOptions.role || "user", - content: chatOptions.prompt, - }, - ] - : chatOptions.messages, - max_tokens: chatOptions.max_tokens || 1024, - temperature: chatOptions.temperature || 0.7, - functions: chatOptions.functions, - function_call: chatOptions.function_call || "auto", - }, - { - headers: { - Authorization: "Bearer " + this.apiKey, - "content-type": "application/json", - "OpenAI-Organization": this.orgId, - }, - } - ) - .then((response) => { - return response.data.choices; - }) - .catch((error) => { - if (error.response) { - console.log("Server responded with status code:", error.response.status); - console.log("Response data:", error.response.data); - } else if (error.request) { - console.log("No response received:", error); - } else { - console.log("Error creating request:", error.message); - } - }); - return response[0].message; - } - - async generateEmbeddings({ input, model }: { input: string[]; model: string }): Promise { - const response = await axios - .post( - "https://api.openai.com/v1/embeddings", - { - model: model, - input, - }, - { - headers: { - Authorization: `Bearer ${this.apiKey}`, - "content-type": "application/json", - "OpenAI-Organization": this.orgId, - }, - } - ) - .then((response) => { - return response.data.data; - }) - .catch((error) => { - if (error.response) { - console.log("Server responded with status code:", error.response.status); - console.log("Response data:", error.response.data); - } else if (error.request) { - console.log("No response received:", error.request); - } else { - console.log("Error creating request:", error.message); - } - }); - return response; - } - - async zodSchemaResponse( - chatOptions: ZodSchemaResponseOptions - ): Promise { - const jsonSchema = zodToJsonSchema(chatOptions.schema, { $refStrategy: "none" }); - const openAIFunctionCallDefinition = { - name: "generateSchema", - description: "Generate a schema based on provided details.", - parameters: jsonSchema, - }; - // Remembrer if any field like url or link is not available please create a dummy link based on the following prompt - const content = ` - You are a Schema generator that can generate answer based on given prompt and then return the response based on the give schema - Remembrer if any field like url or link is not available please create a dummy link based on the following prompt - - prompt: - ${chatOptions.prompt || ""} - `; - - const response = await axios - .post( - openAI_url, - { - model: chatOptions.model || "gpt-3.5-turbo-16k", - messages: [ - { - role: chatOptions.role || "user", - content, - }, - ], - functions: [openAIFunctionCallDefinition], - function_call: "auto", - max_tokens: chatOptions.max_tokens || 1000, - temperature: chatOptions.temperature || 0.7, - }, - { - headers: { - Authorization: "Bearer " + this.apiKey, - "content-type": "application/json", - "OpenAI-Organization": this.orgId, - }, - } - ) - .then((response) => { - return response.data.choices[0].message; - }) - .catch((error) => { - if (error.response) { - console.log("Server responded with status code:", error.response.status); - console.log("Response data:", error.response.data); - } else if (error.request) { - console.log("No response received:", error); - } else { - console.log("Error creating request:", error.message); - } - }); - if (response) { - if (response.content) return response.content; - return chatOptions.schema.parse(JSON.parse(response.function_call.arguments)); - } else { - throw new Error("Response did not contain valid JSON."); - } - } -} diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/router/callbacks.ts b/JS/edgechains/arakoodev/src/ai/src/lib/router/callbacks.ts new file mode 100644 index 000000000..ddcb4e6bc --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/lib/router/callbacks.ts @@ -0,0 +1,75 @@ +// Observability callbacks (litellm parity: sentry + posthog). Each factory takes +// an already-constructed client so neither SDK becomes a dependency of this +// package — the caller injects `@sentry/node` / `posthog-node` (or any compatible +// stub) they already use. + +import { RouterCallback, RouterEvent } from "./types.js"; + +export interface SentryLike { + captureException(error: unknown): void; + captureMessage?(message: string, level?: string): void; + addBreadcrumb?(breadcrumb: Record): void; +} + +export interface PostHogLike { + capture(payload: { + distinctId: string; + event: string; + properties?: Record; + }): void; +} + +/** Reports completion failures to Sentry and leaves a breadcrumb on success. */ +export function sentryCallback(client: SentryLike): RouterCallback { + return { + onSuccess(event: RouterEvent) { + client.addBreadcrumb?.({ + category: "edgechains.router", + message: `completion ${event.provider}/${event.model}`, + level: "info", + data: { ...event.usage, durationMs: event.durationMs }, + }); + }, + onError(event: RouterEvent) { + client.captureException(event.error); + }, + }; +} + +/** Emits a PostHog event per completion success/failure. */ +export function posthogCallback( + client: PostHogLike, + options: { distinctId?: string } = {} +): RouterCallback { + const distinctId = options.distinctId ?? "edgechains"; + return { + onSuccess(event: RouterEvent) { + client.capture({ + distinctId, + event: "edgechains_completion", + properties: { + provider: event.provider, + model: event.model, + deploymentIndex: event.deploymentIndex, + durationMs: event.durationMs, + promptTokens: event.usage?.promptTokens, + completionTokens: event.usage?.completionTokens, + totalTokens: event.usage?.totalTokens, + }, + }); + }, + onError(event: RouterEvent) { + client.capture({ + distinctId, + event: "edgechains_completion_error", + properties: { + provider: event.provider, + model: event.model, + deploymentIndex: event.deploymentIndex, + durationMs: event.durationMs, + error: event.error instanceof Error ? event.error.message : String(event.error), + }, + }); + }, + }; +} diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/router/httpClient.ts b/JS/edgechains/arakoodev/src/ai/src/lib/router/httpClient.ts new file mode 100644 index 000000000..0ea882785 --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/lib/router/httpClient.ts @@ -0,0 +1,95 @@ +// Default axios-backed HttpClient. Reliability (timeouts + retries on transient +// failures) is implemented with `axios.interceptors.response.use`, exactly as +// the issue suggests. The Router talks to the `HttpClient` interface, so tests +// inject a fake and never hit the network. + +import axios, { AxiosInstance } from "axios"; +import { HttpClient, HttpRequestConfig, HttpResponse } from "./types.js"; + +export class HttpError extends Error { + status?: number; + constructor(message: string, status?: number) { + super(message); + this.name = "HttpError"; + this.status = status; + } +} + +const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +function isTransient(error: any): boolean { + if (!error?.response) return true; // network error / timeout + const status = error.response.status; + return status === 429 || (status >= 500 && status <= 599); +} + +export interface AxiosHttpClientOptions { + numRetries?: number; + retryDelay?: number; + timeout?: number; +} + +export class AxiosHttpClient implements HttpClient { + private instance: AxiosInstance; + + constructor(options: AxiosHttpClientOptions = {}) { + const numRetries = options.numRetries ?? 2; + const retryDelay = options.retryDelay ?? 200; + this.instance = axios.create({ timeout: options.timeout ?? 600000 }); + + // Transient-failure retry with exponential backoff. + this.instance.interceptors.response.use( + (response) => response, + async (error) => { + const config: any = error?.config; + if (!config || !isTransient(error)) return Promise.reject(error); + config.__retryCount = config.__retryCount ?? 0; + if (config.__retryCount >= numRetries) return Promise.reject(error); + config.__retryCount += 1; + await sleep(retryDelay * 2 ** (config.__retryCount - 1)); + return this.instance(config); + } + ); + } + + async post(url: string, body: unknown, config: HttpRequestConfig): Promise { + try { + const res = await this.instance.post(url, body, { + headers: config.headers, + timeout: config.timeout, + params: config.params, + }); + return { status: res.status, data: res.data }; + } catch (error: any) { + throw new HttpError( + error?.response?.data?.error?.message || error?.message || "request failed", + error?.response?.status + ); + } + } + + async *stream( + url: string, + body: unknown, + config: HttpRequestConfig + ): AsyncGenerator { + const res = await this.instance.post(url, body, { + headers: config.headers, + timeout: config.timeout, + params: config.params, + responseType: "stream", + }); + let buffer = ""; + for await (const chunk of res.data as AsyncIterable) { + buffer += chunk.toString("utf-8"); + let newlineIndex = buffer.indexOf("\n"); + while (newlineIndex !== -1) { + const line = buffer.slice(0, newlineIndex); + buffer = buffer.slice(newlineIndex + 1); + if (line.trim()) yield line; + newlineIndex = buffer.indexOf("\n"); + } + } + if (buffer.trim()) yield buffer; + } +} diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/router/index.ts b/JS/edgechains/arakoodev/src/ai/src/lib/router/index.ts new file mode 100644 index 000000000..4b22228e8 --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/lib/router/index.ts @@ -0,0 +1,36 @@ +export { + Router, + NoMatchingDeploymentError, + NoDeploymentsAvailableError, +} from "./router.js"; +export { AxiosHttpClient, HttpError } from "./httpClient.js"; +export { sentryCallback, posthogCallback } from "./callbacks.js"; +export type { SentryLike, PostHogLike } from "./callbacks.js"; +export { + openAIAdapter, + palmAdapter, + cohereAdapter, + getAdapter, + estimateTokens, +} from "./providers.js"; +export type { ProviderAdapter, NormalizedResponse, NormalizedEmbedding } from "./providers.js"; +export type { + Deployment, + RouterProvider, + RouterRole, + RouterMessage, + RoutingStrategy, + TokenUsage, + CompletionRequest, + CompletionResult, + FunctionCall, + EmbeddingRequest, + EmbeddingResult, + DeploymentUsage, + RouterEvent, + RouterCallback, + RouterOptions, + HttpClient, + HttpRequestConfig, + HttpResponse, +} from "./types.js"; diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/router/providers.ts b/JS/edgechains/arakoodev/src/ai/src/lib/router/providers.ts new file mode 100644 index 000000000..69e0ec02f --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/lib/router/providers.ts @@ -0,0 +1,293 @@ +// Provider adapters: normalize the three supported backends (openai, palm, +// cohere) onto one request/response shape so the Router stays provider-agnostic. + +import { + Deployment, + CompletionRequest, + EmbeddingRequest, + FunctionCall, + HttpRequestConfig, + RouterMessage, + TokenUsage, +} from "./types.js"; + +export interface NormalizedResponse { + content: string; + usage: TokenUsage; + /** Present when the model returned a function call. */ + functionCall?: FunctionCall; +} + +export interface NormalizedEmbedding { + data: any[]; + usage: TokenUsage; +} + +export interface ProviderRequest { + url: string; + body: unknown; + config: HttpRequestConfig; +} + +export interface StreamDelta { + content?: string; + usage?: TokenUsage; + done?: boolean; +} + +export interface ProviderAdapter { + /** Whether the backend exposes a native token stream. */ + supportsStreaming: boolean; + /** Whether the backend exposes an embeddings endpoint. */ + supportsEmbeddings?: boolean; + buildRequest( + deployment: Deployment, + req: CompletionRequest, + opts: { stream: boolean } + ): ProviderRequest; + parseResponse(data: any): NormalizedResponse; + /** Parse one raw stream line into a delta, or null to skip the line. */ + parseStreamLine(line: string): StreamDelta | null; + buildEmbeddingRequest?(deployment: Deployment, req: EmbeddingRequest): ProviderRequest; + parseEmbeddingResponse?(data: any): NormalizedEmbedding; +} + +/** + * Rough token estimate (~4 chars/token) used for tpm accounting when a backend + * does not report usage. This is intentionally an estimate, like litellm's + * fallback token_counter — never a substitute for real counts when present. + */ +export function estimateTokens(text: string): number { + if (!text) return 0; + return Math.max(1, Math.ceil(text.length / 4)); +} + +function toMessages(req: CompletionRequest): RouterMessage[] { + if (req.messages && req.messages.length) return req.messages; + if (req.prompt != null) return [{ role: "user", content: req.prompt }]; + throw new Error("completion request must include `prompt` or `messages`"); +} + +function joinMessages(messages: RouterMessage[]): string { + return messages.map((m) => m.content).join("\n"); +} + +function trimBase(base: string | undefined, fallback: string): string { + return (base ?? fallback).replace(/\/$/, ""); +} + +const OPENAI_DONE = "[DONE]"; + +export const openAIAdapter: ProviderAdapter = { + supportsStreaming: true, + supportsEmbeddings: true, + buildRequest(deployment, req, opts) { + const messages = toMessages(req); + return { + url: `${trimBase(deployment.apiBase, "https://api.openai.com/v1")}/chat/completions`, + body: { + model: deployment.providerModel || deployment.model, + messages, + max_tokens: req.max_tokens ?? 256, + temperature: req.temperature ?? 0.7, + stream: opts.stream, + ...(req.functions ? { functions: req.functions } : {}), + ...(req.function_call ? { function_call: req.function_call } : {}), + ...(opts.stream ? { stream_options: { include_usage: true } } : {}), + }, + config: { + headers: { + Authorization: `Bearer ${deployment.apiKey}`, + "content-type": "application/json", + }, + }, + }; + }, + parseResponse(data) { + const message = data?.choices?.[0]?.message; + const usage = data?.usage ?? {}; + const fc = message?.function_call; + return { + content: message?.content ?? "", + functionCall: + fc && typeof fc.name === "string" + ? { name: fc.name, arguments: fc.arguments ?? "" } + : undefined, + usage: { + promptTokens: usage.prompt_tokens ?? 0, + completionTokens: usage.completion_tokens ?? 0, + totalTokens: + usage.total_tokens ?? + (usage.prompt_tokens ?? 0) + (usage.completion_tokens ?? 0), + }, + }; + }, + buildEmbeddingRequest(deployment, req) { + return { + url: `${trimBase(deployment.apiBase, "https://api.openai.com/v1")}/embeddings`, + body: { + model: deployment.providerModel || deployment.model, + input: req.input, + }, + config: { + headers: { + Authorization: `Bearer ${deployment.apiKey}`, + "content-type": "application/json", + }, + }, + }; + }, + parseEmbeddingResponse(data) { + const usage = data?.usage ?? {}; + const promptTokens = usage.prompt_tokens ?? 0; + return { + data: data?.data ?? [], + usage: { + promptTokens, + completionTokens: 0, + totalTokens: usage.total_tokens ?? promptTokens, + }, + }; + }, + parseStreamLine(line) { + const trimmed = line.replace(/^data:\s*/, "").trim(); + if (!trimmed) return null; + if (trimmed === OPENAI_DONE) return { done: true }; + let parsed: any; + try { + parsed = JSON.parse(trimmed); + } catch { + return null; + } + const delta: StreamDelta = {}; + const content = parsed?.choices?.[0]?.delta?.content; + if (typeof content === "string") delta.content = content; + if (parsed?.usage) { + delta.usage = { + promptTokens: parsed.usage.prompt_tokens ?? 0, + completionTokens: parsed.usage.completion_tokens ?? 0, + totalTokens: parsed.usage.total_tokens ?? 0, + }; + } + return delta; + }, +}; + +export const palmAdapter: ProviderAdapter = { + // Google PaLM generateText has no streaming endpoint; the Router falls back + // to a single-shot completion emitted as one chunk. + supportsStreaming: false, + buildRequest(deployment, req) { + const text = joinMessages(toMessages(req)); + const model = deployment.providerModel || deployment.model; + return { + url: `${trimBase( + deployment.apiBase, + "https://generativelanguage.googleapis.com/v1beta2" + )}/models/${model}:generateText`, + body: { + prompt: { text }, + temperature: req.temperature ?? 0.7, + maxOutputTokens: req.max_tokens ?? 256, + }, + config: { + headers: { "content-type": "application/json" }, + params: { key: deployment.apiKey }, + }, + }; + }, + parseResponse(data) { + // PaLM generateText omits token counts; the Router backfills an estimate + // (from request + content) whenever a provider reports zero usage. + return { + content: data?.candidates?.[0]?.output ?? "", + usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 }, + }; + }, + parseStreamLine() { + return null; + }, +}; + +export const cohereAdapter: ProviderAdapter = { + supportsStreaming: true, + buildRequest(deployment, req, opts) { + const messages = toMessages(req); + const last = messages[messages.length - 1]; + const history = messages.slice(0, -1).map((m) => ({ + role: m.role === "assistant" ? "CHATBOT" : "USER", + message: m.content, + })); + return { + url: `${trimBase(deployment.apiBase, "https://api.cohere.ai/v1")}/chat`, + body: { + model: deployment.providerModel || deployment.model, + message: last.content, + ...(history.length ? { chat_history: history } : {}), + max_tokens: req.max_tokens ?? 256, + temperature: req.temperature ?? 0.7, + stream: opts.stream, + }, + config: { + headers: { + Authorization: `Bearer ${deployment.apiKey}`, + "content-type": "application/json", + }, + }, + }; + }, + parseResponse(data) { + const tokens = data?.meta?.tokens ?? data?.meta?.billed_units ?? {}; + const promptTokens = tokens.input_tokens ?? 0; + const completionTokens = tokens.output_tokens ?? 0; + return { + content: data?.text ?? "", + usage: { + promptTokens, + completionTokens, + totalTokens: promptTokens + completionTokens, + }, + }; + }, + parseStreamLine(line) { + const trimmed = line.trim(); + if (!trimmed) return null; + let parsed: any; + try { + parsed = JSON.parse(trimmed); + } catch { + return null; + } + if (parsed?.event_type === "text-generation") { + return { content: parsed.text ?? "" }; + } + if (parsed?.event_type === "stream-end") { + const tokens = parsed?.response?.meta?.tokens ?? {}; + return { + done: true, + usage: { + promptTokens: tokens.input_tokens ?? 0, + completionTokens: tokens.output_tokens ?? 0, + totalTokens: (tokens.input_tokens ?? 0) + (tokens.output_tokens ?? 0), + }, + }; + } + return null; + }, +}; + +export const ADAPTERS: Record = { + openai: openAIAdapter, + palm: palmAdapter, + cohere: cohereAdapter, +}; + +export function getAdapter(provider: string): ProviderAdapter { + const adapter = ADAPTERS[provider]; + if (!adapter) { + throw new Error( + `unsupported provider "${provider}" (expected one of openai, palm, cohere)` + ); + } + return adapter; +} diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/router/router.ts b/JS/edgechains/arakoodev/src/ai/src/lib/router/router.ts new file mode 100644 index 000000000..c26e49f56 --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/lib/router/router.ts @@ -0,0 +1,531 @@ +// Litellm-style smart Router for Edgechain.js. +// +// Features: +// 1. Load balancing across multiple deployments (openai, palm, cohere). The +// "usage-based" strategy picks the deployment that is below its rate limit +// and has used the fewest tokens this minute. Failed / rate-limited +// deployments are benched (cooldown) and routing fails over to the next. +// 2. Streaming completions (`router.stream`). +// 3. Token usage accounting (`router.getUsage` / per-result `usage`). +// 4. Logging callbacks (see ./callbacks for sentry + posthog adapters). +// +// Configuration is a plain object, which makes it natural to drive from a +// jsonnet file (see examples/react-chain/jsonnet/router.jsonnet). + +import { z } from "zod"; +import { zodToJsonSchema } from "zod-to-json-schema"; +import { + CompletionRequest, + CompletionResult, + Deployment, + DeploymentUsage, + EmbeddingRequest, + EmbeddingResult, + HttpClient, + RouterCallback, + RouterEvent, + RouterOptions, + RouterRole, + RoutingStrategy, + TokenUsage, +} from "./types.js"; +import { AxiosHttpClient } from "./httpClient.js"; +import { estimateTokens, getAdapter, NormalizedResponse } from "./providers.js"; + +const WINDOW_MS = 60000; + +export class NoMatchingDeploymentError extends Error { + constructor(model?: string) { + super( + model + ? `no deployment is configured for model "${model}"` + : "router has no deployments" + ); + this.name = "NoMatchingDeploymentError"; + } +} + +export class NoDeploymentsAvailableError extends Error { + constructor(model?: string) { + super( + `all deployments${ + model ? ` for model "${model}"` : "" + } are rate-limited or in cooldown` + ); + this.name = "NoDeploymentsAvailableError"; + } +} + +interface DeploymentState { + windowStart: number; + requestsInWindow: number; + tokensInWindow: number; + cooldownUntil: number; + totalRequests: number; + promptTokens: number; + completionTokens: number; + totalTokens: number; +} + +export class Router { + readonly deployments: Deployment[]; + private states: DeploymentState[]; + private strategy: RoutingStrategy; + private numRetries: number; + private retryDelay: number; + private cooldownMs: number; + private timeout: number; + private callbacks: RouterCallback[]; + private http: HttpClient; + private now: () => number; + private rrCounter = 0; + + constructor(options: RouterOptions) { + if (!options || !Array.isArray(options.deployments) || options.deployments.length === 0) { + throw new NoMatchingDeploymentError(); + } + options.deployments.forEach((d, i) => this.validateDeployment(d, i)); + + this.deployments = options.deployments; + this.strategy = options.strategy ?? "usage-based"; + this.numRetries = options.numRetries ?? 2; + this.retryDelay = options.retryDelay ?? 200; + this.cooldownMs = (options.cooldownSeconds ?? 60) * 1000; + this.timeout = options.timeout ?? 600000; + this.callbacks = options.callbacks ?? []; + this.now = options.now ?? Date.now; + this.http = + options.httpClient ?? + new AxiosHttpClient({ + numRetries: this.numRetries, + retryDelay: this.retryDelay, + timeout: this.timeout, + }); + + const start = this.now(); + this.states = this.deployments.map(() => ({ + windowStart: start, + requestsInWindow: 0, + tokensInWindow: 0, + cooldownUntil: 0, + totalRequests: 0, + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + })); + } + + private validateDeployment(d: Deployment, index: number): void { + if (!d || typeof d.model !== "string" || !d.model) { + throw new Error(`deployment[${index}] is missing a "model"`); + } + if (typeof d.apiKey !== "string" || !d.apiKey) { + throw new Error(`deployment[${index}] (${d.model}) is missing an "apiKey"`); + } + // Throws for unknown providers. + getAdapter(d.provider); + } + + /** Token estimate (litellm token_counter parity). */ + static tokenCount(text: string): number { + return estimateTokens(text); + } + + /** Estimate the prompt token cost of a request without sending it. */ + tokenCount(req: CompletionRequest): number { + return estimateTokens(this.promptText(req)); + } + + private promptText(req: CompletionRequest): string { + if (req.prompt != null) return req.prompt; + if (req.messages) return req.messages.map((m) => m.content).join("\n"); + return ""; + } + + private rollWindow(state: DeploymentState, now: number): void { + if (now - state.windowStart >= WINDOW_MS) { + state.windowStart = now; + state.requestsInWindow = 0; + state.tokensInWindow = 0; + } + } + + private isEligible(index: number, now: number): boolean { + const state = this.states[index]; + const dep = this.deployments[index]; + if (now < state.cooldownUntil) return false; + this.rollWindow(state, now); + if (dep.rpm != null && state.requestsInWindow >= dep.rpm) return false; + if (dep.tpm != null && state.tokensInWindow >= dep.tpm) return false; + return true; + } + + /** Indices of deployments serving the requested model group. */ + private matching(model?: string): number[] { + const indices: number[] = []; + this.deployments.forEach((d, i) => { + if (!model || d.model === model) indices.push(i); + }); + return indices; + } + + /** + * Pick the best eligible deployment for the model, or null if every + * matching deployment is currently benched. Throws if none match at all. + */ + private pick(model?: string): number | null { + const matching = this.matching(model); + if (matching.length === 0) throw new NoMatchingDeploymentError(model); + const now = this.now(); + const eligible = matching.filter((i) => this.isEligible(i, now)); + if (eligible.length === 0) return null; + + if (this.strategy === "round-robin") { + const choice = eligible[this.rrCounter % eligible.length]; + this.rrCounter += 1; + return choice; + } + // usage-based: fewest tokens used this minute, then fewest requests. + return eligible.reduce((best, i) => { + const a = this.states[i]; + const b = this.states[best]; + if (a.tokensInWindow !== b.tokensInWindow) + return a.tokensInWindow < b.tokensInWindow ? i : best; + return a.requestsInWindow < b.requestsInWindow ? i : best; + }, eligible[0]); + } + + private backfillUsage(normalized: NormalizedResponse, req: CompletionRequest): TokenUsage { + if (normalized.usage.totalTokens > 0) return normalized.usage; + const promptTokens = estimateTokens(this.promptText(req)); + const completionTokens = estimateTokens(normalized.content); + return { promptTokens, completionTokens, totalTokens: promptTokens + completionTokens }; + } + + private recordSuccess(index: number, usage: TokenUsage): void { + const state = this.states[index]; + this.rollWindow(state, this.now()); + state.requestsInWindow += 1; + state.tokensInWindow += usage.totalTokens; + state.totalRequests += 1; + state.promptTokens += usage.promptTokens; + state.completionTokens += usage.completionTokens; + state.totalTokens += usage.totalTokens; + } + + private recordFailure(index: number): void { + this.states[index].cooldownUntil = this.now() + this.cooldownMs; + } + + private async notify( + kind: "onSuccess" | "onError", + event: RouterEvent + ): Promise { + for (const cb of this.callbacks) { + const handler = cb[kind]; + if (!handler) continue; + try { + await handler.call(cb, event); + } catch { + // Observability must never break the completion. + } + } + } + + /** Non-streaming completion with usage-aware load balancing + failover. */ + async completion(req: CompletionRequest): Promise { + const maxAttempts = this.numRetries + 1; + let lastError: unknown; + for (let attempt = 0; attempt < maxAttempts; attempt++) { + const index = this.pick(req.model); + if (index == null) break; + const dep = this.deployments[index]; + const adapter = getAdapter(dep.provider); + const start = this.now(); + try { + const built = adapter.buildRequest(dep, req, { stream: false }); + const res = await this.http.post(built.url, built.body, { + ...built.config, + timeout: dep.timeout ?? this.timeout, + }); + const normalized = adapter.parseResponse(res.data); + const usage = this.backfillUsage(normalized, req); + this.recordSuccess(index, usage); + const result: CompletionResult = { + content: normalized.content, + usage, + provider: dep.provider, + model: dep.providerModel || dep.model, + deploymentIndex: index, + functionCall: normalized.functionCall, + raw: res.data, + }; + await this.notify("onSuccess", { + provider: dep.provider, + model: result.model, + deploymentIndex: index, + request: req, + usage, + response: result, + durationMs: this.now() - start, + }); + return result; + } catch (error) { + lastError = error; + this.recordFailure(index); + await this.notify("onError", { + provider: dep.provider, + model: dep.providerModel || dep.model, + deploymentIndex: index, + request: req, + error, + durationMs: this.now() - start, + }); + } + } + if (lastError) throw lastError; + throw new NoDeploymentsAvailableError(req.model); + } + + /** Embeddings with the same usage-aware load balancing + failover. */ + async embedding(req: EmbeddingRequest): Promise { + const maxAttempts = this.numRetries + 1; + let lastError: unknown; + for (let attempt = 0; attempt < maxAttempts; attempt++) { + const index = this.pick(req.model); + if (index == null) break; + const dep = this.deployments[index]; + const adapter = getAdapter(dep.provider); + if ( + !adapter.supportsEmbeddings || + !adapter.buildEmbeddingRequest || + !adapter.parseEmbeddingResponse + ) { + throw new Error(`provider "${dep.provider}" does not support embeddings`); + } + const start = this.now(); + try { + const built = adapter.buildEmbeddingRequest(dep, req); + const res = await this.http.post(built.url, built.body, { + ...built.config, + timeout: dep.timeout ?? this.timeout, + }); + const parsed = adapter.parseEmbeddingResponse(res.data); + const usage = + parsed.usage.totalTokens > 0 + ? parsed.usage + : this.estimateEmbeddingUsage(req); + this.recordSuccess(index, usage); + const result: EmbeddingResult = { + data: parsed.data, + usage, + provider: dep.provider, + model: dep.providerModel || dep.model, + deploymentIndex: index, + raw: res.data, + }; + await this.notify("onSuccess", { + provider: dep.provider, + model: result.model, + deploymentIndex: index, + request: { model: req.model }, + usage, + durationMs: this.now() - start, + }); + return result; + } catch (error) { + lastError = error; + this.recordFailure(index); + await this.notify("onError", { + provider: dep.provider, + model: dep.providerModel || dep.model, + deploymentIndex: index, + request: { model: req.model }, + error, + durationMs: this.now() - start, + }); + } + } + if (lastError) throw lastError; + throw new NoDeploymentsAvailableError(req.model); + } + + private estimateEmbeddingUsage(req: EmbeddingRequest): TokenUsage { + const text = Array.isArray(req.input) ? req.input.join("\n") : req.input; + const promptTokens = estimateTokens(text); + return { promptTokens, completionTokens: 0, totalTokens: promptTokens }; + } + + /** + * Structured output via a zod schema. The schema is sent as an OpenAI + * function definition and the returned arguments are parsed back through + * zod, mirroring the litellm `response_format` / function-calling pattern. + */ + async zodSchemaResponse(opts: { + model?: string; + role?: RouterRole; + max_tokens?: number; + temperature?: number; + prompt: string; + schema: S; + }): Promise> { + const jsonSchema = zodToJsonSchema(opts.schema, { $refStrategy: "none" }); + const functionDefinition = { + name: "generateSchema", + description: "Generate a schema based on provided details.", + parameters: jsonSchema, + }; + const content = ` + You are a Schema generator that can generate answer based on given prompt and then return the response based on the give schema + Remembrer if any field like url or link is not available please create a dummy link based on the following prompt + + prompt: + ${opts.prompt || ""} + `; + const result = await this.completion({ + model: opts.model, + max_tokens: opts.max_tokens ?? 1000, + temperature: opts.temperature, + messages: [{ role: opts.role ?? "user", content }], + functions: [functionDefinition], + function_call: "auto", + }); + if (result.content) return result.content as unknown as z.infer; + if (result.functionCall) { + return opts.schema.parse(JSON.parse(result.functionCall.arguments)); + } + throw new Error("Response did not contain valid JSON."); + } + + /** + * Streaming completion. Yields content deltas. Providers without a native + * stream (palm) fall back to a single-shot completion emitted as one chunk. + * Fail-over only happens before the first chunk is delivered. + */ + async *stream(req: CompletionRequest): AsyncGenerator { + const maxAttempts = this.numRetries + 1; + let lastError: unknown; + for (let attempt = 0; attempt < maxAttempts; attempt++) { + const index = this.pick(req.model); + if (index == null) break; + const dep = this.deployments[index]; + const adapter = getAdapter(dep.provider); + const start = this.now(); + let started = false; + try { + if (!adapter.supportsStreaming) { + const built = adapter.buildRequest(dep, req, { stream: false }); + const res = await this.http.post(built.url, built.body, { + ...built.config, + timeout: dep.timeout ?? this.timeout, + }); + const normalized = adapter.parseResponse(res.data); + const usage = this.backfillUsage(normalized, req); + this.recordSuccess(index, usage); + await this.notifyStreamSuccess(index, dep, req, normalized.content, usage, start); + if (normalized.content) { + started = true; + yield normalized.content; + } + return; + } + + const built = adapter.buildRequest(dep, req, { stream: true }); + const lines = this.http.stream(built.url, built.body, { + ...built.config, + timeout: dep.timeout ?? this.timeout, + }); + let assembled = ""; + let streamUsage: TokenUsage | undefined; + for await (const line of lines) { + const delta = adapter.parseStreamLine(line); + if (!delta) continue; + if (delta.usage) streamUsage = delta.usage; + if (delta.content) { + assembled += delta.content; + started = true; + yield delta.content; + } + if (delta.done) break; + } + const usage = + streamUsage ?? + this.backfillUsage({ content: assembled, usage: zeroUsage() }, req); + this.recordSuccess(index, usage); + await this.notifyStreamSuccess(index, dep, req, assembled, usage, start); + return; + } catch (error) { + lastError = error; + this.recordFailure(index); + await this.notify("onError", { + provider: dep.provider, + model: dep.providerModel || dep.model, + deploymentIndex: index, + request: req, + error, + durationMs: this.now() - start, + }); + if (started) throw error; // cannot fail over mid-stream + } + } + if (lastError) throw lastError; + throw new NoDeploymentsAvailableError(req.model); + } + + private async notifyStreamSuccess( + index: number, + dep: Deployment, + req: CompletionRequest, + content: string, + usage: TokenUsage, + start: number + ): Promise { + const model = dep.providerModel || dep.model; + await this.notify("onSuccess", { + provider: dep.provider, + model, + deploymentIndex: index, + request: req, + usage, + response: { + content, + usage, + provider: dep.provider, + model, + deploymentIndex: index, + raw: null, + }, + durationMs: this.now() - start, + }); + } + + /** Per-deployment cumulative usage. */ + getUsage(): DeploymentUsage[] { + return this.deployments.map((d, i) => ({ + deploymentIndex: i, + provider: d.provider, + model: d.model, + requests: this.states[i].totalRequests, + promptTokens: this.states[i].promptTokens, + completionTokens: this.states[i].completionTokens, + totalTokens: this.states[i].totalTokens, + })); + } + + /** Aggregate usage across all deployments. */ + getTotalUsage(): TokenUsage & { requests: number } { + return this.states.reduce( + (acc, s) => ({ + requests: acc.requests + s.totalRequests, + promptTokens: acc.promptTokens + s.promptTokens, + completionTokens: acc.completionTokens + s.completionTokens, + totalTokens: acc.totalTokens + s.totalTokens, + }), + { requests: 0, promptTokens: 0, completionTokens: 0, totalTokens: 0 } + ); + } +} + +function zeroUsage(): TokenUsage { + return { promptTokens: 0, completionTokens: 0, totalTokens: 0 }; +} diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/router/types.ts b/JS/edgechains/arakoodev/src/ai/src/lib/router/types.ts new file mode 100644 index 000000000..5acf9f0a2 --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/lib/router/types.ts @@ -0,0 +1,156 @@ +// Shared types for the litellm-style smart Router. +// The Router load-balances chat completions across multiple provider +// deployments (openai, palm, cohere), tracks token usage, supports +// streaming and pluggable logging callbacks. + +export type RouterProvider = "openai" | "palm" | "cohere"; + +export type RouterRole = "system" | "user" | "assistant"; + +export interface RouterMessage { + role: RouterRole; + content: string; +} + +export type RoutingStrategy = "usage-based" | "round-robin"; + +/** A single provider deployment the Router can route a request to. */ +export interface Deployment { + /** Logical model/group name a caller asks for, e.g. "gpt-3.5-turbo". */ + model: string; + provider: RouterProvider; + apiKey: string; + /** Provider model id sent on the wire. Defaults to `model`. */ + providerModel?: string; + /** Override the provider base url (host + version), no trailing slash. */ + apiBase?: string; + /** Requests-per-minute budget. Deployment is skipped once exceeded. */ + rpm?: number; + /** Tokens-per-minute budget. Deployment is skipped once exceeded. */ + tpm?: number; + /** Per-request timeout in ms. Defaults to the Router timeout. */ + timeout?: number; +} + +export interface TokenUsage { + promptTokens: number; + completionTokens: number; + totalTokens: number; +} + +export interface CompletionRequest { + /** Restrict routing to deployments declaring this model group. */ + model?: string; + prompt?: string; + messages?: RouterMessage[]; + max_tokens?: number; + temperature?: number; + /** OpenAI-style function definitions enabling function calling. */ + functions?: object | object[]; + /** "auto" / "none" / { name } to force a specific function. */ + function_call?: string | { name: string }; +} + +/** An OpenAI-style function call the model decided to make. */ +export interface FunctionCall { + name: string; + arguments: string; +} + +export interface CompletionResult { + content: string; + usage: TokenUsage; + provider: RouterProvider; + /** Provider model id that actually served the request. */ + model: string; + deploymentIndex: number; + /** Present when the model returned a function call. */ + functionCall?: FunctionCall; + raw: unknown; +} + +export interface EmbeddingRequest { + /** Restrict routing to deployments declaring this model group. */ + model?: string; + input: string | string[]; +} + +export interface EmbeddingResult { + /** Provider embedding objects, e.g. [{ embedding: number[], index }]. */ + data: any[]; + usage: TokenUsage; + provider: RouterProvider; + /** Provider model id that actually served the request. */ + model: string; + deploymentIndex: number; + raw: unknown; +} + +/** Cumulative usage book-kept by the Router, per deployment. */ +export interface DeploymentUsage { + deploymentIndex: number; + provider: RouterProvider; + model: string; + requests: number; + promptTokens: number; + completionTokens: number; + totalTokens: number; +} + +export interface RouterEvent { + provider: RouterProvider; + model: string; + deploymentIndex: number; + request: CompletionRequest; + usage?: TokenUsage; + response?: CompletionResult; + error?: unknown; + durationMs: number; +} + +/** + * Observability hook. Mirrors litellm's success/failure callbacks. Concrete + * Sentry / PostHog adapters live in ./callbacks and receive an injected client + * so no SDK becomes a hard dependency of this package. + */ +export interface RouterCallback { + onSuccess?(event: RouterEvent): void | Promise; + onError?(event: RouterEvent): void | Promise; +} + +/** Minimal HTTP seam so tests run without real network calls. */ +export interface HttpRequestConfig { + headers?: Record; + timeout?: number; + params?: Record; +} + +export interface HttpResponse { + status: number; + data: any; +} + +export interface HttpClient { + post(url: string, body: unknown, config: HttpRequestConfig): Promise; + /** Yields raw text lines/chunks from a streamed response body. */ + stream(url: string, body: unknown, config: HttpRequestConfig): AsyncIterable; +} + +export interface RouterOptions { + deployments: Deployment[]; + /** Load-balancing strategy. Defaults to "usage-based" (least tokens used). */ + strategy?: RoutingStrategy; + /** Number of fail-over retries across deployments. Defaults to 2. */ + numRetries?: number; + /** Base backoff between Router-level retries, in ms. Defaults to 200. */ + retryDelay?: number; + /** How long a failing/rate-limited deployment is benched, in seconds. Defaults to 60. */ + cooldownSeconds?: number; + /** Default per-request timeout in ms. Defaults to 600000 (litellm default). */ + timeout?: number; + callbacks?: RouterCallback[]; + /** Injectable HTTP client (defaults to an axios-backed one). */ + httpClient?: HttpClient; + /** Injectable clock, in ms. Defaults to Date.now. Used for windows/cooldown. */ + now?: () => number; +} diff --git a/JS/edgechains/arakoodev/src/ai/src/tests/openAiEndpoints.test.ts b/JS/edgechains/arakoodev/src/ai/src/tests/openAiEndpoints.test.ts deleted file mode 100644 index b6ab6b654..000000000 --- a/JS/edgechains/arakoodev/src/ai/src/tests/openAiEndpoints.test.ts +++ /dev/null @@ -1,81 +0,0 @@ -import axios from "axios"; -import { OpenAI } from "../../../../dist/openai/src/lib/endpoints/OpenAiEndpoint.js"; - -jest.mock("axios"); - -describe("ChatOpenAi", () => { - describe("generateResponse", () => { - test("should generate response from OpenAI", async () => { - const mockResponse = [ - { - message: { - content: "Test response", - }, - }, - ]; - - axios.post = jest.fn().mockResolvedValueOnce({ data: { choices: mockResponse } }); - const chatOpenAi = new OpenAI({ apiKey: "test_api_key" }); - const response = await chatOpenAi.chat({ prompt: "test prompt" }); - expect(response).toEqual("Test response"); - }); - }); - - describe("generateEmbeddings", () => { - test("should generate embeddings from OpenAI", async () => { - const mockResponse = { embeddings: "Test embeddings" }; - axios.post = jest.fn().mockResolvedValue({ data: { data: { choices: mockResponse } } }); - const chatOpenAi = new OpenAI({ apiKey: "test_api_key" }); - const res = await chatOpenAi.generateEmbeddings("test prompt"); - expect(res.choices.embeddings).toEqual("Test embeddings"); - }); - }); - - describe("chatWithAI", () => { - test("should chat with AI using multiple messages", async () => { - const mockResponse = [ - { - message: { - content: "Test response 1", - }, - }, - { - message: { - content: "Test response 2", - }, - }, - ]; - axios.post = jest.fn().mockResolvedValueOnce({ data: { choices: mockResponse } }); - const chatOpenAi = new OpenAI({ apiKey: "test_api_key" }); - const chatMessages = [ - { - role: "user", - content: "message 1", - }, - { - role: "agent", - content: "message 2", - }, - ]; - //@ts-ignore - const responses = await chatOpenAi.chat({ messages: chatMessages }); - expect(responses).toEqual(mockResponse); - }); - }); - - describe("testResponseGeneration", () => { - test("should generate test response from OpenAI", async () => { - const mockResponse = [ - { - message: { - content: "Test response", - }, - }, - ]; - axios.post = jest.fn().mockResolvedValueOnce({ data: { choices: mockResponse } }); - const chatOpenAi = new OpenAI({ apiKey: "test_api_key" }); - const response = await chatOpenAi.chat({ prompt: "test prompt" }); - expect(response).toEqual("Test response"); - }); - }); -}); diff --git a/JS/edgechains/arakoodev/src/ai/src/tests/router.test.ts b/JS/edgechains/arakoodev/src/ai/src/tests/router.test.ts new file mode 100644 index 000000000..d684fbf6f --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/tests/router.test.ts @@ -0,0 +1,432 @@ +import { describe, it, expect, vi } from "vitest"; +import { z } from "zod"; +import { + Router, + NoMatchingDeploymentError, + NoDeploymentsAvailableError, +} from "../lib/router/router.js"; +import { sentryCallback, posthogCallback } from "../lib/router/callbacks.js"; +import { HttpError } from "../lib/router/httpClient.js"; +import { HttpClient, HttpRequestConfig, HttpResponse } from "../lib/router/types.js"; + +// ---- Mocked provider endpoints (openai / palm / cohere) ----------------------- + +type PostFn = (url: string, body: any) => HttpResponse; + +class FakeHttp implements HttpClient { + posts: { url: string; body: any }[] = []; + constructor( + private postFn: PostFn, + private lines: string[] = [] + ) {} + + async post(url: string, body: any, _config: HttpRequestConfig): Promise { + this.posts.push({ url, body }); + return this.postFn(url, body); // postFn may throw to simulate failures + } + + async *stream(url: string, body: any, _config: HttpRequestConfig): AsyncIterable { + this.posts.push({ url, body }); + for (const line of this.lines) yield line; + } +} + +function openAiData(content: string, total = 10): HttpResponse { + return { + status: 200, + data: { + choices: [{ message: { content } }], + usage: { prompt_tokens: 4, completion_tokens: total - 4, total_tokens: total }, + }, + }; +} + +function palmData(content: string): HttpResponse { + return { status: 200, data: { candidates: [{ output: content }] } }; +} + +function cohereData(content: string): HttpResponse { + return { + status: 200, + data: { text: content, meta: { tokens: { input_tokens: 4, output_tokens: 6 } } }, + }; +} + +function openAiFunctionData(name: string, args: object): HttpResponse { + return { + status: 200, + data: { + choices: [ + { message: { content: null, function_call: { name, arguments: JSON.stringify(args) } } }, + ], + usage: { prompt_tokens: 5, completion_tokens: 5, total_tokens: 10 }, + }, + }; +} + +function openAiEmbeddingData(): HttpResponse { + return { + status: 200, + data: { + data: [{ embedding: [0.1, 0.2, 0.3], index: 0 }], + usage: { prompt_tokens: 8, total_tokens: 8 }, + }, + }; +} + +async function collect(gen: AsyncGenerator): Promise { + const out: string[] = []; + for await (const chunk of gen) out.push(chunk); + return out; +} + +const openai = (extra: object = {}) => ({ + model: "gpt", + provider: "openai" as const, + apiKey: "test-key", + ...extra, +}); + +// ------------------------------------------------------------------------------ + +describe("Router", () => { + describe("construction & validation", () => { + it("throws when no deployments are configured", () => { + expect(() => new Router({ deployments: [] })).toThrow(NoMatchingDeploymentError); + }); + + it("fails loud when a deployment is missing its apiKey", () => { + expect( + () => + new Router({ + // @ts-expect-error apiKey intentionally omitted + deployments: [{ model: "gpt", provider: "openai" }], + }) + ).toThrow(/apiKey/); + }); + + it("rejects unknown providers", () => { + expect( + () => + new Router({ + // @ts-expect-error provider intentionally invalid + deployments: [{ model: "x", provider: "anthropic", apiKey: "k" }], + }) + ).toThrow(/unsupported provider/); + }); + }); + + describe("load balancing", () => { + it("usage-based picks the deployment with the fewest tokens used", async () => { + const http = new FakeHttp(() => openAiData("ok", 10)); + const router = new Router({ + deployments: [openai(), openai()], + httpClient: http, + }); + + const r1 = await router.completion({ prompt: "hi" }); + const r2 = await router.completion({ prompt: "hi" }); + const r3 = await router.completion({ prompt: "hi" }); + + // 0 -> 1 (0 had tokens) -> 0 (tie broken by lowest index) + expect([r1.deploymentIndex, r2.deploymentIndex, r3.deploymentIndex]).toEqual([0, 1, 0]); + }); + + it("round-robin alternates across deployments", async () => { + const http = new FakeHttp(() => openAiData("ok")); + const router = new Router({ + deployments: [openai(), openai()], + strategy: "round-robin", + httpClient: http, + }); + + const seq: number[] = []; + for (let i = 0; i < 4; i++) + seq.push((await router.completion({ prompt: "hi" })).deploymentIndex); + expect(seq).toEqual([0, 1, 0, 1]); + }); + }); + + describe("reliability", () => { + it("fails over to the next deployment and benches the rate-limited one", async () => { + let clock = 1000; + const http = new FakeHttp((_url, body) => { + if (body.model === "m0") throw new HttpError("rate limit", 429); + return openAiData("served by m1"); + }); + const router = new Router({ + deployments: [ + openai({ providerModel: "m0" }), + openai({ providerModel: "m1" }), + ], + httpClient: http, + now: () => clock, + }); + + const r1 = await router.completion({ prompt: "hi" }); + expect(r1.deploymentIndex).toBe(1); + expect(r1.content).toBe("served by m1"); + + // m0 is cooling down, so the next call goes straight to m1 (1 post, not 2). + http.posts.length = 0; + const r2 = await router.completion({ prompt: "hi" }); + expect(r2.deploymentIndex).toBe(1); + expect(http.posts.length).toBe(1); + }); + + it("throws NoDeploymentsAvailableError when every deployment is over budget", async () => { + const http = new FakeHttp(() => openAiData("ok")); + const router = new Router({ + deployments: [openai({ rpm: 0 })], + httpClient: http, + }); + await expect(router.completion({ prompt: "hi" })).rejects.toBeInstanceOf( + NoDeploymentsAvailableError + ); + }); + + it("throws NoMatchingDeploymentError for an unknown model group", async () => { + const http = new FakeHttp(() => openAiData("ok")); + const router = new Router({ deployments: [openai()], httpClient: http }); + await expect(router.completion({ model: "does-not-exist", prompt: "hi" })).rejects.toBeInstanceOf( + NoMatchingDeploymentError + ); + }); + }); + + describe("token usage", () => { + it("accumulates per-deployment and aggregate usage", async () => { + const http = new FakeHttp(() => openAiData("ok", 10)); + const router = new Router({ deployments: [openai(), openai()], httpClient: http }); + + await router.completion({ prompt: "hi" }); // dep 0 + await router.completion({ prompt: "hi" }); // dep 1 + + const usage = router.getUsage(); + expect(usage[0].totalTokens).toBe(10); + expect(usage[0].requests).toBe(1); + expect(usage[1].totalTokens).toBe(10); + + const total = router.getTotalUsage(); + expect(total.totalTokens).toBe(20); + expect(total.requests).toBe(2); + }); + }); + + describe("provider normalization (mocked endpoints)", () => { + it("normalizes openai, palm and cohere responses to {content, usage}", async () => { + const http = new FakeHttp((url) => { + if (url.includes("/chat/completions")) return openAiData("openai answer", 10); + if (url.includes(":generateText")) return palmData("palm answer"); + if (url.endsWith("/chat")) return cohereData("cohere answer"); + throw new Error(`unexpected url ${url}`); + }); + const router = new Router({ + deployments: [ + { model: "gpt", provider: "openai", apiKey: "k" }, + { model: "chat-bison", provider: "palm", apiKey: "k" }, + { model: "command", provider: "cohere", apiKey: "k" }, + ], + httpClient: http, + }); + + const o = await router.completion({ model: "gpt", prompt: "hi" }); + expect(o.provider).toBe("openai"); + expect(o.content).toBe("openai answer"); + expect(o.usage.totalTokens).toBe(10); + + const p = await router.completion({ model: "chat-bison", prompt: "hello there" }); + expect(p.provider).toBe("palm"); + expect(p.content).toBe("palm answer"); + // PaLM reports no usage -> Router backfills an estimate (> 0). + expect(p.usage.totalTokens).toBeGreaterThan(0); + + const c = await router.completion({ model: "command", prompt: "hi" }); + expect(c.provider).toBe("cohere"); + expect(c.content).toBe("cohere answer"); + expect(c.usage.totalTokens).toBe(10); + }); + }); + + describe("streaming", () => { + it("assembles openai stream deltas and records reported usage", async () => { + const lines = [ + 'data: {"choices":[{"delta":{"content":"Hello"}}]}', + 'data: {"choices":[{"delta":{"content":" world"}}]}', + 'data: {"choices":[],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}', + "data: [DONE]", + ]; + const http = new FakeHttp(() => openAiData("unused"), lines); + const router = new Router({ deployments: [openai()], httpClient: http }); + + const chunks = await collect(router.stream({ prompt: "hi" })); + expect(chunks).toEqual(["Hello", " world"]); + expect(router.getTotalUsage().totalTokens).toBe(5); + }); + + it("falls back to a single chunk for non-streaming providers (palm)", async () => { + const http = new FakeHttp(() => palmData("palm streamed")); + const router = new Router({ + deployments: [{ model: "chat-bison", provider: "palm", apiKey: "k" }], + httpClient: http, + }); + + const chunks = await collect(router.stream({ prompt: "hi" })); + expect(chunks).toEqual(["palm streamed"]); + expect(router.getTotalUsage().totalTokens).toBeGreaterThan(0); + }); + }); + + describe("logging callbacks", () => { + it("invokes sentry + posthog on success and failure", async () => { + const sentry = { captureException: vi.fn(), addBreadcrumb: vi.fn() }; + const posthog = { capture: vi.fn() }; + + const http = new FakeHttp((_url, body) => { + if (body.model === "bad") throw new HttpError("boom", 500); + return openAiData("ok"); + }); + const router = new Router({ + deployments: [ + openai({ providerModel: "bad" }), + openai({ providerModel: "good" }), + ], + httpClient: http, + callbacks: [sentryCallback(sentry), posthogCallback(posthog, { distinctId: "u1" })], + }); + + await router.completion({ prompt: "hi" }); // bad fails over to good + + expect(sentry.captureException).toHaveBeenCalledTimes(1); + expect(sentry.addBreadcrumb).toHaveBeenCalledTimes(1); + expect(posthog.capture).toHaveBeenCalledWith( + expect.objectContaining({ distinctId: "u1", event: "edgechains_completion_error" }) + ); + expect(posthog.capture).toHaveBeenCalledWith( + expect.objectContaining({ distinctId: "u1", event: "edgechains_completion" }) + ); + }); + + it("never lets a throwing callback break the completion", async () => { + const http = new FakeHttp(() => openAiData("ok")); + const router = new Router({ + deployments: [openai()], + httpClient: http, + callbacks: [ + { + onSuccess() { + throw new Error("callback exploded"); + }, + }, + ], + }); + const res = await router.completion({ prompt: "hi" }); + expect(res.content).toBe("ok"); + }); + }); + + describe("token counting", () => { + it("estimates prompt tokens without sending a request", () => { + const http = new FakeHttp(() => openAiData("ok")); + const router = new Router({ deployments: [openai()], httpClient: http }); + expect(router.tokenCount({ prompt: "a".repeat(40) })).toBe(10); + expect(Router.tokenCount("")).toBe(0); + }); + }); + + describe("function calling", () => { + it("passes functions through and surfaces the returned function call", async () => { + const http = new FakeHttp(() => + openAiFunctionData("get_weather", { location: "Paris" }) + ); + const router = new Router({ deployments: [openai()], httpClient: http }); + + const res = await router.completion({ + prompt: "weather in paris", + functions: [{ name: "get_weather", parameters: {} }], + function_call: "auto", + }); + + expect(res.functionCall?.name).toBe("get_weather"); + expect(JSON.parse(res.functionCall!.arguments)).toEqual({ location: "Paris" }); + // The function definition + forced-call mode reach the provider. + expect(http.posts[0].body.functions).toEqual([{ name: "get_weather", parameters: {} }]); + expect(http.posts[0].body.function_call).toBe("auto"); + }); + + it("leaves functionCall undefined for a plain completion", async () => { + const http = new FakeHttp(() => openAiData("just text")); + const router = new Router({ deployments: [openai()], httpClient: http }); + const res = await router.completion({ prompt: "hi" }); + expect(res.functionCall).toBeUndefined(); + expect(http.posts[0].body.functions).toBeUndefined(); + }); + }); + + describe("embeddings", () => { + it("routes to the embedding deployment and returns vectors + usage", async () => { + const http = new FakeHttp((url) => { + if (url.endsWith("/embeddings")) return openAiEmbeddingData(); + throw new Error(`unexpected url ${url}`); + }); + const router = new Router({ + deployments: [ + { model: "gpt-3.5-turbo", provider: "openai", apiKey: "k" }, + { model: "text-embedding-ada-002", provider: "openai", apiKey: "k" }, + ], + httpClient: http, + }); + + const res = await router.embedding({ + input: ["hello", "world"], + model: "text-embedding-ada-002", + }); + + expect(res.deploymentIndex).toBe(1); + expect(res.data[0].embedding).toEqual([0.1, 0.2, 0.3]); + expect(res.usage.totalTokens).toBe(8); + expect(http.posts[0].url).toContain("/embeddings"); + expect(http.posts[0].body.model).toBe("text-embedding-ada-002"); + expect(router.getUsage()[1].totalTokens).toBe(8); + }); + + it("throws when the selected provider has no embeddings endpoint", async () => { + const http = new FakeHttp(() => palmData("unused")); + const router = new Router({ + deployments: [{ model: "chat-bison", provider: "palm", apiKey: "k" }], + httpClient: http, + }); + await expect(router.embedding({ input: "hi" })).rejects.toThrow( + /does not support embeddings/ + ); + }); + }); + + describe("zodSchemaResponse", () => { + it("parses the function-call arguments back through the zod schema", async () => { + const schema = z.object({ answer: z.string() }); + const http = new FakeHttp(() => + openAiFunctionData("generateSchema", { answer: "42" }) + ); + const router = new Router({ deployments: [openai()], httpClient: http }); + + const res = await router.zodSchemaResponse({ prompt: "the answer?", schema }); + + expect(res).toEqual({ answer: "42" }); + // The schema is sent to the provider as an OpenAI function definition. + expect(http.posts[0].body.functions[0].name).toBe("generateSchema"); + }); + + it("rejects when the model returns neither content nor a function call", async () => { + const schema = z.object({ answer: z.string() }); + const http = new FakeHttp(() => ({ + status: 200, + data: { choices: [{ message: { content: null } }], usage: {} }, + })); + const router = new Router({ deployments: [openai()], httpClient: http }); + await expect(router.zodSchemaResponse({ prompt: "x", schema })).rejects.toThrow( + /did not contain valid JSON/ + ); + }); + }); +}); diff --git a/JS/edgechains/arakoodev/src/ai/src/tests/streaming.test.ts b/JS/edgechains/arakoodev/src/ai/src/tests/streaming.test.ts deleted file mode 100644 index 603353dd7..000000000 --- a/JS/edgechains/arakoodev/src/ai/src/tests/streaming.test.ts +++ /dev/null @@ -1,79 +0,0 @@ -import { Stream } from "../../../../dist/openai/src/lib/streaming/OpenAiStreaming.js"; -import { TextEncoder, TextDecoder } from "text-decoding"; -jest.mock("../lib/streaming/OpenAiStreaming.ts", () => { - return { - Stream: jest.fn().mockImplementation(() => ({ - OpenAIStream: jest.fn().mockImplementation((prompt) => { - return new ReadableStream({ - start(controller) { - controller.enqueue( - new TextEncoder().encode('[{"choices":[{"delta":{"content":"Hi! "}}]}]') - ); - controller.enqueue( - new TextEncoder().encode('[{"choices":[{"delta":{"content":"How "}}]}]') - ); - controller.enqueue( - new TextEncoder().encode('[{"choices":[{"delta":{"content":"can "}}]}]') - ); - controller.enqueue( - new TextEncoder().encode('[{"choices":[{"delta":{"content":"I "}}]}]') - ); - controller.enqueue( - new TextEncoder().encode( - '[{"choices":[{"delta":{"content":"help "}}]}]' - ) - ); - controller.enqueue( - new TextEncoder().encode( - '[{"choices":[{"delta":{"content":"you?."}}]}]' - ) - ); - controller.enqueue(new TextEncoder().encode("[DONE]")); - controller.close(); - }, - }); - }), - })), - }; -}); - -describe("Streaming", () => { - afterEach(() => { - jest.clearAllMocks(); // Clear mock function calls after each test - }); - - test("OpenAIStream should return expected text", async () => { - const options = { - model: "test_model", - OpenApiKey: "test_api_key", - temperature: 0.7, - top_p: 1, - frequency_penalty: 0, - presence_penalty: 0, - max_tokens: 500, - stream: true, - n: 1, - }; - - const stream = new Stream(options); - - //@ts-ignore - const streamReader = await stream.OpenAIStream("hi").getReader(); - const text = await readStreamToString(streamReader); - - expect(text).toBe("Hi! How can I help you?."); - }); -}); - -async function readStreamToString(reader) { - const decoder = new TextDecoder(); - let text = ""; - while (true) { - const { value, done } = await reader.read(); - if (done) break; - const decodedValue = decoder.decode(value); - if (decodedValue.includes("DONE")) break; - text += JSON.parse(decodedValue)[0]["choices"][0]["delta"]["content"]; - } - return text; -} diff --git a/JS/edgechains/examples/chat-with-llm/jsonnet/router.jsonnet b/JS/edgechains/examples/chat-with-llm/jsonnet/router.jsonnet new file mode 100644 index 000000000..641d20fc7 --- /dev/null +++ b/JS/edgechains/examples/chat-with-llm/jsonnet/router.jsonnet @@ -0,0 +1,21 @@ +// Routing config for the litellm-style Router, kept in jsonnet (the way +// Edgechains prefers to manage configuration). Add more deployments to load +// balance across providers/keys; the Router picks the one below its rate limit +// with the fewest tokens used. + +local openai_api_key = std.extVar("openai_api_key"); + +{ + strategy: "usage-based", + numRetries: 2, + cooldownSeconds: 60, + deployments: [ + { + model: "gpt-3.5-turbo", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 100000, + }, + ], +} diff --git a/JS/edgechains/examples/chat-with-llm/src/lib/generateResponse.cts b/JS/edgechains/examples/chat-with-llm/src/lib/generateResponse.cts index a33426693..e94f6faa6 100644 --- a/JS/edgechains/examples/chat-with-llm/src/lib/generateResponse.cts +++ b/JS/edgechains/examples/chat-with-llm/src/lib/generateResponse.cts @@ -1,4 +1,6 @@ -const { OpenAI } = require("@arakoodev/edgechains.js/ai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); +const path = require("path"); +const Jsonnet = require("@arakoodev/jsonnet"); import { z } from "zod"; const schema = z.object({ @@ -30,8 +32,13 @@ const schema = z.object({ async function openAICall({ prompt, openAIApiKey }: any) { try { - const openai = new OpenAI({ apiKey: openAIApiKey }); - let res = await openai.zodSchemaResponse({ prompt, schema: schema }); + const jsonnet = new Jsonnet(); + jsonnet.extString("openai_api_key", openAIApiKey || ""); + const config = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) + ); + const router = new Router(config); + let res = await router.zodSchemaResponse({ prompt, schema }); return JSON.stringify(res); } catch (error) { return error; diff --git a/JS/edgechains/examples/chat-with-pdf/jsonnet/router.jsonnet b/JS/edgechains/examples/chat-with-pdf/jsonnet/router.jsonnet new file mode 100644 index 000000000..4dda9dc9e --- /dev/null +++ b/JS/edgechains/examples/chat-with-pdf/jsonnet/router.jsonnet @@ -0,0 +1,29 @@ +// Routing config for the litellm-style Router, kept in jsonnet (the way +// Edgechains prefers to manage configuration). Add more deployments to load +// balance across providers/keys; the Router picks the one below its rate limit +// with the fewest tokens used. This example needs both a chat model (answers) +// and an embeddings model (vector search), so it declares one deployment each. + +local openai_api_key = std.extVar("openai_api_key"); + +{ + strategy: "usage-based", + numRetries: 2, + cooldownSeconds: 60, + deployments: [ + { + model: "gpt-3.5-turbo", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 100000, + }, + { + model: "text-embedding-ada-002", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 1000000, + }, + ], +} diff --git a/JS/edgechains/examples/chat-with-pdf/src/lib/generateResponse.cts b/JS/edgechains/examples/chat-with-pdf/src/lib/generateResponse.cts index 601a040b2..00ed52ec7 100644 --- a/JS/edgechains/examples/chat-with-pdf/src/lib/generateResponse.cts +++ b/JS/edgechains/examples/chat-with-pdf/src/lib/generateResponse.cts @@ -1,5 +1,5 @@ const path = require("path"); -const { OpenAI } = require("@arakoodev/edgechains.js/openai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); const z = require("zod"); const Jsonnet = require("@arakoodev/jsonnet"); const jsonnet = new Jsonnet(); @@ -7,7 +7,11 @@ const jsonnet = new Jsonnet(); const secretsPath = path.join(__dirname, "../../jsonnet/secrets.jsonnet"); const openAIApiKey = JSON.parse(jsonnet.evaluateFile(secretsPath)).openai_api_key; -const openai = new OpenAI({ apiKey: openAIApiKey }); +jsonnet.extString("openai_api_key", openAIApiKey || ""); +const routerConfig = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) +); +const router = new Router(routerConfig); const schema = z.object({ answer: z.string().describe("The answer to the question"), @@ -16,9 +20,11 @@ const schema = z.object({ function openAICall() { return function (prompt: string) { try { - return openai.zodSchemaResponse({ prompt, schema }).then((res: any) => { - return JSON.stringify(res); - }); + return router + .zodSchemaResponse({ prompt, schema, model: "gpt-3.5-turbo" }) + .then((res: any) => { + return JSON.stringify(res); + }); } catch (error) { return error; } diff --git a/JS/edgechains/examples/chat-with-pdf/src/lib/getEmbeddings.cts b/JS/edgechains/examples/chat-with-pdf/src/lib/getEmbeddings.cts index 680e87aa8..7e60077ef 100644 --- a/JS/edgechains/examples/chat-with-pdf/src/lib/getEmbeddings.cts +++ b/JS/edgechains/examples/chat-with-pdf/src/lib/getEmbeddings.cts @@ -1,4 +1,4 @@ -const { OpenAI } = require("@arakoodev/edgechains.js/openai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); const path = require("path"); const Jsonnet = require("@arakoodev/jsonnet"); @@ -7,15 +7,19 @@ const jsonnet = new Jsonnet(); const secretsPath = path.join(__dirname, "../../jsonnet/secrets.jsonnet"); const openAIApiKey = JSON.parse(jsonnet.evaluateFile(secretsPath)).openai_api_key; -const llm = new OpenAI({ - apiKey: openAIApiKey, -}); +jsonnet.extString("openai_api_key", openAIApiKey || ""); +const routerConfig = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) +); +const router = new Router(routerConfig); function getEmbeddings() { return (content: any) => { - const embeddings = llm.generateEmbeddings(content).then((res: any) => { - return JSON.stringify(res); - }); + const embeddings = router + .embedding({ input: content, model: "text-embedding-ada-002" }) + .then((res: any) => { + return JSON.stringify(res.data); + }); return embeddings; }; } diff --git a/JS/edgechains/examples/getWeather-or-time-function-calling/jsonnet/router.jsonnet b/JS/edgechains/examples/getWeather-or-time-function-calling/jsonnet/router.jsonnet new file mode 100644 index 000000000..0927d68c2 --- /dev/null +++ b/JS/edgechains/examples/getWeather-or-time-function-calling/jsonnet/router.jsonnet @@ -0,0 +1,21 @@ +// Routing config for the litellm-style Router, kept in jsonnet (the way +// Edgechains prefers to manage configuration). Add more deployments to load +// balance across providers/keys; the Router picks the one below its rate limit +// with the fewest tokens used. + +local openai_api_key = std.extVar("openai_api_key"); + +{ + strategy: "usage-based", + numRetries: 2, + cooldownSeconds: 60, + deployments: [ + { + model: "gpt-3.5-turbo-0613", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 100000, + }, + ], +} diff --git a/JS/edgechains/examples/getWeather-or-time-function-calling/src/lib/openAIChat.cts b/JS/edgechains/examples/getWeather-or-time-function-calling/src/lib/openAIChat.cts index 088f675c7..6750344f3 100644 --- a/JS/edgechains/examples/getWeather-or-time-function-calling/src/lib/openAIChat.cts +++ b/JS/edgechains/examples/getWeather-or-time-function-calling/src/lib/openAIChat.cts @@ -1,4 +1,4 @@ -const { OpenAI } = require("@arakoodev/edgechains.js/openai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); const path = require("path"); const Jsonnet = require("@arakoodev/jsonnet"); @@ -7,22 +7,23 @@ const jsonnet = new Jsonnet(); const secretsPath = path.join(__dirname, "../../jsonnet/secrets.jsonnet"); const apiKey = JSON.parse(jsonnet.evaluateFile(secretsPath)).openai_api_key; -const openai = new OpenAI({ - apiKey, -}); +jsonnet.extString("openai_api_key", apiKey || ""); +const routerConfig = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) +); +const router = new Router(routerConfig); function openAIChat() { return (prompt: string) => { try { - const completion = openai - .chat({ - model: "gpt-3.5-turbo-0613", + const completion = router + .completion({ messages: [ { role: "user", content: "Summarize the following input." + prompt }, ], }) - .then((completion: any) => { - return JSON.stringify(completion); + .then((res: any) => { + return JSON.stringify({ content: res.content }); }) .catch((error: any) => { console.error(error); diff --git a/JS/edgechains/examples/getWeather-or-time-function-calling/src/lib/openAIFunction.cts b/JS/edgechains/examples/getWeather-or-time-function-calling/src/lib/openAIFunction.cts index 96ae6f7ec..8bb3261d4 100644 --- a/JS/edgechains/examples/getWeather-or-time-function-calling/src/lib/openAIFunction.cts +++ b/JS/edgechains/examples/getWeather-or-time-function-calling/src/lib/openAIFunction.cts @@ -1,4 +1,4 @@ -const { OpenAI } = require("@arakoodev/edgechains.js/openai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); const path = require("path"); const Jsonnet = require("@arakoodev/jsonnet"); @@ -7,9 +7,11 @@ const jsonnet = new Jsonnet(); const secretsPath = path.join(__dirname, "../../jsonnet/secrets.jsonnet"); const apiKey = JSON.parse(jsonnet.evaluateFile(secretsPath)).openai_api_key; -const openai = new OpenAI({ - apiKey, -}); +jsonnet.extString("openai_api_key", apiKey || ""); +const routerConfig = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) +); +const router = new Router(routerConfig); interface messageOption { prompt: string; @@ -19,15 +21,18 @@ interface messageOption { function openAIFunction() { return ({ prompt, functions }: messageOption) => { try { - const completion = openai - .chatWithFunction({ - model: "gpt-3.5-turbo-0613", + const completion = router + .completion({ messages: [{ role: "user", content: prompt }], + max_tokens: 1024, functions, function_call: "auto", }) - .then((completion: any) => { - return JSON.stringify(completion); + .then((res: any) => { + return JSON.stringify({ + content: res.content || null, + function_call: res.functionCall, + }); }) .catch((error: any) => { console.error(error); diff --git a/JS/edgechains/examples/language-translater/jsonnet/router.jsonnet b/JS/edgechains/examples/language-translater/jsonnet/router.jsonnet new file mode 100644 index 000000000..641d20fc7 --- /dev/null +++ b/JS/edgechains/examples/language-translater/jsonnet/router.jsonnet @@ -0,0 +1,21 @@ +// Routing config for the litellm-style Router, kept in jsonnet (the way +// Edgechains prefers to manage configuration). Add more deployments to load +// balance across providers/keys; the Router picks the one below its rate limit +// with the fewest tokens used. + +local openai_api_key = std.extVar("openai_api_key"); + +{ + strategy: "usage-based", + numRetries: 2, + cooldownSeconds: 60, + deployments: [ + { + model: "gpt-3.5-turbo", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 100000, + }, + ], +} diff --git a/JS/edgechains/examples/language-translater/src/lib/generateResponse.cts b/JS/edgechains/examples/language-translater/src/lib/generateResponse.cts index bf37e5159..e1398ee6f 100644 --- a/JS/edgechains/examples/language-translater/src/lib/generateResponse.cts +++ b/JS/edgechains/examples/language-translater/src/lib/generateResponse.cts @@ -1,13 +1,17 @@ import { z } from "zod"; const path = require("path"); -const { OpenAI } = require("@arakoodev/edgechains.js/openai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); const Jsonnet = require("@arakoodev/jsonnet"); const jsonnet = new Jsonnet(); const secretsPath = path.join(__dirname, "../../jsonnet/secrets.jsonnet"); const openAIApiKey = JSON.parse(jsonnet.evaluateFile(secretsPath)).openai_api_key; -const openai = new OpenAI({ apiKey: openAIApiKey }); +jsonnet.extString("openai_api_key", openAIApiKey || ""); +const routerConfig = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) +); +const router = new Router(routerConfig); const schema = z.object({ answer: z.string().describe("The answer to the question"), @@ -16,7 +20,7 @@ const schema = z.object({ function openAICall() { return function (prompt: string) { try { - return openai.zodSchemaResponse({ prompt, schema: schema }).then((res: any) => { + return router.zodSchemaResponse({ prompt, schema: schema }).then((res: any) => { return JSON.stringify(res); }); } catch (error) { diff --git a/JS/edgechains/examples/react-chain/jsonnet/router.jsonnet b/JS/edgechains/examples/react-chain/jsonnet/router.jsonnet new file mode 100644 index 000000000..641d20fc7 --- /dev/null +++ b/JS/edgechains/examples/react-chain/jsonnet/router.jsonnet @@ -0,0 +1,21 @@ +// Routing config for the litellm-style Router, kept in jsonnet (the way +// Edgechains prefers to manage configuration). Add more deployments to load +// balance across providers/keys; the Router picks the one below its rate limit +// with the fewest tokens used. + +local openai_api_key = std.extVar("openai_api_key"); + +{ + strategy: "usage-based", + numRetries: 2, + cooldownSeconds: 60, + deployments: [ + { + model: "gpt-3.5-turbo", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 100000, + }, + ], +} diff --git a/JS/edgechains/examples/react-chain/src/lib/generateResponse.cts b/JS/edgechains/examples/react-chain/src/lib/generateResponse.cts index a55647a86..bf1d82a5a 100644 --- a/JS/edgechains/examples/react-chain/src/lib/generateResponse.cts +++ b/JS/edgechains/examples/react-chain/src/lib/generateResponse.cts @@ -1,14 +1,17 @@ -const { OpenAI } = require("@arakoodev/edgechains.js/openai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); +const path = require("path"); +const Jsonnet = require("@arakoodev/jsonnet"); async function openAICall({ prompt, apiKey }: any) { try { - const openai = new OpenAI({ - apiKey: apiKey, - temperature: 0, - }); - return openai.chat({ prompt }).then((res: any) => { - return res.content; - }); + const jsonnet = new Jsonnet(); + jsonnet.extString("openai_api_key", apiKey || ""); + const config = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) + ); + const router = new Router(config); + const res = await router.completion({ prompt }); + return res.content; } catch (error) { return error; } diff --git a/JS/edgechains/examples/research-agent/jsonnet/router.jsonnet b/JS/edgechains/examples/research-agent/jsonnet/router.jsonnet new file mode 100644 index 000000000..641d20fc7 --- /dev/null +++ b/JS/edgechains/examples/research-agent/jsonnet/router.jsonnet @@ -0,0 +1,21 @@ +// Routing config for the litellm-style Router, kept in jsonnet (the way +// Edgechains prefers to manage configuration). Add more deployments to load +// balance across providers/keys; the Router picks the one below its rate limit +// with the fewest tokens used. + +local openai_api_key = std.extVar("openai_api_key"); + +{ + strategy: "usage-based", + numRetries: 2, + cooldownSeconds: 60, + deployments: [ + { + model: "gpt-3.5-turbo", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 100000, + }, + ], +} diff --git a/JS/edgechains/examples/research-agent/src/lib/generateResponse.cts b/JS/edgechains/examples/research-agent/src/lib/generateResponse.cts index 6301a0fc4..654442584 100644 --- a/JS/edgechains/examples/research-agent/src/lib/generateResponse.cts +++ b/JS/edgechains/examples/research-agent/src/lib/generateResponse.cts @@ -1,9 +1,16 @@ -import { OpenAI } from "@arakoodev/edgechains.js/openai"; +const { Router } = require("@arakoodev/edgechains.js/ai"); +const path = require("path"); +const Jsonnet = require("@arakoodev/jsonnet"); async function openAICall({ prompt, openAIApiKey }: { prompt: string; openAIApiKey: string }) { try { - const openai = new OpenAI({ apiKey: openAIApiKey }); - const response = await openai.chat({ prompt, max_tokens: 2000 }); + const jsonnet = new Jsonnet(); + jsonnet.extString("openai_api_key", openAIApiKey || ""); + const config = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) + ); + const router = new Router(config); + const response = await router.completion({ prompt, max_tokens: 2000 }); return JSON.stringify(response.content); } catch (error) { return error; diff --git a/JS/edgechains/examples/resume-reviewer/jsonnet/router.jsonnet b/JS/edgechains/examples/resume-reviewer/jsonnet/router.jsonnet new file mode 100644 index 000000000..641d20fc7 --- /dev/null +++ b/JS/edgechains/examples/resume-reviewer/jsonnet/router.jsonnet @@ -0,0 +1,21 @@ +// Routing config for the litellm-style Router, kept in jsonnet (the way +// Edgechains prefers to manage configuration). Add more deployments to load +// balance across providers/keys; the Router picks the one below its rate limit +// with the fewest tokens used. + +local openai_api_key = std.extVar("openai_api_key"); + +{ + strategy: "usage-based", + numRetries: 2, + cooldownSeconds: 60, + deployments: [ + { + model: "gpt-3.5-turbo", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 100000, + }, + ], +} diff --git a/JS/edgechains/examples/resume-reviewer/src/lib/generateResponse.cts b/JS/edgechains/examples/resume-reviewer/src/lib/generateResponse.cts index c8290b006..ec1fe398f 100644 --- a/JS/edgechains/examples/resume-reviewer/src/lib/generateResponse.cts +++ b/JS/edgechains/examples/resume-reviewer/src/lib/generateResponse.cts @@ -1,5 +1,5 @@ const path = require("path"); -const { OpenAI } = require("@arakoodev/edgechains.js/openai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); import { z } from "zod"; const Jsonnet = require("@arakoodev/jsonnet"); const jsonnet = new Jsonnet(); @@ -7,7 +7,11 @@ const jsonnet = new Jsonnet(); const secretsPath = path.join(__dirname, "../../jsonnet/secrets.jsonnet"); const openAIApiKey = JSON.parse(jsonnet.evaluateFile(secretsPath)).openai_api_key; -const openai = new OpenAI({ apiKey: openAIApiKey }); +jsonnet.extString("openai_api_key", openAIApiKey || ""); +const routerConfig = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) +); +const router = new Router(routerConfig); const profileSchema = z.object({ fullTimeExperienceInMonths: z @@ -96,7 +100,7 @@ const profileSchema = z.object({ function openAICall() { return function (prompt: string) { try { - return openai.zodSchemaResponse({ prompt, schema: profileSchema }).then((res: any) => { + return router.zodSchemaResponse({ prompt, schema: profileSchema }).then((res: any) => { return JSON.stringify(res); }); } catch (error) { diff --git a/JS/edgechains/examples/rpa-challenge/jsonnet/router.jsonnet b/JS/edgechains/examples/rpa-challenge/jsonnet/router.jsonnet new file mode 100644 index 000000000..0927d68c2 --- /dev/null +++ b/JS/edgechains/examples/rpa-challenge/jsonnet/router.jsonnet @@ -0,0 +1,21 @@ +// Routing config for the litellm-style Router, kept in jsonnet (the way +// Edgechains prefers to manage configuration). Add more deployments to load +// balance across providers/keys; the Router picks the one below its rate limit +// with the fewest tokens used. + +local openai_api_key = std.extVar("openai_api_key"); + +{ + strategy: "usage-based", + numRetries: 2, + cooldownSeconds: 60, + deployments: [ + { + model: "gpt-3.5-turbo-0613", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 100000, + }, + ], +} diff --git a/JS/edgechains/examples/rpa-challenge/src/lib/generateTask.cts b/JS/edgechains/examples/rpa-challenge/src/lib/generateTask.cts index 7127883f0..6244a6bbf 100644 --- a/JS/edgechains/examples/rpa-challenge/src/lib/generateTask.cts +++ b/JS/edgechains/examples/rpa-challenge/src/lib/generateTask.cts @@ -1,4 +1,6 @@ -const { OpenAI } = require("@arakoodev/edgechains.js/ai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); +const path = require("path"); +const Jsonnet = require("@arakoodev/jsonnet"); function openAICall({ prompt, @@ -10,16 +12,21 @@ function openAICall({ openAIKey: string; }) { try { - const openai = new OpenAI({ apiKey: openAIKey }); - const completion = openai - .chatWithFunction({ - model: "gpt-3.5-turbo-0613", + const jsonnet = new Jsonnet(); + jsonnet.extString("openai_api_key", openAIKey || ""); + const config = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) + ); + const router = new Router(config); + const completion = router + .completion({ messages: [{ role: "user", content: prompt }], + max_tokens: 1024, functions, function_call: { name: functions[0].name }, }) - .then((completion: any) => { - return JSON.stringify(JSON.parse(completion.function_call.arguments).tasks); + .then((res: any) => { + return JSON.stringify(JSON.parse(res.functionCall.arguments).tasks); }) .catch((error: any) => { console.error(error); diff --git a/JS/edgechains/examples/summarize-page/jsonnet/router.jsonnet b/JS/edgechains/examples/summarize-page/jsonnet/router.jsonnet new file mode 100644 index 000000000..641d20fc7 --- /dev/null +++ b/JS/edgechains/examples/summarize-page/jsonnet/router.jsonnet @@ -0,0 +1,21 @@ +// Routing config for the litellm-style Router, kept in jsonnet (the way +// Edgechains prefers to manage configuration). Add more deployments to load +// balance across providers/keys; the Router picks the one below its rate limit +// with the fewest tokens used. + +local openai_api_key = std.extVar("openai_api_key"); + +{ + strategy: "usage-based", + numRetries: 2, + cooldownSeconds: 60, + deployments: [ + { + model: "gpt-3.5-turbo", + provider: "openai", + apiKey: openai_api_key, + rpm: 1000, + tpm: 100000, + }, + ], +} diff --git a/JS/edgechains/examples/summarize-page/src/lib/generateResponse.cts b/JS/edgechains/examples/summarize-page/src/lib/generateResponse.cts index ce20fb710..c5c5b1014 100644 --- a/JS/edgechains/examples/summarize-page/src/lib/generateResponse.cts +++ b/JS/edgechains/examples/summarize-page/src/lib/generateResponse.cts @@ -1,4 +1,6 @@ -const { OpenAI } = require("@arakoodev/edgechains.js/openai"); +const { Router } = require("@arakoodev/edgechains.js/ai"); +const path = require("path"); +const Jsonnet = require("@arakoodev/jsonnet"); import { z } from "zod"; const schema = z.object({ @@ -7,8 +9,13 @@ const schema = z.object({ async function openAICall({ prompt, openAIApiKey }: any) { try { - const openai = new OpenAI({ apiKey: openAIApiKey }); - let res = await openai.zodSchemaResponse({ prompt, schema: schema }); + const jsonnet = new Jsonnet(); + jsonnet.extString("openai_api_key", openAIApiKey || ""); + const config = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../../jsonnet/router.jsonnet")) + ); + const router = new Router(config); + let res = await router.zodSchemaResponse({ prompt, schema: schema }); return JSON.stringify(res); } catch (error) { return error;