llm-ie 1.2.1__tar.gz → 1.2.2__tar.gz

{llm_ie-1.2.1 → llm_ie-1.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-ie
-Version: 1.2.1
+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.1 → llm_ie-1.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "llm-ie"
-version = "1.2.1"
+version = "1.2.2"
 description = "A comprehensive toolkit that provides building blocks for LLM-based named entity recognition, attribute extraction, and relation extraction pipelines."
 authors = ["Enshuo (David) Hsu"]
 license = "MIT"

{llm_ie-1.2.1 → llm_ie-1.2.2}/src/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-1.2.1 → llm_ie-1.2.2}/src/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-1.2.1 → llm_ie-1.2.2}/src/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,24 +157,66 @@ 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):
+        return new_messages
+
+
+class OpenAIReasoningLLMConfig(ReasoningLLMConfig):
     def __init__(self, reasoning_effort:str=None, **kwargs):
         """
         The OpenAI "o" series configuration.
@@ -246,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
@@ -286,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.
 
@@ -298,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
 
@@ -363,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.
 
@@ -436,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.
 
@@ -448,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)
 
@@ -497,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.
         """
@@ -558,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.
 
@@ -570,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)
 
@@ -611,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.
         """
@@ -662,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.
 
@@ -674,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)
 
@@ -723,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.
         """
@@ -813,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.
 
@@ -825,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)
         
@@ -877,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-1.2.1 → llm_ie-1.2.2}/src/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-1.2.1 → llm_ie-1.2.2}/src/llm_ie/prompt_editor.py

@@ -85,7 +85,7 @@ 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:
         """
@@ -101,7 +101,7 @@ 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):
         """
@@ -175,7 +175,7 @@ class PromptEditor:
             self.messages.append({"role": "user", "content": user_input})
             print(f"{Fore.BLUE}Assistant: {Style.RESET_ALL}", end="")
             response = self.inference_engine.chat(self.messages, verbose=True)
-            self.messages.append({"role": "assistant", "content": response})
+            self.messages.append({"role": "assistant", "content": response["response"]})
             
 
     def _IPython_chat(self):
@@ -229,7 +229,7 @@ class PromptEditor:
             # Get assistant's response and append it to conversation
             print("Assistant: ", end="")
             response = self.inference_engine.chat(self.messages, verbose=True)
-            self.messages.append({"role": "assistant", "content": response})
+            self.messages.append({"role": "assistant", "content": response["response"]})
 
             # Display the assistant's response
             with output_area: