llm-ie 1.2.0__py3-none-any.whl → 1.2.2__py3-none-any.whl

llm_ie/__init__.py

@@ -1,11 +1,11 @@
 from .data_types import LLMInformationExtractionFrame, LLMInformationExtractionDocument
-from .engines import BasicLLMConfig, Qwen3LLMConfig, OpenAIReasoningLLMConfig, LlamaCppInferenceEngine, OllamaInferenceEngine, HuggingFaceHubInferenceEngine, OpenAIInferenceEngine, AzureOpenAIInferenceEngine, LiteLLMInferenceEngine
+from .engines import BasicLLMConfig, ReasoningLLMConfig, Qwen3LLMConfig, OpenAIReasoningLLMConfig, LlamaCppInferenceEngine, OllamaInferenceEngine, HuggingFaceHubInferenceEngine, OpenAIInferenceEngine, AzureOpenAIInferenceEngine, LiteLLMInferenceEngine
 from .extractors import DirectFrameExtractor, ReviewFrameExtractor, BasicFrameExtractor, BasicReviewFrameExtractor, SentenceFrameExtractor, SentenceReviewFrameExtractor, AttributeExtractor, BinaryRelationExtractor, MultiClassRelationExtractor
 from .chunkers import UnitChunker, WholeDocumentUnitChunker, SentenceUnitChunker, TextLineUnitChunker, ContextChunker, NoContextChunker, WholeDocumentContextChunker, SlideWindowContextChunker
 from .prompt_editor import PromptEditor
 
 __all__ = ["LLMInformationExtractionFrame", "LLMInformationExtractionDocument",
-           "BasicLLMConfig", "Qwen3LLMConfig", "OpenAIReasoningLLMConfig", "LlamaCppInferenceEngine", "OllamaInferenceEngine", "HuggingFaceHubInferenceEngine", "OpenAIInferenceEngine", "AzureOpenAIInferenceEngine", "LiteLLMInferenceEngine",
+           "BasicLLMConfig", "ReasoningLLMConfig", "Qwen3LLMConfig", "OpenAIReasoningLLMConfig", "LlamaCppInferenceEngine", "OllamaInferenceEngine", "HuggingFaceHubInferenceEngine", "OpenAIInferenceEngine", "AzureOpenAIInferenceEngine", "LiteLLMInferenceEngine",
            "DirectFrameExtractor", "ReviewFrameExtractor", "BasicFrameExtractor", "BasicReviewFrameExtractor", "SentenceFrameExtractor", "SentenceReviewFrameExtractor", "AttributeExtractor", "BinaryRelationExtractor", "MultiClassRelationExtractor",
            "UnitChunker", "WholeDocumentUnitChunker", "SentenceUnitChunker", "TextLineUnitChunker", "ContextChunker", "NoContextChunker", "WholeDocumentContextChunker", "SlideWindowContextChunker",
            "PromptEditor"]

llm_ie/chunkers.py

@@ -41,6 +41,41 @@ class WholeDocumentUnitChunker(UnitChunker):
             text=text
         )]
     
+class SeparatorUnitChunker(UnitChunker):
+    def __init__(self, sep:str):
+        """
+        This class chunks a document by separator provided.
+
+        Parameters:
+        ----------
+        sep : str
+            a separator string.
+        """
+        super().__init__()
+        if not isinstance(sep, str):
+            raise ValueError("sep must be a string")
+
+        self.sep = sep
+
+    def chunk(self, text:str) -> List[FrameExtractionUnit]:
+        """
+        Parameters:
+        ----------
+        text : str
+            The document text.
+        """
+        paragraphs = text.split(self.sep)
+        paragraph_units = []
+        start = 0
+        for paragraph in paragraphs:
+            end = start + len(paragraph)
+            paragraph_units.append(FrameExtractionUnit(
+                start=start,
+                end=end,
+                text=paragraph
+            ))
+            start = end + len(self.sep)
+        return paragraph_units
 
 class SentenceUnitChunker(UnitChunker):
     from nltk.tokenize.punkt import PunktSentenceTokenizer

llm_ie/engines.py

@@ -33,18 +33,18 @@ class LLMConfig(abc.ABC):
         return NotImplemented
 
     @abc.abstractmethod
-    def postprocess_response(self, response:Union[str, Generator[str, None, None]]) -> Union[str, Generator[str, None, None]]:
+    def postprocess_response(self, response:Union[str, Generator[str, None, None]]) -> Union[Dict[str,str], Generator[Dict[str, str], None, None]]:
         """
         This method postprocesses the LLM response after it is generated.
 
         Parameters:
         ----------
-        response : Union[str, Generator[str, None, None]]
-            the LLM response. Can be a string or a generator.
+        response : Union[str, Generator[Dict[str, str], None, None]]
+            the LLM response. Can be a dict or a generator. 
         
         Returns:
         -------
-        response : str
+        response : Union[Dict[str,str], Generator[Dict[str, str], None, None]]
             the postprocessed LLM response
         """
         return NotImplemented
@@ -77,7 +77,7 @@ class BasicLLMConfig(LLMConfig):
         """
         return messages
 
-    def postprocess_response(self, response:Union[str, Generator[str, None, None]]) -> Union[str, Generator[Dict[str, str], None, None]]:
+    def postprocess_response(self, response:Union[str, Generator[str, None, None]]) -> Union[Dict[str,str], Generator[Dict[str, str], None, None]]:
         """
         This method postprocesses the LLM response after it is generated.
 
@@ -86,12 +86,13 @@ class BasicLLMConfig(LLMConfig):
         response : Union[str, Generator[str, None, None]]
             the LLM response. Can be a string or a generator.
         
