nvtitan/graphRAG
0
1"""2Gemini-based Knowledge Graph Extraction3Simple LLM-powered extraction using Google Gemini (cheapest option)4"""5from typing import List, Dict, Any, Optional6from loguru import logger7from models import Chunk, CanonicalTriple, RelationType8from config import settings9import json10import asyncio11 12 13class GeminiExtractor:14 """15 Extract key nodes and relationships using Gemini LLM16 Simple, cost-effective approach for knowledge graph generation17 """18 19 def __init__(self, llm_service=None):20 """Initialize Gemini extractor"""21 logger.info("Initializing GeminiExtractor")22 23 # Import litellm for API calls24 try:25 import litellm26 self.litellm = litellm27 28 # Configure litellm for Gemini29 self.model_name = f"gemini/{settings.gemini_model}"30 self.api_key = settings.gemini_api_key31 32 logger.info(f"✓ GeminiExtractor initialized with model: {self.model_name}")33 34 except ImportError as e:35 logger.error("litellm not installed. Install with: pip install litellm")36 raise RuntimeError("litellm required for Gemini") from e37 38 # Comprehensive list of generic terms to REJECT39 self.generic_stopwords = {40 # Generic nouns41 'system', 'systems', 'data', 'information', 'value', 'values',42 'method', 'methods', 'approach', 'approaches', 'technique', 'techniques',43 'result', 'results', 'study', 'studies', 'paper', 'papers',44 'section', 'sections', 'figure', 'figures', 'table', 'tables',45 'example', 'examples', 'case', 'cases', 'type', 'types',46 'way', 'ways', 'thing', 'things', 'part', 'parts',47 'model', 'models', 'framework', 'frameworks', # Too generic unless specific48 'process', 'processes', 'analysis', 'problem', 'problems',49 'solution', 'solutions', 'set', 'sets', 'group', 'groups',50 'element', 'elements', 'component', 'components',51 'feature', 'features', 'property', 'properties',52 'aspect', 'aspects', 'factor', 'factors', 'parameter', 'parameters',53 'concept', 'concepts', 'idea', 'ideas', 'theory', 'theories',54 'field', 'fields', 'area', 'areas', 'domain', 'domains',55 'task', 'tasks', 'goal', 'goals', 'objective', 'objectives',56 'input', 'inputs', 'output', 'outputs', 'function', 'functions',57 'operation', 'operations', 'step', 'steps', 'stage', 'stages',58 'phase', 'phases', 'level', 'levels', 'layer', 'layers',59 'number', 'numbers', 'amount', 'amounts', 'size', 'sizes',60 'performance', 'accuracy', 'quality', 'efficiency',61 'document', 'documents', 'text', 'texts', 'word', 'words',62 'sentence', 'sentences', 'paragraph', 'paragraphs',63 'item', 'items', 'object', 'objects', 'entity', 'entities',64 'relation', 'relations', 'relationship', 'relationships',65 66 # Generic verbs/actions67 'use', 'uses', 'using', 'used', 'usage',68 'apply', 'applies', 'applying', 'applied', 'application', 'applications',69 'work', 'works', 'working', 'worked',70 'provide', 'provides', 'providing', 'provided',71 'show', 'shows', 'showing', 'shown',72 'present', 'presents', 'presenting', 'presented', 'presentation',73 74 # Generic adjectives75 'new', 'novel', 'existing', 'current', 'previous',76 'different', 'similar', 'same', 'other', 'another',77 'various', 'several', 'multiple', 'single',78 'important', 'significant', 'main', 'key', 'major',79 'good', 'better', 'best', 'high', 'low',80 'large', 'small', 'big', 'little',81 82 # Research-specific generic terms83 'experiment', 'experiments', 'evaluation', 'evaluations',84 'test', 'tests', 'testing', 'validation',85 'comparison', 'comparisons', 'benchmark', 'benchmarks',86 'baseline', 'baselines', 'metric', 'metrics',87 'dataset', 'datasets', 'corpus', 'corpora',88 89 # Time/sequence terms90 'time', 'times', 'period', 'periods', 'year', 'years',91 'first', 'second', 'third', 'last', 'final',92 'next', 'previous', 'current', 'recent',93 94 # Common prepositions/articles (shouldn't appear but just in case)95 'the', 'a', 'an', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by',96 97 # Additional generic ML/AI terms (too broad)98 'neural network', 'deep learning', 'machine learning',99 'training', 'testing', 'prediction', 'classification',100 'regression', 'clustering', 'optimization',101 'network', 'networks', 'algorithm', 'algorithms',102 'learning', 'training data', 'test data',103 'feature extraction', 'preprocessing',104 'hyperparameter', 'hyperparameters',105 'loss', 'error', 'gradient',106 }107 108 async def extract_from_chunks(109 self,110 chunks: List[Chunk],111 use_llm: bool = True112 ) -> List[CanonicalTriple]:113 """114 Extract knowledge graph - PER PAGE with HARD CAP of 2 concepts per page115 116 Args:117 chunks: List of text chunks118 use_llm: Always True for Gemini extraction119 120 Returns:121 List of canonical triples122 """123 logger.info(f"\n{'='*80}")124 logger.info(f"{'GEMINI PER-PAGE EXTRACTION - 2 CONCEPTS MAX PER PAGE':^80}")125 logger.info(f"{'='*80}")126 127 all_triples = []128 129 # Filter text chunks130 text_chunks = [c for c in chunks if c.type.value in ["paragraph", "code"]]131 132 if not text_chunks:133 logger.warning("No text chunks to process")134 return []135 136 # GROUP CHUNKS BY PAGE137 from collections import defaultdict138 chunks_by_page = defaultdict(list)139 for chunk in text_chunks:140 page_num = chunk.page_number or 0141 chunks_by_page[page_num].append(chunk)142 143 logger.info(f"Processing {len(chunks_by_page)} pages in PARALLEL")144 145 # ⚡ PARALLEL PROCESSING: Create tasks for all pages146 tasks = []147 page_numbers = []148 for page_num in sorted(chunks_by_page.keys()):149 page_chunks = chunks_by_page[page_num]150 combined_text = "\n\n".join([chunk.text for chunk in page_chunks])151 152 logger.info(f"📄 PAGE {page_num}: {len(page_chunks)} chunks, {len(combined_text)} chars")153 154 # Create async task for this page155 tasks.append(self._extract_with_gemini(combined_text, page_num))156 page_numbers.append(page_num)157 158 # Execute all Gemini calls in parallel159 logger.info(f"\n🚀 Launching {len(tasks)} parallel Gemini API calls...")160 import time161 start_time = time.time()162 163 results = await asyncio.gather(*tasks, return_exceptions=True)164 165 elapsed = time.time() - start_time166 logger.info(f"✓ All {len(tasks)} Gemini calls completed in {elapsed:.2f}s (parallel)")167 logger.info(f" Average: {elapsed/len(tasks):.2f}s per page (would be {elapsed*len(tasks):.2f}s sequential)")168 169 # Process results170 for page_num, page_triples in zip(page_numbers, results):171 if isinstance(page_triples, Exception):172 logger.error(f" ❌ Page {page_num} failed: {page_triples}")173 continue174 175 if page_triples:176 all_triples.extend(page_triples)177 logger.info(f" ✓ Page {page_num}: Extracted {len(page_triples)} triples")178 for t in page_triples:179 relation_value = t.relation.value if hasattr(t.relation, 'value') else t.relation180 logger.info(f" → {t.subject_label} --[{relation_value}]--> {t.object_label}")181 else:182 logger.warning(f" ⚠️ Page {page_num}: NO TRIPLES EXTRACTED!")183 184 # Summary185 unique_concepts = set()186 concepts_by_page = {}187 for triple in all_triples:188 unique_concepts.add(triple.subject_label)189 unique_concepts.add(triple.object_label)190 page = triple.page_number191 if page not in concepts_by_page:192 concepts_by_page[page] = set()193 concepts_by_page[page].add(triple.subject_label)194 concepts_by_page[page].add(triple.object_label)195 196 logger.info(f"\n{'='*80}")197 logger.info(f"{'EXTRACTION SUMMARY':^80}")198 logger.info(f"{'='*80}")199 logger.info(f"Pages processed: {len(chunks_by_page)}")200 logger.info(f"Total triples: {len(all_triples)}")201 logger.info(f"Unique concepts: {len(unique_concepts)} (max {len(chunks_by_page) * 2})")202 203 if len(all_triples) == 0:204 logger.error(f"\n❌❌❌ CRITICAL ERROR: ZERO TRIPLES EXTRACTED! ❌❌❌")205 logger.error(f"This means:")206 logger.error(f" - Either Gemini returned no concepts")207 logger.error(f" - Or all concepts were rejected by filters")208 logger.error(f" - Or there was an API error")209 logger.error(f"Check the logs above for details!")210 else:211 logger.info(f"\nConcepts per page:")212 for page in sorted(concepts_by_page.keys()):213 logger.info(f" Page {page}: {list(concepts_by_page[page])}")214 215 logger.info(f"{'='*80}\n")216 217 return all_triples218 219 async def _extract_with_gemini(self, text: str, page_number: int) -> List[CanonicalTriple]:220 """221 Call Gemini API to extract technical concepts (nodes) from THIS PAGE222 223 Args:224 text: Text from single page225 page_number: Page number226 227 Returns:228 List of canonical triples229 """230 # Specialized technical concept extraction prompt231 prompt = f"""You are an expert in technical information extraction and knowledge graph construction.232Your task is to identify only the most meaningful *technical concepts* from the given text.233Concepts must represent scientific, mathematical, algorithmic, or methodological entities234that could exist as standalone nodes in a knowledge graph.235Ignore generic words, section titles, variable names, and everyday terms.236Focus on high-value, domain-specific terminology relevant to the text.237 238Extract all important technical concepts from the following text that would form the239nodes of a knowledge graph.240 241⚙️ Rules:242• Each concept should represent a self-contained technical idea, model, method, metric, loss, theorem, or process243• Keep only multi-word phrases when possible ("gradient descent", "convolutional neural network", "cross-entropy loss")244• Skip single, contextless nouns ("data", "model", "value", "equation", "result")245• Merge synonymous terms (e.g., "SGD", "stochastic gradient descent" → one entry)246• Do not include equations, numeric values, figure names, or symbols247• Do not repeat concepts248• Maintain consistent naming conventions (lowercase, hyphen-separated words)249• Extract MAXIMUM 4-5 concepts from this page (quality over quantity)250 251Return output strictly as JSON with "nodes" key:252{{253 "nodes": [254 "gradient descent",255 "neural network",256 "cross entropy loss"257 ]258}}259 260PAGE {page_number} TEXT:261{text}262 263CRITICAL: Return ONLY the JSON. If no technical concepts found, return {{"nodes": []}}"""264 265 logger.info(f" 🚀 Starting Gemini extraction for page {page_number}...")266 logger.info(f" Text length: {len(text)} characters")267 268 try:269 # Call Gemini via litellm270 logger.info(f" 📡 Calling Gemini API for page {page_number}...")271 272 response = await asyncio.to_thread(273 self.litellm.completion,274 model=self.model_name,275 api_key=self.api_key,276 messages=[{277 "role": "user",278 "content": prompt279 }],280 temperature=0.0, 281 max_tokens=settings.llm_max_tokens,282 timeout=settings.llm_timeout283 )284 285 # Extract response text286 response_text = response.choices[0].message.content.strip()287 logger.info(f" 📥 Gemini response ({len(response_text)} chars):")288 logger.info(f" {response_text[:500]}") 289 290 291 if "```json" in response_text:292 response_text = response_text.split("```json")[1].split("```")[0].strip()293 elif "```" in response_text:294 response_text = response_text.split("```")[1].split("```")[0].strip()295 296 data = json.loads(response_text)297 298 299 if isinstance(data, dict) and "nodes" in data:300 nodes = data["nodes"]301 elif isinstance(data, list):302 # Fallback: if Gemini returned a list directly303 nodes = data304 else:305 logger.warning(f" ❌ Gemini returned unexpected format: {type(data)}")306 return []307 308 if not isinstance(nodes, list):309 logger.warning(f" ❌ Nodes is not a list, got: {type(nodes)}")310 return []311 312 logger.info(f" ✓ Gemini extracted {len(nodes)} nodes from page {page_number}")313 logger.info(f" Raw nodes: {nodes}")314 315 # Validate and filter nodes316 valid_nodes = []317 rejected_nodes = []318 319 for node in nodes:320 if not isinstance(node, str):321 logger.warning(f" ⚠️ Skipping non-string node: {node}")322 continue323 324 node = node.strip()325 if not node:326 continue327 328 logger.info(f" Validating node: '{node}'")329 330 # FILTER: Validate node is a technical concept331 if not self._is_technical_concept(node):332 rejected_nodes.append(node)333 logger.warning(f" ✗ REJECTED node '{node}' - not technical enough")334 continue335 336 logger.info(f" ✅ ACCEPTED node: '{node}'")337 valid_nodes.append(node.lower())338 339 # Summary of rejections340 if rejected_nodes:341 logger.warning(f" 📊 Rejected {len(rejected_nodes)} nodes: {rejected_nodes}")342 343 if not valid_nodes:344 logger.warning(f" ⚠️ ALL {len(nodes)} NODES REJECTED for page {page_number}")345 logger.warning(f" No valid technical concepts found. Returning empty list.")346 return []347 348 349 selected_nodes = valid_nodes[:2] #350 logger.info(f" 🎯 Selected {len(selected_nodes)} nodes (hard cap = 2): {selected_nodes}")351 352 353 page_triples = []354 355 if len(selected_nodes) == 1:356 # Only one node - create self-referencing relationship or skip357 logger.info(f" ℹ️ Only 1 node on page {page_number}, cannot create relationships")358 359 return []360 361 elif len(selected_nodes) == 2:362 # Use LLM to determine actual relationship between nodes363 node1, node2 = selected_nodes[0], selected_nodes[1]364 365 # Extract relationship using LLM with page context366 logger.info(f" 🔍 Extracting relationship between: {node1} ↔ {node2}")367 relationship_triple = await self._extract_relationship_with_gemini(368 text=text,369 node1=node1,370 node2=node2,371 page_number=page_number372 )373 374 if relationship_triple:375 page_triples.append(relationship_triple)376 logger.info(f" ✅ Created directed edge:")377 logger.info(f" → {relationship_triple.subject_label} --[{relationship_triple.relation.value}]--> {relationship_triple.object_label}")378 logger.info(f" Justification: {relationship_triple.justification}")379 else:380 logger.warning(f" ⚠️ Could not extract relationship for {node1} ↔ {node2}")381 382 logger.info(f" ✅ Returning {len(page_triples)} triples for page {page_number}")383 return page_triples384 385 except json.JSONDecodeError as e:386 logger.error(f" ❌ JSON PARSE ERROR for page {page_number}: {e}")387 logger.error(f" Response was: {response_text[:500]}")388 return []389 390 except Exception as e:391 logger.error(f" ❌ GEMINI API FAILED for page {page_number}: {e}")392 logger.error(f" Exception type: {type(e).__name__}")393 logger.error(f" Full trace:", exc_info=True)394 return []395 396 async def _extract_relationship_with_gemini(self, text: str, node1: str, node2: str, page_number: int) -> Optional[CanonicalTriple]:397 """398 Use Gemini to determine the actual relationship between two nodes based on page context399 400 Args:401 text: Full page text for context402 node1: First node/concept403 node2: Second node/concept404 page_number: Page number405 406 Returns:407 CanonicalTriple with proper relationship, or None if extraction fails408 """409 # List all available relation types for the LLM410 available_relations = [r.value for r in RelationType]411 412 prompt = f"""You are an expert at extracting knowledge graph relationships from technical text.413 414Given two concepts and the text they appear in, determine the most accurate relationship between them.415 416**Concepts:**417- Concept A: "{node1}"418- Concept B: "{node2}"419 420**Context (page {page_number}):**421{text[:3000]}422 423**Available Relationship Types:**424{', '.join(available_relations)}425 426**Instructions:**4271. Analyze how these two concepts relate in the given context4282. Choose the MOST SPECIFIC relationship type from the list above4293. Determine the direction: which concept is the subject and which is the object4304. Provide a brief justification from the text431 432**Output Format (JSON):**433{{434 "subject": "<node1 or node2>",435 "object": "<node1 or node2>",436 "relation": "<one of the available relationship types>",437 "confidence": <0.0-1.0>,438 "justification": "<brief explanation from text>"439}}440 441**Rules:**442- Use the exact concept names provided443- Choose only ONE relation type from the available list444- If no clear relationship exists, use "related_to"445- Direction matters: subject performs/has the relation to the object446"""447 448 try:449 # Call Gemini API450 response_text = await self.litellm.acompletion(451 model=self.model_name,452 messages=[453 {"role": "system", "content": "You are an expert at knowledge graph relationship extraction. Always output valid JSON."},454 {"role": "user", "content": prompt}455 ],456 api_key=self.api_key,457 temperature=0.1, # Low temperature for consistent relationship extraction458 response_format={"type": "json_object"}459 )460 461 response_content = response_text.choices[0].message.content462 data = json.loads(response_content)463 464 # Validate response465 subject = data.get("subject", "").strip()466 obj = data.get("object", "").strip()467 relation_str = data.get("relation", "related_to").lower().strip().replace(" ", "_")468 confidence = float(data.get("confidence", 0.7))469 justification = data.get("justification", f"Relationship extracted from page {page_number}")470 471 # Map relation string to enum472 try:473 relation = RelationType(relation_str)474 except ValueError:475 logger.warning(f" ⚠️ Invalid relation '{relation_str}', defaulting to RELATED_TO")476 relation = RelationType.RELATED_TO477 478 # Create triple479 triple = CanonicalTriple(480 subject_label=subject,481 object_label=obj,482 relation=relation,483 confidence=confidence,484 justification=justification,485 page_number=page_number486 )487 488 return triple489 490 except json.JSONDecodeError as e:491 logger.error(f" ❌ JSON parse error in relationship extraction: {e}")492 return None493 except Exception as e:494 logger.error(f" ❌ Relationship extraction failed: {e}")495 return None496 497 def _is_technical_concept(self, concept: str) -> bool:498 """499 500 Args:501 concept: Concept string to validate502 503 Returns:504 True if highly technical/specific, False otherwise505 """506 concept_lower = concept.lower().strip()507 508 # RULE 1: Reject if in stopwords509 if concept_lower in self.generic_stopwords:510 logger.debug(f"Rejected '{concept}' - in stopword list")511 return False512 513 # RULE 2: Reject if any word is a generic stopword (stricter)514 words = concept_lower.split()515 for word in words:516 if word in self.generic_stopwords:517 # Allow if it's part of a specific multi-word technical term518 # e.g., "convolutional neural network" has "network" but is specific519 if len(words) < 2:520 logger.debug(f"Rejected '{concept}' - contains generic word '{word}'")521 return False522 523 # RULE 3: Single-word concepts must have SOME specificity (RELAXED)524 if len(words) == 1:525 # Accept if ANY of these are true:526 # - Has uppercase (BERT, Adam, PyTorch)527 # - Has numbers (VGG16, GPT3)528 # - Has special chars (t-SNE, bi-LSTM)529 # - Longish word (8+ chars like "backpropagation")530 has_uppercase = any(c.isupper() for c in concept)531 has_numbers = any(c.isdigit() for c in concept)532 has_special = '-' in concept or '_' in concept533 is_longish = len(concept) >= 8 # RELAXED from 10534 535 if not (has_uppercase or has_numbers or has_special or is_longish):536 logger.debug(f"Rejected '{concept}' - single word not specific enough")537 return False538 539 # RULE 4: Multi-word phrases - very lenient540 if len(words) >= 2:541 # Just check that it's not ALL generic words542 # At least one word should be non-generic or have caps/numbers543 has_caps = any(c.isupper() for c in concept)544 has_numbers = any(c.isdigit() for c in concept)545 has_hyphen = '-' in concept546 547 # Count non-generic words548 non_generic_count = sum(1 for w in words if w not in self.generic_stopwords)549 550 # Accept if ANY of these:551 # - Has caps/numbers/hyphen552 # - At least one word is non-generic553 # - 3+ words (likely specific enough)554 if not (has_caps or has_numbers or has_hyphen or non_generic_count > 0 or len(words) >= 3):555 logger.debug(f"Rejected '{concept}' - multi-word phrase too generic")556 return False557 558 # RULE 5: Reject very short terms (1-2 chars) unless they're known acronyms (all caps)559 if len(concept) <= 2 and concept.upper() != concept:560 logger.debug(f"Rejected '{concept}' - too short")561 return False562 563 # RULE 6: Must contain at least one alphanumeric character564 if not any(c.isalnum() for c in concept):565 logger.debug(f"Rejected '{concept}' - no alphanumeric chars")566 return False567 568 # RULE 7: Reject if it's just a generic category with a modifier569 # e.g., "new algorithm", "proposed method", "our model"570 generic_patterns = [571 'new ', 'novel ', 'proposed ', 'our ', 'this ', 'that ',572 'these ', 'those ', 'such ', 'other ', 'another ',573 'existing ', 'current ', 'previous ', 'standard '574 ]575 for pattern in generic_patterns:576 if concept_lower.startswith(pattern):577 logger.debug(f"Rejected '{concept}' - generic pattern")578 return False579 580 # Passed all strict filters581 return True582 583 def _map_relation(self, relation_str: str) -> RelationType:584 """Map relation string to RelationType enum"""585 relation_lower = relation_str.lower().strip()586 587 # Direct mapping588 mapping = {589 "uses": RelationType.USES,590 "implements": RelationType.IMPLEMENTS,591 "is_a": RelationType.IS_A,592 "is a": RelationType.IS_A,593 "part_of": RelationType.PART_OF,594 "part of": RelationType.PART_OF,595 "requires": RelationType.REQUIRES,596 "produces": RelationType.PRODUCES,597 "enables": RelationType.ENABLES,598 "improves": RelationType.IMPROVES,599 "enhances": RelationType.ENHANCES,600 "contains": RelationType.CONTAINS,601 "depends_on": RelationType.DEPENDS_ON,602 "depends on": RelationType.DEPENDS_ON,603 "related_to": RelationType.RELATED_TO,604 "related to": RelationType.RELATED_TO,605 }606 607 if relation_lower in mapping:608 return mapping[relation_lower]609 610 # Fallback611 logger.debug(f"Unknown relation '{relation_str}', using 'related_to'")612 return RelationType.RELATED_TO613