File size: 5,268 Bytes
55a76a3 | 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 | import asyncio
import json
import re
from app.core.logging import logger
from app.infrastructure.llm.base import LLMClientBase
from app.workflow.prompts import BATCH_GRADING_PROMPT, GRADING_PROMPT
from app.workflow.state import GradedDoc, RAGState
def parse_grade(response: str) -> str:
"""Parse the relevance grade JSON response from the LLM."""
cleaned = response.strip()
if not cleaned:
return "irrelevant"
try:
# Strip potential markdown formatting wraps
if cleaned.startswith("```"):
cleaned = re.sub(r"^```(?:json)?\n", "", cleaned)
cleaned = re.sub(r"\n```$", "", cleaned)
cleaned = cleaned.strip()
# Parse JSON
parsed = json.loads(cleaned)
grade = parsed.get("grade", "irrelevant").strip().lower()
if grade in ("relevant", "irrelevant"):
return grade
return "irrelevant"
except Exception as e:
logger.warning(f"Relevance grade JSON parse failed: {e}. Raw response: {response}")
return "irrelevant"
def parse_batch_grades(response: str, num_chunks: int) -> dict[int, str]:
"""Parse the relevance grades JSON response from the LLM for batch grading."""
cleaned = response.strip()
# Default to irrelevant for all indices
default_grades = {i: "irrelevant" for i in range(num_chunks)}
if not cleaned:
return default_grades
try:
# Strip potential markdown formatting wraps
if cleaned.startswith("```"):
cleaned = re.sub(r"^```(?:json)?\n", "", cleaned)
cleaned = re.sub(r"\n```$", "", cleaned)
cleaned = cleaned.strip()
# Parse JSON
parsed = json.loads(cleaned)
# We expect a list of dicts in "grades" or direct list
grades_list = []
if isinstance(parsed, dict):
if "grades" in parsed:
grades_list = parsed["grades"]
elif "grade" in parsed:
# Mock fallback or single grade format: apply to all chunks
single_grade = parsed.get("grade", "irrelevant").strip().lower()
if single_grade in ("relevant", "irrelevant"):
return {i: single_grade for i in range(num_chunks)}
elif isinstance(parsed, list):
grades_list = parsed
grades_map = {}
for item in grades_list:
if isinstance(item, dict):
idx = item.get("chunk_index")
grade = item.get("grade", "irrelevant").strip().lower()
if idx is not None and grade in ("relevant", "irrelevant"):
grades_map[int(idx)] = grade
# Merge with defaults to ensure all indices are covered
for i in range(num_chunks):
if i not in grades_map:
grades_map[i] = "irrelevant"
return grades_map
except Exception as e:
logger.warning(f"Batch relevance grade JSON parse failed: {e}. Raw response: {response}")
return default_grades
def document_grading_node(llm: LLMClientBase):
"""Factory that creates document grading node function."""
async def node(state: RAGState) -> dict:
logger.info("Executing Node 3: Document Relevance Grading (Batch)...")
question = state["question"]
retrieved_docs = state.get("retrieved_docs", [])
if not retrieved_docs:
logger.info("No documents to grade.")
return {
"graded_docs": [],
"relevant_docs": [],
"should_fallback": state.get("retry_count", 0) >= state.get("max_retries", 2)
}
# Format the retrieved documents for batch grading
chunks_formatted = ""
for idx, doc in enumerate(retrieved_docs):
chunks_formatted += f"Chunk {idx}:\n---\n{doc.content}\n---\n\n"
prompt = BATCH_GRADING_PROMPT.format(question=question, chunks_formatted=chunks_formatted)
try:
response = await llm.ainvoke([{"role": "user", "content": prompt}], temperature=0.0)
grades_map = parse_batch_grades(response, len(retrieved_docs))
except Exception as e:
logger.error(f"Batch grading request failed: {e}")
grades_map = {i: "irrelevant" for i in range(len(retrieved_docs))}
graded_docs = []
for idx, doc in enumerate(retrieved_docs):
grade = grades_map.get(idx, "irrelevant")
graded_docs.append(GradedDoc(chunk=doc, grade=grade))
relevant_docs = [gd.chunk for gd in graded_docs if gd.grade == "relevant"]
logger.info(f"Grading complete: {len(relevant_docs)} out of {len(retrieved_docs)} chunks relevant.")
# Check if fallback is triggered
should_fallback = False
if len(relevant_docs) == 0 and state.get("retry_count", 0) >= state.get("max_retries", 2):
should_fallback = True
return {
"graded_docs": graded_docs,
"relevant_docs": relevant_docs,
"should_fallback": should_fallback
}
return node
|