File size: 3,847 Bytes
0d72ed4 732c7d6 0d72ed4 6022cc9 0d72ed4 732c7d6 0d72ed4 732c7d6 0d72ed4 6022cc9 0d72ed4 732c7d6 0d72ed4 6022cc9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """LLM provider registry and non-streaming call helpers."""
from __future__ import annotations
import httpx
class RateLimitError(Exception):
"""Raised when all attempted providers return HTTP 429."""
_PROVIDERS: dict[str, dict] = {
"groq": {"env": "GROQ_API_KEY", "model": "llama-3.3-70b-versatile"},
"mistral": {"env": "MISTRAL_API_KEY", "model": "codestral-latest"},
"gemini": {"env": "GEMINI_API_KEY", "model": "gemini-2.0-flash"},
"cohere": {"env": "COHERE_API_KEY", "model": "command-r-plus-08-2024"},
}
def get_provider_cfg(provider: str) -> dict:
return _PROVIDERS.get(provider, _PROVIDERS["groq"])
async def _call_groq(prompt: str, model: str, key: str) -> str:
async with httpx.AsyncClient(timeout=30) as client:
resp = await client.post(
"https://api.groq.com/openai/v1/chat/completions",
headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
json={
"model": model,
"messages": [{"role": "user", "content": prompt}],
"stream": False,
"max_tokens": 512,
"temperature": 0.1,
},
)
if resp.status_code == 429: raise RateLimitError("groq")
resp.raise_for_status()
return resp.json()["choices"][0]["message"]["content"]
async def _call_gemini(prompt: str, model: str, key: str) -> str:
url = (
f"https://generativelanguage.googleapis.com/v1beta/models/"
f"{model}:generateContent?key={key}"
)
async with httpx.AsyncClient(timeout=30) as client:
resp = await client.post(
url,
json={
"contents": [{"role": "user", "parts": [{"text": prompt}]}],
"generationConfig": {"maxOutputTokens": 512, "temperature": 0.1},
},
)
if resp.status_code == 429: raise RateLimitError("gemini")
resp.raise_for_status()
return resp.json()["candidates"][0]["content"]["parts"][0]["text"]
async def _call_mistral(prompt: str, model: str, key: str) -> str:
async with httpx.AsyncClient(timeout=30) as client:
resp = await client.post(
"https://api.mistral.ai/v1/chat/completions",
headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
json={
"model": model,
"messages": [{"role": "user", "content": prompt}],
"stream": False,
"max_tokens": 512,
"temperature": 0.1,
},
)
if resp.status_code == 429: raise RateLimitError("mistral")
resp.raise_for_status()
return resp.json()["choices"][0]["message"]["content"]
async def _call_cohere(prompt: str, model: str, key: str) -> str:
async with httpx.AsyncClient(timeout=30) as client:
resp = await client.post(
"https://api.cohere.ai/v2/chat",
headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
json={
"model": model,
"messages": [{"role": "user", "content": prompt}],
"stream": False,
"max_tokens": 512,
"temperature": 0.1,
},
)
if resp.status_code == 429: raise RateLimitError("cohere")
resp.raise_for_status()
return resp.json()["message"]["content"][0]["text"]
_CALLERS = {
"groq": _call_groq,
"mistral": _call_mistral,
"gemini": _call_gemini,
"cohere": _call_cohere,
}
async def call_provider(provider: str, prompt: str, model: str, key: str) -> str:
"""Dispatch a prompt to the named provider (unknown names fall back to groq)."""
fn = _CALLERS.get(provider, _call_groq)
return await fn(prompt, model, key)
|