-        Returns: Union[str, Generator[Dict[str, str], None, None]]
+        Returns: Union[Dict[str,str], Generator[Dict[str, str], None, None]]
             the postprocessed LLM response. 
-            if input is a generator, the output will be a generator {"data": <content>}.
+            If input is a string, the output will be a dict {"response": <response>}. 
+            if input is a generator, the output will be a generator {"type": "response", "data": <content>}.
         """
         if isinstance(response, str):
-            return response
+            return {"response": response}
 
         def _process_stream():
             for chunk in response:
@@ -99,23 +100,19 @@ class BasicLLMConfig(LLMConfig):
 
         return _process_stream()
 
-class Qwen3LLMConfig(LLMConfig):
-    def __init__(self, thinking_mode:bool=True, **kwargs):
-        """
-        The Qwen3 LLM configuration for reasoning models.
 
-        Parameters:
-        ----------
-        thinking_mode : bool, Optional
-            if True, a special token "/think" will be placed after each system and user prompt. Otherwise, "/no_think" will be placed.
+class ReasoningLLMConfig(LLMConfig):
+    def __init__(self, thinking_token_start="<think>", thinking_token_end="</think>", **kwargs):
+        """
+        The general LLM configuration for reasoning models.
         """
         super().__init__(**kwargs)
-        self.thinking_mode = thinking_mode
+        self.thinking_token_start = thinking_token_start
+        self.thinking_token_end = thinking_token_end
 
     def preprocess_messages(self, messages:List[Dict[str,str]]) -> List[Dict[str,str]]:
         """
-        Append a special token to the system and user prompts.
-        The token is "/think" if thinking_mode is True, otherwise "/no_think".
+        This method preprocesses the input messages before passing them to the LLM.
 
         Parameters:
         ----------
@@ -127,23 +124,11 @@ class Qwen3LLMConfig(LLMConfig):
         messages : List[Dict[str,str]]
             a list of dict with role and content. role must be one of {"system", "user", "assistant"}
         """
-        thinking_token = "/think" if self.thinking_mode else "/no_think"
-        new_messages = []
-        for message in messages:
-            if message['role'] in ['system', 'user']:
-                new_message = {'role': message['role'], 'content': f"{message['content']} {thinking_token}"}
-            else:
-                new_message = {'role': message['role'], 'content': message['content']}
-
-            new_messages.append(new_message)
-
-        return new_messages
+        return messages
 
-    def postprocess_response(self, response:Union[str, Generator[str, None, None]]) -> Union[str, Generator[Dict[str,str], None, None]]:
+    def postprocess_response(self, response:Union[str, Generator[str, None, None]]) -> Union[Dict[str,str], Generator[Dict[str,str], None, None]]:
         """
-        If input is a generator, tag contents in <think> and </think> as {"type": "reasoning", "data": <content>},
-        and the rest as {"type": "response", "data": <content>}.
-        If input is a string, drop contents in <think> and </think>.
+        This method postprocesses the LLM response after it is generated.
 
         Parameters:
         ----------
@@ -153,11 +138,16 @@ class Qwen3LLMConfig(LLMConfig):
         Returns:
         -------
         response : Union[str, Generator[str, None, None]]
-            the postprocessed LLM response.
+            the postprocessed LLM response as a dict {"reasoning": <reasoning>, "response": <content>}
             if input is a generator, the output will be a generator {"type": <reasoning or response>, "data": <content>}.
         """
         if isinstance(response, str):
-            return re.sub(r"<think>.*?</think>\s*", "", response, flags=re.DOTALL).strip()
+            # get contents between thinking_token_start and thinking_token_end
+            match = re.search(f"{self.thinking_token_start}.*?{self.thinking_token_end}", response, re.DOTALL)
+            reasoning = match.group(0) if match else ""
+            # get response AFTER thinking_token_end
+            response = re.sub(f".*?{self.thinking_token_end}", "", response, flags=re.DOTALL).strip()
+            return {"reasoning": reasoning, "response": response}
 
         if isinstance(response, Generator):
             def _process_stream():
@@ -167,28 +157,71 @@ class Qwen3LLMConfig(LLMConfig):
                     if isinstance(chunk, str):
                         buffer += chunk
                         # switch between reasoning and response
-                        if "<think>" in buffer:
+                        if self.thinking_token_start in buffer:
                             think_flag = True
-                            buffer = buffer.replace("<think>", "")
-                        elif "</think>" in buffer:
+                            buffer = buffer.replace(self.thinking_token_start, "")
+                        elif self.thinking_token_end in buffer:
                             think_flag = False
-                            buffer = buffer.replace("</think>", "")
+                            buffer = buffer.replace(self.thinking_token_end, "")
                         
                         # if chunk is in thinking block, tag it as reasoning; else tag it as response
-                        if chunk not in ["<think>", "</think>"]:
+                        if chunk not in [self.thinking_token_start, self.thinking_token_end]:
                             if think_flag:
                                 yield {"type": "reasoning", "data": chunk}
                             else:
                                 yield {"type": "response", "data": chunk}
 
             return _process_stream()
