11from abc import ABC , abstractmethod
22from typing import List , Tuple , Dict , Any
33from .primitives import Entity , Relation
4+ from .parsers .base_parser import BaseParser
5+ from .parsers .prompts import DELETE_PROMPT , UPDATE_ENTITIES_PROMPT
6+ import requests
7+ import json
48
59class Extractor (ABC ):
610 @abstractmethod
@@ -15,16 +19,111 @@ def extract_relations(self) -> List[Relation]:
1519 def entities_json_schema (self ) -> Dict [str , Any ]:
1620 pass
1721
22+ @abstractmethod
23+ def update_entities (self , new_entities : List [Entity ]) -> List [Entity ]:
24+ pass
25+
1826class FileExtractor (Extractor ):
19- def __init__ (self , file_path : str , parser ):
27+ def __init__ (self , file_path : str , parser : BaseParser ):
2028 self .file_path = file_path
2129 self .parser = parser
2230
2331 def extract_entities (self ) -> List [Entity ]:
24- return self .parser .extract_entities (self .file_path )
32+ new_entities = self .parser .extract_entities (self .file_path )
33+ return self .update_entities (new_entities )
2534
2635 def extract_relations (self ) -> List [Relation ]:
2736 return self .parser .extract_relations (self .file_path )
2837
2938 def entities_json_schema (self ) -> Dict [str , Any ]:
30- return self .parser .entities_json_schema (self .file_path )
39+ return self .parser .entities_json_schema (self .file_path )
40+
41+ def delete_entity_or_relation (self , item_description : str ) -> None :
42+ """
43+ Delete an entity or relation based on user description.
44+
45+ :param item_description: User's description of the entity or relation to delete
46+ """
47+ entities_ids = [e .id for e in self .parser .get_entities ()]
48+ relations_ids = [(r .source , r .target , r .name ) for r in self .parser .get_relations ()]
49+ prompt = DELETE_PROMPT .format (
50+ entities = entities_ids ,
51+ relations = relations_ids ,
52+ item_description = item_description
53+ )
54+
55+ response = self ._get_llm_response (prompt )[8 :- 3 ]
56+ response_dict = json .loads (response )
57+
58+ for key , value in response_dict .items ():
59+ if key == 'Type' :
60+ if value == 'Entity' :
61+ self ._delete_entity (response_dict ['ID' ])
62+ elif value == 'Relation' :
63+ self ._delete_relation (response_dict ['ID' ])
64+
65+
66+ def _delete_entity (self , entity_id : str ) -> None :
67+ """Delete an entity and its related relations."""
68+ entities = self .parser .get_entities ()
69+ relations = self .parser .get_relations ()
70+
71+ entities = [e for e in entities if e .id != entity_id ]
72+ relations = [r for r in relations if r .source != entity_id and r .target != entity_id ]
73+
74+ self .parser .set_entities (entities )
75+ self .parser .set_relations (relations )
76+ print (f"Entity '{ entity_id } )
77+
78+ def _delete_relation (self , relation_id : str ) -> None :
79+ """Delete a relation."""
80+ relations = self .parser .get_relations ()
81+
82+ source , target , name = eval (relation_id )
83+ relations = [r for r in relations if not (r .source == source and r .target == target and r .name == name )]
84+
85+ self .parser .set_relations (relations )
86+ print (f"Relation '{ name } { source } { target } )
87+
88+ def _get_llm_response (self , prompt : str ) -> str :
89+ """Get a response from the language model."""
90+ payload = {
91+ "model" : self .parser .get_model (),
92+ "temperature" : self .parser .get_temperature (),
93+ "messages" : [
94+ {"role" : "user" , "content" : prompt }
95+ ],
96+ }
97+ response = requests .post (self .parser .get_inference_base_url (), headers = self .parser .get_headers (), json = payload )
98+ return response .json ()['choices' ][0 ]['message' ]['content' ]
99+
100+ def update_entities (self , new_entities : List [Entity ]) -> List [Entity ]:
101+ """
102+ Update the existing entities with new entities, integrating and deduplicating as necessary.
103+
104+ :param new_entities: List of new entities to be integrated
105+ :return: Updated list of entities
106+ """
107+ existing_entities = self .parser .get_entities ()
108+
109+ # Prepare the prompt for the LLM
110+ prompt = UPDATE_ENTITIES_PROMPT .format (
111+ existing_entities = json .dumps ([e .__dict__ for e in existing_entities ], indent = 2 ),
112+ new_entities = json .dumps ([e .__dict__ for e in new_entities ], indent = 2 )
113+ )
114+
115+ # Get the LLM response
116+ response = self ._get_llm_response (prompt )
117+
118+ try :
119+ updated_entities_data = json .loads (response )
120+ updated_entities = [Entity (** entity_data ) for entity_data in updated_entities_data ]
121+
122+ # Update the parser's entities
123+ self .parser .set_entities (updated_entities )
124+
125+ print (f"Entities updated. New count: { len (updated_entities )} )
126+ return updated_entities
127+ except json .JSONDecodeError :
128+ print ("Error: Unable to parse the LLM response." )
129+ return existing_entities
0 commit comments