OutstandingOm/knowledge-graph-env

1"""2LLM-as-a-Judge graders for Knowledge Graph Environment tasks.3 4Instead of brittle keyword matching, we send the agent's response to an LLM5and ask it to judge whether the issue was correctly addressed.  Falls back to6a simple keyword heuristic when no API key is available.7 8IMPORTANT: All scores MUST be strictly between 0 and 1 (exclusive).9"""10 11import os12 13# ── LLM client (lazy-initialised) ─────────────────────────────────────────────14_client = None15 16API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")17API_KEY = os.getenv("HF_TOKEN", "")18MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")19 20 21def _get_client():22    global _client23    if _client is None:24        from openai import OpenAI25        _client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)26    return _client27 28 29def _clamp(score: float) -> float:30    """Ensure score is strictly within (0, 1) — never 0.0 or 1.0."""31    return max(0.01, min(0.99, score))32 33 34def _llm_judge(agent_response: str, task_description: str, keywords: list) -> float:35    """36    Ask an LLM to score the agent's response on a 0-1 scale.37    Falls back to keyword matching if LLM is unavailable.38    Always returns a value strictly in (0, 1).39    """40    if not isinstance(agent_response, str) or not agent_response.strip():41        return _clamp(0.01)42 43    # Try LLM judge first44    if API_KEY:45        try:46            prompt = (47                "You are an expert evaluator for a customer support AI.\n"48                f"The customer's issue: {task_description}\n"49                f"The agent responded: {agent_response}\n\n"50                "Did the agent correctly identify and address the issue?\n"51                "Score the response from 0.0 (completely wrong) to 1.0 (perfect).\n"52                "Reply with ONLY a single decimal number, nothing else."53            )54            client = _get_client()55            resp = client.chat.completions.create(56                model=MODEL_NAME,57                messages=[58                    {"role": "system", "content": "You are an evaluation judge. Respond with only a number between 0.0 and 1.0."},59                    {"role": "user", "content": prompt},60                ],61                temperature=0.0,62                max_tokens=10,63            )64            raw = (resp.choices[0].message.content or "").strip()65            score = float(raw)66            return _clamp(score)67        except Exception:68            pass  # Fall through to keyword fallback69 70    # Keyword fallback71    return _keyword_fallback(agent_response, keywords)72 73 74def _keyword_fallback(text: str, keywords: list) -> float:75    """Simple keyword matching fallback. Always returns strictly in (0, 1)."""76    if not isinstance(text, str) or not text.strip():77        return _clamp(0.01)78 79    text_lower = text.lower().strip()80    if not keywords:81        return _clamp(0.01)82 83    matches = sum(1 for kw in keywords if kw.lower() in text_lower)84    score = 0.01 + (matches / len(keywords)) * 0.9885    return _clamp(score)86 87 88# ── Task definitions ──────────────────────────────────────────────────────────89 90TASK_DESCRIPTIONS = {91    "task_easy": "The user cannot log in to their account. Their password is not working and they keep getting locked out.",92    "task_medium": "The user's bill shows a double charge for their subscription. They need a refund for the extra payment.",93    "task_hard": "The user's account is locked after multiple failed password attempts. They suspect a security breach.",94}95 96# Keywords for fallback grading — short tokens that an LLM response would97# naturally contain if it addresses the issue correctly.98TASK_KEYWORDS = {99    "task_easy":   ["login", "account", "password", "access", "sign in", "authentication", "credential", "reset"],100    "task_medium": ["bill", "payment", "charge", "invoice", "refund", "subscription", "double", "overcharge"],101    "task_hard":   ["locked", "security", "breach", "blocked", "verify", "critical", "password", "unauthorized"],102}103 104 105def task_easy(input_text: str) -> float:106    return _llm_judge(input_text, TASK_DESCRIPTIONS["task_easy"], TASK_KEYWORDS["task_easy"])107 108 109def task_medium(input_text: str) -> float:110    return _llm_judge(input_text, TASK_DESCRIPTIONS["task_medium"], TASK_KEYWORDS["task_medium"])111 112 113def task_hard(input_text: str) -> float:114    return _llm_judge(input_text, TASK_DESCRIPTIONS["task_hard"], TASK_KEYWORDS["task_hard"])115 116 117TASKS = ["task_easy", "task_medium", "task_hard"]118GRADERS = {119    "task_easy": task_easy,120    "task_medium": task_medium,121    "task_hard": task_hard,122}123