+        
+
+class Qwen3LLMConfig(ReasoningLLMConfig):
+    def __init__(self, thinking_mode:bool=True, **kwargs):
+        """
+        The Qwen3 **hybrid thinking** LLM configuration. 
+        For Qwen3 thinking 2507, use ReasoningLLMConfig instead; for Qwen3 Instruct, use BasicLLMConfig instead.
+
+        Parameters:
+        ----------
+        thinking_mode : bool, Optional
+            if True, a special token "/think" will be placed after each system and user prompt. Otherwise, "/no_think" will be placed.
+        """
+        super().__init__(**kwargs)
+        self.thinking_mode = thinking_mode
+
+    def preprocess_messages(self, messages:List[Dict[str,str]]) -> List[Dict[str,str]]:
+        """
+        Append a special token to the system and user prompts.
+        The token is "/think" if thinking_mode is True, otherwise "/no_think".
+
+        Parameters:
+        ----------
+        messages : List[Dict[str,str]]
+            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
+        
+        Returns:
+        -------
+        messages : List[Dict[str,str]]
+            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
+        """
+        thinking_token = "/think" if self.thinking_mode else "/no_think"
+        new_messages = []
+        for message in messages:
+            if message['role'] in ['system', 'user']:
+                new_message = {'role': message['role'], 'content': f"{message['content']} {thinking_token}"}
+            else:
+                new_message = {'role': message['role'], 'content': message['content']}
 
+            new_messages.append(new_message)
 
-class OpenAIReasoningLLMConfig(LLMConfig):
-    def __init__(self, reasoning_effort:str="low", **kwargs):
+        return new_messages
+
+
+class OpenAIReasoningLLMConfig(ReasoningLLMConfig):
+    def __init__(self, reasoning_effort:str=None, **kwargs):
         """
         The OpenAI "o" series configuration.
-        1. The reasoning effort is set to "low" by default.
+        1. The reasoning effort as one of {"low", "medium", "high"}.
+            For models that do not support setting reasoning effort (e.g., o1-mini, o1-preview), set to None.
         2. The temperature parameter is not supported and will be ignored.
         3. The system prompt is not supported and will be concatenated to the next user prompt.
 
@@ -198,11 +231,12 @@ class OpenAIReasoningLLMConfig(LLMConfig):
             the reasoning effort. Must be one of {"low", "medium", "high"}. Default is "low".
         """
         super().__init__(**kwargs)
-        if reasoning_effort not in ["low", "medium", "high"]:
-            raise ValueError("reasoning_effort must be one of {'low', 'medium', 'high'}.")
+        if reasoning_effort is not None:
+            if reasoning_effort not in ["low", "medium", "high"]:
+                raise ValueError("reasoning_effort must be one of {'low', 'medium', 'high'}.")
 
-        self.reasoning_effort = reasoning_effort
-        self.params["reasoning_effort"] = self.reasoning_effort
+            self.reasoning_effort = reasoning_effort
+            self.params["reasoning_effort"] = self.reasoning_effort
 
         if "temperature" in self.params:
             warnings.warn("Reasoning models do not support temperature parameter. Will be ignored.", UserWarning)
@@ -244,28 +278,6 @@ class OpenAIReasoningLLMConfig(LLMConfig):
 
         return new_messages
 
-    def postprocess_response(self, response:Union[str, Generator[str, None, None]]) -> Union[str, Generator[Dict[str, str], None, None]]:
-        """
-        This method postprocesses the LLM response after it is generated.
-
-        Parameters:
-        ----------
-        response : Union[str, Generator[str, None, None]]
-            the LLM response. Can be a string or a generator.
-        
-        Returns: Union[str, Generator[Dict[str, str], None, None]]
-            the postprocessed LLM response. 
-            if input is a generator, the output will be a generator {"type": "response", "data": <content>}.
-        """
-        if isinstance(response, str):
-            return response
-
-        def _process_stream():
-            for chunk in response:
-                yield {"type": "response", "data": chunk}
-
-        return _process_stream()
-
 
 class InferenceEngine:
     @abc.abstractmethod
@@ -284,7 +296,7 @@ class InferenceEngine:
 
     @abc.abstractmethod
     def chat(self, messages:List[Dict[str,str]], 
-             verbose:bool=False, stream:bool=False) -> Union[str, Generator[Dict[str, str], None, None]]:
+             verbose:bool=False, stream:bool=False) -> Union[Dict[str,str], Generator[Dict[str, str], None, None]]:
         """
         This method inputs chat messages and outputs LLM generated text.
 
@@ -296,6 +308,11 @@ class InferenceEngine:
             if True, LLM generated text will be printed in terminal in real-time.
         stream : bool, Optional
             if True, returns a generator that yields the output in real-time.  
+
+        Returns:
+        -------
+        response : Union[Dict[str,str], Generator[Dict[str, str], None, None]]
+            a dict {"reasoning": <reasoning>, "response": <response>} or Generator {"type": <reasoning or response>, "data": <content>}
         """
         return NotImplemented
 
@@ -361,7 +378,7 @@ class LlamaCppInferenceEngine(InferenceEngine):
 
         return formatted_params
 
-    def chat(self, messages:List[Dict[str,str]], verbose:bool=False) -> str:
+    def chat(self, messages:List[Dict[str,str]], verbose:bool=False) -> Dict[str,str]:
         """
         This method inputs chat messages and outputs LLM generated text.
 
@@ -434,7 +451,7 @@ class OllamaInferenceEngine(InferenceEngine):
         return formatted_params
 
     def chat(self, messages:List[Dict[str,str]], 
-             verbose:bool=False, stream:bool=False) -> Union[str, Generator[Dict[str, str], None, None]]:
+             verbose:bool=False, stream:bool=False) -> Union[Dict[str,str], Generator[Dict[str, str], None, None]]:
         """
         This method inputs chat messages and outputs VLM generated text.
 
@@ -446,6 +463,11 @@ class OllamaInferenceEngine(InferenceEngine):
             if True, VLM generated text will be printed in terminal in real-time.
         stream : bool, Optional
             if True, returns a generator that yields the output in real-time.
