Spaces:
Running
Running
| """ | |
| Groq API Key Rotation Utility | |
| ────────────────────────────────────────────────────────────────────────────── | |
| Loads ALL env vars named GROQ_API_KEY or GROQ_API_KEY_* at startup. | |
| When any route hits a 429 rate-limit error, calling rotate() silently | |
| advances to the next key. Users never see a rate-limit error as long as | |
| at least one key still has remaining quota. | |
| Rotation is round-robin and thread-safe (uses a threading.Lock so concurrent | |
| FastAPI workers don't race on the index). | |
| Usage: | |
| from utils.groq_client import call_groq_json, HEAVY_MODEL, LIGHT_MODEL | |
| data = call_groq_json(system_prompt, user_prompt, json_mode=True) | |
| """ | |
| import os | |
| import json | |
| import time | |
| import threading | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| from groq import Groq | |
| # ─── E2E fixture mode ───────────────────────────────────────────────────────── | |
| # When E2E_FIXTURE_MODE=1, every call below returns a canned fixture instead of | |
| # hitting the real Groq API — no cost, deterministic, works offline. Used by | |
| # the E2E/surface-discovery test suite (see docs/qa/test-strategy.md and the | |
| # project-e2e-testing-design memory), never by production or by the unit tests | |
| # in backend/tests/ (those stay pure-function tests with no I/O). | |
| # Each call site passes a `fixture` key identifying which canned response it | |
| # wants; fixture mode with no key given is a bug at the call site, not a | |
| # silent fallback, so it raises loudly instead of guessing. | |
| FIXTURE_MODE = os.environ.get("E2E_FIXTURE_MODE") == "1" | |
| _FIXTURES_DIR = Path(__file__).resolve().parent.parent / "tests" / "fixtures" / "groq" | |
| # DEMO_MODE=1 (with fixture mode, used by `npm run demo`): every AI answer reads | |
| # "Test Environment". Only fields the app needs to keep working are left alone. | |
| DEMO_MODE = FIXTURE_MODE and os.environ.get("DEMO_MODE") == "1" | |
| DEMO_TEXT = "Test Environment" | |
| _DEMO_KEEP_KEYS = {"id", "type", "column", "targetDate", "date", "status"} | |
| def _demo_value(value, key=None): | |
| if isinstance(value, dict): | |
| return {k: _demo_value(v, k) for k, v in value.items()} | |
| if isinstance(value, list): | |
| strings = [v for v in value if isinstance(v, str)] | |
| if strings and len(strings) == len(value): | |
| # Numbered so choice lists and React keys stay distinct | |
| return [f"{DEMO_TEXT} {i + 1}" for i in range(len(value))] if len(value) > 1 else [DEMO_TEXT] | |
| return [_demo_value(v, key) for v in value] | |
| if isinstance(value, str) and key not in _DEMO_KEEP_KEYS: | |
| return DEMO_TEXT | |
| return value | |
| def _load_json_fixture(name: str) -> dict: | |
| path = _FIXTURES_DIR / f"{name}.json" | |
| if not path.exists(): | |
| raise RuntimeError( | |
| f"[E2E_FIXTURE_MODE] no fixture file for '{name}' at {path} — " | |
| "add one under backend/tests/fixtures/groq/" | |
| ) | |
| data = json.loads(path.read_text(encoding="utf-8")) | |
| return _demo_value(data) if DEMO_MODE else data | |
| def _load_text_fixture(name: str) -> str: | |
| path = _FIXTURES_DIR / f"{name}.txt" | |
| if not path.exists(): | |
| raise RuntimeError( | |
| f"[E2E_FIXTURE_MODE] no fixture file for '{name}' at {path} — " | |
| "add one under backend/tests/fixtures/groq/" | |
| ) | |
| return DEMO_TEXT if DEMO_MODE else path.read_text(encoding="utf-8") | |
| # ─── Load all keys ──────────────────────────────────────────────────────────── | |
| # Accepts: GROQ_API_KEY, GROQ_API_KEY_1 ... GROQ_API_KEY_N (any suffix) | |
| _KEYS: list[str] = [] | |
| for k in sorted(os.environ.keys()): | |
| if k == 'GROQ_API_KEY' or k.startswith('GROQ_API_KEY_'): | |
| v = os.environ[k].strip() | |
| if v: | |
| _KEYS.append(v) | |
| if not _KEYS: | |
| raise RuntimeError( | |
| "No Groq API keys found. " | |
| "Set GROQ_API_KEY (and optionally GROQ_API_KEY_7, etc.) in your .env file." | |
| ) | |
| print(f"[groq] {len(_KEYS)} API key(s) loaded for rotation") | |
| # ─── Model names ────────────────────────────────────────────────────────────── | |
| # llama-3.3-70b-versatile / llama-3.1-8b-instant were retired by Groq — replaced | |
| # with the current openai/gpt-oss lineup (see Groq's models page for the live list). | |
| HEAVY_MODEL = "openai/gpt-oss-120b" # reasoning, profile generation, consultation | |
| LIGHT_MODEL = "openai/gpt-oss-20b" # chat, quick responses | |
| TRANSCRIBE_MODEL = "whisper-large-v3-turbo" | |
| # ─── Round-robin state ──────────────────────────────────────────────────────── | |
| _idx = 0 | |
| _lock = threading.Lock() | |
| def _current_key() -> str: | |
| return _KEYS[_idx % len(_KEYS)] | |
| def _rotate() -> str: | |
| """Advance to the next key and return it. Thread-safe.""" | |
| global _idx | |
| with _lock: | |
| _idx = (_idx + 1) % len(_KEYS) | |
| new_key = _KEYS[_idx] | |
| print(f"[groq] Rotated to key index {_idx % len(_KEYS)} of {len(_KEYS)}") | |
| return new_key | |
| # ─── Core call ──────────────────────────────────────────────────────────────── | |
| def call_groq( | |
| system_prompt: str, | |
| user_prompt: str, | |
| model: str = HEAVY_MODEL, | |
| json_mode: bool = False, | |
| max_tokens: int = 1500, | |
| temperature: float = 0.7, | |
| ) -> str: | |
| """ | |
| Calls Groq with automatic key rotation on 429 errors. | |
| Tries every available key before giving up. | |
| Raises the last error if all keys are exhausted. | |
| Returns the raw string content of the first choice message. | |
| """ | |
| last_error: Exception | None = None | |
| # Try each key in the rotation at most once, plus one extra attempt | |
| # in case the first key has recovered by the time we loop around. | |
| for attempt in range(len(_KEYS) + 1): | |
| try: | |
| client = Groq(api_key=_current_key()) | |
| kwargs: dict = dict( | |
| model=model, | |
| messages=[ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": user_prompt}, | |
| ], | |
| temperature=temperature, | |
| max_tokens=max_tokens, | |
| ) | |
| if json_mode: | |
| kwargs["response_format"] = {"type": "json_object"} | |
| response = client.chat.completions.create(**kwargs) | |
| return response.choices[0].message.content # type: ignore | |
| except Exception as e: | |
| last_error = e | |
| err_str = str(e) | |
| is_rate_limit = ( | |
| '429' in err_str | |
| or 'rate_limit' in err_str.lower() | |
| or 'rate limit' in err_str.lower() | |
| ) | |
| if is_rate_limit and attempt < len(_KEYS): | |
| # Rotate and pause briefly before retrying | |
| _rotate() | |
| time.sleep(1) | |
| continue | |
| # Non-rate-limit errors bubble up immediately | |
| raise | |
| # All keys exhausted | |
| raise last_error or RuntimeError("All Groq API keys are rate-limited") | |
| def call_groq_transcribe( | |
| file: tuple, | |
| model: str = TRANSCRIBE_MODEL, | |
| response_format: str = "json", | |
| language: str | None = None, | |
| fixture: str | None = None, | |
| ): | |
| """ | |
| Calls Groq's Whisper transcription endpoint with the same key-rotation | |
| on 429 errors as the chat completion calls above. | |
| `file` is the (filename, bytes, content_type) tuple the Groq SDK expects. | |
| """ | |
| if FIXTURE_MODE: | |
| if not fixture: | |
| raise RuntimeError( | |
| "[E2E_FIXTURE_MODE] call_groq_transcribe was called with no fixture key — " | |
| "every call site must pass one when fixture mode is on" | |
| ) | |
| data = _load_json_fixture(fixture) | |
| return SimpleNamespace(text=data.get("text", "")) | |
| last_error: Exception | None = None | |
| for attempt in range(len(_KEYS) + 1): | |
| try: | |
| client = Groq(api_key=_current_key()) | |
| kwargs: dict = dict(model=model, file=file, response_format=response_format) | |
| if language: | |
| kwargs["language"] = language | |
| return client.audio.transcriptions.create(**kwargs) | |
| except Exception as e: | |
| last_error = e | |
| err_str = str(e) | |
| is_rate_limit = ( | |
| '429' in err_str | |
| or 'rate_limit' in err_str.lower() | |
| or 'rate limit' in err_str.lower() | |
| ) | |
| if is_rate_limit and attempt < len(_KEYS): | |
| _rotate() | |
| time.sleep(1) | |
| continue | |
| raise | |
| raise last_error or RuntimeError("All Groq API keys are rate-limited") | |
| def call_groq_json( | |
| system_prompt: str, | |
| user_prompt: str, | |
| model: str = HEAVY_MODEL, | |
| max_tokens: int = 1500, | |
| fixture: str | None = None, | |
| ) -> dict: | |
| """Convenience wrapper: calls Groq in JSON mode and returns parsed dict.""" | |
| if FIXTURE_MODE: | |
| if not fixture: | |
| raise RuntimeError( | |
| "[E2E_FIXTURE_MODE] call_groq_json was called with no fixture key — " | |
| "every call site must pass one when fixture mode is on" | |
| ) | |
| return _load_json_fixture(fixture) | |
| raw = call_groq(system_prompt, user_prompt, model=model, | |
| json_mode=True, max_tokens=max_tokens) | |
| if not raw: | |
| # Reasoning models (e.g. openai/gpt-oss-*) can spend the token budget on | |
| # internal reasoning and return empty content if max_tokens is too tight. | |
| raise RuntimeError( | |
| "Groq returned empty content — likely truncated by max_tokens " | |
| f"(model={model}, max_tokens={max_tokens})" | |
| ) | |
| try: | |
| return json.loads(raw) | |
| except json.JSONDecodeError as e: | |
| raise RuntimeError(f"Groq returned invalid JSON (likely truncated): {e}") from e | |
| def _fixture_stream(text: str): | |
| """Yields fake chunks shaped like the Groq SDK's streaming response — | |
| chunk.choices[0].delta.content — so callers don't need to know they're | |
| in fixture mode.""" | |
| for word in text.split(" "): | |
| yield SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=word + " "))]) | |
| def call_groq_stream( | |
| system_prompt: str, | |
| user_prompt: str, | |
| model: str = LIGHT_MODEL, | |
| max_tokens: int = 512, | |
| fixture: str | None = None, | |
| messages: list[dict] | None = None, | |
| ): | |
| """ | |
| Returns a Groq streaming completion object. | |
| Rotates key on 429 like all other calls. | |
| `messages`, when given, is the full chat list (system prompt first) and | |
| replaces system_prompt/user_prompt — used by chat to send conversation | |
| history. Callers without history keep passing the two prompt strings. | |
| """ | |
| if FIXTURE_MODE: | |
| if not fixture: | |
| raise RuntimeError( | |
| "[E2E_FIXTURE_MODE] call_groq_stream was called with no fixture key — " | |
| "every call site must pass one when fixture mode is on" | |
| ) | |
| return _fixture_stream(_load_text_fixture(fixture)) | |
| last_error: Exception | None = None | |
| for attempt in range(len(_KEYS) + 1): | |
| try: | |
| client = Groq(api_key=_current_key()) | |
| return client.chat.completions.create( | |
| model=model, | |
| messages=messages or [ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": user_prompt}, | |
| ], | |
| temperature=0.7, | |
| max_tokens=max_tokens, | |
| stream=True, | |
| ) | |
| except Exception as e: | |
| last_error = e | |
| if ('429' in str(e) or 'rate_limit' in str(e).lower()) and attempt < len(_KEYS): | |
| _rotate() | |
| time.sleep(1) | |
| continue | |
| raise | |
| raise last_error or RuntimeError("All Groq API keys are rate-limited") | |