import json import asyncio import re import time as _time from typing import Dict, List, Any, Optional from ..base import BaseGraphStorage, BaseVectorStorage from ..utils import logger, compute_mdhash_id from ..prompts import get_prompts PROMPTS = get_prompts("text2sql") import os def detect_schema_type(schema_data: Dict[str, Any]) -> str: """ Detect if schema has pre-built relationships or needs inference Args: schema_data: Parsed JSON schema data Returns: "db_summary" if has explicit relationships, "reconstruct" if needs inference """ # Check for explicit relationship indicators if "relationships" in schema_data: return "db_summary" if "foreign_keys" in schema_data: return "db_summary" # Check if tables have relationship information tables = schema_data.get("tables", {}) for table_data in tables.values(): if "foreign_keys" in table_data: return "db_summary" if "relationships" in table_data: return "db_summary" # Check if columns are in DB summary format (list with detailed FK info) columns = table_data.get("columns", []) if columns and isinstance(columns, list): for col in columns[:3]: # Check first few columns if isinstance(col, dict) and ("references_table" in col or "is_foreign_key" in col): return "db_summary" # If no explicit relationships found, assume reconstruct data return "reconstruct" class DatabaseSchemaBuilder: """ Database Schema Builder for QAFD_RAG with LLM Enhancement This class handles the manual construction of knowledge graphs from database schema JSON files, with LLM-powered description enhancement and relationship weight optimization. """ def __init__(self, graph_storage: BaseGraphStorage, entities_vdb: BaseVectorStorage, relationships_vdb: BaseVectorStorage, llm_model_func: callable, schema_file_path: Optional[str] = None): self.graph_storage = graph_storage self.entities_vdb = entities_vdb self.relationships_vdb = relationships_vdb self.llm_model_func = llm_model_func self.schema_file_path = schema_file_path self._tables_info_cache = {} # Cache for tables_info async def build_from_json_schema(self, schema_file_path: str, metadata_file_path: Optional[str] = None, language: str = "English") -> Dict[str, Any]: """ Build knowledge graph from JSON schema file with LLM enhancement Args: schema_file_path: Path to the JSON schema file metadata_file_path: Optional path to metadata file language: Output language for descriptions Returns: Dictionary containing build statistics """ logger.info(f"Building knowledge graph from schema: {schema_file_path}") # Update instance schema file path for validation methods self.schema_file_path = schema_file_path # Load JSON schema with open(schema_file_path, 'r', encoding='utf-8') as f: schema_data = json.load(f) # Detect schema type schema_type = detect_schema_type(schema_data) logger.info(f"Detected schema type: {schema_type}") # Load metadata if provided metadata_content = None if metadata_file_path: with open(metadata_file_path, 'r', encoding='utf-8') as f: metadata_content = f.read() # Step 1: Extract tables and columns based on schema type if schema_type == "db_summary": logger.info("Processing DB summary with existing relationships") tables_info = self._extract_tables_from_schema(schema_data) else: logger.info("Processing reconstruct data - will infer relationships from schema structure") tables_info = self._extract_tables_from_reconstruct_schema(schema_data) self._tables_info_cache = tables_info # Cache for weight enhancement # Step 2: Insert tables and columns as entities WITH MINIMAL DESCRIPTIONS _t0 = _time.time() _build_start = _t0 print(f" [1/4] Inserting entities for {len(tables_info)} tables...", flush=True) entities_added = await self._insert_schema_entities_minimal(tables_info) print(f" [1/4] Done ({entities_added} entities, {_time.time()-_t0:.1f}s)", flush=True) # Step 3: Create relationships between tables and columns _t0 = _time.time() print(f" [2/4] Creating relationships...", flush=True) if schema_type == "db_summary": relationships_added = await self._create_schema_relationships(tables_info) else: relationships_added = await self._create_inferred_relationships(tables_info, schema_data) print(f" [2/4] Done ({relationships_added} relationships, {_time.time()-_t0:.1f}s)", flush=True) # Step 3.5: Clean up any duplicate nodes that might have been created duplicates_removed = await self.graph_storage.remove_duplicate_nodes() if duplicates_removed > 0: logger.info(f"Cleaned up {duplicates_removed} duplicate nodes") # Step 3.6: Log graph statistics graph_stats = await self.graph_storage.get_graph_stats() logger.info(f"Graph statistics: {graph_stats['total_nodes']} nodes, {graph_stats['total_edges']} edges") logger.info(f"Node types: {graph_stats['node_types']}") # Step 4: LLM ENHANCEMENT (MOVED EARLIER AND IMPROVED) if self.llm_model_func: logger.info("Starting LLM enhancement phase...") # by Wenjun self.edge_descriptions = {} _t0 = _time.time() print(f" [3/4] Enhancing descriptions with LLM...", flush=True) await self._enhance_descriptions_with_llm_chunked(tables_info, metadata_content, language) print(f" [3/4] Done ({_time.time()-_t0:.1f}s)", flush=True) _t0 = _time.time() print(f" [4/4] Enhancing relationship weights with LLM...", flush=True) await self._enhance_relationship_weights_with_llm_chunked(metadata_content, language) print(f" [4/4] Done ({_time.time()-_t0:.1f}s, total: {_time.time()-_build_start:.1f}s)", flush=True) else: logger.warning("No LLM function provided - skipping description enhancement") return { "schema_type": schema_type, "tables_added": len(tables_info), "entities_added": entities_added, "relationships_added": relationships_added, "duplicates_removed": duplicates_removed, "graph_stats": graph_stats } def _extract_tables_from_reconstruct_schema(self, schema_data: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: """Extract table and column information from reconstruct-style JSON schema""" tables_info = {} # Extract metadata metadata = schema_data.get("metadata", {}) # Extract tables - reconstruct format has tables as dict with column lists tables = schema_data.get("tables", {}) for table_name, table_data in tables.items(): table_info = { "name": table_name, "column_count": len(table_data.get("columns", [])), "row_count": len(table_data.get("sample_data", [])), "columns": [] } # Extract columns - reconstruct format has columns as list of dicts columns = table_data.get("columns", []) for col in columns: column_info = { "name": col.get("name", ""), "type": col.get("type", ""), "description": col.get("description", ""), "is_primary_key": False, # Will be inferred "is_foreign_key": False, # Will be inferred "not_null": False, "default": None, "references_table": None, # Will be inferred "references_column": None # Will be inferred } # Infer primary key from naming patterns if col.get("name", "").endswith("_id") and "id" in col.get("name", ""): if col.get("name", "") == f"{table_name.split('.')[-1]}_id": column_info["is_primary_key"] = True # Infer foreign keys from naming patterns if col.get("name", "").endswith("_id") and not column_info["is_primary_key"]: potential_ref_table = col.get("name", "").replace("_id", "") # Look for matching table for other_table in tables.keys(): if potential_ref_table in other_table: column_info["is_foreign_key"] = True column_info["references_table"] = other_table column_info["references_column"] = col.get("name", "") break table_info["columns"].append(column_info) # Add sample data info if "sample_data" in table_data: table_info["sample_rows"] = table_data["sample_data"] # Add project structure info if available if "project" in table_data: table_info["project"] = table_data["project"] if "dataset" in table_data: table_info["dataset"] = table_data["dataset"] tables_info[table_name] = table_info return tables_info async def _create_inferred_relationships(self, tables_info: Dict[str, Dict[str, Any]], schema_data: Dict[str, Any]) -> int: """Create inferred relationships for reconstruct data based on schema analysis""" relationships_added = 0 # Create table-column relationships (same as existing) for table_name, table_data in tables_info.items(): table_id = f'"{table_name}"' for column in table_data.get("columns", []): column_name = column.get("name", "") if not column_name: continue col_id = f'"{table_name}.{column_name}"' # Table contains column relationship edge_data = { "weight": 8.0, "description": f"Table '{table_name}' contains column '{column_name}'", "keywords": "table_structure, contains_column", "source_id": "reconstruct_inference" } await self.graph_storage.upsert_edge(table_id, col_id, edge_data=edge_data) relationships_added += 1 # Create inferred table-to-table relationships table_names = list(tables_info.keys()) for i, table1 in enumerate(table_names): for j, table2 in enumerate(table_names): if i >= j: # Avoid duplicates and self-references continue table1_data = tables_info[table1] table2_data = tables_info[table2] relationship_weight = 0.0 relationship_type = "related_to" # Check for same project/dataset (for BigQuery style) if (table1_data.get("project") == table2_data.get("project") and table1_data.get("dataset") == table2_data.get("dataset")): relationship_weight = 6.0 relationship_type = "same_dataset" # Check for naming similarity table1_base = table1.split('.')[-1] table2_base = table2.split('.')[-1] if any(word in table1_base for word in table2_base.split('_')) or \ any(word in table2_base for word in table1_base.split('_')): if relationship_weight < 5.0: relationship_weight = 5.0 relationship_type = "semantically_related" # Add relationship if weight is significant if relationship_weight > 4.0: table1_id = f'"{table1}"' table2_id = f'"{table2}"' edge_data = { "weight": relationship_weight, "description": f"Inferred {relationship_type} relationship between '{table1}' and '{table2}'", "keywords": f"inferred_relationship, {relationship_type}", "source_id": "reconstruct_inference" } await self.graph_storage.upsert_edge(table1_id, table2_id, edge_data=edge_data) relationships_added += 1 logger.info(f"Created {relationships_added} inferred relationships for reconstruct data") return relationships_added def _extract_tables_from_schema(self, schema_data: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: """Extract table and column information from JSON schema""" tables_info = {} # Extract tables tables = schema_data.get("tables", {}) for table_name, table_data in tables.items(): table_info = { "name": table_name, "column_count": table_data.get("column_count", 0), "row_count": table_data.get("row_count", 0), "columns": [] } # Extract columns columns = table_data.get("columns", []) for col in columns: column_info = { "name": col.get("name", ""), "type": col.get("type", ""), "is_primary_key": col.get("is_primary_key", False), "is_foreign_key": col.get("is_foreign_key", False), "not_null": col.get("not_null", False), "default": col.get("default", None), "references_table": col.get("references_table", None), "references_column": col.get("references_column", None) } table_info["columns"].append(column_info) # ✅ NEW: Read table-level foreign_keys array and merge into columns table_level_fks = table_data.get("foreign_keys", []) if table_level_fks: logger.info(f"Processing {len(table_level_fks)} table-level FKs for {table_name}") for fk in table_level_fks: col_name = fk.get("column") ref_info = fk.get("references", {}) ref_table = ref_info.get("table") ref_column = ref_info.get("column") if not col_name or not ref_table or not ref_column: continue # Find the matching column and update its FK info for col_info in table_info["columns"]: if col_info["name"] == col_name: col_info["is_foreign_key"] = True # Only set if not already populated from column-level data if not col_info["references_table"]: col_info["references_table"] = ref_table logger.info(f"Set FK: {table_name}.{col_name} -> {ref_table}.{ref_column}") if not col_info["references_column"]: col_info["references_column"] = ref_column break tables_info[table_name] = table_info return tables_info async def _insert_schema_entities_minimal(self, tables_info: Dict[str, Dict[str, Any]]) -> int: """ Insert tables and columns as entities with MINIMAL descriptions for LLM enhancement Args: tables_info: Dictionary of table information Returns: Number of entities added """ entities_added = 0 entities_for_vdb = {} for table_name, table_data in tables_info.items(): # Add table entity with MINIMAL description table_id = f'"{table_name}"' table_node_data = { "entity_type": "complete_table", "description": f"Table: {table_name}", # MINIMAL - will be enhanced by LLM "source_id": "schema_extraction", "table_name": table_name, "column_count": table_data["column_count"], "row_count": table_data["row_count"] } await self.graph_storage.upsert_node(table_id, node_data=table_node_data) # Prepare for vector database ???????????????????????????????? table_vdb_id = compute_mdhash_id(table_id, prefix="ent-") entities_for_vdb[table_vdb_id] = { "content": table_id + " " + table_node_data["description"], "entity_name": table_id } entities_added += 1 # Add column entities with MINIMAL descriptions for col in table_data["columns"]: col_name = col["name"] if not col_name: continue col_id = f'"{table_name}.{col_name}"' col_node_data = { "entity_type": "column", "description": f"Column: {table_name}.{col_name}", # MINIMAL - will be enhanced by LLM "source_id": "schema_extraction", "table_name": table_name, "column_name": col_name, "data_type": col["type"], "is_primary_key": col["is_primary_key"], "is_foreign_key": col["is_foreign_key"], "not_null": col["not_null"] } # Add optional fields only if they are not None if col["default"] is not None: col_node_data["default"] = col["default"] if col["references_table"] is not None: col_node_data["references_table"] = col["references_table"] if col["references_column"] is not None: col_node_data["references_column"] = col["references_column"] await self.graph_storage.upsert_node(col_id, node_data=col_node_data) # Prepare for vector database col_vdb_id = compute_mdhash_id(col_id, prefix="ent-") entities_for_vdb[col_vdb_id] = { "content": col_id + " " + col_node_data["description"], "entity_name": col_id } entities_added += 1 # Insert entities into vector database if entities_for_vdb and self.entities_vdb: await self.entities_vdb.upsert(entities_for_vdb) logger.info(f"Inserted {len(entities_for_vdb)} entities into vector database") return entities_added async def _create_schema_relationships(self, tables_info: Dict[str, Dict[str, Any]]) -> int: """Create relationships between tables and columns with fully qualified names""" relationships_added = 0 relationships_for_vdb = {} # Add debug logging logger.info(f"Creating relationships for {len(tables_info)} tables") available_tables = set(tables_info.keys()) logger.info(f"Available tables: {sorted(available_tables)}") for table_name, table_data in tables_info.items(): table_id = f'"{table_name}"' logger.info(f"Processing table: {table_name}") for col in table_data["columns"]: col_name = col["name"] if not col_name: continue col_id = f'"{table_name}.{col_name}"' # Create table-to-column relationship edge_data = { "weight": 10.0, "description": f"Table '{table_name}' contains column '{col_name}'", "keywords": "table_structure, contains_column", "source_id": "schema_extraction" } await self.graph_storage.upsert_edge(table_id, col_id, edge_data=edge_data) # Prepare for vector database ??????????????????????????????? edge_vdb_id = compute_mdhash_id(f"{table_id}->{col_id}", prefix="rel-") relationships_for_vdb[edge_vdb_id] = { "src_id": table_id, "tgt_id": col_id, "content": edge_data["keywords"] + " " + table_id + " " + col_id + " " + edge_data["description"] } relationships_added += 1 # Create foreign key relationships if applicable if col["is_foreign_key"] and col["references_table"] and col["references_column"]: ref_table_full = col["references_table"] ref_column = col["references_column"] logger.info(f"Processing FK: {col_name} -> {ref_table_full}.{ref_column}") # ✅ Use fully qualified names consistently if ref_table_full in available_tables: ref_table_id = f'"{ref_table_full}"' ref_col_id = f'"{ref_table_full}.{ref_column}"' logger.info(f"Creating FK edge: {col_id} -> {ref_col_id}") fk_edge_data = { "weight": 15.0, "description": f"Foreign key: '{col_name}' in '{table_name}' " f"references '{ref_column}' in '{ref_table_full}'", "keywords": "foreign_key, references, data_integrity", "source_id": "schema_extraction" } await self.graph_storage.upsert_edge(col_id, ref_col_id, edge_data=fk_edge_data) # Prepare for vector database fk_edge_vdb_id = compute_mdhash_id(f"{col_id}->{ref_col_id}", prefix="rel-") relationships_for_vdb[fk_edge_vdb_id] = { "src_id": col_id, "tgt_id": ref_col_id, "content": fk_edge_data["keywords"] + " " + col_id + " " + ref_col_id + " " + fk_edge_data["description"] } relationships_added += 1 logger.info(f"Successfully created FK relationship: {col_id} -> {ref_col_id}") else: logger.warning(f"Referenced table '{ref_table_full}' not found in schema for FK {col_name}") logger.warning(f"Available tables: {sorted(available_tables)}") elif col["is_foreign_key"]: logger.warning( f"FK column {col_name} missing reference info: " f"references_table={col.get('references_table')}, " f"references_column={col.get('references_column')}" ) # ✅ Insert into vector DB once, after all loops if relationships_for_vdb and self.relationships_vdb: await self.relationships_vdb.upsert(relationships_for_vdb) logger.info(f"Inserted {len(relationships_for_vdb)} relationships into vector database") logger.info(f"Total relationships created: {relationships_added}") return relationships_added async def _enhance_descriptions_with_llm_chunked(self, tables_info: Dict[str, Dict[str, Any]], metadata_content: Optional[str], language: str) -> None: """ Use LLM to enhance entity descriptions with chunking to avoid token limits """ logger.info("Enhancing descriptions with LLM (chunked)...") from ..utils import encode_string_by_tiktoken # Configuration that actually works max_tokens_per_chunk = 8000 # Reasonable chunk size base_prompt_tokens = 1000 # Reduce prompt overhead response_buffer_tokens = 1000 # Reduce response buffer available_tokens = max_tokens_per_chunk - base_prompt_tokens - response_buffer_tokens # Split tables into chunks table_chunks = self._chunk_tables_by_tokens(tables_info, available_tokens) logger.info(f"Processing {len(tables_info)} tables in {len(table_chunks)} chunks") _chunk_start = _time.time() # Process each chunk for i, chunk_tables in enumerate(table_chunks, 1): logger.info(f"Processing chunk {i}/{len(table_chunks)} with {len(chunk_tables)} tables") elapsed = _time.time() - _chunk_start if i > 1: eta = elapsed / (i - 1) * (len(table_chunks) - i + 1) print(f"\r [3/4] Enhancing descriptions: chunk {i}/{len(table_chunks)} (ETA: {eta:.0f}s)", end='', flush=True) else: print(f"\r [3/4] Enhancing descriptions: chunk {i}/{len(table_chunks)}", end='', flush=True) try: # Create prompt for this chunk chunk_schema_text = self._format_schema_for_llm({"tables": chunk_tables}) prompt = self._create_description_enhancement_prompt(chunk_schema_text, metadata_content, language) # Verify token count prompt_tokens = len(encode_string_by_tiktoken(prompt)) if prompt_tokens > max_tokens_per_chunk: logger.warning(f"Chunk {i} still exceeds token limit ({prompt_tokens} tokens), skipping") continue logger.info(f"Chunk {i}: Sending {prompt_tokens} tokens to LLM...") # Get enhanced descriptions from LLM enhanced_result = await self.llm_model_func(prompt) logger.info(f"Chunk {i}: Received LLM response, applying enhancements...") # Parse and apply enhanced descriptions await self._apply_enhanced_descriptions(enhanced_result, chunk_tables) except Exception as e: logger.error(f"Error enhancing descriptions for chunk {i}: {e}") continue print(flush=True) async def _enhance_relationship_weights_with_llm_chunked(self, metadata_content: Optional[str], language: str = "English") -> None: """ Use LLM to enhance relationship weights with chunking to avoid token limits """ logger.info("Enhancing relationship weights with LLM (chunked)...") from ..utils import encode_string_by_tiktoken # Configuration max_tokens_per_chunk = 100000 base_prompt_tokens = 3000 response_buffer_tokens = 2000 available_tokens = max_tokens_per_chunk - base_prompt_tokens - response_buffer_tokens # Get all relationships relationships_list = await self._get_all_relationships() if not relationships_list.strip(): logger.warning("No relationships found for weight enhancement") return # Split relationships into chunks relationship_chunks = self._chunk_relationships_by_tokens(relationships_list, available_tokens) logger.info(f"Processing relationships in {len(relationship_chunks)} chunks") _chunk_start = _time.time() # Process each chunk for i, chunk_relationships in enumerate(relationship_chunks, 1): logger.info(f"Processing relationship chunk {i}/{len(relationship_chunks)}") elapsed = _time.time() - _chunk_start if i > 1: eta = elapsed / (i - 1) * (len(relationship_chunks) - i + 1) print(f"\r [4/4] Enhancing weights: chunk {i}/{len(relationship_chunks)} (ETA: {eta:.0f}s)", end='', flush=True) else: print(f"\r [4/4] Enhancing weights: chunk {i}/{len(relationship_chunks)}", end='', flush=True) try: # Create minimal schema for this chunk involved_tables = self._extract_tables_from_relationships(chunk_relationships) minimal_schema = self._create_minimal_schema_for_tables(involved_tables) # Create prompt for this chunk prompt = self._create_weight_enhancement_prompt(minimal_schema, metadata_content, chunk_relationships, language) # Verify token count prompt_tokens = len(encode_string_by_tiktoken(prompt)) if prompt_tokens > max_tokens_per_chunk: logger.warning(f"Relationship chunk {i} exceeds token limit ({prompt_tokens} tokens), skipping") continue logger.info(f"Relationship chunk {i}: Sending {prompt_tokens} tokens to LLM...") # Get enhanced weights from LLM enhanced_result = await self.llm_model_func(prompt) logger.info(f"Relationship chunk {i}: Received LLM response, applying weight enhancements...") # Parse and apply enhanced weights await self._apply_enhanced_weights(enhanced_result) except Exception as e: logger.error(f"Error enhancing weights for chunk {i}: {e}") continue async def _get_all_relationships(self) -> str: """Get all relationships from the graph for LLM processing""" relationships = [] try: # Get all edges from the actual graph all_edges = await self.graph_storage.edges() if all_edges: # Convert EdgeView to list for easier handling edges_list = list(all_edges) logger.info(f"Found {len(edges_list)} edges in graph") for source, target in edges_list: # Get edge data to get current weight and description edge_data = await self.graph_storage.get_edge(source, target) if edge_data: current_weight = edge_data.get('weight', 1.0) description = edge_data.get('description', 'unknown') relationship_info = f"- {source} -> {target} (current_weight: {current_weight}, description: {description})" relationships.append(relationship_info) else: logger.warning(f"Could not get edge data for {source} -> {target}") else: # Fallback to schema-based relationships if graph is empty logger.warning("Graph is empty, falling back to schema-based relationships") for table_name, table_data in self._tables_info_cache.items(): table_id = f'"{table_name}"' for col in table_data.get("columns", []): col_name = col.get("name", "") if col_name: col_id = f'"{table_name}.{col_name}"' # Table-to-column relationship relationship_info = f"- {table_id} -> {col_id} (current_weight: 10.0, description: table_structure)" relationships.append(relationship_info) # Foreign key relationships if applicable if col.get("is_foreign_key") and col.get("references_table") and col.get("references_column"): ref_table_id = f'"{col["references_table"]}"' ref_col_id = f'"{col["references_table"]}.{col["references_column"]}"' fk_relationship_info = f"- {col_id} -> {ref_col_id} (current_weight: 9.0, description: foreign_key)" relationships.append(fk_relationship_info) except Exception as e: logger.error(f"Error getting relationships from graph: {e}") # Fallback to empty list relationships = [] return "\n".join(relationships) def _create_weight_enhancement_prompt(self, schema_text: str, metadata_content: Optional[str], relationships_list: str, language: str) -> str: """Create prompt for LLM weight enhancement""" # Load the complete JSON file to include sample data try: # Use the instance schema file path if available if self.schema_file_path and os.path.exists(self.schema_file_path): with open(self.schema_file_path, 'r', encoding='utf-8') as f: complete_schema_data = json.load(f) complete_schema_text = json.dumps(complete_schema_data, indent=2, ensure_ascii=False) else: complete_schema_text = schema_text except Exception as e: logger.warning(f"Could not load complete schema file: {e}") complete_schema_text = schema_text # Use the enhanced weight assignment prompt template prompt_template = PROMPTS["enhanced_graph_weight_assignment"] # Format the prompt with the complete schema data prompt = prompt_template.format( language=language, schema_text=complete_schema_text, metadata_content=metadata_content or "No additional metadata provided.", edge_list=relationships_list ) return prompt async def _apply_enhanced_weights(self, enhanced_result: str) -> None: """Apply enhanced weights from LLM to graph relationships with improved debugging""" try: logger.info(f"Raw LLM weight response (first 500 chars): {enhanced_result[:500]}...") # Extract JSON from the response with improved parsing import re # Try multiple patterns to find JSON json_patterns = [ r'\{.*\}', # Basic JSON object r'```json\s*(\{.*?\})\s*```', # JSON in code blocks r'```\s*(\{.*?\})\s*```', # JSON in generic code blocks r'```\s*(\{.*\})\s*```', # JSON in code blocks (non-greedy) ] enhanced_data = None for pattern in json_patterns: json_match = re.search(pattern, enhanced_result, re.DOTALL) if json_match: try: json_str = json_match.group(1) if len(json_match.groups()) > 0 else json_match.group(0) enhanced_data = json.loads(json_str) logger.info(f"Successfully parsed JSON with keys: {list(enhanced_data.keys())}") break except json.JSONDecodeError: continue if not enhanced_data: logger.warning("No valid JSON found in enhanced result, skipping weight enhancement") return # Apply relationship weights (CoFD-style multiplication) relationship_weights = enhanced_data.get("relationship_weights", {}) weighting_rationale = enhanced_data.get("weighting_rationale", {}) logger.info(f"Found {len(relationship_weights)} relationship weights to apply") weights_updated = 0 for relationship_key, llm_score in relationship_weights.items(): if '->' in relationship_key: source, target = relationship_key.split('->', 1) # Clean up the source and target IDs (remove quotes if present) source = source.strip().strip('"') target = target.strip().strip('"') # Add quotes to match the graph format source_with_quotes = f'"{source}"' target_with_quotes = f'"{target}"' # Try both formats: with and without quotes edge_data = await self.graph_storage.get_edge(source_with_quotes, target_with_quotes) actual_source = source_with_quotes actual_target = target_with_quotes if not edge_data: # Try without quotes as fallback edge_data = await self.graph_storage.get_edge(source, target) actual_source = source actual_target = target # ✅ FIXED: Only proceed if edge exists if edge_data: # Apply description from self.edge_descriptions if available # (This is redundant since Phase 3 already applied it, but kept for backward compatibility) if relationship_key in self.edge_descriptions: if 'original_description' not in edge_data: edge_data['original_description'] = edge_data.get('description', '') edge_data['description'] = self.edge_descriptions[relationship_key] # Get original weight (CoFD-style) original_weight = edge_data.get('weight', 1.0) # Multiply original weight with LLM score (CoFD approach) enhanced_weight = original_weight * llm_score # Update edge with new weight edge_data['weight'] = enhanced_weight edge_data['llm_enhanced'] = True edge_data['llm_score'] = llm_score edge_data['original_weight'] = original_weight # Add rationale if available if relationship_key in weighting_rationale: edge_data['weighting_rationale'] = weighting_rationale[relationship_key] await self.graph_storage.upsert_edge(actual_source, actual_target, edge_data=edge_data) weights_updated += 1 logger.info(f"Updated weight for {relationship_key}: {original_weight} * {llm_score} = {enhanced_weight}") else: logger.warning(f"Edge not found in graph: {source} -> {target} (tried both quoted and unquoted formats)") logger.info(f"Successfully updated {weights_updated} relationship weights with LLM enhancement") except json.JSONDecodeError as e: logger.error(f"Error parsing JSON from enhanced result: {e}") except Exception as e: logger.error(f"Error applying enhanced weights: {e}") def _format_schema_for_llm(self, tables_info: Dict[str, Dict[str, Any]]) -> str: """Format schema information for LLM processing""" schema_text = "JSON Schema:\n" schema_text += json.dumps({"tables": tables_info}, indent=2, ensure_ascii=False) return schema_text def _create_description_enhancement_prompt(self, schema_text: str, metadata_content: Optional[str], language: str) -> str: """Create prompt for LLM description enhancement using only chunked data""" # Use the provided chunked schema_text instead of loading complete file prompt_template = PROMPTS["enhanced_graph_description"] prompt = prompt_template.format( language=language, schema_text=schema_text, # Use chunked data passed in metadata_content=metadata_content or "No additional metadata provided." ) return prompt async def _apply_enhanced_descriptions(self, enhanced_result: str, tables_info: Dict[str, Dict[str, Any]]) -> None: """Apply enhanced descriptions from LLM to graph entities and edges""" try: logger.info(f"Raw LLM description response (first 500 chars): {enhanced_result[:500]}...") # Also check what node IDs actually exist: all_nodes = await self.graph_storage.nodes() existing_node_ids = list(all_nodes) if all_nodes else [] logger.info(f"Total nodes in graph: {len(existing_node_ids)}") # Extract JSON from the response with improved parsing import re # Try multiple patterns to find JSON json_patterns = [ r'\{.*\}', # Basic JSON object r'```json\s*(\{.*?\})\s*```', # JSON in code blocks r'```\s*(\{.*?\})\s*```', # JSON in generic code blocks r'```\s*(\{.*\})\s*```', # JSON in code blocks (non-greedy) ] enhanced_data = None for pattern in json_patterns: json_match = re.search(pattern, enhanced_result, re.DOTALL) if json_match: try: json_str = json_match.group(1) if len(json_match.groups()) > 0 else json_match.group(0) enhanced_data = json.loads(json_str) logger.info(f"Successfully parsed JSON with keys: {list(enhanced_data.keys())}") break except json.JSONDecodeError as e: logger.warning(f"JSON parse failed for pattern {pattern}: {e}") continue if not enhanced_data: logger.warning("No valid JSON found in enhanced result, skipping description enhancement") return # ==================================================================== # PHASE 1: ENHANCE NODES (Tables and Columns) # ==================================================================== # Apply table descriptions table_descriptions = enhanced_data.get("table_descriptions", {}) logger.info(f"Found {len(table_descriptions)} table descriptions to apply") for table_name, description in table_descriptions.items(): # Try different table name formats table_ids_to_try = [ f'"{table_name}"', f'"{table_name.lower()}"', f'"{table_name.upper()}"' ] description_applied = False for table_id in table_ids_to_try: if await self.graph_storage.has_node(table_id): # Get existing node data and update description existing_data = await self.graph_storage.get_node(table_id) if existing_data: existing_data["description"] = description existing_data["llm_enhanced"] = True await self.graph_storage.upsert_node(table_id, existing_data) logger.info(f"Updated table description for {table_name} (ID: {table_id})") description_applied = True break if not description_applied: logger.warning(f"Could not find table node for: {table_name}") # Apply column descriptions column_descriptions = enhanced_data.get("column_descriptions", {}) logger.info(f"Found {len(column_descriptions)} column descriptions to apply") for column_name, description in column_descriptions.items(): # Try different column name formats column_ids_to_try = [ f'"{column_name}"', f'"{column_name.lower()}"', f'"{column_name.upper()}"' ] description_applied = False for column_id in column_ids_to_try: if await self.graph_storage.has_node(column_id): # Get existing node data and update description existing_data = await self.graph_storage.get_node(column_id) if existing_data: existing_data["description"] = description existing_data["llm_enhanced"] = True await self.graph_storage.upsert_node(column_id, existing_data) logger.info(f"Updated column description for {column_name} (ID: {column_id})") description_applied = True break if not description_applied: logger.warning(f"Could not find column node for: {column_name}") # ==================================================================== # PHASE 2: UPDATE VECTOR DATABASE WITH ENHANCED NODE DESCRIPTIONS # ==================================================================== entities_added = 0 entities_for_vdb = {} for table_name, description in table_descriptions.items(): table_id = f'"{table_name}"' # Prepare for vector database table_vdb_id = compute_mdhash_id(table_id, prefix="ent-") entities_for_vdb[table_vdb_id] = { "content": table_id + " " + description, "entity_name": table_id } entities_added += 1 for col_name, description in column_descriptions.items(): col_id = f'"{col_name}"' # Prepare for vector database col_vdb_id = compute_mdhash_id(col_id, prefix="ent-") entities_for_vdb[col_vdb_id] = { "content": col_id + " " + description, "entity_name": col_id } entities_added += 1 if entities_for_vdb and self.entities_vdb: await self.entities_vdb.upsert(entities_for_vdb) logger.info(f"Inserted llm enhanced {len(entities_for_vdb)} entities into vector database") # ==================================================================== # PHASE 3: ENHANCE GRAPH EDGES (Table-Column & Column-Column) # ==================================================================== # Get relationship descriptions table_column_descriptions = enhanced_data.get("table_column_relationship_descriptions", {}) column_column_descriptions = enhanced_data.get("column_relationship_descriptions", {}) logger.info(f"Found {len(table_column_descriptions)} table-column relationship descriptions") logger.info(f"Found {len(column_column_descriptions)} column-column relationship descriptions") # Apply table-column relationship descriptions to GRAPH EDGES edges_enhanced = 0 for rel_key, description in table_column_descriptions.items(): if '->' in rel_key: source, target = rel_key.split('->', 1) source = source.strip().strip('"') target = target.strip().strip('"') source_quoted = f'"{source}"' target_quoted = f'"{target}"' # Try to find and update the edge in the graph edge_data = await self.graph_storage.get_edge(source_quoted, target_quoted) actual_source = source_quoted actual_target = target_quoted if not edge_data: edge_data = await self.graph_storage.get_edge(source, target) actual_source = source actual_target = target if edge_data: # Store original description before overwriting if 'description' in edge_data and 'original_description' not in edge_data: edge_data['original_description'] = edge_data['description'] edge_data['description'] = description edge_data['llm_enhanced'] = True await self.graph_storage.upsert_edge(actual_source, actual_target, edge_data=edge_data) edges_enhanced += 1 logger.info(f"Enhanced table-column edge: {source} -> {target}") else: logger.warning(f"Could not find graph edge for: {source} -> {target}") # Apply column-column relationship descriptions to GRAPH EDGES for rel_key, description in column_column_descriptions.items(): if '->' in rel_key: source, target = rel_key.split('->', 1) source = source.strip().strip('"') target = target.strip().strip('"') source_quoted = f'"{source}"' target_quoted = f'"{target}"' edge_data = await self.graph_storage.get_edge(source_quoted, target_quoted) actual_source = source_quoted actual_target = target_quoted if not edge_data: edge_data = await self.graph_storage.get_edge(source, target) actual_source = source actual_target = target if edge_data: # Store original description before overwriting if 'description' in edge_data and 'original_description' not in edge_data: edge_data['original_description'] = edge_data['description'] edge_data['description'] = description edge_data['llm_enhanced'] = True await self.graph_storage.upsert_edge(actual_source, actual_target, edge_data=edge_data) edges_enhanced += 1 logger.info(f"Enhanced col-col edge: {source} -> {target}") else: logger.warning(f"Could not find graph edge for: {source} -> {target}") logger.info(f"Enhanced {edges_enhanced} graph edges with LLM descriptions") # ==================================================================== # PHASE 4: UPDATE VECTOR DATABASE WITH ENHANCED EDGE DESCRIPTIONS # ==================================================================== relationships_added = 0 relationships_for_vdb = {} for relationship_key, description in table_column_descriptions.items(): if '->' in relationship_key: table_name, column_name = relationship_key.split('->') table_name = table_name.strip() column_name = column_name.strip() table_id = f'"{table_name}"' col_id = f'"{column_name}"' # Prepare for vector database edge_vdb_id = compute_mdhash_id(f"{table_id}->{col_id}", prefix="rel-") relationships_for_vdb[edge_vdb_id] = { "src_id": table_id, "tgt_id": col_id, "content": "table_structure, contains_column" + " " + table_id + " " + col_id + " " + description } relationships_added += 1 for relationship_key, description in column_column_descriptions.items(): if '->' in relationship_key: col_name, ref_col_name = relationship_key.split('->') col_name = col_name.strip() ref_col_name = ref_col_name.strip() col_id = f'"{col_name}"' ref_col_id = f'"{ref_col_name}"' # Prepare for vector database fk_edge_vdb_id = compute_mdhash_id(f"{col_id}->{ref_col_id}", prefix="rel-") relationships_for_vdb[fk_edge_vdb_id] = { "src_id": col_id, "tgt_id": ref_col_id, "content": "foreign_key, references, data_integrity" + " " + col_id + " " + ref_col_id + " " + description } relationships_added += 1 if relationships_for_vdb and self.relationships_vdb: await self.relationships_vdb.upsert(relationships_for_vdb) logger.info(f"Inserted llm enhanced {len(relationships_for_vdb)} relationships into vector database") # ==================================================================== # PHASE 5: STORE EDGE DESCRIPTIONS FOR WEIGHT ENHANCEMENT PHASE # ==================================================================== # Store descriptions from the LLM response for use in weight enhancement temp_d = enhanced_data.get("relationship_descriptions", {}) if len(temp_d) > 0: self.edge_descriptions = self.edge_descriptions | temp_d # Merge all edge descriptions for weight phase self.edge_descriptions = self.edge_descriptions | table_column_descriptions self.edge_descriptions = self.edge_descriptions | column_column_descriptions logger.info(f"Stored {len(self.edge_descriptions)} edge descriptions for weight enhancement phase") # ==================================================================== # PHASE 6: LOG DATA INSIGHTS IF AVAILABLE # ==================================================================== # Log data insights if available data_insights = enhanced_data.get("data_insights", {}) if data_insights: logger.info(f"Data insights extracted: {data_insights}") logger.info("Finished applying enhanced descriptions") except json.JSONDecodeError as e: logger.error(f"Error parsing JSON from enhanced result: {e}") except Exception as e: logger.error(f"Error applying enhanced descriptions: {e}") import traceback logger.error(f"Full traceback: {traceback.format_exc()}") # Utility methods for chunking def _chunk_tables_by_tokens(self, tables_info: Dict[str, Dict[str, Any]], max_tokens: int) -> List[Dict[str, Dict[str, Any]]]: """Split tables into chunks based on token limits, with column-level chunking for large tables""" from ..utils import encode_string_by_tiktoken all_chunks = [] current_chunk = {} current_tokens = 0 for table_name, table_data in tables_info.items(): # Estimate tokens for this table table_text = json.dumps({table_name: table_data}, indent=2) table_tokens = len(encode_string_by_tiktoken(table_text)) # If single table exceeds limit, split by columns if table_tokens > max_tokens: logger.info(f"Table {table_name} too large ({table_tokens} tokens), splitting by columns") # Split this large table by columns columns = table_data.get("columns", []) max_columns_per_chunk = 15 # Start with 15 columns per chunk for i in range(0, len(columns), max_columns_per_chunk): chunk_columns = columns[i:i + max_columns_per_chunk] # Create table data for this column chunk chunk_table_data = table_data.copy() chunk_table_data["columns"] = chunk_columns chunk_table_data["column_count"] = len(chunk_columns) # Remove sample data to reduce size chunk_table_data.pop("sample_data", None) chunk_table_data.pop("sample_rows", None) # Test if this chunk is small enough chunk_text = json.dumps({table_name: chunk_table_data}, indent=2) chunk_tokens = len(encode_string_by_tiktoken(chunk_text)) if chunk_tokens > max_tokens: # Still too large, try smaller chunks smaller_max = 5 for j in range(i, min(i + max_columns_per_chunk, len(columns)), smaller_max): smaller_chunk_columns = columns[j:j + smaller_max] smaller_chunk_data = table_data.copy() smaller_chunk_data["columns"] = smaller_chunk_columns smaller_chunk_data["column_count"] = len(smaller_chunk_columns) smaller_chunk_data.pop("sample_data", None) smaller_chunk_data.pop("sample_rows", None) all_chunks.append({table_name: smaller_chunk_data}) else: # Chunk is acceptable size all_chunks.append({table_name: chunk_table_data}) else: # Normal chunking logic for manageable tables if current_tokens + table_tokens > max_tokens and current_chunk: all_chunks.append(current_chunk) current_chunk = {table_name: table_data} current_tokens = table_tokens else: current_chunk[table_name] = table_data current_tokens += table_tokens # Add the last normal chunk if it has content if current_chunk: all_chunks.append(current_chunk) return all_chunks def _chunk_relationships_by_tokens(self, relationships_list: str, max_tokens: int) -> List[str]: """Split relationships into chunks based on token limits""" from ..utils import encode_string_by_tiktoken relationships = relationships_list.strip().split('\n') chunks = [] current_chunk_lines = [] current_tokens = 0 for relationship in relationships: if not relationship.strip(): continue relationship_tokens = len(encode_string_by_tiktoken(relationship)) # If adding this relationship would exceed the limit and current chunk is not empty if current_tokens + relationship_tokens > max_tokens and current_chunk_lines: chunks.append('\n'.join(current_chunk_lines)) current_chunk_lines = [relationship] current_tokens = relationship_tokens else: current_chunk_lines.append(relationship) current_tokens += relationship_tokens # Add the last chunk if it has content if current_chunk_lines: chunks.append('\n'.join(current_chunk_lines)) return chunks def _extract_tables_from_relationships(self, relationships_chunk: str) -> List[str]: """Extract fully qualified table names from this chunk of relationships""" import re tables = set() for line in relationships_chunk.split('\n'): # Find quoted strings (tables or columns) matches = re.findall(r'"([^"]+)"', line) for match in matches: if '.' in match: parts = match.split('.') if len(parts) > 1: # ✅ Keep all but the last as table (e.g., project.dataset.table) table_name = '.'.join(parts[:-1]) tables.add(table_name) else: tables.add(match) else: tables.add(match) return list(tables) def _create_minimal_schema_for_tables(self, table_names: List[str]) -> str: """Create minimal schema containing only the specified tables""" minimal_tables = {} # Get table info from cache for table_name in table_names: if table_name in self._tables_info_cache: # Include only essential info to minimize tokens table_data = self._tables_info_cache[table_name] minimal_tables[table_name] = { "name": table_data["name"], "column_count": table_data["column_count"], "columns": [ { "name": col["name"], "type": col["type"], "is_primary_key": col.get("is_primary_key", False), "is_foreign_key": col.get("is_foreign_key", False) } for col in table_data.get("columns", []) ] } return json.dumps({"tables": minimal_tables}, indent=2, ensure_ascii=False)