Files
diegosouzapw--omniroute/src/lib/a2a/skills/quotaManagement.ts
T
2026-07-13 13:39:12 +08:00

122 lines
4.5 KiB
TypeScript

/**
* A2A Skill: Quota Management
*
* Handles natural-language queries about provider quotas and
* returns structured responses with actionable data.
*/
import type { A2ATask, TaskArtifact } from "../taskManager";
import { normalizeQuotaResponse } from "@/shared/contracts/quota";
import { resolveOmniRouteBaseUrl } from "@/shared/utils/resolveOmniRouteBaseUrl";
const OMNIROUTE_BASE_URL = resolveOmniRouteBaseUrl();
const OMNIROUTE_API_KEY = process.env.OMNIROUTE_API_KEY || "";
async function quotaFetch(path: string): Promise<any> {
const url = `${OMNIROUTE_BASE_URL}${path}`;
const headers: Record<string, string> = {
"Content-Type": "application/json",
...(OMNIROUTE_API_KEY ? { Authorization: `Bearer ${OMNIROUTE_API_KEY}` } : {}),
};
const res = await fetch(url, { headers, signal: AbortSignal.timeout(10000) });
if (!res.ok) throw new Error(`API [${res.status}]`);
return res.json();
}
export interface QuotaManagementResult {
artifacts: TaskArtifact[];
metadata: Record<string, unknown>;
}
export async function executeQuotaManagement(task: A2ATask): Promise<QuotaManagementResult> {
const query = task.input.messages[task.input.messages.length - 1]?.content?.toLowerCase() || "";
const [quotaRaw, combosRaw] = await Promise.allSettled([
quotaFetch("/api/usage/quota"),
quotaFetch("/api/combos"),
]);
const quota =
quotaRaw.status === "fulfilled"
? normalizeQuotaResponse(quotaRaw.value)
: normalizeQuotaResponse({});
const combos =
combosRaw.status === "fulfilled"
? Array.isArray((combosRaw.value as any)?.combos)
? (combosRaw.value as any).combos
: Array.isArray(combosRaw.value)
? (combosRaw.value as any[])
: []
: [];
const providers = quota.providers;
const availableQuota = (provider: { quotaTotal: number | null; quotaUsed: number }) => {
if (provider.quotaTotal === null) return Number.POSITIVE_INFINITY;
return provider.quotaTotal - provider.quotaUsed;
};
// Query classification
if (query.includes("ranking") || query.includes("most quota") || query.includes("best")) {
const sorted = [...providers].sort((a, b) => availableQuota(b) - availableQuota(a));
return {
artifacts: [
{
type: "text",
content: `**Quota Ranking (most available first):**\n${sorted
.map((p, i) => {
const remaining = availableQuota(p);
const remainingLabel =
remaining === Number.POSITIVE_INFINITY ? "unlimited" : remaining.toLocaleString();
const percentLabel =
p.quotaTotal === null ? "n/a" : `${Math.round(p.percentRemaining)}%`;
return `${i + 1}. **${p.provider}** — ${remainingLabel} remaining (${percentLabel})`;
})
.join("\n")}`,
},
],
metadata: { queryType: "ranking", providers: sorted.map((p) => p.provider) },
};
}
if (query.includes("free") || query.includes("suggest")) {
const freeCombos = combos.filter((c: any) => {
const name = (c.name || "").toLowerCase();
return name.includes("free") || name.includes("gratis");
});
return {
artifacts: [
{
type: "text",
content:
freeCombos.length > 0
? `**Free combos available:**\n${freeCombos.map((c: any) => `- **${c.name}** (ID: ${c.id})`).join("\n")}`
: "No free combos configured. Consider adding providers with free tiers (Gemini, Groq, etc.).",
},
],
metadata: { queryType: "free_suggestion", freeCombos: freeCombos.length },
};
}
// Default: general quota summary
const totalUsed = providers.reduce((sum, p) => sum + (p.quotaUsed || 0), 0);
const totalAvailable = providers.reduce((sum, p) => sum + (p.quotaTotal || 0), 0);
const warnings = providers.filter((p) => p.percentRemaining <= 10);
return {
artifacts: [
{
type: "text",
content: `**Quota Summary (${providers.length} providers):**\n- Total used: ${totalUsed.toLocaleString()} / ${totalAvailable.toLocaleString()}\n${providers
.map(
(p) =>
`- **${p.provider}:** ${p.quotaUsed.toLocaleString()} / ${p.quotaTotal?.toLocaleString() || "∞"} (${p.tokenStatus})`
)
.join(
"\n"
)}${warnings.length > 0 ? `\n\n⚠️ **Warning:** ${warnings.map((w) => w.provider).join(", ")} at or below 10% remaining` : ""}`,
},
],
metadata: { queryType: "summary", providerCount: providers.length, warnings: warnings.length },
};
}