File size: 4,175 Bytes
eaae571 | 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 107 108 109 | from __future__ import annotations
import json
import os
import re
import time
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
from dotenv import load_dotenv
from prompt_injection_framework.models.base import BaseModelAdapter, ModelRequest
load_dotenv()
def _extract_retry_delay_seconds(detail: str, attempt: int) -> float:
match = re.search(r"Please try again in ([0-9]+(?:\.[0-9]+)?)(ms|s)", detail)
if match:
value = float(match.group(1))
unit = match.group(2)
seconds = value / 1000.0 if unit == "ms" else value
# compound routing models share TPM pools — enforce a longer minimum wait
return max(15.0, seconds + 2.0)
return min(10.0 * attempt, 60.0)
class GroqModelAdapter(BaseModelAdapter):
provider = "groq"
api_url = "https://api.groq.com/openai/v1/chat/completions"
max_retries = 8
def __init__(self, model_name: str, api_key: str | None = None) -> None:
super().__init__(model_name=model_name)
self.api_key = api_key or os.getenv("GROQ_APIKEY") or os.getenv("GROQ_API_KEY")
if not self.api_key:
raise ValueError(
"Missing Groq API key. Set GROQ_APIKEY or GROQ_API_KEY in the environment."
)
def _generate_text(self, request: ModelRequest) -> tuple[str, dict[str, object]]:
messages = []
if request.context:
messages.append({"role": "system", "content": request.context})
for turn in request.conversation_history:
messages.append({"role": turn.role, "content": turn.content})
user_parts = [request.prompt]
if request.task_input:
user_parts.append(request.task_input)
user_content = "\n\n".join(part for part in user_parts if part)
# groq/compound has a strict payload limit; truncate to ~8 000 chars to stay under 413
if len(user_content) > 8000:
user_content = user_content[:8000]
messages.append({"role": "user", "content": user_content})
payload = {
"model": self.model_name,
"messages": messages,
"temperature": request.temperature,
"max_completion_tokens": request.max_tokens,
}
if request.seed is not None:
payload["seed"] = request.seed
req = Request(
self.api_url,
data=json.dumps(payload).encode("utf-8"),
headers={
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
"Accept": "application/json",
"User-Agent": "llm-prompt-injection-security-eval-framework/0.1",
},
method="POST",
)
for attempt in range(1, self.max_retries + 1):
try:
with urlopen(req, timeout=60) as response:
raw_response = json.loads(response.read().decode("utf-8"))
break
except HTTPError as exc:
detail = exc.read().decode("utf-8", errors="replace")
if exc.code == 429 and attempt < self.max_retries:
time.sleep(_extract_retry_delay_seconds(detail, attempt))
continue
if exc.code == 413:
# Payload too large for this model — return a sentinel so the
# experiment records the case as failed rather than crashing.
return "[SKIPPED: payload_too_large]", {"error": "413_payload_too_large"}
raise RuntimeError(
f"Groq API request failed with status {exc.code}: {detail}"
) from exc
except URLError as exc:
raise RuntimeError(f"Groq API request failed: {exc.reason}") from exc
else:
raise RuntimeError("Groq API request failed after retry attempts.")
choices = raw_response.get("choices", [])
if not choices:
raise RuntimeError("Groq API returned no choices.")
text = choices[0].get("message", {}).get("content", "")
return text, raw_response
|