Skip to content

Commit 6c1a954

Browse files
Merge pull request #71 from Android-PowerUser/tackle-mistral-ai-rate-limit-issue
Introduce MistralRequestCoordinator and queueing for Mistral auto-screenshots
2 parents c04e2c2 + df82714 commit 6c1a954

4 files changed

Lines changed: 220 additions & 156 deletions

File tree

app/src/main/kotlin/com/google/ai/sample/ScreenCaptureApiClients.kt

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import kotlinx.serialization.json.JsonClassDiscriminator
1515
import kotlinx.serialization.modules.SerializersModule
1616
import kotlinx.serialization.modules.polymorphic
1717
import kotlinx.serialization.modules.subclass
18+
import com.google.ai.sample.network.MistralRequestCoordinator
1819
import okhttp3.MediaType.Companion.toMediaType
1920
import okhttp3.OkHttpClient
2021
import okhttp3.Request
@@ -70,7 +71,7 @@ data class ServiceMistralResponseMessage(
7071
val content: String
7172
)
7273

73-
internal suspend fun callMistralApi(modelName: String, apiKey: String, chatHistory: List<Content>, inputContent: Content): Pair<String?, String?> {
74+
internal suspend fun callMistralApi(modelName: String, apiKeys: List<String>, chatHistory: List<Content>, inputContent: Content): Pair<String?, String?> {
7475
var responseText: String? = null
7576
var errorMessage: String? = null
7677

@@ -126,10 +127,18 @@ internal suspend fun callMistralApi(modelName: String, apiKey: String, chatHisto
126127
.url("https://api.mistral.ai/v1/chat/completions")
127128
.post(jsonBody.toRequestBody(mediaType))
128129
.addHeader("Content-Type", "application/json")
129-
.addHeader("Authorization", "Bearer $apiKey")
130+
.addHeader("Authorization", "Bearer ${apiKeys.first()}")
130131
.build()
131132

132-
client.newCall(request).execute().use { response ->
133+
val coordinated = MistralRequestCoordinator.execute(apiKeys = apiKeys, maxAttempts = apiKeys.size * 4 + 8) { key ->
134+
client.newCall(
135+
request.newBuilder()
136+
.header("Authorization", "Bearer $key")
137+
.build()
138+
).execute()
139+
}
140+
141+
coordinated.response.use { response ->
133142
val responseBody = response.body?.string()
134143
if (!response.isSuccessful) {
135144
Log.e("ScreenCaptureService", "Mistral API Error ($response.code): $responseBody")

app/src/main/kotlin/com/google/ai/sample/ScreenCaptureService.kt

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -297,7 +297,11 @@ class ScreenCaptureService : Service() {
297297
if (apiProvider == ApiProvider.VERCEL) {
298298
responseText = callVercelApi(applicationContext, modelName, apiKey, chatHistoryDtos, inputContentDto)
299299
} else if (apiProvider == ApiProvider.MISTRAL) {
300-
val result = callMistralApi(modelName, apiKey, chatHistory, inputContent)
300+
val apiKeyManager = ApiKeyManager.getInstance(applicationContext)
301+
val availableKeys = apiKeyManager.getApiKeys(ApiProvider.MISTRAL)
302+
.filter { it.isNotBlank() }
303+
.distinct()
304+
val result = callMistralApi(modelName, availableKeys, chatHistory, inputContent)
301305
responseText = result.first
302306
errorMessage = result.second
303307
} else if (apiProvider == ApiProvider.PUTER) {

app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningViewModel.kt

Lines changed: 86 additions & 152 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import com.google.ai.sample.feature.multimodal.ModelDownloadManager
3434
import com.google.ai.sample.ModelOption
3535
import com.google.ai.sample.GenerativeAiViewModelFactory
3636
import com.google.ai.sample.InferenceBackend
37+
import com.google.ai.sample.network.MistralRequestCoordinator
3738
import com.google.ai.sample.feature.multimodal.dtos.toDto
3839
import com.google.ai.sample.feature.multimodal.dtos.TempFilePathCollector
3940
import kotlinx.coroutines.Dispatchers
@@ -70,8 +71,6 @@ import kotlinx.serialization.modules.subclass
7071
import com.google.ai.sample.webrtc.WebRTCSender
7172
import com.google.ai.sample.webrtc.SignalingClient
7273
import org.webrtc.IceCandidate
73-
import kotlin.math.max
74-
import kotlin.math.roundToLong
7574

7675
class PhotoReasoningViewModel(
7776
application: Application,
@@ -184,11 +183,14 @@ class PhotoReasoningViewModel(
184183
// to avoid re-executing already-executed commands
185184
private var incrementalCommandCount = 0
186185

187-
// Mistral rate limiting per API key (1.5 seconds between requests with same key)
188-
private val mistralNextAllowedRequestAtMsByKey = mutableMapOf<String, Long>()
189-
private var lastMistralTokenTimeMs = 0L
190-
private var lastMistralTokenKey: String? = null
191-
private val MISTRAL_MIN_INTERVAL_MS = 1500L
186+
private data class QueuedMistralScreenshotRequest(
187+
val bitmap: Bitmap,
188+
val screenshotUri: String,
189+
val screenInfo: String?
190+
)
191+
private val mistralAutoScreenshotQueueLock = Any()
192+
private var mistralAutoScreenshotInFlight = false
193+
private var queuedMistralScreenshotRequest: QueuedMistralScreenshotRequest? = null
192194

193195
// Accumulated full text during streaming for incremental command parsing
194196
private var streamingAccumulatedText = StringBuilder()
@@ -1136,129 +1138,17 @@ class PhotoReasoningViewModel(
11361138

11371139
// Validate that we have at least one key before proceeding
11381140
require(availableKeys.isNotEmpty()) { "No valid Mistral API keys available after filtering" }
1139-
1140-
fun markKeyCooldown(key: String, referenceTimeMs: Long) {
1141-
val nextAllowedAt = referenceTimeMs + MISTRAL_MIN_INTERVAL_MS
1142-
val existing = mistralNextAllowedRequestAtMsByKey[key] ?: 0L
1143-
mistralNextAllowedRequestAtMsByKey[key] = max(existing, nextAllowedAt)
1144-
}
1145-
1146-
fun markKeyCooldown(key: String, referenceTimeMs: Long, extraDelayMs: Long) {
1147-
val normalizedExtraDelay = extraDelayMs.coerceAtLeast(0L)
1148-
val nextAllowedAt = referenceTimeMs + max(MISTRAL_MIN_INTERVAL_MS, normalizedExtraDelay)
1149-
val existing = mistralNextAllowedRequestAtMsByKey[key] ?: 0L
1150-
mistralNextAllowedRequestAtMsByKey[key] = max(existing, nextAllowedAt)
1151-
}
1152-
1153-
fun remainingWaitForKeyMs(key: String, nowMs: Long): Long {
1154-
val nextAllowedAt = mistralNextAllowedRequestAtMsByKey[key] ?: 0L
1155-
return (nextAllowedAt - nowMs).coerceAtLeast(0L)
1156-
}
1157-
1158-
fun parseRetryAfterMs(headerValue: String?): Long? {
1159-
if (headerValue.isNullOrBlank()) return null
1160-
val seconds = headerValue.trim().toDoubleOrNull() ?: return null
1161-
return (seconds * 1000.0).roundToLong().coerceAtLeast(0L)
1162-
}
1163-
1164-
fun parseRateLimitResetDelayMs(response: okhttp3.Response, nowMs: Long): Long? {
1165-
val resetHeader = response.header("x-ratelimit-reset") ?: return null
1166-
val resetEpochSeconds = resetHeader.trim().toLongOrNull() ?: return null
1167-
val resetMs = resetEpochSeconds * 1000L
1168-
return (resetMs - nowMs).coerceAtLeast(0L)
1169-
}
1170-
1171-
fun adaptiveRetryDelayMs(failureCount: Int): Long {
1172-
val cappedExponent = (failureCount - 1).coerceIn(0, 5)
1173-
return 1000L shl cappedExponent // 1s, 2s, 4s, 8s, 16s, 32s
1174-
}
1175-
1176-
fun isRetryableMistralFailure(code: Int): Boolean {
1177-
return code == 429 || code >= 500
1178-
}
1179-
1180-
var response: okhttp3.Response? = null
1181-
var selectedKeyForResponse: String? = null
1182-
var consecutiveFailures = 0
1183-
var blockedKeysThisRound = mutableSetOf<String>()
1184-
11851141
val maxAttempts = availableKeys.size * 4 + 8
1186-
while (response == null && consecutiveFailures < maxAttempts) {
1187-
if (stopExecutionFlag.get()) break
1188-
1189-
val now = System.currentTimeMillis()
1190-
val keyPool = availableKeys.filter { it !in blockedKeysThisRound }.ifEmpty {
1191-
blockedKeysThisRound.clear()
1192-
availableKeys
1193-
}
1194-
1195-
val keyWithLeastWait = keyPool.minByOrNull { remainingWaitForKeyMs(it, now) } ?: availableKeys.first()
1196-
val waitMs = remainingWaitForKeyMs(keyWithLeastWait, now)
1197-
if (waitMs > 0L) {
1198-
delay(waitMs)
1142+
val coordinated = MistralRequestCoordinator.execute(
1143+
apiKeys = availableKeys,
1144+
maxAttempts = maxAttempts
1145+
) { selectedKey ->
1146+
if (stopExecutionFlag.get()) {
1147+
throw IOException("Mistral request aborted.")
11991148
}
1200-
1201-
val selectedKey = keyWithLeastWait
1202-
selectedKeyForResponse = selectedKey
1203-
1204-
try {
1205-
val attemptResponse = client.newCall(buildRequest(selectedKey)).execute()
1206-
val requestEndMs = System.currentTimeMillis()
1207-
val retryAfterMs = parseRetryAfterMs(attemptResponse.header("Retry-After"))
1208-
val resetDelayMs = parseRateLimitResetDelayMs(attemptResponse, requestEndMs)
1209-
val serverRequestedDelayMs = max(retryAfterMs ?: 0L, resetDelayMs ?: 0L)
1210-
markKeyCooldown(selectedKey, requestEndMs, serverRequestedDelayMs)
1211-
1212-
if (attemptResponse.isSuccessful) {
1213-
response = attemptResponse
1214-
break
1215-
}
1216-
1217-
val isRetryable = isRetryableMistralFailure(attemptResponse.code)
1218-
if (!isRetryable) {
1219-
val errBody = attemptResponse.body?.string()
1220-
attemptResponse.close()
1221-
throw IllegalStateException("Mistral Error ${attemptResponse.code}: $errBody")
1222-
}
1223-
1224-
attemptResponse.close()
1225-
blockedKeysThisRound.add(selectedKey)
1226-
consecutiveFailures++
1227-
val adaptiveDelay = adaptiveRetryDelayMs(consecutiveFailures)
1228-
markKeyCooldown(
1229-
selectedKey,
1230-
requestEndMs,
1231-
max(serverRequestedDelayMs, adaptiveDelay)
1232-
)
1233-
withContext(Dispatchers.Main) {
1234-
replaceAiMessageText(
1235-
"Mistral temporär nicht verfügbar (Versuch $consecutiveFailures/$maxAttempts). Warte auf Server-Rate-Limit und wiederhole...",
1236-
isPending = true
1237-
)
1238-
}
1239-
} catch (e: IOException) {
1240-
val requestEndMs = System.currentTimeMillis()
1241-
val adaptiveDelay = adaptiveRetryDelayMs(consecutiveFailures + 1)
1242-
markKeyCooldown(selectedKey, requestEndMs, adaptiveDelay)
1243-
blockedKeysThisRound.add(selectedKey)
1244-
consecutiveFailures++
1245-
if (consecutiveFailures >= maxAttempts) {
1246-
throw IOException("Mistral request failed after $maxAttempts attempts: ${e.message}", e)
1247-
}
1248-
withContext(Dispatchers.Main) {
1249-
replaceAiMessageText(
1250-
"Mistral Netzwerkfehler (Versuch $consecutiveFailures/$maxAttempts). Wiederhole...",
1251-
isPending = true
1252-
)
1253-
}
1254-
}
1255-
}
1256-
1257-
if (stopExecutionFlag.get()) {
1258-
throw IOException("Mistral request aborted.")
1149+
client.newCall(buildRequest(selectedKey)).execute()
12591150
}
1260-
1261-
val finalResponse = response ?: throw IOException("Mistral request failed after $maxAttempts attempts.")
1151+
val finalResponse = coordinated.response
12621152

12631153
if (!finalResponse.isSuccessful) {
12641154
val errBody = finalResponse.body?.string()
@@ -1268,27 +1158,12 @@ class PhotoReasoningViewModel(
12681158

12691159
val body = finalResponse.body ?: throw IOException("Empty response body from Mistral")
12701160
val aiResponseText = openAiStreamParser.parse(body) { accText ->
1271-
selectedKeyForResponse?.let { key ->
1272-
lastMistralTokenKey = key
1273-
lastMistralTokenTimeMs = System.currentTimeMillis()
1274-
markKeyCooldown(key, lastMistralTokenTimeMs)
1275-
} ?: run {
1276-
Log.w(TAG, "selectedKeyForResponse is null during streaming callback")
1277-
}
12781161
withContext(Dispatchers.Main) {
12791162
replaceAiMessageText(accText, isPending = true)
12801163
processCommandsIncrementally(accText)
12811164
}
12821165
}
12831166
finalResponse.close()
1284-
selectedKeyForResponse?.let { key ->
1285-
val reference = if (lastMistralTokenKey == key && lastMistralTokenTimeMs > 0L) {
1286-
lastMistralTokenTimeMs
1287-
} else {
1288-
System.currentTimeMillis()
1289-
}
1290-
markKeyCooldown(key, reference)
1291-
}
12921167

12931168
withContext(Dispatchers.Main) {
12941169
_uiState.value = PhotoReasoningUiState.Success(aiResponseText)
@@ -1306,6 +1181,7 @@ class PhotoReasoningViewModel(
13061181
}
13071182
} finally {
13081183
withContext(Dispatchers.Main) {
1184+
releaseAndDrainMistralAutoScreenshotQueue()
13091185
refreshStopButtonState()
13101186
}
13111187
}
@@ -2404,16 +2280,22 @@ private fun processCommands(text: String) {
24042280
_commandExecutionStatus.value = status
24052281
}
24062282

2407-
// Create prompt with screen information if available
2408-
val genericAnalysisPrompt = createGenericScreenshotPrompt()
2409-
2410-
// Re-send the query with only the latest screenshot
2411-
reason(
2412-
userInput = genericAnalysisPrompt,
2413-
selectedImages = listOf(bitmap),
2414-
screenInfoForPrompt = screenInfo,
2415-
imageUrisForChat = listOf(screenshotUri.toString()) // Add this argument
2416-
)
2283+
val currentModel = GenerativeAiViewModelFactory.getCurrentModel()
2284+
if (currentModel.apiProvider == ApiProvider.MISTRAL) {
2285+
enqueueMistralAutoScreenshotRequest(
2286+
bitmap = bitmap,
2287+
screenshotUri = screenshotUri.toString(),
2288+
screenInfo = screenInfo
2289+
)
2290+
} else {
2291+
// Re-send the query with only the latest screenshot
2292+
reason(
2293+
userInput = createGenericScreenshotPrompt(),
2294+
selectedImages = listOf(bitmap),
2295+
screenInfoForPrompt = screenInfo,
2296+
imageUrisForChat = listOf(screenshotUri.toString())
2297+
)
2298+
}
24172299

24182300
PhotoReasoningScreenshotUiNotifier.showAddedToConversation(context)
24192301
} else {
@@ -2436,5 +2318,57 @@ private fun processCommands(text: String) {
24362318
}
24372319
}
24382320
}
2321+
2322+
private fun enqueueMistralAutoScreenshotRequest(
2323+
bitmap: Bitmap,
2324+
screenshotUri: String,
2325+
screenInfo: String?
2326+
) {
2327+
val request = QueuedMistralScreenshotRequest(
2328+
bitmap = bitmap,
2329+
screenshotUri = screenshotUri,
2330+
screenInfo = screenInfo
2331+
)
2332+
var shouldStartNow = false
2333+
synchronized(mistralAutoScreenshotQueueLock) {
2334+
if (mistralAutoScreenshotInFlight) {
2335+
queuedMistralScreenshotRequest = request
2336+
Log.d(TAG, "Mistral auto screenshot request queued (latest wins).")
2337+
} else {
2338+
mistralAutoScreenshotInFlight = true
2339+
shouldStartNow = true
2340+
}
2341+
}
2342+
if (shouldStartNow) {
2343+
dispatchMistralAutoScreenshotRequest(request)
2344+
}
2345+
}
2346+
2347+
private fun dispatchMistralAutoScreenshotRequest(request: QueuedMistralScreenshotRequest) {
2348+
val genericAnalysisPrompt = createGenericScreenshotPrompt()
2349+
reasonWithMistral(
2350+
userInput = genericAnalysisPrompt,
2351+
selectedImages = listOf(request.bitmap),
2352+
screenInfoForPrompt = request.screenInfo,
2353+
imageUrisForChat = listOf(request.screenshotUri)
2354+
)
2355+
}
2356+
2357+
private fun releaseAndDrainMistralAutoScreenshotQueue() {
2358+
val nextRequest: QueuedMistralScreenshotRequest? = synchronized(mistralAutoScreenshotQueueLock) {
2359+
val queued = queuedMistralScreenshotRequest
2360+
if (queued == null) {
2361+
mistralAutoScreenshotInFlight = false
2362+
null
2363+
} else {
2364+
queuedMistralScreenshotRequest = null
2365+
queued
2366+
}
2367+
}
2368+
if (nextRequest != null) {
2369+
Log.d(TAG, "Draining queued Mistral auto screenshot request.")
2370+
dispatchMistralAutoScreenshotRequest(nextRequest)
2371+
}
2372+
}
24392373

24402374
}

0 commit comments

Comments
 (0)