Coverage for collectra/tasks/llms.py: 78.23%

1"""Large Language Model (LLM) integration tasks for Collectra workflows. 

2 

3This module provides integration with various Large Language Models through 

4the LangChain framework. It supports both text and multimodal (text+image) 

5processing tasks with configurable prompts and model parameters. 

6 

7The module includes: 

8 - LLM task class with configurable model backends 

9 - Support for text and image inputs 

10 - Template-based prompt formatting 

11 - Integration with various LLM providers via llmloader 

12 

13Classes: 

14 LLM: Task for Large Language Model inference operations 

15""" 

16 

17__all__ = ["LLM"] 

18 

19import os 

20import re 

21from functools import cached_property 

22from pathlib import Path 

23 

24import yaml 

25from dotenv import load_dotenv 

26from langchain_core.messages import HumanMessage, SystemMessage 

27from langchain_core.output_parsers import StrOutputParser 

28 

29from collectra.tasks.base import Task 

30from collectra.types.images import Image 

31from collectra.types.texts import Text 

32 

33from ..logger import get_logger 

34 

35 

36def get_llmloader(): 

37 try: 

38 import llmloader 

39 except ModuleNotFoundError as error: 

40 if error.name != "llmloader": 

41 raise 

42 raise ImportError( 

43 'LLM execution requires optional dependencies. Install with: pip install "collectra[ml]"' 

44 ) from error 

45 return llmloader 

46 

47load_dotenv() 

48 

49logger = get_logger(__name__) 

50 

51 

52class LLM(Task): 

53 """Task for Large Language Model inference operations. 

54 

55 Integrates with various LLM providers to perform text generation and 

56 multimodal processing tasks. Supports configurable prompts, temperature, 

57 and token limits with template-based input formatting. 

58 

59 Attributes: 

60 name (str): Name identifier for the LLM task. 

61 model (str): Model identifier or path for the LLM. 

62 template (str | Path): Prompt template string or path to template file. 

63 temperature (float): Sampling temperature for generation (0.0 to 1.0). 

64 max_tokens (int): Maximum number of tokens to generate. 

65 variables (dict): Additional variables for template substitution. 

66 """ 

67 

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

69 super().__init__(name, **kwargs) 

70 self.template: str = kwargs.get("template", "") 

71 self.preamble: str = kwargs.get("preamble", "") 

72 self.temperature = kwargs.get("temperature", 0.8) 

73 self.max_tokens = kwargs.get("max_tokens", None) 

74 self.model = model 

75 self.parser = StrOutputParser() 

76 

77 init_messages = ( 

78 SystemMessage(content=kwargs.get("system", "")) 

79 if kwargs.get("system", "") 

80 else SystemMessage(content="You are a helpful assistant.") 

81 ) 

82 self.messages: list[SystemMessage | HumanMessage] = [init_messages] 

83 

84 @cached_property 

85 def llm(self): 

86 """Load the LLM lazily so that it is only instantiated when first needed.""" 

87 return get_llmloader().load( 

88 self.model, temperature=self.temperature, max_tokens=self.max_tokens 

89 ) 

90 

91 @property 

92 def chain(self): 

93 return self.llm 

94 

95 def get_pattern_matches(self) -> tuple: 

96 prompt = f"{self.preamble}\n\n{self.template}".strip() 

97 pattern = r"\{(.*?)\}" 

98 return prompt, pattern 

99 

100 def image_content(self, image: Image): 

101 return get_llmloader().LLMWrapper.format( 

102 self.llm, 

103 data_type="image", 

104 data={ 

105 "data": image.get_encoding(), 

106 "mime_type": image.mime(), 

107 }, 

108 ) 

109 

110 def _add_content(self, value: Image | Text) -> dict: 

111 if isinstance(value, Image): 

112 return self.image_content(value) 

113 return self._add_text(value()) 

114 

115 def _add_text(self, text: str) -> dict: 

116 text = ( 

117 "''" if len(text.strip()) == 0 else text.strip() 

118 ) # Avoid empty string issues in some LLMs 

