entity_extractor.py

import sys, os 

sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from datetime import datetime
from typing import Dict, Any, Optional
from uuid import uuid4
from langchain_community.graphs.graph_document import GraphDocument, Node, Relationship
from langchain_core.documents import Document
from utils.llm import LLM
from utils.output_parser import GraphOutputParser
import re 
from prompt.graph_extractor import get_graph_extractor_prompt


import logging 

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class EntityExtractor:

    def __init__(self, graph):
        self.llm = LLM.get_gpt_model()
        self.graph = graph  


    def extract(self, text: str, source_id: Optional[str] = None) -> Optional[GraphDocument]:
        """
        Extract entities and relationships from the given text using LLM.
        If source_id is not provided, a new one will be generated.
        """

        if not source_id:
            source_id = f"source_{uuid4().hex}"
        
        source_node = self._create_or_get_source_node(source_id, text)
            
        prompt = get_graph_extractor_prompt(text)
        
        try:
            chain = self.llm | GraphOutputParser() 
            data = chain.invoke(prompt)
            graph_doc = self._build_graph_document(data, source_node, text)
            return graph_doc
        except Exception as e:
            logger.error(f"Error extracting entities and relations: {e}")
            return None
    

    def _create_or_get_source_node(self, source_id: str, text: str) -> Node:

        if self.graph:
            query = """
            MATCH (s:Source {id: $source_id})
            RETURN s.id as id, s.text as text, s.created_at as created_at
            """

            try:
                result = self.graph.query(query, {"source_id": source_id})
                        
                if result and len(result) > 0:
                    # Source node exists, return it
                    source_data = result[0]
                    return Node(
                        id=source_id,
                        type="Source",
                        properties={
                            "id": source_id,
                            "text": source_data.get("text", "") + text,
                            "created_at": source_data.get("created_at", str(datetime.now()))
                        }
                    )
            except Exception as e:
                print(f"Error checking for existing source node: {e}")
        

        return Node(
            id=source_id,
            type="Source",
            properties={
                "id": source_id,
                "text": text,
                "created_at": str(datetime.now())
            }
        )

    def _build_graph_document(self, data: Dict[str, Any], source_node: Node, text: str) -> GraphDocument:
    
        graph_doc = GraphDocument(
            nodes=[source_node],
            relationships=[],
            source=Document(page_content=text, metadata={"source_id": source_node.id})
        )
    
        entity_map = {source_node.id: source_node}

        # Process entities 
        for entity in data.get("entities", []):
        
            entity_id = entity.get("id", f"entity_{uuid4().hex}")
            entity_type = entity.get("type", "Entity")
            
            properties = {
                "id": entity_id,
                "name": entity.get("name", "Unknown"),
                "description": entity.get("description", ""),
                "type": entity_type
            }
   
            for key, value in entity.get("attributes", {}).items():
                properties[key] = value
                
            for key, value in entity.get("temporal_attributes", {}).items():
                properties[key] = value
            
            entity_node = self._create_or_get_entity_node(entity_id, entity_type, properties)
            
            graph_doc.nodes.append(entity_node)
            entity_map[entity_id] = entity_node
            
            # Connect to source node
            source_rel = Relationship(
                source=entity_node,
                target=source_node,
                type="SOURCED_FROM",
                properties={"description": f"Entity extracted from {source_node.id}"}
            )
            graph_doc.relationships.append(source_rel)
            
            temporal_attrs = entity.get("temporal_attributes", {})
            self._process_timepoints(temporal_attrs, entity_node, graph_doc, entity_map)
            self._process_timeranges(temporal_attrs, entity_node, graph_doc, entity_map)
        
        # Process relationships
        for rel in data.get("relationships", []):
            source_id = rel.get("source_id")
            target_id = rel.get("target_id")
            
            if source_id not in entity_map or target_id not in entity_map:
                print(f"Warning: Relationship refers to missing entities: {source_id} -> {target_id}")
                continue
            
            rel_type = rel.get("relationship_type", "RELATES_TO")
            properties = {
                "description": rel.get("description", ""),
                "source_id": source_node.id
            }
            
            relationship = Relationship(
                source=entity_map[source_id],
                target=entity_map[target_id],
                type=rel_type,
                properties=properties
            )
            graph_doc.relationships.append(relationship)
            
        return graph_doc
    
    def _create_or_get_entity_node(self, entity_id: str, entity_type: str, properties: Dict[str, Any]) -> Dict[str, Any]:
        """
        Create a new entity node or retrieve an existing one.
        """
        if self.graph: 
            
            try:
                query = """
                MATCH (e {id: $entity_id})
                RETURN e
                """
                result = self.graph.query(query, {"entity_id": entity_id})
                
                if result and len(result) > 0:
                   
                    existing_entity = result[0]["e"]
                    merged_properties = {**properties, **existing_entity}
                    merged_properties["id"] = entity_id 
                    
                    return Node(
                        id=entity_id,
                        type=entity_type,
                        properties=merged_properties
                    )
            except Exception as e:
                logger.error(f"Error checking for existing entity node {entity_id}: {e}")

        return Node(
            id=entity_id,
            type=entity_type,
            properties=properties
        )

    def _process_timepoints(self, temporal_attrs, entity_node, graph_doc, entity_map):
        """Process time point information."""

        date = temporal_attrs.get("date", "")
        date_pattern = re.compile(r"^\d{4}-\d{2}-\d{2}$")
      
        if not date or not date_pattern.match(date):
            return
            
        calendar_system = temporal_attrs.get("calendar_system", "gregorian")
        day_of_week = temporal_attrs.get("day_of_week", "")
        
        tp_id = f"timepoint_{date}_{calendar_system}"
    
        tp_node = entity_map.get(tp_id)
        
        if not tp_node:
            tp_node = Node(
                id=tp_id,
                type="TimePoint",
                properties={
                    "id": tp_id,
                    "date": date,
                    "calendar_system": calendar_system,
                    "day_of_week": day_of_week
                }
            )
            graph_doc.nodes.append(tp_node)
            entity_map[tp_id] = tp_node
        
        # Create relationship from entity to timepoint
        rel = Relationship(
            source=entity_node,
            target=tp_node,
            type="OCCURRED_AT",
            properties={"description": f"{entity_node.id} occurred at {date}"}
        )
        graph_doc.relationships.append(rel)
    
    def _process_timeranges(self, temporal_attrs, entity_node, graph_doc, entity_map):
        """Process time range information."""

        start_date = temporal_attrs.get("start_date", "")
        end_date = temporal_attrs.get("end_date", "")
        
        if not start_date or not end_date:
            return
            
        duration_days = temporal_attrs.get("duration_days", 0)
        tr_id = f"timerange_{start_date}_{end_date}"
        
        # Check if timerange already exists
        tr_node = entity_map.get(tr_id)
        
        if not tr_node:
            # Create new timerange node
            tr_node = Node(
                id=tr_id,
                type="TimeRange",
                properties={
                    "id": tr_id,
                    "start_date": start_date,
                    "end_date": end_date,
                    "duration_days": duration_days
                }
            )
            graph_doc.nodes.append(tr_node)
            entity_map[tr_id] = tr_node
        
        # Create relationship from entity to timerange
        rel = Relationship(
            source=entity_node,
            target=tr_node,
            type="HAPPENED_DURING",
            properties={"description": f"{entity_node.id} occurred during {start_date} to {end_date}"}
        )
        graph_doc.relationships.append(rel)