api-key-manager 2.1.0__py3-none-any.whl

key_manager/providers/mistral.py

@@ -0,0 +1,109 @@
+import asyncio
+import time
+from .base import ProviderBase, CheckResult, TestResult
+
+
+class MistralProvider(ProviderBase):
+    name = "mistral"
+    base_url = "https://api.mistral.ai"
+    check_endpoint = "/v1/models"
+
+    def build_headers(self, key: str) -> dict:
+        return {"Authorization": f"Bearer {key}"}
+
+    async def get_models(self, client, key: str) -> list[str]:
+        headers = self.build_headers(key)
+        try:
+            resp = await client.get(
+                f"{self.get_base_url()}{self.check_endpoint}",
+                headers=headers
+            )
+            if resp.status_code == 200:
+                data = resp.json()
+                if "data" in data:
+                    return [m["id"] for m in data["data"] if "id" in m]
+            return []
+        except Exception:
+            return []
+
+    async def check(self, client, key: str) -> CheckResult:
+        """Real usage test - try to make a minimal chat completion request."""
+        headers = self.build_headers(key)
+        headers["Content-Type"] = "application/json"
+        start = time.monotonic()
+        try:
+            resp = await client.post(
+                f"{self.get_base_url()}/v1/chat/completions",
+                headers=headers,
+                json={
+                    "model": "mistral-tiny",
+                    "messages": [{"role": "user", "content": "hi"}],
+                    "max_tokens": 5
+                }
+            )
+            latency = (time.monotonic() - start) * 1000
+
+            if resp.status_code == 200:
+                return CheckResult(True, 200, latency, None)
+            elif resp.status_code in (401, 403):
+                return CheckResult(False, resp.status_code, latency, "invalid key or forbidden")
+            elif resp.status_code == 429:
+                return CheckResult(False, 429, latency, "rate limited")
+            else:
+                try:
+                    data = resp.json()
+                    error_msg = data.get("error", {}).get("message", f"status {resp.status_code}")
+                except:
+                    error_msg = f"status {resp.status_code}"
+                return CheckResult(False, resp.status_code, latency, error_msg)
+        except Exception as e:
+            return CheckResult(False, None, (time.monotonic() - start) * 1000, str(e))
+
+    async def test_token_limit(self, client, key: str, token_steps: list[int]) -> TestResult:
+        headers = self.build_headers(key)
+        last_success = None
+        for step in token_steps:
+            try:
+                resp = await client.post(
+                    f"{self.get_base_url()}/v1/chat/completions",
+                    headers=headers,
+                    json={
+                        "model": "mistral-tiny",
+                        "messages": [{"role": "user", "content": "hi"}],
+                        "max_tokens": step
+                    }
+                )
+                if resp.status_code == 200:
+                    last_success = step
+                elif resp.status_code in (400, 413):
+                    break
+                elif resp.status_code == 429:
+                    await asyncio.sleep(1)
+                    continue
+                else:
+                    break
+            except Exception:
+                break
+        return TestResult(max_tokens=last_success)
+
+
+    async def check_real(self, client, key: str) -> CheckResult:
+        return await self.check(client, key)
+    async def test_concurrency(self, client, key: str, concurrency_steps: list[int]) -> TestResult:
+        headers = self.build_headers(key)
+        last_success = None
+        for step in concurrency_steps:
+            tasks = [self._probe(client, headers) for _ in range(step)]
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+            rate_limited = sum(1 for r in results if not isinstance(r, Exception) and not r)
+            if rate_limited / step >= 0.3:
+                break
+            last_success = step
+        return TestResult(max_concurrency=last_success)
+
+    async def _probe(self, client, headers) -> bool:
+        try:
+            resp = await client.get(f"{self.get_base_url()}{self.check_endpoint}", headers=headers)
+            return resp.status_code == 200
+        except Exception:
+            return False