feat: SmartRouter — load balancing, streaming, token tracking, logging (#286)

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

@@ -3,3 +3,5 @@ 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 { SmartRouter, OpenAIProvider, GeminiProvider, CohereProvider, TokenTracker, Logger } from "./lib/router/index.js";
+export type { DeploymentConfig, RouterChatOptions, RouterStreamChunk, TokenUsage, LogEvent, BaseProvider } from "./lib/router/index.js";

JS/edgechains/arakoodev/src/ai/src/lib/router/SmartRouter.ts

@@ -0,0 +1,169 @@
+import { DeploymentConfig, RouterChatOptions, RouterStreamChunk, TokenUsage, LogEvent, BaseProvider } from "./types.js";
+import { OpenAIProvider } from "./providers/openai.js";
+import { GeminiProvider } from "./providers/gemini.js";
+import { CohereProvider } from "./providers/cohere.js";
+import { TokenTracker } from "./middleware/tokenTracker.js";
+import { Logger } from "./middleware/logger.js";
+
+export class SmartRouter {
+  private deployments: BaseProvider[] = [];
+  private tokenTracker: TokenTracker;
+  private logger: Logger;
+  private rateLimitMap: Map<string, { count: number; resetAt: number }> = new Map();
+
+  constructor(configs: DeploymentConfig[]) {
+    this.tokenTracker = new TokenTracker();
+    this.logger = new Logger();
+
+    for (const config of configs) {
+      const provider = this.createProvider(config);
+      if (provider) this.deployments.push(provider);
+    }
+  }
+
+  private createProvider(config: DeploymentConfig): BaseProvider | null {
+    switch (config.provider) {
+      case "openai": return new OpenAIProvider(config);
+      case "gemini": return new GeminiProvider(config);
+      case "cohere": return new CohereProvider(config);
+      default: return null;
+    }
+  }
+
+  private selectDeployment(): BaseProvider {
+    const leastTokens = this.tokenTracker.getDeploymentWithLeastTokens();
+    const eligible = this.deployments.filter(d => {
+      const id = d.getDeploymentId();
+      const rl = this.rateLimitMap.get(id);
+      if (!rl) return true;
+      if (Date.now() > rl.resetAt) {
+        this.rateLimitMap.delete(id);
+        return true;
+      }
+      return rl.count < 60;
+    });
+
+    if (eligible.length === 0) {
+      const fallback = this.deployments[0];
+      this.logger.log({ type: "rate_limit", provider: fallback.name, model: "", timestamp: new Date().toISOString(), error: "All deployments rate-limited, using first" });
+      return fallback;
+    }
+
+    if (leastTokens) {
+      const best = eligible.find(d => d.getDeploymentId() === leastTokens);
+      if (best) {
+        this.logger.log({ type: "deployment_switch", provider: best.name, model: "", timestamp: new Date().toISOString(), deploymentId: best.getDeploymentId() });
+        return best;
+      }
+    }
+
+    return eligible[0];
+  }
+
+  private trackRateLimit(deploymentId: string): void {
+    const rl = this.rateLimitMap.get(deploymentId) || { count: 0, resetAt: Date.now() + 60000 };
+    rl.count++;
+    this.rateLimitMap.set(deploymentId, rl);
+  }
+
+  useLogger(cb: (event: LogEvent) => void | Promise<void>): void {
+    this.logger.use(cb);
+  }
+
+  getLogger(): Logger {
+    return this.logger;
+  }
+
+  getTokenTracker(): TokenTracker {
+    return this.tokenTracker;
+  }
+
+  async chat(options: RouterChatOptions): Promise<{ content: string; usage?: TokenUsage }> {
+    const deployment = this.selectDeployment();
+    const start = Date.now();
+    const deploymentId = deployment.getDeploymentId();
+
+    try {
+      this.trackRateLimit(deploymentId);
+      const result = await deployment.chat(options);
+      const duration = Date.now() - start;
+
+      if (result.usage) {
+        this.tokenTracker.record(deploymentId, result.usage);
+      }
+
+      this.logger.log({
+        type: "completion",
+        provider: deployment.name,
+        model: this.getModelName(options),
+        timestamp: new Date().toISOString(),
+        durationMs: duration,
+        tokens: result.usage,
+        deploymentId,
+      });
+
+      return result;
+    } catch (error: any) {
+      this.logger.log({
+        type: "error",
+        provider: deployment.name,
+        model: this.getModelName(options),
+        timestamp: new Date().toISOString(),
+        error: error.message,
+        deploymentId,
+      });
+
+      const fallback = this.deployments.find(d => d.getDeploymentId() !== deploymentId);
+      if (fallback) {
+        this.logger.log({ type: "deployment_switch", provider: fallback.name, model: this.getModelName(options), timestamp: new Date().toISOString(), deploymentId: fallback.getDeploymentId() });
+        return fallback.chat(options);
+      }
+      throw error;
+    }
+  }
+
+  async *streamChat(options: RouterChatOptions): AsyncGenerator<RouterStreamChunk> {
+    const deployment = this.selectDeployment();
+    const deploymentId = deployment.getDeploymentId();
+
+    try {
+      this.trackRateLimit(deploymentId);
+      const stream = deployment.streamChat(options);
+      let fullContent = "";
+      for await (const chunk of stream) {
+        fullContent += chunk.content;
+        yield chunk;
+      }
+
+      const usage: TokenUsage = { promptTokens: 0, completionTokens: fullContent.length / 4, totalTokens: fullContent.length / 4 };
+      this.tokenTracker.record(deploymentId, usage);
+
+      this.logger.log({
+        type: "completion",
+        provider: deployment.name,
+        model: this.getModelName(options),
+        timestamp: new Date().toISOString(),
+        tokens: usage,
+        deploymentId,
+      });
+    } catch (error: any) {
+      this.logger.log({
+        type: "error",
+        provider: deployment.name,
+        model: this.getModelName(options),
+        timestamp: new Date().toISOString(),
+        error: error.message,
+        deploymentId,
+      });
+      throw error;
+    }
+  }
+
+  getDeployments(): BaseProvider[] {
+    return this.deployments;
+  }
+
+  private getModelName(options: RouterChatOptions): string {
+    return options.model || "default";
+  }
+}

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

@@ -0,0 +1,10 @@
+import { SmartRouter } from "./SmartRouter.js";
+import { DeploymentConfig, RouterChatOptions, RouterStreamChunk, TokenUsage, LogEvent, BaseProvider } from "./types.js";
+import { OpenAIProvider } from "./providers/openai.js";
+import { GeminiProvider } from "./providers/gemini.js";
+import { CohereProvider } from "./providers/cohere.js";
+import { TokenTracker } from "./middleware/tokenTracker.js";
+import { Logger } from "./middleware/logger.js";
+
+export { SmartRouter, OpenAIProvider, GeminiProvider, CohereProvider, TokenTracker, Logger };
+export type { DeploymentConfig, RouterChatOptions, RouterStreamChunk, TokenUsage, LogEvent, BaseProvider };

JS/edgechains/arakoodev/src/ai/src/lib/router/middleware/logger.ts

@@ -0,0 +1,36 @@
+import { LogCallback, LogEvent } from "../types.js";
+
+export class Logger {
+  private callbacks: LogCallback[] = [];
+
+  constructor() {
+    this.callbacks = [];
+  }
+
+  use(cb: LogCallback): void {
+    this.callbacks.push(cb);
+  }
+
+  async log(event: LogEvent): Promise<void> {
+    for (const cb of this.callbacks) {
+      try {
+        await cb(event);
+      } catch {}
+    }
+  }
+
+  sentryLog(dsn?: string): LogCallback {
+    return (event: LogEvent) => {
+      if (event.type === "error" || event.type === "rate_limit") {
+        console.error(`[sentry] ${event.type}:`, event.provider, event.error || "");
+      }
+    };
+  }
+
+  posthogLog(apiKey?: string, host?: string): LogCallback {
+    return (event: LogEvent) => {
+      const level = event.type === "error" ? "error" : "info";
+      console.log(`[posthog] ${level}:`, event.provider, event.model, event.type);
+    };
+  }
+}

JS/edgechains/arakoodev/src/ai/src/lib/router/middleware/tokenTracker.ts

@@ -0,0 +1,41 @@
+import { TokenUsage } from "../types.js";
+
+export class TokenTracker {
+  private usage: Map<string, TokenUsage> = new Map();
+  private totalCost: number = 0;
+
+  record(deploymentId: string, usage: TokenUsage): void {
+    const existing = this.usage.get(deploymentId) || { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
+    existing.promptTokens += usage.promptTokens;
+    existing.completionTokens += usage.completionTokens;
+    existing.totalTokens += usage.totalTokens;
+    this.usage.set(deploymentId, existing);
+
+    const rate = 0.002 / 1000;
+    this.totalCost += usage.totalTokens * rate;
+  }
+
+  getUsage(deploymentId: string): TokenUsage {
+    return this.usage.get(deploymentId) || { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
+  }
+
+  getAllUsage(): Record<string, TokenUsage> {
+    return Object.fromEntries(this.usage);
+  }
+
+  getTotalCost(): number {
+    return this.totalCost;
+  }
+
+  getDeploymentWithLeastTokens(): string | null {
+    let min: string | null = null;
+    let minTokens = Infinity;
+    for (const [id, u] of this.usage) {
+      if (u.totalTokens < minTokens) {
+        minTokens = u.totalTokens;
+        min = id;
+      }
+    }
+    return min;
+  }
+}

JS/edgechains/arakoodev/src/ai/src/lib/router/providers/cohere.ts

@@ -0,0 +1,86 @@
+import axios from "axios";
+import { BaseProvider, DeploymentConfig, RouterChatOptions, RouterStreamChunk, TokenUsage } from "../types.js";
+
+export class CohereProvider extends BaseProvider {
+  readonly name = "cohere";
+  private config: DeploymentConfig;
+  private baseURL = "https://api.cohere.ai/v1";
+
+  constructor(config: DeploymentConfig) {
+    super();
+    this.config = config;
+  }
+
+  getDeploymentId(): string {
+    return `cohere:${this.config.model || "command-r"}`;
+  }
+
+  async chat(options: RouterChatOptions): Promise<{ content: string; usage?: TokenUsage }> {
+    const apiKey = this.config.apiKey || process.env.COHERE_API_KEY;
+    const response = await axios.post(
+      `${this.baseURL}/chat`,
+      {
+        model: this.config.model || "command-r",
+        message: options.prompt || options.messages?.map(m => m.content).join("\n") || "",
+        max_tokens: options.maxTokens || 256,
+        temperature: options.temperature ?? 0.7,
+      },
+      {
+        headers: {
+          Authorization: `Bearer ${apiKey}`,
+          "Content-Type": "application/json",
+        },
+        timeout: 30000,
+      }
+    );
+    const data = response.data;
+    return {
+      content: data.text ?? data.generations?.[0]?.text ?? "",
+      usage: data.meta?.billed_units
+        ? { promptTokens: data.meta.billed_units.input_tokens || 0, completionTokens: data.meta.billed_units.output_tokens || 0, totalTokens: (data.meta.billed_units.input_tokens || 0) + (data.meta.billed_units.output_tokens || 0) }
+        : undefined,
+    };
+  }
+
+  async *streamChat(options: RouterChatOptions): AsyncGenerator<RouterStreamChunk> {
+    const apiKey = this.config.apiKey || process.env.COHERE_API_KEY;
+    const response = await axios.post(
+      `${this.baseURL}/chat`,
+      {
+        model: this.config.model || "command-r",
+        message: options.prompt || options.messages?.map(m => m.content).join("\n") || "",
+        max_tokens: options.maxTokens || 256,
+        temperature: options.temperature ?? 0.7,
+        stream: true,
+      },
+      {
+        headers: {
+          Authorization: `Bearer ${apiKey}`,
+          "Content-Type": "application/json",
+        },
+        responseType: "stream",
+        timeout: 60000,
+        adapter: "fetch",
+      }
+    );
+
+    const stream = response.data;
+    let buffer = "";
+    for await (const chunk of stream) {
+      buffer += chunk.toString();
+      const lines = buffer.split("\n");
+      buffer = lines.pop() || "";
+      for (const line of lines) {
+        const trimmed = line.trim();
+        if (!trimmed) continue;
+        try {
+          const json = JSON.parse(trimmed);
+          const text = json.text || json.event?.text || "";
+          if (text) yield { content: text, done: false };
+          if (json.is_finished) yield { content: "", done: true };
+        } catch {}
+      }
+    }
+    yield { content: "", done: true };
+  }
+}