Merge pull request #501 from punkpeye/add-polling

feat: poll usage

Author: mrubens

src/api/providers/glama.ts

@@ -110,26 +110,38 @@ export class GlamaHandler implements ApiHandler, SingleCompletionHandler {
 		}
 
 		try {
-			const response = await axios.get(
-				`https://glama.ai/api/gateway/v1/completion-requests/${completionRequestId}`,
-				{
-					headers: {
-						Authorization: `Bearer ${this.options.glamaApiKey}`,
-					},
-				},
-			)
+			let attempt = 0
 
-			const completionRequest = response.data
+			const delay = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms))
 
-			if (completionRequest.tokenUsage) {
-				yield {
-					type: "usage",
-					cacheWriteTokens: completionRequest.tokenUsage.cacheCreationInputTokens,
-					cacheReadTokens: completionRequest.tokenUsage.cacheReadInputTokens,
-					inputTokens: completionRequest.tokenUsage.promptTokens,
-					outputTokens: completionRequest.tokenUsage.completionTokens,
-					totalCost: parseFloat(completionRequest.totalCostUsd),
+			while (attempt++ < 10) {
+				// In case of an interrupted request, we need to wait for the upstream API to finish processing the request
+				// before we can fetch information about the token usage and cost.
+				const response = await axios.get(
+					`https://glama.ai/api/gateway/v1/completion-requests/${completionRequestId}`,
+					{
+						headers: {
+							Authorization: `Bearer ${this.options.glamaApiKey}`,
+						},
+					},
+				)
+
+				const completionRequest = response.data
+
+				if (completionRequest.tokenUsage && completionRequest.totalCostUsd) {
+					yield {
+						type: "usage",
+						cacheWriteTokens: completionRequest.tokenUsage.cacheCreationInputTokens,
+						cacheReadTokens: completionRequest.tokenUsage.cacheReadInputTokens,
+						inputTokens: completionRequest.tokenUsage.promptTokens,
+						outputTokens: completionRequest.tokenUsage.completionTokens,
+						totalCost: parseFloat(completionRequest.totalCostUsd),
+					}
+
+					break
 				}
+
+				await delay(200)
 			}
 		} catch (error) {
 			console.error("Error fetching Glama completion details", error)