import json import os import re import time import textwrap from dataclasses import dataclass, field from typing import Optional # for static type hints import numpy as np from openai import OpenAI from dotenv import load_dotenv from sentence_transformers import SentenceTransformer load_dotenv() GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY") JUDGE_MODEL_ID = "gemma-3-4b-it" # Lightweight bi-encoder for relevancy cosine similarity. # all-MiniLM-L6-v2 is 80MB, fast on CPU, well-suited for short sentence similarity. SIM_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" @dataclass class ClaimVerification: claim: str supported: bool reasoning: str = "" # brief LLM explanation @dataclass class FaithfulnessResult: score: float # 0–1: fraction of claims supported claims: list[ClaimVerification] = field(default_factory=list) total: int = 0 supported_count: int = 0 @dataclass class RelevancyResult: score: float # 0–1: mean cosine sim of 3 generated Qs vs original Q generated_questions: list[str] = field(default_factory=list) similarities: list[float] = field(default_factory=list) @dataclass class EvalResult: query: str answer: str faithfulness: FaithfulnessResult relevancy: RelevancyResult @property def faithfulness_score(self) -> float: return round(self.faithfulness.score, 4) @property def relevancy_score(self) -> float: return round(self.relevancy.score, 4) # Sliding-window rate limiter for the Gemini API class RateLimiter: """ Tracks API calls and token usage in a rolling 60-second window and sleeps the minimum required time before any call that would breach either limit. Why sliding window instead of a fixed "reset every minute" bucket? A fixed bucket is fragile — if you fire 14 calls at t=0:59 and 14 more at t=1:01 you've used 28 calls in 2 seconds while technically staying inside two separate minute buckets. A sliding window always looks at the last 60 seconds, so the constraint is continuously enforced regardless of where you are in the clock. Gemini free tier hard limits: 30 RPM, 15 000 TPM. We set effective limits slightly below to absorb timing jitter. """ EFFECTIVE_RPM: int = 28 # hard limit: 30 EFFECTIVE_TPM: int = 14_000 # hard limit: 15 000 WINDOW: float = 60.0 # seconds def __init__(self) -> None: # Each entry: (unix_timestamp_of_call, tokens_consumed_by_that_call) self._history: list[tuple[float, int]] = [] # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ def _prune(self) -> None: """Drop entries older than WINDOW seconds.""" cutoff = time.time() - self.WINDOW self._history = [(t, tok) for t, tok in self._history if t > cutoff] def _window_stats(self) -> tuple[int, int]: """Return (call_count, token_count) within the current window.""" self._prune() calls = len(self._history) tokens = sum(tok for _, tok in self._history) return calls, tokens # ------------------------------------------------------------------ # Public interface # ------------------------------------------------------------------ def wait_if_needed(self, estimated_tokens: int) -> None: """ Block until there is headroom for one more call that will consume approximately *estimated_tokens* tokens. Called BEFORE every _llm() invocation. """ while True: calls, tokens = self._window_stats() rpm_ok = calls < self.EFFECTIVE_RPM tpm_ok = tokens + estimated_tokens <= self.EFFECTIVE_TPM if rpm_ok and tpm_ok: return # safe to proceed # Sleep until the oldest entry slides out of the window. # Adding a 300 ms buffer avoids immediately re-entering the loop # due to floating-point timing imprecision. if self._history: oldest_ts = self._history[0][0] sleep_secs = (oldest_ts + self.WINDOW) - time.time() + 0.3 sleep_secs = max(0.3, sleep_secs) else: sleep_secs = 1.0 print(f"[RateLimiter] Window full ({calls} calls, {tokens} tokens). " f"Sleeping {sleep_secs:.1f}s...") time.sleep(sleep_secs) def record(self, tokens_used: int) -> None: """ Record a completed call. Called AFTER every _llm() invocation with the actual token count from the API response (or the pre-call estimate if the response did not include usage data). """ self._history.append((time.time(), tokens_used)) @staticmethod def estimate_tokens(prompt: str, max_response_tokens: int) -> int: """ Rough upper-bound estimate of total tokens for a call. Rule of thumb: ~4 chars per token for English text. We add max_response_tokens to account for the output side of the TPM budget, since Gemini counts both input and output tokens. """ input_estimate = max(1, len(prompt) // 4) return input_estimate + max_response_tokens class RAGJudge: def __init__(self, verbose: bool = True): if not GOOGLE_API_KEY: raise ValueError("GOOGLE_API_KEY not found.") self.verbose = verbose self._log("Loading judge LLM client...") self.client = OpenAI( base_url="https://generativelanguage.googleapis.com/v1beta/openai/", api_key=GOOGLE_API_KEY, ) self._log(f"Loading sentence similarity model: {SIM_MODEL_NAME}...") self.sim_model = SentenceTransformer(SIM_MODEL_NAME) # One shared rate limiter for all LLM calls made by this judge instance self.rate_limiter = RateLimiter() self._log("Judge ready.") def _log(self, msg: str) -> None: if self.verbose: print(f"[Judge] {msg}") def _llm(self, prompt: str, max_tokens: int = 512, temperature: float = 0.0) -> str: """ Single-turn LLM call with integrated rate limiting. Flow: 1. Estimate token cost of this call (prompt tokens + max response tokens). 2. Ask the rate limiter to sleep if the window is too full. 3. Make the API call. 4. Record the actual token count (from response.usage if present, otherwise fall back to our pre-call estimate). temperature=0.0 for deterministic judge outputs. System role not supported by Gemma chat template — merged into user turn. """ estimated_tokens = RateLimiter.estimate_tokens(prompt, max_tokens) # Block here if needed — this is the only sleep in the whole eval loop self.rate_limiter.wait_if_needed(estimated_tokens) try: response = self.client.chat.completions.create( model=JUDGE_MODEL_ID, messages=[{"role": "user", "content": prompt}], max_tokens=max_tokens, temperature=temperature, top_p=1.0, ) # Use actual token count when the API returns it; fall back to estimate. # Gemini's OpenAI-compat endpoint populates response.usage reliably. actual_tokens = ( response.usage.total_tokens if response.usage and response.usage.total_tokens else estimated_tokens ) self.rate_limiter.record(actual_tokens) return (response.choices[0].message.content or "").strip() except Exception as e: # Still record estimated usage so the limiter doesn't think the # window is emptier than it really is after a failed call. self.rate_limiter.record(estimated_tokens) self._log(f"LLM call failed: {e}") return "" def _extract_claims(self, answer: str) -> list[str]: """ Step 1 of faithfulness: extract atomic factual claims from the answer, in JSON output """ prompt = textwrap.dedent(f""" You are a precise fact-extractor. Read the following answer and extract every distinct factual claim it makes. A claim is a short declarative sentence that asserts one specific fact. Do NOT include opinions, hedges ("may", "could"), or meta-statements. Return ONLY a valid JSON array of strings. No preamble, no explanation. Example output: ["The shot clock is 24 seconds in the NBA.", "A player must release the ball before the buzzer sounds."] ANSWER: {answer} JSON array of claims: """).strip() raw = self._llm(prompt, max_tokens=400) # Parse JSON, falling back gracefully try: # Strip markdown code fences if present clean = re.sub(r"```(?:json)?|```", "", raw).strip() claims = json.loads(clean) if isinstance(claims, list): return [str(c).strip() for c in claims if str(c).strip()] except (json.JSONDecodeError, ValueError): pass # Fallback: split on newlines / numbered lines lines = [re.sub(r"^\s*[\d\.\-\*]+\s*", "", line).strip() for line in raw.split("\n")] return [l for l in lines if len(l) > 15] def _verify_claim(self, claim: str, context: str) -> ClaimVerification: """ Step 2 of faithfulness: verify one claim against the retrieved context. Returns a ClaimVerification with supported=True/False + a short reasoning. """ prompt = textwrap.dedent(f""" You are a reasonable and accurate fact-checker. CONTEXT (retrieved from a basketball rulebook / encyclopedia): {context} CLAIM TO VERIFY: {claim} TASK: Does the CONTEXT support this claim? Rules: - Answer "SUPPORTED" if the claim is explicitly stated OR reasonably implied by the context. - Allow for synonyms and paraphrasing (e.g., if the text says someone is a "post player", it supports the claim that they "play in the post"). - Answer "NOT SUPPORTED" only if the claim is entirely missing, directly contradicted, or requires outside knowledge. # - Do not use your own basketball knowledge — judge only from the CONTEXT above. Respond in this exact format (two lines): VERDICT: SUPPORTED or VERDICT: NOT SUPPORTED REASON: """).strip() raw = self._llm(prompt, max_tokens=120) verdict_match = re.search(r"VERDICT:\s*(SUPPORTED|NOT SUPPORTED)", raw, re.IGNORECASE) reason_match = re.search(r"REASON:\s*(.+)", raw, re.IGNORECASE) supported = bool(verdict_match and "NOT" not in verdict_match.group(1).upper()) reasoning = reason_match.group(1).strip() if reason_match else raw[:120] return ClaimVerification(claim=claim, supported=supported, reasoning=reasoning) def faithfulness(self, answer: str, context_chunks: list[dict]) -> FaithfulnessResult: """ Full faithfulness pipeline for one query. context_chunks: the list of retrieved chunk dicts The old time.sleep(1) between claim verifications has been removed — the rate limiter in _llm() now handles all pacing adaptively. """ # Concatenate all chunk texts into one context string for the verifier context_text = "\n\n---\n\n".join( c.get("text", "") for c in context_chunks if c.get("text") ) self._log("Extracting claims from answer...") claims = self._extract_claims(answer) self._log(f" Found {len(claims)} claims.") if not claims: return FaithfulnessResult(score=0.0, claims=[], total=0, supported_count=0) verifications = [] for i, claim in enumerate(claims): self._log(f" Verifying claim {i+1}/{len(claims)}: {claim[:60]}...") result = self._verify_claim(claim, context_text) # No sleep here — _llm() calls wait_if_needed() before every call verifications.append(result) supported = sum(1 for v in verifications if v.supported) score = supported / len(verifications) return FaithfulnessResult( score=score, claims=verifications, total=len(verifications), supported_count=supported, ) def _generate_questions(self, answer: str) -> list[str]: """ Step 1 of relevancy: generate 3 questions that the answer could be answering. These are then compared (via cosine similarity) to the original query. """ prompt = textwrap.dedent(f""" You are a question generator. Read the following answer about basketball rules or strategy. Generate exactly 3 different questions that this answer could be a response to. The questions should be specific, natural, and varied. Return ONLY a valid JSON array of exactly 3 question strings. No preamble. Example: ["What is the shot clock rule in the NBA?", "How long does a team have to shoot?", "What happens if the shot clock expires?"] ANSWER: {answer} JSON array of 3 questions: """).strip() raw = self._llm(prompt, max_tokens=200) try: clean = re.sub(r"```(?:json)?|```", "", raw).strip() questions = json.loads(clean) if isinstance(questions, list): qs = [str(q).strip() for q in questions if str(q).strip()] return qs[:3] except (json.JSONDecodeError, ValueError): pass # Fallback: line-by-line lines = [re.sub(r"^\s*[\d\.\-\*]+\s*", "", line).strip() for line in raw.split("\n")] qs = [l for l in lines if len(l) > 10 and "?" in l] return qs[:3] def relevancy(self, query: str, answer: str) -> RelevancyResult: """ Full relevancy pipeline for one query. Generates 3 questions from the answer, then measures cosine similarity between each generated question and the original query. """ self._log("Generating alternate questions for relevancy...") questions = self._generate_questions(answer) self._log(f" Generated {len(questions)} questions.") if not questions: return RelevancyResult(score=0.0, generated_questions=[], similarities=[]) # Embed original query + all generated questions in one batch all_texts = [query] + questions embeddings = self.sim_model.encode(all_texts, normalize_embeddings=True) query_vec = embeddings[0] # shape (384,) question_vecs = embeddings[1:] # shape (n, 384) # Cosine similarity — since embeddings are L2-normalised, sim = dot product similarities = [float(query_vec @ qv) for qv in question_vecs] mean_score = float(np.mean(similarities)) if similarities else 0.0 return RelevancyResult( score=mean_score, generated_questions=questions, similarities=similarities, ) # Combining it all: def evaluate(self, query: str, answer: str, context_chunks: list[dict]) -> EvalResult: self._log(f"\nEvaluating: {query[:60]}...") faith = self.faithfulness(answer, context_chunks) rel = self.relevancy(query, answer) self._log( f" Faithfulness: {faith.score:.2%} ({faith.supported_count}/{faith.total} claims supported)" ) self._log(f" Relevancy: {rel.score:.4f} (mean cosine sim)") return EvalResult( query=query, answer=answer, faithfulness=faith, relevancy=rel, ) TEST_QUERIES = [ # NBA rules "What is the defensive three-second rule in the NBA?", "What is the 'Restricted Area' on an NBA court and how far is it from the basket?", "What are the rules for a technical foul in the NBA?", "How does the shot clock reset work after an offensive rebound in the NBA?", # FIBA rules "What is the rule for a backcourt violation in FIBA?", "How many seconds does a player have to inbound the ball in FIBA?", # Strategy / tactics "Coach's Challenge: We have a very athletic team but we are undersized. Which press defense (1-2-1-1 or 1-2-2) should we use to maximize our speed for turnovers?", "How does a zone defense differ from man-to-man defense?", "What is the purpose of a box-and-one defense?", "Player Scenario: I'm guarding a 'slasher' in a 1v1 game. Should I play 'tight pressure' or 'sagging defense', and why?", ] _judge_singleton: Optional[RAGJudge] = None def get_judge() -> RAGJudge: global _judge_singleton if _judge_singleton is None: _judge_singleton = RAGJudge(verbose=True) return _judge_singleton # --------------------------------------------------------------------------- # Smoke test # --------------------------------------------------------------------------- # if __name__ == "__main__": # judge = RAGJudge(verbose=True) # test_answer = ( # "In the NBA, a player in the paint on defense cannot remain there for " # "more than three consecutive seconds while their team is in control of " # "the ball [Source 1]. This rule is called the defensive three-second " # "violation and results in a technical foul [Source 2]." # ) # test_query = "What is the defensive three seconds rule in the NBA?" # test_chunks = [ # {"text": "NBA Rule No. 26—Defensive Three-Second Rule: No defensive player may " # "remain in the lane for more than three consecutive seconds while the " # "ball is in play. Violation results in a technical foul."} # ] # result = judge.evaluate(test_query, test_answer, test_chunks) # print(f"\nFaithfulness: {result.faithfulness_score:.2%}") # for cv in result.faithfulness.claims: # mark = "✓" if cv.supported else "✗" # print(f" {mark} {cv.claim}") # print(f"\nRelevancy: {result.relevancy_score:.4f}") # for q, s in zip(result.relevancy.generated_questions, result.relevancy.similarities): # print(f" {s:.4f} — {q}")