| 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: |
| |
| if cleaned.startswith("```"): |
| cleaned = re.sub(r"^```(?:json)?\n", "", cleaned) |
| cleaned = re.sub(r"\n```$", "", cleaned) |
| cleaned = cleaned.strip() |
| |
| |
| 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_grades = {i: "irrelevant" for i in range(num_chunks)} |
| if not cleaned: |
| return default_grades |
| |
| try: |
| |
| if cleaned.startswith("```"): |
| cleaned = re.sub(r"^```(?:json)?\n", "", cleaned) |
| cleaned = re.sub(r"\n```$", "", cleaned) |
| cleaned = cleaned.strip() |
| |
| |
| parsed = json.loads(cleaned) |
| |
| |
| grades_list = [] |
| if isinstance(parsed, dict): |
| if "grades" in parsed: |
| grades_list = parsed["grades"] |
| elif "grade" in parsed: |
| |
| 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 |
| |
| |
| 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) |
| } |
| |
| |
| 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.") |
| |
| |
| 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 |
|
|
|
|
