""" backend/rag/agent_workflow.py 4-Stage Agentic RAG Workflow ============================ Stage 1 — Planner : Expands 1 question into 3-4 targeted search queries Stage 2 — Searcher : Runs each query against FAISS, merges + deduplicates chunks Stage 3 — Validator : Scores & filters chunks, keeps only the most relevant ones Stage 4 — Synthesizer: Handled by existing generate_answer_stream (no extra call) Rate-limit strategy: - Stages 1 & 3 use llama-3.1-8b-instant (FREE tier, 6000 TPM, super fast) - Stage 4 (synthesis) uses llama-3.3-70b-versatile (main model, existing limit) - Total extra tokens per query: ~400 (planner) + ~300 (validator) = ~700 tokens - This is well within free Groq limits (~6000 tokens/min on 8b model) """ from __future__ import annotations import json import re import time from typing import Generator from groq import Groq from backend.config import GROQ_API_KEY, GROQ_MODEL from backend.rag.multi_retriever import MultiSourceResult, _build_source_groups from backend.rag.retriever import RetrievedChunk # ── Small fast model for planner & validator (higher rate limits, cheaper) ────── _FAST_MODEL = "llama-3.1-8b-instant" # ── Max chunks to keep after validation ───────────────────────────────────────── # Reduced from 12 to 8 to increase context density and reduce "lost in the middle" _MAX_FINAL_CHUNKS = 8 def _get_groq(api_key: str) -> Groq: return Groq(api_key=api_key) # ───────────────────────────────────────────────────────────────────────────── # STAGE 1 — PLANNER # ───────────────────────────────────────────────────────────────────────────── def _stage1_planner(question: str, is_legal: bool) -> list[str]: """ Uses a few-shot approach to generate 3 targeted queries. """ if is_legal: domain_examples = ( "Example 1: 'Does the IT Act 2000 apply to crypto?' -> ['Section 66 IT Act crypto', 'Indian crypto legality IPC', 'RBI digital asset circulars']\n" "Example 2: 'Bail for non-bailable offense' -> ['Section 437 CrPC conditions', 'Supreme Court bail guidelines', 'Anticipatory bail landmark cases']" ) system_prompt = ( "You are a Senior Legal Research Planner. Decompose the user question into 3 precise legal search queries. " "Focus on: Statutory sections, Case Law, and Procedural guidelines.\n" f"{domain_examples}\n" "Return ONLY a JSON array of 3 strings." ) else: domain_examples = ( "Example 1: 'How does BERT work?' -> ['BERT architecture transformer', 'Self-attention mechanism BERT', 'BERT pre-training MLM NSP']\n" "Example 2: 'Climate change impact on farming' -> ['Climate change crop yield statistics', 'Sustainable farming adaptations', 'Soil degradation carbon cycles']" ) system_prompt = ( "You are an expert Research Planner. Decompose the user question into 3 targeted search queries. " "Focus on: Technical foundations, Methodology, and Current Benchmarks.\n" f"{domain_examples}\n" "Return ONLY a JSON array of 3 strings." ) prompt = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": f"Question: {question}"}, ] try: client = _get_groq(GROQ_API_KEY) resp = client.chat.completions.create( model=_FAST_MODEL, messages=prompt, temperature=0.1, # Lower for consistency max_tokens=250, ) raw = resp.choices[0].message.content.strip() match = re.search(r"\[.*\]", raw, re.DOTALL) if match: queries = json.loads(match.group()) if isinstance(queries, list): # We always prioritize the original question all_queries = [question] + [str(q) for q in queries[:3]] return list(dict.fromkeys(all_queries))[:4] except Exception as e: print(f"[AgentWorkflow] Planner error: {e}") return [question] # ───────────────────────────────────────────────────────────────────────────── # STAGE 2 — SEARCHER # ───────────────────────────────────────────────────────────────────────────── def _stage2_searcher( queries: list[str], retriever_fn, source_ids: list[str] | None, ) -> list[RetrievedChunk]: seen_ids: set[str] = set() all_chunks: list[RetrievedChunk] = [] for i, q in enumerate(queries): try: result: MultiSourceResult = retriever_fn(q, source_ids) for chunk in result.all_chunks: if chunk.chunk_id not in seen_ids: seen_ids.add(chunk.chunk_id) # We slightly boost the original question's chunks if i == 0: chunk.score *= 1.1 all_chunks.append(chunk) except Exception: pass time.sleep(0.02) return all_chunks # ───────────────────────────────────────────────────────────────────────────── # STAGE 3 — VALIDATOR (HYBRID) # ───────────────────────────────────────────────────────────────────────────── def _stage3_validator( question: str, chunks: list[RetrievedChunk], max_chunks: int = _MAX_FINAL_CHUNKS, ) -> list[RetrievedChunk]: """ Combines LLM relevance scoring with Vector similarity (Hybrid). """ if not chunks: return [] if len(chunks) <= 3: return chunks # Too few to filter # Limit chunks to validate to avoid token limits (top 15 by vector score) candidate_chunks = sorted(chunks, key=lambda c: c.score, reverse=True)[:15] snippets = "\n".join( f"ID {i}: {c.chunk_text[:350].strip()}" for i, c in enumerate(candidate_chunks) ) prompt = [ { "role": "system", "content": ( "You are an Elite Document Validator. Rate the relevance of each document snippet to the question. " "Return a JSON object where keys are IDs and values are integer scores from 0 (useless) to 10 (perfect answer).\n" "Example: {'0': 9, '1': 2, '2': 7}" ), }, { "role": "user", "content": f"Question: {question}\n\nSnippets:\n{snippets}", }, ] try: client = _get_groq(GROQ_API_KEY) resp = client.chat.completions.create( model=_FAST_MODEL, messages=prompt, temperature=0.0, max_tokens=200, response_format={"type": "json_object"} ) llm_scores = json.loads(resp.choices[0].message.content) # Apply Hybrid Scoring: 70% LLM + 30% FAISS for i, chunk in enumerate(candidate_chunks): llm_val = float(llm_scores.get(str(i), llm_scores.get(i, 5))) # Normalize vector score (usually -10 to 10 or 0 to 1) # We assume a base score of 5 if LLM fails chunk.score = (llm_val * 1.5) + (chunk.score * 0.5) # Sort by hybrid score and take top N final = sorted(candidate_chunks, key=lambda c: c.score, reverse=True)[:max_chunks] return final except Exception as e: print(f"[AgentWorkflow] Validator error: {e}") return sorted(chunks, key=lambda c: c.score, reverse=True)[:max_chunks] # ───────────────────────────────────────────────────────────────────────────── # MAIN ENTRY POINT # ───────────────────────────────────────────────────────────────────────────── def run_agentic_workflow( question: str, retriever_fn, source_ids: list[str] | None = None, is_legal: bool = False, ) -> tuple[MultiSourceResult, list[str]]: status_log: list[str] = [] # Planner status_log.append("🧠 Planner: Thinking across multiple research dimensions...") queries = _stage1_planner(question, is_legal) status_log.append(f"🎯 Planner: Targeting {len(queries)} specific data angles") time.sleep(1) # Searcher status_log.append("📚 Searcher: Parallel retrieval in progress...") all_chunks = _stage2_searcher(queries, retriever_fn, source_ids) status_log.append(f"🔎 Searcher: Found {len(all_chunks)} potential evidence blocks") time.sleep(1) # Validator status_log.append("⚖️ Validator: Hybrid re-ranking for maximum precision...") validated_chunks = _stage3_validator(question, all_chunks) # Calculate a simple "Confidence" based on top chunk score confidence = "High" if len(validated_chunks) > 0 and validated_chunks[0].score > 12 else "Medium" status_log.append(f"✨ Validator: {confidence} confidence context finalized ({len(validated_chunks)} blocks)") source_groups = _build_source_groups(validated_chunks) result = MultiSourceResult( query_intent="synthesis", source_groups=source_groups, all_chunks=validated_chunks, source_count=len(source_groups), ) status_log.append("🖋️ Synthesizer: Drafting Professional Research Memorandum...") return result, status_log