+
+        Returns:
+        -------
+        response : Union[Dict[str,str], Generator[Dict[str, str], None, None]]
+            a dict {"reasoning": <reasoning>, "response": <response>} or Generator {"type": <reasoning or response>, "data": <content>}
         """
         processed_messages = self.config.preprocess_messages(messages)
 
@@ -495,7 +517,7 @@ class OllamaInferenceEngine(InferenceEngine):
             return self.config.postprocess_response(res)
         
 
-    async def chat_async(self, messages:List[Dict[str,str]]) -> str:
+    async def chat_async(self, messages:List[Dict[str,str]]) -> Dict[str,str]:
         """
         Async version of chat method. Streaming is not supported.
         """
@@ -556,7 +578,7 @@ class HuggingFaceHubInferenceEngine(InferenceEngine):
 
 
     def chat(self, messages:List[Dict[str,str]], 
-             verbose:bool=False, stream:bool=False) -> Union[str, Generator[Dict[str, str], None, None]]:
+             verbose:bool=False, stream:bool=False) -> Union[Dict[str,str], Generator[Dict[str, str], None, None]]:
         """
         This method inputs chat messages and outputs LLM generated text.
 
@@ -568,6 +590,11 @@ class HuggingFaceHubInferenceEngine(InferenceEngine):
             if True, VLM generated text will be printed in terminal in real-time.
         stream : bool, Optional
             if True, returns a generator that yields the output in real-time.
+
+        Returns:
+        -------
+        response : Union[Dict[str,str], Generator[Dict[str, str], None, None]]
+            a dict {"reasoning": <reasoning>, "response": <response>} or Generator {"type": <reasoning or response>, "data": <content>}
         """
         processed_messages = self.config.preprocess_messages(messages)
 
@@ -609,7 +636,7 @@ class HuggingFaceHubInferenceEngine(InferenceEngine):
             res = response.choices[0].message.content
             return self.config.postprocess_response(res)
     
-    async def chat_async(self, messages:List[Dict[str,str]]) -> str:
+    async def chat_async(self, messages:List[Dict[str,str]]) -> Dict[str,str]:
         """
         Async version of chat method. Streaming is not supported.
         """
@@ -660,7 +687,7 @@ class OpenAIInferenceEngine(InferenceEngine):
 
         return formatted_params
 
-    def chat(self, messages:List[Dict[str,str]], verbose:bool=False, stream:bool=False) -> Union[str, Generator[Dict[str, str], None, None]]:
+    def chat(self, messages:List[Dict[str,str]], verbose:bool=False, stream:bool=False) -> Union[Dict[str, str], Generator[Dict[str, str], None, None]]:
         """
         This method inputs chat messages and outputs LLM generated text.
 
@@ -672,6 +699,11 @@ class OpenAIInferenceEngine(InferenceEngine):
             if True, VLM generated text will be printed in terminal in real-time.
         stream : bool, Optional
             if True, returns a generator that yields the output in real-time.
+
+        Returns:
+        -------
+        response : Union[Dict[str,str], Generator[Dict[str, str], None, None]]
+            a dict {"reasoning": <reasoning>, "response": <response>} or Generator {"type": <reasoning or response>, "data": <content>}
         """
         processed_messages = self.config.preprocess_messages(messages)
 
@@ -721,7 +753,7 @@ class OpenAIInferenceEngine(InferenceEngine):
             return self.config.postprocess_response(res)
     
 
-    async def chat_async(self, messages:List[Dict[str,str]]) -> str:
+    async def chat_async(self, messages:List[Dict[str,str]]) -> Dict[str,str]:
         """
         Async version of chat method. Streaming is not supported.
         """
@@ -811,7 +843,7 @@ class LiteLLMInferenceEngine(InferenceEngine):
 
         return formatted_params
 
-    def chat(self, messages:List[Dict[str,str]], verbose:bool=False, stream:bool=False) -> Union[str, Generator[Dict[str, str], None, None]]:
+    def chat(self, messages:List[Dict[str,str]], verbose:bool=False, stream:bool=False) -> Union[Dict[str,str], Generator[Dict[str, str], None, None]]:
         """
         This method inputs chat messages and outputs LLM generated text.
 
@@ -823,6 +855,11 @@ class LiteLLMInferenceEngine(InferenceEngine):
             if True, VLM generated text will be printed in terminal in real-time.
         stream : bool, Optional
             if True, returns a generator that yields the output in real-time.
