Spaces:
Running
Running
File size: 5,700 Bytes
c650d43 | 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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | import json
import logging
from typing import List, Dict, Any
from src.utils.config import Config
from src.llm.llm_client import _get_llm_clients
LOGGER = logging.getLogger("grader")
def _run_json_grader(prompt: str, user_content: str, fallback_value: bool = True) -> bool:
"""Helper to run a fast JSON-mode LLM call for grading."""
last_exc = None
try:
clients = _get_llm_clients()
for client in clients:
try:
response = client.chat.completions.create(
model=Config.GRADER_LLM_MODEL,
response_format={"type": "json_object"},
messages=[
{"role": "system", "content": prompt},
{"role": "user", "content": user_content},
],
temperature=0.0,
)
content = response.choices[0].message.content
parsed = json.loads(content)
return bool(parsed.get("valid", fallback_value))
except Exception as e:
last_exc = e
continue
except Exception as e:
last_exc = e
LOGGER.error(f"Grader LLM failed: {last_exc}")
return fallback_value # Fail-open to avoid breaking the pipeline entirely
def grade_documents(query: str, docs: List[Dict[str, Any]]) -> bool:
"""Checks if the retrieved documents are relevant to the query."""
if not docs:
return False
# We can use the fast reranker score as a primary filter to save LLM calls
# Threshold lowered to 0.05: heuristic reranker scores are in 0.0–0.5 range,
# so 0.2 was incorrectly discarding valid docs before the LLM grader could evaluate them.
max_score = max([float(d.get("rerank_score") or d.get("score") or 0) for d in docs])
if max_score < 0.05:
return False
prompt = (
"You are a grader assessing relevance of retrieved documents to a user question. "
"Return JSON with a single key 'valid' set to true or false. "
"It is true if ANY document contains keywords or semantic meaning related to the user question "
"(e.g., if the user asks for 'revenue', documents mentioning 'sales' are highly relevant). "
"Ignore all markdown formatting and focus on the text content."
)
doc_text = "\n\n".join([
str(d.get("payload", {}).get("text") or d.get("snippet", ""))[:1200]
for d in docs
])
user_content = f"Question: {query}\n\nDocuments:\n{doc_text}"
return _run_json_grader(prompt, user_content)
def grade_hallucination(answer: str, docs: List[Dict[str, Any]]) -> bool:
"""Checks if the generated answer is grounded in the documents (no hallucinations)."""
if not docs or not answer:
return True # Nothing to grade
prompt = (
"You are a grader assessing whether an AI generation is grounded in a set of retrieved facts. "
"Return JSON with a single key 'valid' set to true or false. "
"It is true if the KEY FACTUAL CLAIMS in the generation are supported by the facts. "
"Ignore introductory phrases, hedging language, or formatting — focus only on whether the "
"core facts and numbers are backed by the retrieved context. "
"If the answer says it cannot find information, that is also valid (return true)."
)
doc_text = "\n\n".join([
str(d.get("payload", {}).get("text") or d.get("snippet", ""))[:1200]
for d in docs
])
user_content = f"Facts:\n{doc_text}\n\nGeneration: {answer}"
return _run_json_grader(prompt, user_content)
def grade_answer_relevance(query: str, answer: str) -> bool:
"""Checks if the generated answer actually resolves the user's question."""
if not answer:
return False
prompt = (
"You are an AI grader. Does the generated answer discuss the same general topic as the question? "
"Return JSON with a single key 'valid' set to true or false. "
"If the question is about Amazon revenue, and the answer discusses Amazon sales/revenue, it is valid (true). "
"If the question is about Meta AI, and the answer discusses Meta AI or infrastructure, it is valid (true). "
"Ignore formatting and charts. Output true if it is on-topic."
)
user_content = f"Question: {query}\n\nAnswer: {answer}"
return _run_json_grader(prompt, user_content)
def rewrite_query(query: str) -> str:
"""Rewrites a query to be better optimized for retrieval after a failure."""
last_exc = None
try:
clients = _get_llm_clients()
prompt = (
"You are a search query optimizer. The user's previous search failed to find relevant information. "
"Rewrite the user's query to be a better semantic search query. "
"Extract key entities and relationships. Do not include introductory text, just return the optimized query string."
)
for client in clients:
try:
response = client.chat.completions.create(
model=Config.GRADER_LLM_MODEL,
messages=[
{"role": "system", "content": prompt},
{"role": "user", "content": f"Original query: {query}"},
],
temperature=0.2,
)
return response.choices[0].message.content.strip()
except Exception as e:
last_exc = e
continue
except Exception as e:
last_exc = e
LOGGER.error(f"Rewrite LLM failed: {last_exc}")
return query
|