agentgraph/extraction/graph_utilities/knowledge_graph_merger.py

"""
Knowledge Graph Merger

This module provides functionality for merging multiple knowledge graphs into a single, coherent graph.
It uses a specialized agent to perform entity resolution and relationship consolidation.
"""

# Suppress all warnings
import warnings
warnings.filterwarnings("ignore")

# Configure logging early
import logging
logger = logging.getLogger(__name__)

# Silence all other loggers
for log_name in ['httpx', 'httpcore', 'LiteLLM', 'litellm', 'openai', 
                'pydantic', 'crewai', 'langchain', 'requests', 'urllib3']:
    logging.getLogger(log_name).setLevel(logging.ERROR)  # Using ERROR instead of WARNING

import json
import os
from typing import List, Dict, Any, Optional, Callable
from datetime import datetime
import uuid

# Import CrewAI for graph merging
from crewai import Agent, Task, Crew, Process
from crewai.memory import LongTermMemory, ShortTermMemory, EntityMemory
from crewai.memory.storage.rag_storage import RAGStorage
from crewai.memory.storage.ltm_sqlite_storage import LTMSQLiteStorage

# Import Pydantic models for structured data
from agentgraph.shared.models.reference_based import KnowledgeGraph

# ImportanceAssessor removed - multi-agent extractor already assigns importance levels

# from utils.config import LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY, LANGFUSE_AUTH, LANGFUSE_HOST

# os.environ["LANGFUSE_PUBLIC_KEY"] = LANGFUSE_PUBLIC_KEY
# os.environ["LANGFUSE_SECRET_KEY"] = LANGFUSE_SECRET_KEY


# if LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY:
#     os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = f"{LANGFUSE_HOST}/api/public/otel"
#     os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"Authorization=Basic {LANGFUSE_AUTH}"
# import openlit

# openlit.init()

# Load OpenAI API key from configuration
from utils.config import OPENAI_API_KEY
# Only set environment variable if OPENAI_API_KEY is not None
if OPENAI_API_KEY:
    os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
# Note: OPENAI_MODEL_NAME will be set dynamically in __init__ method






# Add asyncio and tqdm imports at the top
import asyncio
from tqdm import tqdm