+
+        Returns:
+        -------
+        response : Union[Dict[str,str], Generator[Dict[str, str], None, None]]
+            a dict {"reasoning": <reasoning>, "response": <response>} or Generator {"type": <reasoning or response>, "data": <content>}
         """
         processed_messages = self.config.preprocess_messages(messages)
         
@@ -875,7 +912,7 @@ class LiteLLMInferenceEngine(InferenceEngine):
             res = response.choices[0].message.content
             return self.config.postprocess_response(res)
     
-    async def chat_async(self, messages:List[Dict[str,str]]) -> str:
+    async def chat_async(self, messages:List[Dict[str,str]]) -> Dict[str,str]:
         """
         Async version of chat method. Streaming is not supported.
         """

llm_ie/extractors.py

@@ -489,7 +489,10 @@ class DirectFrameExtractor(FrameExtractor):
                         )
 
             if return_messages_log:
-                messages.append({"role": "assistant", "content": gen_text})
+                message = {"role": "assistant", "content": gen_text["response"]}
+                if "reasoning" in gen_text:
+                    message["reasoning"] = gen_text["reasoning"]
+                messages.append(message)
                 messages_log.append(messages)
 
             # add to output
@@ -497,7 +500,7 @@ class DirectFrameExtractor(FrameExtractor):
                             start=unit.start,
                             end=unit.end,
                             text=unit.text,
-                            gen_text=gen_text)
+                            gen_text=gen_text["response"])
             output.append(result)
             
         if return_messages_log:
@@ -581,7 +584,8 @@ class DirectFrameExtractor(FrameExtractor):
             )
             for chunk in response_stream:
                 yield chunk
-                current_gen_text += chunk
+                if chunk["type"] == "response":
+                    current_gen_text += chunk["data"]
            
             # Store the result for this unit
             result_for_unit = FrameExtractionUnitResult(
@@ -679,7 +683,11 @@ class DirectFrameExtractor(FrameExtractor):
                 gen_text = await self.inference_engine.chat_async(
                     messages=messages
                 )
-            return {"original_index": original_index, "unit": unit, "gen_text": gen_text, "messages": messages}
+
+            out = {"original_index": original_index, "unit": unit, "gen_text": gen_text["response"], "messages": messages}
+            if "reasoning" in gen_text:
+                out["reasoning"] = gen_text["reasoning"]
+            return out
            
         # Create and gather tasks
         tasks = []
@@ -713,7 +721,10 @@ class DirectFrameExtractor(FrameExtractor):
 
             # Append to messages log if requested
             if return_messages_log:
-                final_messages = result_data["messages"] + [{"role": "assistant", "content": gen_text}]
+                message = {"role": "assistant", "content": gen_text}
+                if "reasoning" in result_data:
+                    message["reasoning"] = result_data["reasoning"]
+                final_messages = result_data["messages"] + [message]
                 messages_log.append(final_messages)
 
         if return_messages_log:
@@ -975,15 +986,11 @@ class ReviewFrameExtractor(DirectFrameExtractor):
                             stream=False
                         )
 
-            if return_messages_log:
-                messages.append({"role": "assistant", "content": initial})
-                messages_log.append(messages)
-
             # <--- Review step --->
             if verbose:
                 print(f"\n{Fore.YELLOW}Review:{Style.RESET_ALL}")
 
-            messages.append({'role': 'assistant', 'content': initial})
+            messages.append({'role': 'assistant', 'content': initial["response"]})
             messages.append({'role': 'user', 'content': self.review_prompt})
             
             review = self.inference_engine.chat(
@@ -994,12 +1001,18 @@ class ReviewFrameExtractor(DirectFrameExtractor):
 
             # Output
             if self.review_mode == "revision":
-                gen_text = review
+                gen_text = review["response"]
             elif self.review_mode == "addition":
-                gen_text = initial + '\n' + review
+                gen_text = initial["response"] + '\n' + review["response"]
 
             if return_messages_log:
-                messages.append({"role": "assistant", "content": review})
+                if "reasoning" in initial:
+                    messages[-2]["reasoning"] = initial["reasoning"]
+
+                message = {"role": "assistant", "content": review["response"]}
+                if "reasoning" in review:
+                    message["reasoning"] = review["reasoning"]
+                messages.append(message)
                 messages_log.append(messages)
 
             # add to output
@@ -1192,7 +1205,10 @@ class ReviewFrameExtractor(DirectFrameExtractor):
                     messages=messages
                 )
             # Return initial generation result along with the messages used and the unit
-            return {"original_index": original_index, "unit": unit, "initial_gen_text": gen_text, "initial_messages": messages}
+            out = {"original_index": original_index, "unit": unit, "initial_gen_text": gen_text["response"], "initial_messages": messages}
+            if "reasoning" in gen_text:
+                out["reasoning"] = gen_text["reasoning"]
+            return out
 
         # Create and gather initial generation tasks
         initial_tasks = [
@@ -1218,28 +1234,35 @@ class ReviewFrameExtractor(DirectFrameExtractor):
                 {'role': 'user', 'content': self.review_prompt}
             ]
             # Store data needed for review task
+            if "reasoning" in result_data:
+                message = {'role': 'assistant', 'content': initial_gen_text, "reasoning": result_data["reasoning"]}
+            else:
+                message = {'role': 'assistant', 'content': initial_gen_text}
+
             review_tasks_input.append({
                 "unit": result_data["unit"],
                 "initial_gen_text": initial_gen_text,
                 "messages": review_messages, 
                 "original_index": result_data["original_index"],
-                "full_initial_log": initial_messages + [{'role': 'assistant', 'content': initial_gen_text}] if return_messages_log else None # Log up to initial generation
+                "full_initial_log": initial_messages + [message] + [{'role': 'user', 'content': self.review_prompt}] if return_messages_log else None
             })
 
 
         async def review_semaphore_helper(task_data: Dict, **kwrs):
             messages = task_data["messages"] 
-            original_index = task_data["original_index"]
 
             async with semaphore:
                 review_gen_text = await self.inference_engine.chat_async(
                     messages=messages
                 )
             # Combine initial and review results
-            task_data["review_gen_text"] = review_gen_text
+            task_data["review_gen_text"] = review_gen_text["response"]
             if return_messages_log:
                 # Log for the review call itself
-                 task_data["full_review_log"] = messages + [{'role': 'assistant', 'content': review_gen_text}]
+                message = {'role': 'assistant', 'content': review_gen_text["response"]}
+                if "reasoning" in review_gen_text:
+                    message["reasoning"] = review_gen_text["reasoning"]
+                task_data["full_review_log"] = task_data["full_initial_log"] + [message]
             return task_data # Return the augmented dictionary
 
         # Create and gather review tasks
@@ -1283,7 +1306,7 @@ class ReviewFrameExtractor(DirectFrameExtractor):
 
             # Append full conversation log if requested
             if return_messages_log:
-                full_log_for_unit = result_data.get("full_initial_log", []) + [{'role': 'user', 'content': self.review_prompt}] + [{'role': 'assistant', 'content': review_gen}]
+                full_log_for_unit = result_data["full_review_log"]
                 messages_log.append(full_log_for_unit)
 
         if return_messages_log:
@@ -1541,15 +1564,18 @@ class AttributeExtractor(Extractor):
             
             print(f"{Fore.BLUE}Extraction:{Style.RESET_ALL}")
 
-        get_text = self.inference_engine.chat(
+        gen_text = self.inference_engine.chat(
                             messages=messages,
                             verbose=verbose,
                             stream=False
                         )
         if return_messages_log:
-            messages.append({"role": "assistant", "content": get_text})
+            message = {"role": "assistant", "content": gen_text["response"]}
+            if "reasoning" in gen_text:
+                message["reasoning"] = gen_text["reasoning"]
+            messages.append(message)
 
-        attribute_list = self._extract_json(gen_text=get_text)
+        attribute_list = self._extract_json(gen_text=gen_text["response"])
         if isinstance(attribute_list, list) and len(attribute_list) > 0:
             attributes = attribute_list[0]
             if return_messages_log:
@@ -1658,9 +1684,12 @@ class AttributeExtractor(Extractor):
                 gen_text = await self.inference_engine.chat_async(messages=messages)
                 
                 if return_messages_log:
-                    messages.append({"role": "assistant", "content": gen_text})
+                    message = {"role": "assistant", "content": gen_text["response"]}
+                    if "reasoning" in gen_text:
+                        message["reasoning"] = gen_text["reasoning"]
+                    messages.append(message)
 
-                attribute_list = self._extract_json(gen_text=gen_text)
+                attribute_list = self._extract_json(gen_text=gen_text["response"])
                 attributes = attribute_list[0] if isinstance(attribute_list, list) and len(attribute_list) > 0 else {}
                 return {"frame": frame, "attributes": attributes, "messages": messages}
 
@@ -1824,12 +1853,15 @@ class RelationExtractor(Extractor):
                     messages=task_payload['messages'],
                     verbose=verbose
                 )
-                relation = self._post_process_result(gen_text, task_payload)
+                relation = self._post_process_result(gen_text["response"], task_payload)
                 if relation:
                     relations.append(relation)
 
                 if return_messages_log:
-                    task_payload['messages'].append({"role": "assistant", "content": gen_text})
+                    message = {"role": "assistant", "content": gen_text["response"]}
+                    if "reasoning" in gen_text:
+                        message["reasoning"] = gen_text["reasoning"]
+                    task_payload['messages'].append(message)
                     messages_log.append(task_payload['messages'])
 
         return (relations, messages_log) if return_messages_log else relations
@@ -1853,12 +1885,15 @@ class RelationExtractor(Extractor):
         results = await asyncio.gather(*tasks)
 
         for gen_text, task_payload in results:
-            relation = self._post_process_result(gen_text, task_payload)
+            relation = self._post_process_result(gen_text["response"], task_payload)
             if relation:
                 relations.append(relation)
 
             if return_messages_log:
-                task_payload['messages'].append({"role": "assistant", "content": gen_text})
+                message = {"role": "assistant", "content": gen_text["response"]}
+                if "reasoning" in gen_text:
+                    message["reasoning"] = gen_text["reasoning"]
+                task_payload['messages'].append(message)
                 messages_log.append(task_payload['messages'])
 
         return (relations, messages_log) if return_messages_log else relations

llm_ie/prompt_editor.py

@@ -1,9 +1,11 @@
 import sys
+import warnings
 from typing import List, Dict, Generator
 import importlib.resources
 from llm_ie.engines import InferenceEngine
 from llm_ie.extractors import FrameExtractor
 import re
+import json
 from colorama import Fore, Style
 
     
@@ -40,7 +42,9 @@ class PromptEditor:
         file_path = importlib.resources.files('llm_ie.asset.PromptEditor_prompts').joinpath('system.txt')
         with open(file_path, 'r') as f:
             self.system_prompt =  f.read()
-        
+
+        # internal memory (history messages) for the `chat` method
+        self.messages = []
 
     def _apply_prompt_template(self, text_content:Dict[str,str], prompt_template:str) -> str:
         """