119 return {"type": "text", "text": text} 

120 

121 def pre_run(self) -> tuple: 

122 self.messages = self.messages[:1] # Reset to initial system message 

123 prompt, pattern = self.get_pattern_matches() 

124 return prompt, pattern 

125 

126 def check_inputs(self, *args: Text | Image) -> None: 

127 """Validate that the messages are not longer than the expected number of inputs.""" 

128 residual_inputs = len(args) 

129 if len(self.messages) != 2: 

130 raise ValueError("Expected one system message and one human message.") 

131 human_message = self.messages[1] 

132 for content in human_message.content: 

133 if isinstance(content, dict) and ( 

134 content.get("type") == "image" or content.get("type") == "text" 

135 ): 

136 residual_inputs -= 1 

137 if residual_inputs > 0: 

138 raise ValueError("Number of inputs does not match the expected count") 

139 

140 def replace_inputs( 

141 self, pattern: str, prompt: str, *args: Text | Image, **kwargs 

142 ) -> list[str | dict]: 

143 messages: list[str | dict] = list() 

144 while re.search(pattern, prompt): 

145 match = next(re.finditer(pattern, prompt)) 

146 start, end = match.span() 

147 item = match[1].strip() 

148 replaced = False 

149 if prompt[:start]: 

150 messages.append(self._add_text(prompt[:start])) 

151 for arg in args: 

152 key = arg.name 

153 if key == item: 

154 messages.append(self._add_content(arg)) 

155 replaced = True 

156 break 

157 if item == "entities" and kwargs.get("entities", None): 

158 replaced = True 

159 messages.append(self._add_text(kwargs["entities"])) 

160 if not replaced: 

161 messages.append( 

162 self._add_text(f"No content provided for {item}. Ignore this part.") 

163 ) 

164 prompt = prompt[end:].strip() 

165 

166 if prompt: 

167 messages.append(self._add_text(prompt)) 

168 

169 return messages 

170 

171 def invoke(self) -> str: 

172 response = self.chain.invoke(self.messages) 

173 

174 # Use context for usage tracking instead of environment variable 

175 if self.context.usage_file: 

176 usage = get_llmloader().LLMWrapper.get_token_count(response) 

177 usage_file = self.context.usage_file 

178 data = {f"{self.name}": usage} 

179 old_data = {} 

180 if usage_file.exists(): 

181 with open(usage_file, "r") as f: 

182 old_data = yaml.safe_load(f) or dict() 

183 for key, value in old_data.items(): 

184 if key not in data: 

185 data[key] = value 

186 else: 

187 for item, item_value in value.items(): 

188 data[key][item] = data.get(key, {}).get(item, 0) + item_value 

189 usage_file.parent.mkdir(parents=True, exist_ok=True) 

190 with open(usage_file, "w") as f: 

191 yaml.dump(data, f) 

192 

193 response = self.parser.invoke(response) 

194 

195 return "" if response in ('""', "''") else response.strip() 

196 

197 def run(self, *args: Text | Image) -> Text | None: 

198 """Execute LLM inference on the provided inputs with template-based prompt generation. 

199 

200 For list inputs, each item is processed individually and results are collected 

201 in a list. The method supports multimodal inputs (text + images) by creating 

202 appropriate message structures for the LangChain conversation format. 

203 

204 Args: 

205 **kwargs: Input data for LLM processing. Keys should match template 

206 placeholders. Values can be text strings, file paths (Path objects), 

207 Image objects, or lists of these types for batch processing. 

208 

209 Side Effects: 

210 Updates self.output dictionary with generated text responses. Output keys 

211 that contain the input key as a substring will be populated with LLM results. 

212 Prints a success message when inference completes. 

213 """ 

214 try: 

215 prompt, pattern = self.pre_run() 

216 

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

218 

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

220 

221 self.check_inputs(*args) 

222 

223 response = self.invoke() 

224 

225 name = self.get_output_name() 

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

227 return output 

228 except Exception as e: 

229 logger.error("LLM task '%s' failed: %s", self.name, e, exc_info=True) 

230 return None