Update ner_tool.py

ner_tool.py

@@ -7,7 +7,7 @@ class NamedEntityRecognitionTool(Tool):
     name = "ner_tool"
     description = "Identifies and labels entities such as persons, organizations, and locations in a given text."
     inputs = ["text"]
-    outputs = ["text"]
+    outputs = ["entities"]
 
     def __call__(self, text: str):
         # Initialize the named entity recognition pipeline
@@ -16,10 +16,13 @@ class NamedEntityRecognitionTool(Tool):
         # Perform named entity recognition on the input text
         entities = ner_analyzer(text)
 
-        # Print the identified entities
-        print(f"Identified Entities: {entities}")
+        # Extract relevant information for each identified entity
+        entity_info = [{"entity": entity.get("entity", "UNKNOWN"), "word": entity.get("word", ""), "start": entity.get("start", -1), "end": entity.get("end", -1)} for entity in entities]
+
+        # Extract the actual text span for each identified location entity
+        location_entities = [text[start:end] for entity in entity_info if entity["entity"] == "I-LOC" for start, end in [(entity["start"], entity["end"])]]
 
-        # Extract entity labels and return as a list
-        entity_labels = [entity.get("entity", "UNKNOWN") for entity in entities]
+        # Print the identified entities
+        print(f"Identified Location Entities: {location_entities}")
 
-        return {"entities": entity_labels}  # Return a dictionary with the specified output component
+        return {"entities": location_entities}  # Return a dictionary with the specified output component