@@ -70,6 +74,7 @@ class PromptEditor:
     def rewrite(self, draft:str) -> str:
         """
         This method inputs a prompt draft and rewrites it following the extractor's guideline.
+        This method is stateless.
         """
         file_path = importlib.resources.files('llm_ie.asset.PromptEditor_prompts').joinpath('rewrite.txt')
         with open(file_path, 'r') as f:
@@ -80,11 +85,12 @@ class PromptEditor:
         messages = [{"role": "system", "content": self.system_prompt},
                     {"role": "user", "content": prompt}]
         res = self.inference_engine.chat(messages, verbose=True)
-        return res
+        return res["response"]
     
     def comment(self, draft:str) -> str:
         """
         This method inputs a prompt draft and comment following the extractor's guideline.
+        This method is stateless.
         """
         file_path = importlib.resources.files('llm_ie.asset.PromptEditor_prompts').joinpath('comment.txt')
         with open(file_path, 'r') as f:
@@ -95,24 +101,66 @@ class PromptEditor:
         messages = [{"role": "system", "content": self.system_prompt},
                     {"role": "user", "content": prompt}]
         res = self.inference_engine.chat(messages, verbose=True)
-        return res
+        return res["response"]
     
+    def clear_messages(self):
+        """
+        Clears the current chat history.
+        """
+        self.messages = []
 
