Coverage for collectra/tasks/canonicalisers.py: 0.00%

1"""Entity canonicalization and text processing tasks for Collectra workflows. 

2 

3This module provides specialized tasks for entity recognition and canonicalization 

4operations. It includes functionality for matching text entities against known 

5canonical forms using fuzzy string matching algorithms. 

6 

7The module supports: 

8 - Entity canonicalization with configurable similarity thresholds 

9 - Fuzzy string matching using difflib algorithms 

10 - Entity lists loaded from files or provided directly 

11 

12Classes: 

13 EntityCanonicalizer: Task for canonicalizing entities in text 

14""" 

15 

16__all__ = ["LLMCanonicaliser"] 

17 

18import re 

19from pathlib import Path 

20 

21from dotenv import load_dotenv 

22from langchain_chroma import Chroma 

23from langchain_core.documents import Document 

24from langchain_core.messages import HumanMessage 

25 

26from collectra.tasks.llms import LLM as CollectraLLM 

27from collectra.types.texts import Text 

28 

29load_dotenv() 

30 

31 

32class LLMCanonicaliser(CollectraLLM): 

33 """Task for canonicalizing entities in text using fuzzy string matching. 

34 

35 Matches input text against a list of known entities and returns the best 

36 canonical match if it exceeds the similarity threshold. Can load entity 

37 lists from files or use provided lists directly. 

38 

39 Attributes: 

40 entities (list[str]|Path): List of canonical entities or path to file containing them. 

41 threshold (float): Minimum similarity threshold for matches (0.0 to 1.0). 

42 """ 

43 

44 def __init__(self, name: str, model: str, **kwargs): 

45 super().__init__(name, model, **kwargs) 

46 self.embedding_model = kwargs.get("embedding_model", "text-embedding-3-large") 

47 self.threshold = kwargs.get("threshold", 0.8) 

48 self.count = kwargs.get("count", 5) 

49 self.entities = kwargs.get("entities", []) 

50 entities_path = Path(self.entities) 

51 if not entities_path.exists() and entities_path.suffix == "": 

52 if isinstance(self.entities, str) or not isinstance(self.entities, list): 

53 self.entities = [entity for entity in str(self.entities).split(",")] 

54 else: 

55 self.entities = entities_path 

56 

57 def retrieve_entities(self, query_data: str) -> dict: 

58 from langchain_openai import OpenAIEmbeddings 

59 

60 index_path = Path(self.entities.stem + "_index.embed") 

61 embeddings = OpenAIEmbeddings(model=self.embedding_model) 

62 if index_path.exists() and index_path.is_dir(): 

63 print(f"Loading existing vector store from {index_path}") 

64 vector_store = Chroma( 

65 collection_name=f"{self.entities.stem}_collection", 

66 embedding_function=embeddings, 

67 persist_directory=str(index_path), 

68 ) 

69 else: 

70 with open(self.entities, "r") as f: 

71 raw_data = [line.strip() for line in f.readlines() if line.strip()] 

72 ids = [str(i) for i in range(len(raw_data))] 

73 documents = [Document(page_content=content) for content in raw_data] 

74 vector_store = Chroma( 

75 collection_name=f"{self.entities.stem}_collection", 

76 embedding_function=embeddings, 

77 persist_directory=str(index_path), 

78 ) 

79 vector_store.add_documents(documents=documents, ids=ids) 

80 

81 entities = dict() 

82 

83 results = vector_store.similarity_search_with_score(query_data, k=self.count) 

84 for result in results: 

85 if result[0].page_content in entities: 

86 entities[result[0].page_content].append(result[0].id) 

87 else: 

88 entities[result[0].page_content] = [result[0].id] 

89 return entities 

90 

91 def run(self, *args: Text) -> Text: 

92 """Canonicalize the input text against known entities. 

93 

94 Args: 

95 text (str): Input text to canonicalize. 

96 

97 Returns: 

98 str: Canonicalized text if a match is found, otherwise original text. 

99 """ 

100 prompt, pattern = self.pre_run() 

101 

102 entity_stem = re.sub("_.*", "", self.name) 

103 task_inputs = list(args) 

104 query_data: str = ", ".join( 

105 [ 

106 str(task_input.data) 

107 for task_input in task_inputs 

108 if entity_stem in task_input.name 

109 ] 

110 ) 

111 entities = "" 

112 

113 if isinstance(self.entities, Path): 

114 entities_dict = self.retrieve_entities(query_data) 

115 entities = ", ".join(list(entities_dict.keys())) 

116 elif isinstance(self.entities, list): 

117 entities = ", ".join(self.entities) 

118 elif isinstance(self.entities, str): 

119 entities = self.entities 

120 

121 messages = self.replace_inputs(pattern, prompt, *args, entities=entities) 

122 

123 self.messages.append(HumanMessage(content=messages)) 

124 

125 self.check_inputs(*args) 

126 

127 response = self.invoke() 

128 

129 name = self.get_output_name() 

130 

131 output = Text(name=name, data=response) 

132 

133 return output