import os import json import asyncio import itertools import time import pandas as pd from dotenv import load_dotenv from google import genai load_dotenv(override=True) # Initialize API client api_key = os.getenv("GOOGLE_API_KEY") client = None if api_key: client = genai.Client(api_key=api_key) else: print("Warning: GOOGLE_API_KEY is not set in environment.") MODELS = ["models/gemma-4-26b-a4b-it", "models/gemma-4-31b-it"] MIN_INTERVAL = 6.0 # 6 seconds between calls to same model (= 10 RPM safety pace) model_next_allowed_time = {m: 0.0 for m in MODELS} model_locks = {m: asyncio.Lock() for m in MODELS} model_cycle = itertools.cycle(MODELS) # Metrics tracking new_calls_count = 0 model_calls_tracker = {m: 0 for m in MODELS} async def call_model(prompt: str) -> str: """ Alternates between Gemma models and enforces 15 RPM per model rate limits using an async leaky bucket slot reservation system. """ global new_calls_count model_name = next(model_cycle) async with model_locks[model_name]: now = asyncio.get_event_loop().time() target_time = max(now, model_next_allowed_time[model_name]) delay = target_time - now model_next_allowed_time[model_name] = target_time + MIN_INTERVAL # Wait outside the lock so other tasks can reserve their slots concurrently if delay > 0: await asyncio.sleep(delay) if not client: raise ValueError("Google GenAI client is not configured (missing GOOGLE_API_KEY).") # Call SDK in executor thread response = await asyncio.to_thread( lambda: client.models.generate_content(model=model_name, contents=prompt) ) model_calls_tracker[model_name] += 1 new_calls_count += 1 return response.text.strip() def extract_final_summary(text: str) -> str: """ Extracts only the final summary paragraph from Gemma's output, ignoring thinking process, scratchpads, and bullet points. """ paragraphs = [p.strip() for p in text.split("\n") if p.strip()] if not paragraphs: return "" for p in reversed(paragraphs): p_clean = p.replace("```", "").strip() if not p_clean: continue # Skip markdown lists or analysis steps if p_clean.startswith(("*", "-", "1.", "2.", "3.", "Sentence ", "Draft ", "Option ", "Check ", "Topic:", "Source:", "Task:", "Constraint ")): continue if len(p_clean) > 80: return p_clean return paragraphs[-1].replace("```", "").strip() async def call_model_with_retry(prompt: str, retries=5) -> str: """ Helper with retry and exponential backoff, handling rate limits specifically. """ for attempt in range(retries): try: raw_text = await call_model(prompt) return extract_final_summary(raw_text) except Exception as e: err_msg = str(e) print(f"Error calling model (attempt {attempt+1}/{retries}): {e}") if "429" in err_msg or "Quota" in err_msg or "ResourceExhausted" in err_msg: backoff = 35 + (attempt * 15) print(f"Rate limit (429) hit. Backing off for {backoff} seconds...") await asyncio.sleep(backoff) else: await asyncio.sleep(2 ** attempt + 5) if attempt == retries - 1: raise e # Cache Management CACHE_FILE = "tree/summary_cache.json" summary_cache = {} def load_cache(): global summary_cache if os.path.exists(CACHE_FILE): try: with open(CACHE_FILE, "r", encoding="utf-8") as f: summary_cache = json.load(f) print(f"Loaded {len(summary_cache)} summaries from cache.") except Exception as e: print(f"Error loading cache: {e}. Starting fresh.") summary_cache = {} else: summary_cache = {} def save_cache(): os.makedirs(os.path.dirname(CACHE_FILE), exist_ok=True) with open(CACHE_FILE, "w", encoding="utf-8") as f: json.dump(summary_cache, f, indent=2, ensure_ascii=False) def format_schedule_row_summary(node): """ Auto-formats a one-line summary for schedule rows without LLM calls. """ content = node.get("content", "") lines = content.split("\n") data = {} for line in lines: if ":" in line: k, v = line.split(":", 1) data[k.strip().lower()] = v.strip() section = data.get("section", "N/A") offence = data.get("offence", "N/A") punishment = data.get("punishment", "N/A") offence_summary = offence[:80] + "..." if len(offence) > 80 else offence punishment_summary = punishment[:60] + "..." if len(punishment) > 60 else punishment return f"Section {section}: {offence_summary} — {punishment_summary}" def chunk_text(text, max_chars=14000): """ Groups paragraphs to stay safely under token limits, preserving paragraph integrity. """ paragraphs = text.split("\n") chunks = [] current_chunk = [] current_len = 0 for para in paragraphs: para_len = len(para) if current_len + para_len > max_chars and current_chunk: chunks.append("\n".join(current_chunk)) current_chunk = [para] current_len = para_len else: current_chunk.append(para) current_len += para_len + 1 if current_chunk: chunks.append("\n".join(current_chunk)) return chunks ACT_FULL_NAMES = { "BNS": "Bharatiya Nyaya Sanhita, 2023", "BNSS": "Bharatiya Nagarik Suraksha Sanhita, 2023", "BSA": "Bharatiya Sakshya Adhiniyam, 2023", "SOP": "Telangana Police Standard Operating Procedures" } def get_leaf_prompt(act_code, section_no, title, content, node_id): act_name = ACT_FULL_NAMES.get(act_code, act_code) special_instruction = "" if node_id == "BNSS_S193": special_instruction = ( "\nNote: This section may reference other BNS sections without incorporating their definitions. " "Do NOT describe referenced BNS sections as if they are defined inside this section.\n" ) if act_code == "SOP": return f"""You are a precise legal procedure summarizer for Indian police officers. Summarize the following Standard Operating Procedure of the {act_name} in 2–4 sentences. Rules: - Make the summary highly keyword-rich and dense, preserving all specific procedural steps, officer roles, timelines, and required actions. - Preserve exact numbers, timelines (e.g. days, hours), and forms verbatim. - If it references specific sections of BNSS, BNS, or BSA, name them explicitly. - Do not add interpretation or inference beyond what is stated. SOP {section_no}. {title}: {content}""" return f"""You are a precise legal summarizer. Summarize the following section of the {act_name} in 2–4 sentences. Rules: - Make the summary highly keyword-rich and dense, preserving all specific legal concepts, offences, actors, and mechanisms mentioned. - Preserve exact numbers, thresholds, punishment durations, and defined terms verbatim. - Do not paraphrase numeric penalties (e.g., "not less than seven years" must not become "several years"). - Do not add interpretation or inference beyond what is stated. - If the section contains Provisos, Exceptions, or Explanations, mention them explicitly.{special_instruction} Section {section_no}. {title}: {content}""" async def summarize_section(node): """ Summarizes a single section (leaf) node, supporting caching and chunked fallback. """ node_id = node["node_id"] stable_hash = node["metadata"]["stable_hash"] cache_key = f"{node_id}:{stable_hash}" if cache_key in summary_cache: node["summary"] = extract_final_summary(summary_cache[cache_key]) return act_code = node["metadata"]["act_code"] title = node["title"] content = node.get("content", "") # Extract section number from title section_no = node_id.split("_S")[-1] if "_S" in node_id else "" # Check if we need chunked fallback est_tokens = node["metadata"]["token_estimate"] try: if est_tokens > 4000: print(f"[{node_id}] Content is very long ({est_tokens} est. tokens). Running chunked summarization fallback...") chunks = chunk_text(content) chunk_summaries = [] for idx, chunk in enumerate(chunks): chunk_prompt = get_leaf_prompt(act_code, f"{section_no} (Part {idx+1})", title, chunk, node_id) summary = await call_model_with_retry(chunk_prompt) chunk_summaries.append(summary) # Merge summaries merged_content = "\n\n".join(chunk_summaries) merge_prompt = f"""You are a precise legal summarizer. The following are summaries of different parts of Section {section_no} of the {act_code} ({title}). Combine them into a single, cohesive, keyword-dense summary of 2-4 sentences. Part summaries: {merged_content}""" final_summary = await call_model_with_retry(merge_prompt) else: prompt = get_leaf_prompt(act_code, section_no, title, content, node_id) final_summary = await call_model_with_retry(prompt) except Exception as e: print(f"\nCRITICAL WARNING: Failed to summarize leaf {node_id}: {e}. Using fallback snippet.") snippet = content[:180].strip() if content else "" final_summary = f"Summary of Section {section_no or title}: {snippet}..." summary_cache[cache_key] = final_summary node["summary"] = final_summary save_cache() async def summarize_schedule_chapter(node): """ Summarizes the First Schedule chapter node (non-leaf description). """ node_id = node["node_id"] stable_hash = node["metadata"]["stable_hash"] cache_key = f"{node_id}:{stable_hash}" if cache_key in summary_cache: node["summary"] = extract_final_summary(summary_cache[cache_key]) return prompt = """Summarize the First Schedule of the Bharatiya Nagarik Suraksha Sanhita, 2023. It is a structured table classifying Bharatiya Nyaya Sanhita offences by section number, offence description, punishment, cognizability, bailability, and triable court. Provide a keyword-dense 2–3 sentence summary covering its overall structure, purpose, and columns.""" summary = await call_model_with_retry(prompt) summary_cache[cache_key] = summary node["summary"] = summary save_cache() async def summarize_chapter(node): """ Summarizes a chapter node from its child section summaries. """ node_id = node["node_id"] stable_hash = node["metadata"]["stable_hash"] cache_key = f"{node_id}:{stable_hash}" if cache_key in summary_cache: node["summary"] = extract_final_summary(summary_cache[cache_key]) return act_code = node["metadata"]["act_code"] act_name = ACT_FULL_NAMES.get(act_code, act_code) children_summaries = [] for child in node["children"]: # Only use sections/schedule/front_matter, or sub-summaries c_title = child["title"] c_sum = child.get("summary", "") if c_sum: children_summaries.append(f"- {c_title}: {c_sum}") summaries_text = "\n".join(children_summaries) # Cap input text length if it gets too large if len(summaries_text) > 40000: summaries_text = summaries_text[:40000] + "\n...[truncated for length]..." prompt = f"""You are a precise legal summarizer. The following are summaries of all sections in {node['title']} of the {act_name}. Provide a keyword-rich and dense 2–3 sentence summary of this chapter's scope, primary offences/topics, and legal subject matter. Preserve key terms. Section summaries: {summaries_text}""" summary = await call_model_with_retry(prompt) summary_cache[cache_key] = summary node["summary"] = summary save_cache() async def summarize_root(node): """ Summarizes the entire Act's root node using its chapter summaries. """ node_id = node["node_id"] stable_hash = node["metadata"]["stable_hash"] cache_key = f"{node_id}:{stable_hash}" if cache_key in summary_cache: node["summary"] = extract_final_summary(summary_cache[cache_key]) return act_code = node["metadata"]["act_code"] act_name = ACT_FULL_NAMES.get(act_code, act_code) chapter_summaries = [] for child in node["children"]: c_title = child["title"] c_sum = child.get("summary", "") if c_sum: chapter_summaries.append(f"- {c_title}: {c_sum}") summaries_text = "\n".join(chapter_summaries) if act_code == "SOP": prompt = f"""You are a precise legal procedure summarizer. The following are summaries of individual Standard Operating Procedures for the {act_name}. Provide a dense, keyword-rich one-paragraph summary (3–5 sentences) of the overall scope, key procedural workflows (such as investigation, arrest, and electronic evidence), and operational guidelines introduced by this manual. SOP summaries: {summaries_text}""" else: prompt = f"""You are a precise legal summarizer. The following are chapter summaries for the {act_name}. Provide a dense, keyword-rich one-paragraph summary (3–5 sentences) of the overall scope, structure, and main legal systems/topics introduced by this Act. Chapter summaries: {summaries_text}""" summary = await call_model_with_retry(prompt) summary_cache[cache_key] = summary node["summary"] = summary save_cache() async def run_summarization_pipeline(trees): """ Runs the bottom-up summarization pipeline: leaf nodes -> chapter nodes -> root nodes. """ load_cache() # 1. Flatten all leaf nodes across the acts print("Flattening and preparing leaf nodes...") sections = [] schedule_rows = [] front_matters = [] schedules = [] chapters = [] roots = [] def traverse(node): nt = node["node_type"] if nt == "root": roots.append(node) elif nt == "chapter": chapters.append(node) elif nt == "front_matter": front_matters.append(node) elif nt == "schedule": schedules.append(node) elif nt in ["section", "sop_procedure", "sop_form", "sop_reference", "sop_table"]: sections.append(node) elif nt == "schedule_row": schedule_rows.append(node) for child in node["children"]: traverse(child) for act_root in trees.values(): traverse(act_root) print(f"Found: {len(sections)} sections, {len(schedule_rows)} schedule rows, {len(front_matters)} front matters, {len(schedules)} schedule chapters, {len(chapters)} chapters, {len(roots)} roots.") # 2. Process schedule rows (free, no LLM calls) print("Processing schedule rows...") cache_hits = 0 new_rows = 0 for row in schedule_rows: node_id = row["node_id"] stable_hash = row["metadata"]["stable_hash"] cache_key = f"{node_id}:{stable_hash}" if cache_key in summary_cache: row["summary"] = summary_cache[cache_key] cache_hits += 1 else: summary = format_schedule_row_summary(row) summary_cache[cache_key] = summary row["summary"] = summary new_rows += 1 save_cache() print(f"Schedule rows formatted: {cache_hits} cached, {new_rows} newly generated.") # 3. Summarize normal section leaves and the schedule chapters # We do sections and front_matters as leaf nodes leaves_to_summarize = sections + front_matters + schedules # Check cache hits for leaves uncached_leaves = [] cached_leaves_count = 0 for leaf in leaves_to_summarize: node_id = leaf["node_id"] stable_hash = leaf["metadata"]["stable_hash"] cache_key = f"{node_id}:{stable_hash}" if cache_key in summary_cache: leaf["summary"] = extract_final_summary(summary_cache[cache_key]) cached_leaves_count += 1 else: uncached_leaves.append(leaf) print(f"Leaves: {cached_leaves_count} cached, {len(uncached_leaves)} to call via API.") # Concurrency limit with locks / pacing if uncached_leaves: print("Starting rate-paced Leaf Summarization LLM calls...") start_time = time.time() async def summarize_task(node, idx, total): if node["node_type"] == "schedule": await summarize_schedule_chapter(node) else: await summarize_section(node) # Print progress every 50 calls if idx > 0 and idx % 50 == 0: elapsed = time.time() - start_time avg_time = elapsed / idx rem_calls = total - idx est_rem = avg_time * rem_calls print(f"Progress: {idx}/{total} leaf calls completed. Elapsed: {elapsed:.1f}s. Est. remaining: {est_rem:.1f}s.") # Fire async tasks tasks = [summarize_task(node, i+1, len(uncached_leaves)) for i, node in enumerate(uncached_leaves)] await asyncio.gather(*tasks) print("Leaf Summarization complete.") else: print("All leaf nodes retrieved from cache. No LLM calls needed.") # 4. Chapters Roll-Up Summarization (Bottom-up level 1) uncached_chapters = [] cached_chapters_count = 0 for chap in chapters: node_id = chap["node_id"] stable_hash = chap["metadata"]["stable_hash"] cache_key = f"{node_id}:{stable_hash}" if cache_key in summary_cache: chap["summary"] = extract_final_summary(summary_cache[cache_key]) cached_chapters_count += 1 else: uncached_chapters.append(chap) print(f"Chapters: {cached_chapters_count} cached, {len(uncached_chapters)} to call via API.") if uncached_chapters: print("Starting Chapter roll-up LLM calls...") tasks = [summarize_chapter(chap) for chap in uncached_chapters] await asyncio.gather(*tasks) print("Chapter summaries complete.") # 5. Roots Roll-Up Summarization (Bottom-up level 0) uncached_roots = [] cached_roots_count = 0 for r in roots: node_id = r["node_id"] stable_hash = r["metadata"]["stable_hash"] cache_key = f"{node_id}:{stable_hash}" if cache_key in summary_cache: r["summary"] = extract_final_summary(summary_cache[cache_key]) cached_roots_count += 1 else: uncached_roots.append(r) print(f"Roots: {cached_roots_count} cached, {len(uncached_roots)} to call via API.") if uncached_roots: print("Starting Root roll-up LLM calls...") tasks = [summarize_root(r) for r in uncached_roots] await asyncio.gather(*tasks) print("Root summaries complete.") # Save final cache save_cache() print("Summarization pipeline completed successfully.")