class KnowledgeGraphMerger:
    """
    Merges multiple knowledge graphs into a single, coherent knowledge graph.
    Uses CrewAI and specialized agents to perform entity resolution and relationship consolidation.
    """
    
    def __init__(self, model: str = "gpt-5-mini"):
        """
        Initialize the knowledge graph merger.
        
        Args:
            model: Model to use for LLM operations
        """
        self.model = model
        # Set the model dynamically for CrewAI agents
        os.environ["OPENAI_MODEL_NAME"] = model
        logger.info(f"KnowledgeGraphMerger initialized with model: {model}")
        self._create_kg_merger_crew()
    
    def _create_kg_merger_crew(self):
        """Create a CrewAI crew for merging knowledge graphs using Pydantic models for structure"""
        # Define the merger agent
        self.kg_merger_agent = Agent(
            role="Knowledge Graph Merger",
            goal="Merge multiple knowledge graphs into a coherent, consolidated graph with strict schema adherence",
            backstory="""You are an expert in graph theory and knowledge representation. 
            You can identify duplicate entities across multiple knowledge graphs and
            merge them into a unified representation while preserving unique relationships.
            Your specialty is resolving conflicts, identifying the most accurate information
            when duplicates differ, and ensuring the final graph is consistent and accurate.
            
            You work with structured knowledge graphs following strict Pydantic schemas to ensure
            data integrity and proper typing throughout the merging process. Your expertise in 
            entity resolution is particularly valuable for maintaining a clean, consolidated graph
            without duplicate or inconsistent information.
            
            You excel at merging knowledge graphs derived from overlapping sliding windows of chronological data.
            You understand that entities may evolve over time across windows, and you can track this evolution
            while maintaining entity identity. You're skilled at resolving temporal inconsistencies and
            preserving the chronological narrative flow in the merged graph. When entities appear in multiple
            windows, you can determine the most current or complete representation while preserving important
            historical information.""",
            verbose=False,
            llm=self.model
        )
        
        # Define the merger task with Pydantic output model
        self.kg_merger_task = Task(
            description="""
            Merge multiple knowledge graphs into a single, coherent knowledge graph.
            These knowledge graphs represent overlapping sliding windows from a chronological trace.
            All merging decisions and resulting data must be based *strictly* on the content of the provided knowledge_graphs input.
            Do not introduce external knowledge or assumptions.
            
            WINDOW CHARACTERISTICS (interpret based on input KGs):
            - Windows have temporal overlaps where the same entities may appear in multiple windows.
            - Later windows generally contain more recent information about entities.
            - Entity evolution should be preserved while maintaining consistent identity, based on evidence in the KGs.
            - Relationships may change over time as the system executes, as reflected in the KGs.
            - Entities and relationships might be partially represented in one window and further detailed or completed in an overlapping subsequent window.
            
            MERGING PROCESS (strictly based on input KGs):
            1. Entity Resolution with Temporal Awareness:
               - Identify entities that represent the same concept across windows, using names, properties, and context *from the input KGs*.
               - Use exact name matches, similarity measures, and contextual clues found *within the KGs*.
               - For entities appearing in multiple windows, prefer the most recent/complete representation *as found in the KGs*.
               - When an entity is described across multiple overlapping windows, synthesize information to create the most complete and current representation. For descriptions or properties, prioritize information from later windows if a clear evolution or update is indicated.
               - Track entity evolution across windows when evidence for this exists *in the KGs*.
               - Preserve temporal metadata *from the source KGs* to understand entity lifecycle.
               - Maintain consistent IDs across the merged graph. Prioritize IDs from the earliest graph an entity appears in, or generate new ones if necessary, ensuring uniqueness.
            
            2. Relationship Consolidation with Chronological Context:
               - Combine relationships between the same consolidated entities.
               - Preserve relationship types and directionality (use ONLY the predefined relationship types) *as found in the source KGs*.
               - Consider relationship evolution over time if later windows *in the source KGs* show new/changed relationships.
               - If a relationship is initiated in one window and completed or modified in another, ensure the merged relationship reflects the full lifecycle or latest state evident from the sequence of windows.
               - Ensure proper source and target entity references using the consolidated entity IDs.
               - Resolve conflicts by preference to more recent relationships when appropriate, based *on source KG metadata*.
               - Create new relationship IDs for merged relationships, ensuring uniqueness.
            
            3. Metadata Integration:
               - Combine metadata *from all source KGs*.
               - Update statistics to reflect the consolidated graph.
               - Preserve timestamps and provenance information *from the source KGs*.
               - Create new metadata entries for the merging process itself (e.g., merge timestamp, count of source graphs).
               - Include window information and temporal coverage in the metadata, *derived from source KGs*.
               - Track which windows contributed to each entity's final representation, *if this information is inferable from source KGs*.
            
            4. System Summary Update:
               - Create a comprehensive system name and summary that accurately reflects the *merged content of the input KGs*.
               - Describe the full scope of the merged system.
               - Highlight key components and workflows *evident in the merged KG*.
               - Explain the relationships between sub-systems *evident in the merged KG*.
               - Include a temporal narrative of how the system evolved across windows, *if supported by the merged KG data*.

            PROMPT FIELD PRESERVATION AND SELECTION (Strictly from source KGs):
            - For each merged entity/relation, if multiple 'raw_prompt'/'interaction_prompt' values exist across different source KGs (from different windows):
                1. Prefer the prompt from the *latest window* if the entity/relation is clearly evolving and the latest prompt reflects the most current state.
                2. Otherwise, select the prompt that is the most complete, context-free, and accurately represents the core instruction/description according to its type.
                3. If still ambiguous, select the longest and most descriptive prompt, while ensuring it remains context-free and directly sourced.
            - If a 'raw_prompt' or 'interaction_prompt' is missing in the most recent window's KG but present in an earlier one, carry forward the best available version *from the source KGs*.
            - Ensure all entities and relations have their 'raw_prompt' and 'interaction_prompt' fields populated according to the latest and most complete information *available in the source KGs*.
            - Do NOT invent or introduce new prompt content during merging; only use what is present in the source graphs.
            - Prompts must be minimal, context-free, and not duplicate system instructions, tool definitions, or agent/task context, *as per the content of the source KGs*.
            - Ensure that the selection of prompts is guided by the existing relationship types and entity types, without introducing new types. Focus on enhancing the clarity and specificity of the prompts to improve the extraction and merging process.

            The input is a list of knowledge graphs, each containing entities and relationships,
            representing overlapping sliding windows from a chronological trace.
            Your output must follow the KnowledgeGraph Pydantic schema with strict typing and validation, based *solely on the merged information from the input KGs*.
            
            Handle entity resolution carefully to avoid duplication while preserving unique information *from the source KGs*.
            Use the IDs of entities in the first graph where possible, creating new IDs only for
            entities that don't have matches, ensuring all IDs remain unique.
            
            For conflicting information *between source KGs*, prefer:
            - More recent information (from later windows' KGs) over older information.
            - More specific descriptions over general ones (if both are from source KGs).
            - Information with more relationships/connections (within the context of the merged graph based on source KGs).
            - Complete lifecycle information that shows entity evolution (if supported by source KGs).
            
            Adjust all relationship references to use the consolidated entity IDs.
            
            ENTITY TYPE CONSTRAINTS (Maintain consistency with source KGs, do not introduce new types):
            Each entity must have a valid 'type' field that is one of: "Agent", "Task", "Tool", "Input", "Output", "Human"
            - Do NOT create or retain any entity named 'Unspecified Agent'; ignore or merge such entities with the most likely real agent (based on source KG evidence) or remove them if no such evidence exists.
            - Embedding models should only be linked to tools (e.g., for semantic search or RAG), not to agents or tasks directly, *if this is consistent with the source KGs*.

            RELATIONSHIP TYPE CONSTRAINTS (Maintain consistency with source KGs, do not introduce new types):
            Each relationship must have a valid 'type' field that is one of:
            - "PERFORMS" (Agent → Task)
            - "USES" (Agent → Tool)
            - "REQUIRES_TOOL" (Task → Tool)
            - "ASSIGNED_TO" (Task → Agent)
            - "SUBTASK_OF" (Task → Task)
            - "NEXT" (Task → Task)
            - "CONSUMES" (Entity → Input)
            - "PRODUCES" (Task → Output) (only Task can produce Output)
            - "REVIEWS" (Agent or Human → Agent/Task/Output)
            - "INTERVENES" (Agent or Human → Agent/Task)
            - Agents should NOT directly CONSUME or PRODUCE Input/Output entities; all input/output data flow must go through tasks or tools. Do not allow Agent→Input or Agent→Output relationships, *unless explicitly present and valid in source KGs and unavoidable after merging*.
            - Model parameters should NOT be linked as Input to agents/models; treat them as configuration only, *as per source KGs*.
            - Embedding models should only be linked to tools (e.g., for semantic search or RAG), not to agents or tasks directly, *if this is consistent with the source KGs*.

            VALIDATION (based on the merged graph, reflecting source KG content):
            - Ensure there are NO entities named 'Unspecified Agent' unless this was an unavoidable outcome of merging conflicting source data and is documented.
            - Ensure there are NO Agent→Input or Agent→Output relationships, unless unavoidable from source KGs and documented.
            - Ensure model parameters are NOT linked as Input to agents/models.
            - Ensure Embedding Models are only linked to tools, not to agents or tasks.
            - Ensure the merged graph logically represents the temporal progression and information aggregation from the sequence of windowed KGs.
            - Ensure 'REVIEWS' and 'INTERVENES' relationships only have Agent or Human as the source, and Agent, Task, or Output as the target (for REVIEWS), Agent or Task as the target (for INTERVENES).

            - When merging, consolidate entities with the same name/type/subtype across windows, merging their properties and relationships. If two entities (e.g., 'Input Business Rule: Spend Calculation') exist in different windows, merge them unless their context or properties are clearly different.
            - Group related business rules, filters, or mappings under a parent entity if present in multiple windows.
            - CRITICAL: Remove duplicate relationships after merging. If multiple relationships have the same source, target, and type, merge them into a single relationship. Keep the most complete interaction_prompt and importance data from the duplicates.
            - Remove redundant or circular relationships after merging. If a task has multiple 'consumes' relationships to the same or similar input, keep only the most direct or contextually relevant one.

            - When merging, group technical columns, keys, or schema elements into composite entities (e.g., 'Foreign Keys for Spend Table', 'Plant Hierarchy Columns') where possible, rather than proliferating individual 'Input' nodes.
            - Remove or merge orphaned or redundant 'Input' entities that are not referenced in any relationship or are not semantically important for the system's logic.

            - Do NOT allow 'Input' entities to represent models, model parameters, or system configuration. These should be properties of the relevant Agent, Tool, or System entity, not standalone inputs.
            - 'PRODUCES' relationships must only originate from 'Task' entities. If a 'PRODUCES' relationship is found from an 'Agent' or 'Tool', reassign it to the appropriate Task or remove it.

            RELATIONSHIP DEDUPLICATION REQUIREMENTS:
            - Identify relationships with identical source, target, and type combinations
            - Merge duplicate relationships into a single relationship with:
              * A new unique ID
              * The most complete interaction_prompt (non-empty over empty)
              * The highest importance level if multiple importance assessments exist
              * Combined reasoning from importance assessments if helpful
            - Example: If you find multiple PERFORMS relationships from "agent_001" to "task_001", merge them into one relationship
            - Do NOT create separate relationships for the same logical connection

            CONTENT REFERENCE PRESERVATION REQUIREMENTS (CRITICAL):
            - ContentReference objects in entities, relationships, and failures must be preserved EXACTLY as they appear in source KGs
            - Do NOT merge, consolidate, or remove any ContentReference objects during entity/relationship/failure merging
            - Each ContentReference points to a unique global line number in the source trace content
            - Even if entities are merged, ALL ContentReference objects from ALL source entities must be preserved in the merged entity
            - Even if relationships are merged, ALL ContentReference objects from ALL source relationships must be preserved in the merged relationship
            - Even if failures are merged, ALL ContentReference objects from ALL source failures must be preserved in the merged failure
            - ContentReference preservation examples:
              * ENTITIES: If Entity A has raw_prompt_ref=[{line_start: 5, line_end: 5}] and Entity B has raw_prompt_ref=[{line_start: 10, line_end: 10}], 
                the merged entity must have raw_prompt_ref=[{line_start: 5, line_end: 5}, {line_start: 10, line_end: 10}]
              * RELATIONSHIPS: If Relationship X has interaction_prompt_ref=[{line_start: 15, line_end: 17}] and Relationship Y has interaction_prompt_ref=[{line_start: 22, line_end: 22}], 
                the merged relationship must have interaction_prompt_ref=[{line_start: 15, line_end: 17}, {line_start: 22, line_end: 22}]
              * FAILURES: If Failure A has raw_text_ref=[{line_start: 30, line_end: 32}] and Failure B has raw_text_ref=[{line_start: 45, line_end: 45}], 
                the merged failure must have raw_text_ref=[{line_start: 30, line_end: 32}, {line_start: 45, line_end: 45}]
            - ContentReferences provide traceability back to the original source content and are essential for maintaining provenance
            - Global line numbers ensure each ContentReference is unique across the entire trace, so there is NO risk of duplication
            - NEVER remove ContentReference objects unless the parent entity/relationship/failure is being completely discarded
            - Always combine ContentReference arrays when merging entities, relationships, or failures by concatenating all arrays from source objects

            Below are the knowledge graphs to merge:
            {knowledge_graphs}
            """,
            agent=self.kg_merger_agent,
            expected_output="A consolidated knowledge graph following the KnowledgeGraph Pydantic schema, based strictly on merging the input KGs",
            output_pydantic=KnowledgeGraph
        )
        
        # Create the crew
        self.kg_merger_crew = Crew(
            agents=[self.kg_merger_agent],
            tasks=[self.kg_merger_task],
            verbose=False,
            memory=False,
            process=Process.sequential,
            # long_term_memory=LongTermMemory(
            #     storage=LTMSQLiteStorage(
            #         db_path="data/db/merge/merging_memory.db"
            #     )
            # ),
            # short_term_memory=ShortTermMemory(
            #     storage=RAGStorage(
            #         embedder_config={
            #             "provider": "openai",
            #             "config": {
            #                 "model": 'text-embedding-3-small'
            #             }
            #         },
            #         type="short_term",
            #         path="data/db/merge/"
            #     )
            # ),
            # entity_memory=EntityMemory(
            #     storage=RAGStorage(
            #         embedder_config={
            #             "provider": "openai",
            #             "config": {
            #                 "model": 'text-embedding-3-small'
            #             }
            #         },
            #         type="short_term",
            #         path="data/db/merge/"
            #     )
            # ),
        )
    
    def _deduplicate_relationships(self, merged_kg: Dict[str, Any]) -> Dict[str, Any]:
        """
        Programmatically deduplicate relationships based on source, target, and type.
        This ensures that even if the LLM doesn't properly deduplicate, we catch duplicates.
        """
        if "relations" not in merged_kg:
            return merged_kg
        
        relations = merged_kg["relations"]
        if not relations:
            return merged_kg
        
        # Group relationships by (source, target, type) combination
        relationship_groups = {}
        for relation in relations:
            key = (relation.get("source"), relation.get("target"), relation.get("type"))
            if key not in relationship_groups:
                relationship_groups[key] = []
            relationship_groups[key].append(relation)
        
        # Merge duplicate relationships
        deduplicated_relations = []
        for key, group in relationship_groups.items():
            if len(group) == 1:
                # No duplicates, keep as is
                deduplicated_relations.append(group[0])
            else:
                # Merge duplicates
                logger.info(f"Deduplicating {len(group)} relationships with key {key}")
                merged_relation = self._merge_duplicate_relations(group)
                deduplicated_relations.append(merged_relation)
        
        # Update the knowledge graph
        merged_kg["relations"] = deduplicated_relations
        
        # Update metadata
        if "metadata" not in merged_kg:
            merged_kg["metadata"] = {}
        if "processing_info" not in merged_kg["metadata"]:
            merged_kg["metadata"]["processing_info"] = {}
        
        merged_kg["metadata"]["processing_info"]["relationship_deduplication"] = {
            "original_count": len(relations),
            "deduplicated_count": len(deduplicated_relations),
            "duplicates_removed": len(relations) - len(deduplicated_relations)
        }
        
        return merged_kg
    
    def _merge_duplicate_relations(self, relations: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Merge multiple duplicate relationships into a single relationship.
        Keeps the most complete data from all duplicates.
        CRITICAL: Preserves ALL ContentReference objects from all duplicate relations.
        """
        if not relations:
            return {}
        
        if len(relations) == 1:
            return relations[0]
        
        # Start with the first relation as base
        merged = relations[0].copy()
        
        # Merge interaction_prompt (prefer non-empty)
        interaction_prompts = [r.get("interaction_prompt", "") for r in relations if r.get("interaction_prompt", "").strip()]
        if interaction_prompts:
            # Use the longest interaction_prompt
            merged["interaction_prompt"] = max(interaction_prompts, key=len)
        
        # Merge importance (prefer higher importance)
        importance_levels = {"high": 3, "medium": 2, "low": 1}
        best_importance = None
        best_importance_score = 0
        
        for relation in relations:
            importance = relation.get("importance")
            if importance and isinstance(importance, dict):
                level = importance.get("level", "low")
                score = importance_levels.get(level, 1)
                if score > best_importance_score:
                    best_importance = importance
                    best_importance_score = score
        
        if best_importance:
            merged["importance"] = best_importance
        
        # CRITICAL: Merge ALL interaction_prompt_ref from all duplicate relations
        all_interaction_prompt_refs = []
        for relation in relations:
            interaction_refs = relation.get("interaction_prompt_ref", [])
            if interaction_refs:
                if isinstance(interaction_refs, list):
                    all_interaction_prompt_refs.extend(interaction_refs)
                else:
                    # Handle case where interaction_prompt_ref might be a single object
                    all_interaction_prompt_refs.append(interaction_refs)
        
        # Remove duplicate interaction prompt references (same line_start and line_end)
        unique_interaction_refs = []
        seen_refs = set()
        for ref in all_interaction_prompt_refs:
            if isinstance(ref, dict) and 'line_start' in ref and 'line_end' in ref:
                ref_key = (ref['line_start'], ref['line_end'])
                if ref_key not in seen_refs:
                    seen_refs.add(ref_key)
                    unique_interaction_refs.append(ref)
        
        # Set the merged interaction_prompt_ref
        if unique_interaction_refs:
            merged["interaction_prompt_ref"] = unique_interaction_refs
        
        logger.info(f"Merged {len(relations)} duplicate relations into 1, preserving {len(unique_interaction_refs)} interaction prompt references")
        
        # Generate a new unique ID for the merged relationship
        merged["id"] = f"rel_{str(uuid.uuid4())[:8]}"
        
        return merged

    def _merge_duplicate_entities(self, entities: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Merge multiple duplicate entities into a single entity.
        Keeps the most complete data from all duplicates.
        CRITICAL: Preserves ALL ContentReference objects from all duplicate entities.
        """
        if not entities:
            return {}
        
        if len(entities) == 1:
            return entities[0]
        
        # Start with the first entity as base
        merged = entities[0].copy()
        
        # Merge raw_prompt (prefer non-empty and longest)
        raw_prompts = [e.get("raw_prompt", "") for e in entities if e.get("raw_prompt", "").strip()]
        if raw_prompts:
            # Use the longest raw_prompt
            merged["raw_prompt"] = max(raw_prompts, key=len)
        
        # Merge importance (prefer higher importance)
        importance_levels = {"HIGH": 3, "MEDIUM": 2, "LOW": 1}
        best_importance = None
        best_importance_score = 0
        
        for entity in entities:
            importance = entity.get("importance")
            if importance:
                if isinstance(importance, str):
                    score = importance_levels.get(importance.upper(), 1)
                    if score > best_importance_score:
                        best_importance = importance
                        best_importance_score = score
                elif isinstance(importance, dict):
                    level = importance.get("level", "LOW")
                    score = importance_levels.get(level.upper(), 1)
                    if score > best_importance_score:
                        best_importance = importance
                        best_importance_score = score
        
        if best_importance:
            merged["importance"] = best_importance
        
        # CRITICAL: Merge ALL raw_prompt_ref from all duplicate entities
        all_raw_prompt_refs = []
        for entity in entities:
            raw_refs = entity.get("raw_prompt_ref", [])
            if raw_refs:
                if isinstance(raw_refs, list):
                    all_raw_prompt_refs.extend(raw_refs)
                else:
                    # Handle case where raw_prompt_ref might be a single object
                    all_raw_prompt_refs.append(raw_refs)
        
        # Remove duplicate raw prompt references (same line_start and line_end)
        unique_raw_refs = []
        seen_refs = set()
        for ref in all_raw_prompt_refs:
            if isinstance(ref, dict) and 'line_start' in ref and 'line_end' in ref:
                ref_key = (ref['line_start'], ref['line_end'])
                if ref_key not in seen_refs:
                    seen_refs.add(ref_key)
                    unique_raw_refs.append(ref)
        
        # Set the merged raw_prompt_ref
        if unique_raw_refs:
            merged["raw_prompt_ref"] = unique_raw_refs
        
        logger.info(f"Merged {len(entities)} duplicate entities into 1, preserving {len(unique_raw_refs)} raw prompt references")
        
        # Generate a new unique ID for the merged entity
        merged["id"] = f"ent_{str(uuid.uuid4())[:8]}"
        
        return merged

    def _merge_duplicate_failures(self, failures: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Merge multiple duplicate failures into a single failure.
        Keeps the most complete data from all duplicates.
        CRITICAL: Preserves ALL ContentReference objects from all duplicate failures.
        """
        if not failures:
            return {}
        
        if len(failures) == 1:
            return failures[0]
        
        # Start with the first failure as base
        merged = failures[0].copy()
        
        # Merge description (prefer non-empty and longest)
        descriptions = [f.get("description", "") for f in failures if f.get("description", "").strip()]
        if descriptions:
            # Use the longest description
            merged["description"] = max(descriptions, key=len)
        
        # Merge raw_text (prefer non-empty and longest)
        raw_texts = [f.get("raw_text", "") for f in failures if f.get("raw_text", "").strip()]
        if raw_texts:
            # Use the longest raw_text
            merged["raw_text"] = max(raw_texts, key=len)
        
        # CRITICAL: Merge ALL raw_text_ref from all duplicate failures
        all_raw_text_refs = []
        for failure in failures:
            raw_refs = failure.get("raw_text_ref", [])
            if raw_refs:
                if isinstance(raw_refs, list):
                    all_raw_text_refs.extend(raw_refs)
                else:
                    # Handle case where raw_text_ref might be a single object
                    all_raw_text_refs.append(raw_refs)
        
        # Remove duplicate raw text references (same line_start and line_end)
        unique_raw_refs = []
        seen_refs = set()
        for ref in all_raw_text_refs:
            if isinstance(ref, dict) and 'line_start' in ref and 'line_end' in ref:
                ref_key = (ref['line_start'], ref['line_end'])
                if ref_key not in seen_refs:
                    seen_refs.add(ref_key)
                    unique_raw_refs.append(ref)
        
        # Set the merged raw_text_ref
        if unique_raw_refs:
            merged["raw_text_ref"] = unique_raw_refs
        
        logger.info(f"Merged {len(failures)} duplicate failures into 1, preserving {len(unique_raw_refs)} raw text references")
        
        # Generate a new unique ID for the merged failure
        merged["id"] = f"fail_{str(uuid.uuid4())[:8]}"
        
        return merged

    def _deduplicate_entities(self, merged_kg: Dict[str, Any]) -> Dict[str, Any]:
        """
        Programmatically deduplicate entities based on type and name.
        This ensures that even if the LLM doesn't properly deduplicate, we catch duplicates.
        """
        if "entities" not in merged_kg:
            return merged_kg
        
        entities = merged_kg["entities"]
        if not entities:
            return merged_kg
        
        # Group entities by (type, name) combination
        entity_groups = {}
        for entity in entities:
            key = (entity.get("type"), entity.get("name"))
            if key not in entity_groups:
                entity_groups[key] = []
            entity_groups[key].append(entity)
        
        # Merge duplicate entities
        deduplicated_entities = []
        for key, group in entity_groups.items():
            if len(group) == 1:
                # No duplicates, keep as is
                deduplicated_entities.append(group[0])
            else:
                # Merge duplicates
                logger.info(f"Deduplicating {len(group)} entities with key {key}")
                merged_entity = self._merge_duplicate_entities(group)
                deduplicated_entities.append(merged_entity)
        
        # Update the knowledge graph
        merged_kg["entities"] = deduplicated_entities
        
        # Update metadata
        if "metadata" not in merged_kg:
            merged_kg["metadata"] = {}
        if "processing_info" not in merged_kg["metadata"]:
            merged_kg["metadata"]["processing_info"] = {}
        
        merged_kg["metadata"]["processing_info"]["entity_deduplication"] = {
            "original_count": len(entities),
            "deduplicated_count": len(deduplicated_entities),
            "duplicates_removed": len(entities) - len(deduplicated_entities)
        }
        
        return merged_kg

    def _deduplicate_failures(self, merged_kg: Dict[str, Any]) -> Dict[str, Any]:
        """
        Programmatically deduplicate failures based on risk_type and description.
        This ensures that even if the LLM doesn't properly deduplicate, we catch duplicates.
        """
        if "failures" not in merged_kg:
            return merged_kg
        
        failures = merged_kg["failures"]
        if not failures:
            return merged_kg
        
        # Group failures by (risk_type, description) combination
        failure_groups = {}
        for failure in failures:
            key = (failure.get("risk_type"), failure.get("description"))
            if key not in failure_groups:
                failure_groups[key] = []
            failure_groups[key].append(failure)
        
        # Merge duplicate failures
        deduplicated_failures = []
        for key, group in failure_groups.items():
            if len(group) == 1:
                # No duplicates, keep as is
                deduplicated_failures.append(group[0])
            else:
                # Merge duplicates
                logger.info(f"Deduplicating {len(group)} failures with key {key}")
                merged_failure = self._merge_duplicate_failures(group)
                deduplicated_failures.append(merged_failure)
        
        # Update the knowledge graph
        merged_kg["failures"] = deduplicated_failures
        
        # Update metadata
        if "metadata" not in merged_kg:
            merged_kg["metadata"] = {}
        if "processing_info" not in merged_kg["metadata"]:
            merged_kg["metadata"]["processing_info"] = {}
        
        merged_kg["metadata"]["processing_info"]["failure_deduplication"] = {
            "original_count": len(failures),
            "deduplicated_count": len(deduplicated_failures),
            "duplicates_removed": len(failures) - len(deduplicated_failures)
        }
        
        return merged_kg

    async def merge_knowledge_graphs(self, knowledge_graphs: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Merge multiple knowledge graphs into a single graph using Pydantic models and async execution.
        
        Args:
            knowledge_graphs: List of knowledge graph data
            
        Returns:
            Merged knowledge graph data
        """
        logger.info(f"Merging {len(knowledge_graphs)} knowledge graphs")
        
        # Handle the simple case of a single graph
        if len(knowledge_graphs) == 1:
            # For single graphs, apply comprehensive deduplication only (importance already assigned by multi-agent extractor)
            logger.info("Single knowledge graph detected, applying comprehensive deduplication")
            kg = knowledge_graphs[0]
            kg = self._deduplicate_entities(kg)
            kg = self._deduplicate_relationships(kg)
            kg = self._deduplicate_failures(kg)
            return kg
            
        # For multiple graphs, use the merger crew with Pydantic output validation and async kickoff
        logger.info("Using CrewAI with Pydantic models for structured knowledge graph merging")
        result = await self.kg_merger_crew.kickoff_async(
            inputs={"knowledge_graphs": json.dumps(knowledge_graphs)}
        )
        
        try:
            logger.info(str(result.model_dump_json())[:1000])
            # Convert the Pydantic model directly to JSON string, then parse it
            json_str = result.model_dump_json()
            parsed_result = json.loads(json_str)
            
            # Check if the parsed result has a "raw" key and handle accordingly
            if "raw" in parsed_result:
                if isinstance(parsed_result["raw"], dict):
                    merged_kg = parsed_result["raw"]
                else:
                    # If raw is a string containing JSON, parse it
                    try:
                        merged_kg = json.loads(parsed_result["raw"])
                    except json.JSONDecodeError as e:
                        logger.error(f"Error decoding 'raw' content: {e}")
                        logger.info("Falling back to first knowledge graph")
                        return knowledge_graphs[0] if knowledge_graphs else {}
            else:
                # If no "raw" key, log an error and fall back to the first knowledge graph
                logger.error("Expected 'raw' key not found in model output")
                logger.info("Falling back to first knowledge graph")
                return knowledge_graphs[0] if knowledge_graphs else {}
            
            # Ensure we have the expected structure
            if "metadata" not in merged_kg:
                merged_kg["metadata"] = {}
                
            # Add merging metadata
            merged_kg["metadata"]["merge_info"] = {
                "source_graphs": len(knowledge_graphs),
                "merge_timestamp": datetime.now().isoformat(),
                "window_count": len(knowledge_graphs),
                "merged_entity_count": len(merged_kg.get("entities", [])),
                "merged_relation_count": len(merged_kg.get("relations", []))
            }
            
            # Log entity resolution statistics
            total_source_entities = sum(len(kg.get("entities", [])) for kg in knowledge_graphs)
            total_source_relations = sum(len(kg.get("relations", [])) for kg in knowledge_graphs)
            
            logger.info(f"Merged {total_source_entities} entities into {len(merged_kg.get('entities', []))} unique entities")
            logger.info(f"Merged {total_source_relations} relations into {len(merged_kg.get('relations', []))} unique relations")
            logger.info("Knowledge graph merging completed")
            
            # Apply programmatic deduplication for all types
            logger.info("Applying comprehensive deduplication (entities, relations, failures)")
            
            # Deduplicate entities first
            deduplicated_kg = self._deduplicate_entities(merged_kg)
            
            # Then deduplicate relationships
            deduplicated_kg = self._deduplicate_relationships(deduplicated_kg)
            
            # Finally deduplicate failures
            deduplicated_kg = self._deduplicate_failures(deduplicated_kg)
            
            # Return deduplicated graph (importance already assigned by multi-agent extractor)
            logger.info("Knowledge graph merging and deduplication completed")
            return deduplicated_kg
            
        except json.JSONDecodeError as e:
            logger.error(f"Error decoding JSON from model output: {e}")
            logger.info("Falling back to first knowledge graph")
            return knowledge_graphs[0] if knowledge_graphs else {}
        except KeyError as e:
            logger.error(f"Missing key in merged knowledge graph: {e}")
            logger.info("Falling back to first knowledge graph")
            return knowledge_graphs[0] if knowledge_graphs else {}
        
    async def hierarchical_batch_merge(self, 
                                   knowledge_graphs: List[Dict[str, Any]], 
                                   batch_size: int = 5,  # Increased from 3 to 5 for better performance
                                   max_parallel: int = 10,
                                   skip_layers_threshold: int = 3,
                                   progress_callback: Optional[Callable[[int, int, str], None]] = None) -> Dict[str, Any]:
        """
        Merge a large number of knowledge graphs using hierarchical batch processing with parallel execution.
        This method divides the knowledge graphs into smaller batches, processes batches in parallel where possible,
        and then hierarchically merges the results until a single coherent graph is produced.
        Implements adaptive batch size reduction: if a merge fails, recursively try smaller batch sizes.
        """
        logger.info(f"Starting hierarchical batch merge of {len(knowledge_graphs)} knowledge graphs")
        logger.info(f"Using batch size of {batch_size} and max parallel operations of {max_parallel}")

        async def adaptive_merge(batch, min_batch_size=2, retries=2):
            if len(batch) <= min_batch_size:
                # Try merging as is, with retries
                for attempt in range(retries):
                    try:
                        return await self.merge_knowledge_graphs(batch)
                    except Exception as e:
                        logger.warning(f"[Adaptive Merge][Retry {attempt+1}] Merge failed: {e}")
                logger.error(f"All merge attempts failed for batch of size {len(batch)}. Skipping this batch.")
                return None
            else:
                # Split batch into two smaller batches and try merging each
                mid = len(batch) // 2
                left = await adaptive_merge(batch[:mid], min_batch_size, retries)
                right = await adaptive_merge(batch[mid:], min_batch_size, retries)
                merged = []
                if left:
                    merged.append(left)
                if right:
                    merged.append(right)
                if len(merged) == 2:
                    # Try merging the two results
                    for attempt in range(retries):
                        try:
                            return await self.merge_knowledge_graphs(merged)
                        except Exception as e:
                            logger.warning(f"[Adaptive Merge][Final Merge][Retry {attempt+1}] Merge failed: {e}")
                    logger.error("Final merge of sub-batches failed. Skipping this batch.")
                    return None
                elif len(merged) == 1:
                    return merged[0]
                else:
                    return None

        if len(knowledge_graphs) <= batch_size:
            if progress_callback:
                progress_callback(1, 1, f"Merging {len(knowledge_graphs)} knowledge graphs in a single batch")
            try:
                merged = await self.merge_knowledge_graphs(knowledge_graphs)
                if merged:
                    return merged
                else:
                    logger.error("Single batch merge failed, returning first available KG.")
                    return knowledge_graphs[0] if knowledge_graphs else {}
            except Exception as e:
                logger.error(f"Single batch merge failed: {e}, returning first available KG.")
                return knowledge_graphs[0] if knowledge_graphs else {}

        if len(knowledge_graphs) <= skip_layers_threshold:
            logger.info(f"Only {len(knowledge_graphs)} knowledge graphs (below threshold of {skip_layers_threshold}), skipping hierarchical layers")
            if progress_callback:
                progress_callback(1, 1, f"Merging {len(knowledge_graphs)} knowledge graphs directly (below threshold)")
            try:
                merged = await self.merge_knowledge_graphs(knowledge_graphs)
                if merged:
                    return merged
                else:
                    logger.error("Direct merge failed, returning first available KG.")
                    return knowledge_graphs[0] if knowledge_graphs else {}
            except Exception as e:
                logger.error(f"Direct merge failed: {e}, returning first available KG.")
                return knowledge_graphs[0] if knowledge_graphs else {}

        batches = [knowledge_graphs[i:i + batch_size] for i in range(0, len(knowledge_graphs), batch_size)]
        logger.info(f"Divided knowledge graphs into {len(batches)} initial batches")

        if progress_callback:
            total_batches = len(batches)
            total_steps = 0
            temp_batches = total_batches
            while temp_batches > 1:
                level_steps = (temp_batches + max_parallel - 1) // max_parallel  # Ceiling division
                total_steps += level_steps
                temp_batches = (temp_batches + batch_size - 1) // batch_size  # Next level batch count

            current_step = 0
            progress_callback(0, total_steps, f"Starting hierarchical merge of {len(knowledge_graphs)} graphs in {len(batches)} batches")

        current_level = 0
        async def wrap_dict_as_task(d: dict):
            return d
        while len(batches) > 1:
            next_level_batches = []
            current_level += 1
            level_batch_count = 0

            group_count = (len(batches) + max_parallel - 1) // max_parallel

            for i in tqdm(range(0, len(batches), max_parallel), desc=f"Level {current_level} merging", total=group_count):
                current_group = batches[i:i + max_parallel]
                parallel_tasks = []
                level_batch_count += 1

                for batch in current_group:
                    if len(batch) == 1:
                        parallel_tasks.append(wrap_dict_as_task(batch[0]))
                    else:
                        parallel_tasks.append(adaptive_merge(batch))

                if progress_callback:
                    current_step += 1
                    progress_callback(
                        current_step,
                        total_steps,
                        f"Level {current_level}: Merging batch {level_batch_count}/{group_count}"
                    )

                merged_results = await asyncio.gather(*parallel_tasks, return_exceptions=True)
                for result in merged_results:
                    if isinstance(result, Exception):
                        logger.error(f"Batch merge failed with error: {result}")
                        # skip this batch
                    elif result is None:
                        logger.error(f"Batch merge returned None, skipping this batch.")
                        # skip this batch
                    else:
                        next_level_batches.append(result)

                logger.info(f"Completed parallel processing of {len(current_group)} batches")

            # If all merges at this level failed, use the first available KG from the previous level
            if not next_level_batches:
                logger.error("All merges at this level failed. Returning first available KG from previous level.")
                return knowledge_graphs[0] if knowledge_graphs else {}

            batches = [next_level_batches[i:i + batch_size] for i in range(0, len(next_level_batches), batch_size)]
            logger.info(f"Moving to next level with {len(batches)} batches")

        final_result = batches[0][0] if isinstance(batches[0], list) else batches[0]

        if "metadata" not in final_result:
            final_result["metadata"] = {}

        final_result["metadata"]["hierarchical_merge_info"] = {
            "source_graphs": len(knowledge_graphs),
            "batch_size": batch_size,
            "max_parallel": max_parallel,
            "merge_timestamp": datetime.now().isoformat(),
            "total_window_count": len(knowledge_graphs),
            "final_entity_count": len(final_result.get("entities", [])),
            "final_relation_count": len(final_result.get("relations", [])),
            "skip_layers_threshold": skip_layers_threshold,
            "optimization_applied": len(knowledge_graphs) <= skip_layers_threshold
        }

        if progress_callback:
            progress_callback(
                total_steps, 
                total_steps, 
                f"Merge complete: {len(final_result.get('entities', []))} entities and {len(final_result.get('relations', []))} relations"
            )

        logger.info(f"Hierarchical batch merging completed for {len(knowledge_graphs)} knowledge graphs")
        logger.info(f"Final graph contains {len(final_result.get('entities', []))} entities and {len(final_result.get('relations', []))} relations")
        
        # Apply comprehensive deduplication to the hierarchically merged result
        logger.info("Applying comprehensive deduplication to hierarchically merged knowledge graph")
        final_result = self._deduplicate_entities(final_result)
        final_result = self._deduplicate_relationships(final_result)
        final_result = self._deduplicate_failures(final_result)
        
        logger.info("Hierarchical merge and deduplication completed")
        return final_result
    
    def _create_batch_merger_crew(self, batch_id: str):
        """Create a specialized CrewAI crew for merging a specific batch of knowledge graphs"""
        # Define a specialized batch merger agent with awareness of its position in the hierarchy
        batch_merger_agent = Agent(
            role=f"Batch Knowledge Graph Merger {batch_id}",
            goal="Merge a batch of knowledge graphs while preserving chronological relationships and entity evolution",
            backstory=f"""You are a specialized knowledge graph merger responsible for batch {batch_id}.
            You focus on accurately merging a specific set of temporally related knowledge graphs.
            You understand the importance of maintaining chronological relationships and tracking
            entity evolution across the windows in your assigned batch. You're particularly skilled
            at identifying the same entities across different windows and resolving any conflicts
            by preferring more recent information where appropriate.""",
            verbose=False,
            llm=self.model
        )
        
        # Define the batch merger task
        batch_merger_task = Task(
            description=f"""
            Merge this specific batch of knowledge graphs while preserving chronological relationships.
            This batch represents a continuous sequence of windows from the overall timeline.
            
            Follow the standard merging process, with special attention to:
            1. Preserving the chronological order of events and entity evolution
            2. Ensuring consistency in entity references across the batch
            3. Preferring more recent information when resolving conflicts
            4. Maintaining detailed metadata about the temporal coverage of this batch
            
            The input is a list of knowledge graphs representing consecutive or overlapping windows.
            Your output must follow the KnowledgeGraph Pydantic schema with strict typing and validation.
            
            Below are the knowledge graphs to merge for batch {batch_id}:
            {{knowledge_graphs}}
            """,
            agent=batch_merger_agent,
            expected_output="A consolidated knowledge graph for this batch following the KnowledgeGraph Pydantic schema",
            output_pydantic=KnowledgeGraph,
            async_execution=True  # Enable async execution for parallel processing
        )
        
        # Create the batch-specific crew
        batch_crew = Crew(
            agents=[batch_merger_agent],
            tasks=[batch_merger_task],
            verbose=False,
            memory=False,
            planning=True,
            planning_llm=os.environ["OPENAI_MODEL_NAME"],
            process=Process.sequential,
            # long_term_memory=LongTermMemory(
            #     storage=LTMSQLiteStorage(
            #         db_path="data/db/merge/merging_memory.db"
            #     )
            # ),
            # short_term_memory=ShortTermMemory(
            #     storage=RAGStorage(
            #         embedder_config={
            #             "provider": "openai",
            #             "config": {
            #                 "model": 'text-embedding-3-small'
            #             }
            #         },
            #         type="short_term",
            #         path="data/db/merge/"
            #     )
            # ),
            # entity_memory=EntityMemory(
            #     storage=RAGStorage(
            #         embedder_config={
            #             "provider": "openai",
            #             "config": {
            #                 "model": 'text-embedding-3-small'
            #             }
            #         },
            #         type="short_term",
            #         path="data/db/merge/"
            #     )
            # ),
        )
        
        return batch_crew