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

key_manager/providers/doubao.py

@@ -0,0 +1,109 @@
+import asyncio
+import time
+from .base import ProviderBase, CheckResult, TestResult
+
+
+class DoubaoProvider(ProviderBase):
+    name = "doubao"
+    base_url = "https://ark.cn-beijing.volces.com/api/v3"
+    check_endpoint = "/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()}/chat/completions",
+                headers=headers,
+                json={
+                    "model": "doubao-pro-32k",
+                    "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()}/chat/completions",
+                    headers=headers,
+                    json={
+                        "model": "doubao-pro-32k",
+                        "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

key_manager/providers/fireworks.py

@@ -0,0 +1,109 @@
+import asyncio
+import time
+from .base import ProviderBase, CheckResult, TestResult
+
+
+class FireworksProvider(ProviderBase):
+    name = "fireworks"
+    base_url = "https://api.fireworks.ai"
+    check_endpoint = "/inference/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()}/inference/v1/chat/completions",
+                headers=headers,
+                json={
+                    "model": "accounts/fireworks/models/llama-v3p1-8b-instruct",
+                    "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()}/inference/v1/chat/completions",
+                    headers=headers,
+                    json={
+                        "model": "accounts/fireworks/models/llama-v3p1-8b-instruct",
+                        "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

key_manager/providers/google.py

@@ -0,0 +1,99 @@
+import asyncio
+import time
+from .base import ProviderBase, CheckResult, TestResult
+
+
+class GoogleProvider(ProviderBase):
+    name = "google"
+    base_url = "https://generativelanguage.googleapis.com"
+    check_endpoint = "/v1beta/models"
+
+    def build_headers(self, key: str) -> dict:
+        return {}
+
+    async def get_models(self, client, key: str) -> list[str]:
+        try:
+            resp = await client.get(
+                f"{self.get_base_url()}{self.check_endpoint}?key={key}"
+            )
+            if resp.status_code == 200:
+                data = resp.json()
+                if "models" in data:
+                    return [m["name"].split("/")[-1] for m in data["models"] if "name" 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."""
+        start = time.monotonic()
+        try:
+            resp = await client.post(
+                f"{self.get_base_url()}/v1beta/models/gemini-1.5-flash:generateContent?key={key}",
+                json={
+                    "contents": [{"parts": [{"text": "hi"}]}],
+                    "generationConfig": {"maxOutputTokens": 5}
+                }
+            )
+            latency = (time.monotonic() - start) * 1000
+
+            if resp.status_code == 200:
+                return CheckResult(True, 200, latency, None)
+            elif resp.status_code in (400, 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:
+        last_success = None
+        for step in token_steps:
+            try:
+                resp = await client.post(
+                    f"{self.get_base_url()}/v1beta/models/gemini-1.5-flash:generateContent?key={key}",
+                    json={
+                        "contents": [{"parts": [{"text": "hi"}]}],
+                        "generationConfig": {"maxOutputTokens": 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:
+        last_success = None
+        for step in concurrency_steps:
+            tasks = [self._probe(client, key) 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, key) -> bool:
+        try:
+            resp = await client.get(f"{self.get_base_url()}{self.check_endpoint}?key={key}")
+            return resp.status_code == 200
+        except Exception:
+            return False

key_manager/providers/grok.py

@@ -0,0 +1,109 @@
+import asyncio
+import time
+from .base import ProviderBase, CheckResult, TestResult
+
+
+class GrokProvider(ProviderBase):
+    name = "grok"
+    base_url = "https://api.x.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": "grok-1",
+                    "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 check_real(self, client, key: str) -> CheckResult:
+        return await self.check(client, key)
+    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": "grok-2-1212",
+                        "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 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

key_manager/providers/groq.py

@@ -0,0 +1,109 @@
+import asyncio
+import time
+from .base import ProviderBase, CheckResult, TestResult
+
+
+class GroqProvider(ProviderBase):
+    name = "groq"
+    base_url = "https://api.groq.com"
+    check_endpoint = "/openai/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()}/openai/v1/chat/completions",
+                headers=headers,
+                json={
+                    "model": "llama3-8b-8192",
+                    "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()}/openai/v1/chat/completions",
+                    headers=headers,
+                    json={
+                        "model": "llama3-8b-8192",
+                        "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

key_manager/providers/huggingface.py

@@ -0,0 +1,54 @@
+import asyncio
+import time
+from .base import ProviderBase, CheckResult, TestResult
+
+
+class HuggingFaceProvider(ProviderBase):
+    name = "huggingface"
+    base_url = "https://huggingface.co"
+    check_endpoint = "/api/whoami-v2"
+
+    def build_headers(self, key: str) -> dict:
+        return {"Authorization": f"Bearer {key}"}
+
+    async def get_models(self, client, key: str) -> list[str]:
+        return []
+
+    async def check(self, client, key: str) -> CheckResult:
+        headers = self.build_headers(key)
+        start = time.monotonic()
+        try:
+            resp = await client.get(f"{self.get_base_url()}{self.check_endpoint}", headers=headers)
+            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")
+            elif resp.status_code == 429:
+                return CheckResult(True, 429, latency, "rate limited")
+            else:
+                return CheckResult(False, resp.status_code, latency, f"status {resp.status_code}")
+        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:
+        return TestResult(max_tokens=None)
+
+    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