-    def _terminal_chat(self):
+    def export_chat(self, file_path: str):
         """
-        This method runs an interactive chat session in the terminal to help users write prompt templates.
+        Exports the current chat history to a JSON file.
+
+        Parameters
+        ----------
+        file_path : str
+            path to the file where the chat history will be saved.
+            Should have a .json extension.
         """
-        file_path = importlib.resources.files('llm_ie.asset.PromptEditor_prompts').joinpath('chat.txt')
-        with open(file_path, 'r') as f:
-            chat_prompt_template = f.read()
+        if not self.messages:
+            raise ValueError("Chat history is empty. Nothing to export.")
 
-        prompt = self._apply_prompt_template(text_content={"prompt_guideline": self.prompt_guide}, 
-                                             prompt_template=chat_prompt_template)
+        with open(file_path, 'w', encoding='utf-8') as f:
+            json.dump(self.messages, f, indent=4)
 
-        messages = [{"role": "system", "content": self.system_prompt},
-                    {"role": "user", "content": prompt}]
+    def import_chat(self, file_path: str):
+        """
+        Imports a chat history from a JSON file, overwriting the current history.
+
+        Parameters
+        ----------
+        file_path : str
+            The path to the .json file containing the chat history.
+        """
+        with open(file_path, 'r', encoding='utf-8') as f:
+            loaded_messages = json.load(f)
+
+        # Validate the loaded messages format.
+        if not isinstance(loaded_messages, list):
+            raise TypeError("Invalid format: The file should contain a JSON list of messages.")
+        for message in loaded_messages:
+            if not (isinstance(message, dict) and 'role' in message and 'content' in message):
+                raise ValueError("Invalid format: Each message must be a dictionary with 'role' and 'content' keys.")
         
+        self.messages = loaded_messages
+
+
+    def _terminal_chat(self):
+        """
+        This method runs an interactive chat session in the terminal to help users write prompt templates.
+        """
         print(f'Welcome to the interactive chat! Type "{Fore.RED}exit{Style.RESET_ALL}" or {Fore.YELLOW}control + C{Style.RESET_ALL} to end the conversation.')
+        if len(self.messages) > 0:
+            print(f"\nPrevious conversation:")
+            for message in self.messages:
+                if message["role"] == "user":
+                    print(f"{Fore.GREEN}\nUser: {Style.RESET_ALL}")
+                    print(message["content"])
+                elif message["role"] == "assistant":
+                    print(f"{Fore.BLUE}Assistant: {Style.RESET_ALL}", end="")
+                    print(message["content"])
 
         while True:
             # Get user input
@@ -124,10 +172,10 @@ class PromptEditor:
                 break
             
             # Chat
-            messages.append({"role": "user", "content": user_input})
+            self.messages.append({"role": "user", "content": user_input})
             print(f"{Fore.BLUE}Assistant: {Style.RESET_ALL}", end="")
-            response = self.inference_engine.chat(messages, verbose=True)
-            messages.append({"role": "assistant", "content": response})
+            response = self.inference_engine.chat(self.messages, verbose=True)
+            self.messages.append({"role": "assistant", "content": response["response"]})
             
 
     def _IPython_chat(self):
@@ -144,19 +192,6 @@ class PromptEditor:
             raise ImportError("IPython not found. Please install IPython (```pip install ipython```).")
         from IPython.display import display, HTML
 
-        # Load the chat prompt template from the resources
-        file_path = importlib.resources.files('llm_ie.asset.PromptEditor_prompts').joinpath('chat.txt')
-        with open(file_path, 'r') as f:
-            chat_prompt_template = f.read()
-
-        # Prepare the initial system message with the prompt guideline
-        prompt = self._apply_prompt_template(text_content={"prompt_guideline": self.prompt_guide}, 
-                                            prompt_template=chat_prompt_template)
-
-        # Initialize conversation messages
-        messages = [{"role": "system", "content": self.system_prompt},
-                    {"role": "user", "content": prompt}]
-
         # Widgets for user input and chat output
         input_box = widgets.Text(placeholder="Type your message here...")
         output_area = widgets.Output()
@@ -164,6 +199,13 @@ class PromptEditor:
         # Display initial instructions
         with output_area:
             display(HTML('Welcome to the interactive chat! Type "<span style="color: red;">exit</span>" to end the conversation.'))
+            if len(self.messages) > 0:
+                display(HTML(f'<p style="color: red;">Previous messages:</p>'))
+                for message in self.messages:
+                    if message["role"] == "user":
+                        display(HTML(f'<p style="color: green;">User: {message["content"]}</p>'))
+                    elif message["role"] == "assistant":
+                        display(HTML(f'<p style="color: blue;">Assistant: {message["content"]}</p>'))
 
         def handle_input(sender):
             user_input = input_box.value
@@ -177,7 +219,7 @@ class PromptEditor:
                 return
 
             # Append user message to conversation
-            messages.append({"role": "user", "content": user_input})
+            self.messages.append({"role": "user", "content": user_input})
             print(f"User: {user_input}")
             
             # Display the user message
@@ -186,8 +228,8 @@ class PromptEditor:
 
             # Get assistant's response and append it to conversation
             print("Assistant: ", end="")
-            response = self.inference_engine.chat(messages, verbose=True)
-            messages.append({"role": "assistant", "content": response})
+            response = self.inference_engine.chat(self.messages, verbose=True)
+            self.messages.append({"role": "assistant", "content": response["response"]})
 
             # Display the assistant's response
             with output_area:
@@ -203,7 +245,20 @@ class PromptEditor:
     def chat(self):
         """
         External method that detects the environment and calls the appropriate chat method.
+        This method use and updates the `messages` list (internal memory).
+        This method is stateful.
         """
