Function prm_judge

scripts/run_memory_ablation.py:106–150

Rate response quality 0-10 via LLM-as-judge.

(context: str, question: str, response: str)

# PRM Judge (single call per response, no max_tokens)                #
# ------------------------------------------------------------------ #
async def prm_judge(context: str, question: str, response: str) -> dict:
    """Rate response quality 0-10 via LLM-as-judge."""
    prompt = f"""Rate this AI response on a scale of 0 to 10.
Context from previous sessions: {context}
User question: {question}
Response: {response[:800]}
Reply ONLY with JSON: {{"score": <0-10>, "reasoning": "<brief>"}}"""

    for attempt in range(3):
        try:
            async with httpx.AsyncClient(timeout=60.0) as client:
                resp = await client.post(
                    f"{AZURE_API_BASE}/chat/completions",
                    json={
                        "model": AZURE_MODEL,
                        "messages": [{"role": "user", "content": prompt}],
                        "temperature": 0.1,
                        "max_completion_tokens": 150,
                    },
                    headers={
                        "Authorization": f"Bearer {AZURE_API_KEY}",
                        "Content-Type": "application/json",
                    },
                )
                resp.raise_for_status()
                text = resp.json()["choices"][0]["message"]["content"].strip()
                if "```" in text:
                    text = text.split("```")[1]
                    if text.startswith("json"):
                        text = text[4:]
                result = json.loads(text)
                return {"score": float(result.get("score", 5)) / 10.0, "reasoning": result.get("reasoning", "")}
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 429:
                wait = 15 * (attempt + 1)
                logger.warning("  PRM judge rate limited, waiting %ds...", wait)
                await asyncio.sleep(wait)
                continue
            logger.warning("  PRM judge HTTP error: %s", e)
            return {"score": 0.5, "reasoning": f"http_error_{e.response.status_code}"}
        except Exception as e:
            logger.warning("  PRM judge error: %s", e)
            return {"score": 0.5, "reasoning": "parse_error"}

    return {"score": 0.5, "reasoning": "rate_limited"}


# ------------------------------------------------------------------ #