Spaces:
Sleeping
Sleeping
File size: 2,189 Bytes
98ee05e e6fb135 98ee05e e6fb135 d0e810a e6fb135 | 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 | """Groq client wrapper and JSON parsing helpers.
One chat() helper is shared by the generator and the judge so model config
lives in one place. No LangChain.
"""
import json
import os
import re
from typing import Any, Optional
from groq import Groq
DEFAULT_MODEL = os.getenv("GROQ_MODEL", "openai/gpt-oss-120b")
_client: Optional[Groq] = None
def get_client() -> Groq:
"""Lazily construct a shared Groq client. Requires GROQ_API_KEY."""
global _client
if _client is None:
api_key = os.getenv("GROQ_API_KEY")
if not api_key:
raise RuntimeError(
"GROQ_API_KEY is not set. Copy .env.example to .env and add your key."
)
_client = Groq(api_key=api_key)
return _client
def chat(
prompt: str,
temperature: float = 0.7,
model: str | None = None,
max_tokens: int = 1024,
) -> str:
"""Single-shot chat completion returning the raw assistant text."""
client = get_client()
resp = client.chat.completions.create(
model=model or DEFAULT_MODEL,
messages=[{"role": "user", "content": prompt}],
temperature=temperature,
max_tokens=max_tokens,
)
return resp.choices[0].message.content or ""
def extract_json(text: str) -> dict[str, Any]:
"""Best-effort parse of a JSON object out of a model response.
Models occasionally wrap JSON in prose or markdown fences despite
instructions; we strip fences and fall back to the first {...} block.
"""
text = text.strip()
# Strip ```json ... ``` or ``` ... ``` fences if present.
fence = re.match(r"^```(?:json)?\s*(.*?)\s*```$", text, re.DOTALL)
if fence:
text = fence.group(1).strip()
try:
return json.loads(text)
except json.JSONDecodeError:
pass
# Fall back to the first balanced-looking object.
match = re.search(r"\{.*\}", text, re.DOTALL)
if match:
try:
return json.loads(match.group(0))
except json.JSONDecodeError as exc:
raise ValueError(f"Could not parse JSON from model output: {exc}\nRaw: {text}")
raise ValueError(f"No JSON object found in model output.\nRaw: {text}")
|