+        # Check if the conversation is empty, if so, load the initial chat prompt template.
+        if len(self.messages) == 0:
+            file_path = importlib.resources.files('llm_ie.asset.PromptEditor_prompts').joinpath('chat.txt')
+            with open(file_path, 'r') as f:
+                chat_prompt_template = f.read()
+
+            guideline = self._apply_prompt_template(text_content={"prompt_guideline": self.prompt_guide}, 
+                                                    prompt_template=chat_prompt_template)
+
+            self.messages = [{"role": "system", "content": self.system_prompt + guideline}]
+
         if 'ipykernel' in sys.modules:
             self._IPython_chat()
         else:
@@ -213,6 +268,7 @@ class PromptEditor:
         """
         This method processes messages and yields response chunks from the inference engine.
         This is for frontend App.
+        This method is stateless.
 
         Parameters:
         ----------
@@ -232,12 +288,10 @@ class PromptEditor:
         with open(file_path, 'r') as f:
             chat_prompt_template = f.read()
 
-        prompt = self._apply_prompt_template(text_content={"prompt_guideline": self.prompt_guide}, 
-                                            prompt_template=chat_prompt_template)
+        guideline = self._apply_prompt_template(text_content={"prompt_guideline": self.prompt_guide}, 
+                                                prompt_template=chat_prompt_template)
 
-        messages = [{"role": "system", "content": self.system_prompt},
-                    {"role": "user", "content": prompt}] + messages
-        
+        messages = [{"role": "system", "content": self.system_prompt + guideline}] + messages
 
         stream_generator = self.inference_engine.chat(messages, stream=True)
         yield from stream_generator

{llm_ie-1.2.0.dist-info → llm_ie-1.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-ie
-Version: 1.2.0
+Version: 1.2.2
 Summary: A comprehensive toolkit that provides building blocks for LLM-based named entity recognition, attribute extraction, and relation extraction pipelines.
 License: MIT
 Author: Enshuo (David) Hsu

{llm_ie-1.2.0.dist-info → llm_ie-1.2.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-llm_ie/__init__.py,sha256=rLP01qXkIisX0WLzZOv6y494Braw89g5JLmA6ZyrGGA,1590
+llm_ie/__init__.py,sha256=wNyek7i90UlQRylV3hSG9RlzMZ4MVzZSe_uhQYTQab4,1632
 llm_ie/asset/PromptEditor_prompts/chat.txt,sha256=Fq62voV0JQ8xBRcxS1Nmdd7DkHs1fGYb-tmNwctZZK0,118
 llm_ie/asset/PromptEditor_prompts/comment.txt,sha256=C_lxx-dlOlFJ__jkHKosZ8HsNAeV1aowh2B36nIipBY,159
 llm_ie/asset/PromptEditor_prompts/rewrite.txt,sha256=JAwY9vm1jSmKf2qcLBYUvrSmME2EJH36bALmkwZDWYQ,178
@@ -18,11 +18,11 @@ llm_ie/asset/prompt_guide/MultiClassRelationExtractor_prompt_guide.txt,sha256=EQ
 llm_ie/asset/prompt_guide/ReviewFrameExtractor_prompt_guide.txt,sha256=rBRIXg8JQWUHTRdoluTS0zkbTkBAacEtHHvr3lZaQCw,10437
 llm_ie/asset/prompt_guide/SentenceFrameExtractor_prompt_guide.txt,sha256=97_-y_vHMLG4Kb8fLsGgibLxB-3mest8k3LHfLo5h-I,10465
 llm_ie/asset/prompt_guide/SentenceReviewFrameExtractor_prompt_guide.txt,sha256=97_-y_vHMLG4Kb8fLsGgibLxB-3mest8k3LHfLo5h-I,10465
-llm_ie/chunkers.py,sha256=24h9l-Ldyx3EgfYicFqGhV_b-XofUS3yovC1nBWdDoo,5143
+llm_ie/chunkers.py,sha256=jXmUk3beF3EZWqDN_ArtoeerXObRKVCDIdUsv3loO80,6100
 llm_ie/data_types.py,sha256=72-3bzzYpo7KZpD9bjoroWT2eiM0zmWyDkBr2nHoBV0,18559
-llm_ie/engines.py,sha256=uE5sag1YeKBYBFF4gY7rYZK9e1ttatf9T7bV_xSg9Pk,36075
-llm_ie/extractors.py,sha256=aCRqKhjSoKTAWZ3WhX_O6V-S_rIvYhPsk78nZLDpQw8,95149
-llm_ie/prompt_editor.py,sha256=zh7Es5Ta2qSTgHtfF9Y9ZKXs4DMue6XlyRt9O6_Uk6c,10962
-llm_ie-1.2.0.dist-info/METADATA,sha256=X9zsMDwBAq1QzIkX8SSbmwLsEFiiAVeNeA0GTiNkAkQ,728
-llm_ie-1.2.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-llm_ie-1.2.0.dist-info/RECORD,,
+llm_ie/engines.py,sha256=Ofrbcu8j2dp2X25oMQ3Xg7FGPynHse_liQ8oFTEdeHA,38418
+llm_ie/extractors.py,sha256=5y4vuB53R2EAyHGH3wVZ3M1DvN3fPJHdypsTbzbK78s,96889
+llm_ie/prompt_editor.py,sha256=nAgCJQY5kVWTAhmrngdWRG-JKxCCPBh0dyaUcIk_-c0,13198
+llm_ie-1.2.2.dist-info/METADATA,sha256=or9H0YdfLVgjqftn3zg4nlRHmGHcK4hxBYR6R-1otuE,728
+llm_ie-1.2.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+llm_ie-1.2.2.dist-info/RECORD,,