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

llm_ie/extractors.py

@@ -8,11 +8,12 @@ import warnings
 import itertools
 import asyncio
 import nest_asyncio
-from typing import Any, Set, List, Dict, Tuple, Union, Callable, Generator, Optional
-from llm_ie.data_types import FrameExtractionUnit, FrameExtractionUnitResult, LLMInformationExtractionFrame, LLMInformationExtractionDocument
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Set, List, Dict, Tuple, Union, Callable, Generator, Optional, AsyncGenerator
+from llm_ie.data_types import FrameExtractionUnit, LLMInformationExtractionFrame, LLMInformationExtractionDocument
 from llm_ie.chunkers import UnitChunker, WholeDocumentUnitChunker, SentenceUnitChunker
 from llm_ie.chunkers import ContextChunker, NoContextChunker, WholeDocumentContextChunker, SlideWindowContextChunker
-from llm_ie.engines import InferenceEngine
+from llm_ie.engines import InferenceEngine, MessagesLogger
 from colorama import Fore, Style        
 
 
@@ -405,7 +406,7 @@ class DirectFrameExtractor(FrameExtractor):
 
 
     def extract(self, text_content:Union[str, Dict[str,str]], 
-                document_key:str=None, verbose:bool=False, return_messages_log:bool=False) -> List[FrameExtractionUnitResult]:
+                document_key:str=None, verbose:bool=False, return_messages_log:bool=False) -> List[FrameExtractionUnit]:
         """
         This method inputs a text and outputs a list of outputs per unit.
 
@@ -423,11 +424,9 @@ class DirectFrameExtractor(FrameExtractor):
         return_messages_log : bool, Optional
             if True, a list of messages will be returned.
 
-        Return : List[FrameExtractionUnitResult]
+        Return : List[FrameExtractionUnit]
             the output from LLM for each unit. Contains the start, end, text, and generated text.
         """
-        # define output
-        output = []
         # unit chunking
         if isinstance(text_content, str):
             doc_text = text_content
@@ -440,76 +439,70 @@ class DirectFrameExtractor(FrameExtractor):
         units = self.unit_chunker.chunk(doc_text)
         # context chunker init
         self.context_chunker.fit(doc_text, units)
+
         # messages log
-        if return_messages_log:
-            messages_log = []
+        messages_logger = MessagesLogger() if return_messages_log else None
 
         # generate unit by unit
         for i, unit in enumerate(units):
-            # construct chat messages
-            messages = []
-            if self.system_prompt:
-                messages.append({'role': 'system', 'content': self.system_prompt})
-
-            context = self.context_chunker.chunk(unit)
-            
-            if context == "":
-                # no context, just place unit in user prompt
-                if isinstance(text_content, str):
-                    messages.append({'role': 'user', 'content': self._get_user_prompt(unit.text)})
-                else:
-                    unit_content = text_content.copy()
-                    unit_content[document_key] = unit.text
-                    messages.append({'role': 'user', 'content': self._get_user_prompt(unit_content)})
-            else:
-                # insert context to user prompt
-                if isinstance(text_content, str):
-                    messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
-                else:
-                    context_content = text_content.copy()
-                    context_content[document_key] = context
-                    messages.append({'role': 'user', 'content': self._get_user_prompt(context_content)})
-                # simulate conversation where assistant confirms
-                messages.append({'role': 'assistant', 'content': 'Sure, please provide the unit text (e.g., sentence, line, chunk) of interest.'})
-                # place unit of interest
-                messages.append({'role': 'user', 'content': unit.text})
+            try:
+                # construct chat messages
+                messages = []
+                if self.system_prompt:
+                    messages.append({'role': 'system', 'content': self.system_prompt})
 
-            if verbose:
-                print(f"\n\n{Fore.GREEN}Unit {i}:{Style.RESET_ALL}\n{unit.text}\n")
-                if context != "":
-                    print(f"{Fore.YELLOW}Context:{Style.RESET_ALL}\n{context}\n")
+                context = self.context_chunker.chunk(unit)
                 
-                print(f"{Fore.BLUE}Extraction:{Style.RESET_ALL}")
+                if context == "":
+                    # no context, just place unit in user prompt
+                    if isinstance(text_content, str):
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(unit.text)})
+                    else:
+                        unit_content = text_content.copy()
+                        unit_content[document_key] = unit.text
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(unit_content)})
+                else:
+                    # insert context to user prompt
+                    if isinstance(text_content, str):
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                    else:
+                        context_content = text_content.copy()
+                        context_content[document_key] = context
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(context_content)})
+                    # simulate conversation where assistant confirms
+                    messages.append({'role': 'assistant', 'content': 'Sure, please provide the unit text (e.g., sentence, line, chunk) of interest.'})
+                    # place unit of interest
+                    messages.append({'role': 'user', 'content': unit.text})
 
-            
-            gen_text = self.inference_engine.chat(
-                            messages=messages, 
-                            verbose=verbose,
-                            stream=False
-                        )
+                if verbose:
+                    print(f"\n\n{Fore.GREEN}Unit {i + 1}/{len(units)}:{Style.RESET_ALL}\n{unit.text}\n")
+                    if context != "":
+                        print(f"{Fore.YELLOW}Context:{Style.RESET_ALL}\n{context}\n")
+                    
+                    print(f"{Fore.BLUE}Extraction:{Style.RESET_ALL}")
 
-            if return_messages_log:
-                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
-            result = FrameExtractionUnitResult(
-                            start=unit.start,
-                            end=unit.end,
-                            text=unit.text,
-                            gen_text=gen_text["response"])
-            output.append(result)
+                
+                gen_text = self.inference_engine.chat(
+                                messages=messages, 
+                                verbose=verbose,
+                                stream=False,
+                                messages_logger=messages_logger
+                            )
+
+                # add generated text to unit
+                unit.set_generated_text(gen_text["response"])
+                unit.set_status("success")
+            except Exception as e:
+                unit.set_status("fail")
+                warnings.warn(f"LLM inference failed for unit {i} ({unit.start}, {unit.end}): {e}", RuntimeWarning)
             
         if return_messages_log:
-            return output, messages_log
+            return units, messages_logger.get_messages_log()
         
-        return output
+        return units
     
     def stream(self, text_content: Union[str, Dict[str, str]], 
-               document_key: str = None) -> Generator[Dict[str, Any], None, List[FrameExtractionUnitResult]]:
+               document_key: str = None) -> Generator[Dict[str, Any], None, List[FrameExtractionUnit]]:
         """
         Streams LLM responses per unit with structured event types,
         and returns collected data for post-processing.
@@ -525,12 +518,10 @@ class DirectFrameExtractor(FrameExtractor):
 
         Returns:
         --------
-        List[FrameExtractionUnitResult]:
-            A list of FrameExtractionUnitResult objects, each containing the
+        List[FrameExtractionUnit]:
+            A list of FrameExtractionUnit objects, each containing the
             original unit details and the fully accumulated 'gen_text' from the LLM.
         """
-        collected_results: List[FrameExtractionUnitResult] = []
-
         if isinstance(text_content, str):
             doc_text = text_content
         elif isinstance(text_content, dict):
@@ -588,19 +579,14 @@ class DirectFrameExtractor(FrameExtractor):
                     current_gen_text += chunk["data"]
            
             # Store the result for this unit
-            result_for_unit = FrameExtractionUnitResult(
-                start=unit.start,
-                end=unit.end,
-                text=unit.text,
-                gen_text=current_gen_text
-            )
-            collected_results.append(result_for_unit)
+            unit.set_generated_text(current_gen_text)
+            unit.set_status("success")
 
         yield {"type": "info", "data": "All units processed by LLM."}
-        return collected_results
+        return units
 
     async def extract_async(self, text_content:Union[str, Dict[str,str]], document_key:str=None, 
-                            concurrent_batch_size:int=32, return_messages_log:bool=False) -> List[FrameExtractionUnitResult]:
+                            concurrent_batch_size:int=32, return_messages_log:bool=False) -> List[FrameExtractionUnit]:
         """
         This is the asynchronous version of the extract() method.
 
@@ -618,7 +604,7 @@ class DirectFrameExtractor(FrameExtractor):
         return_messages_log : bool, Optional
             if True, a list of messages will be returned.
 
-        Return : List[FrameExtractionUnitResult]
+        Return : List[FrameExtractionUnit]
             the output from LLM for each unit. Contains the start, end, text, and generated text.
         """
         if isinstance(text_content, str):
@@ -637,6 +623,9 @@ class DirectFrameExtractor(FrameExtractor):
         # context chunker init 
         self.context_chunker.fit(doc_text, units)
 
+        # messages logger init
+        messages_logger = MessagesLogger() if return_messages_log else None
+
         # Prepare inputs for all units first
         tasks_input = []
         for i, unit in enumerate(units):
@@ -677,17 +666,15 @@ class DirectFrameExtractor(FrameExtractor):
         async def semaphore_helper(task_data: Dict, **kwrs):
             unit = task_data["unit"]
             messages = task_data["messages"]
-            original_index = task_data["original_index"]
 
             async with semaphore:
                 gen_text = await self.inference_engine.chat_async(
-                    messages=messages
+                    messages=messages,
+                    messages_logger=messages_logger
                 )
 
-            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
+            unit.set_generated_text(gen_text["response"])
+            unit.set_status("success")
            
         # Create and gather tasks
         tasks = []
@@ -697,40 +684,13 @@ class DirectFrameExtractor(FrameExtractor):
             ))
             tasks.append(task)
 
-        results_raw = await asyncio.gather(*tasks)
-
-        # Sort results back into original order using the index stored
-        results_raw.sort(key=lambda x: x["original_index"])
-
-        # Restructure the results
-        output: List[FrameExtractionUnitResult] = []
-        messages_log: Optional[List[List[Dict[str, str]]]] = [] if return_messages_log else None
-
-        for result_data in results_raw:
-            unit = result_data["unit"]
-            gen_text = result_data["gen_text"]
-
-            # Create result object
-            result = FrameExtractionUnitResult(
-                start=unit.start,
-                end=unit.end,
-                text=unit.text,
-                gen_text=gen_text
-            )
-            output.append(result)
-
-            # Append to messages log if requested
-            if return_messages_log:
-                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)
+        await asyncio.gather(*tasks)
 
+        # Return units
         if return_messages_log:
-            return output, messages_log
+            return units, messages_logger.get_messages_log()
         else:
-            return output
+            return units
 
 
     def extract_frames(self, text_content:Union[str, Dict[str,str]], document_key:str=None, 
@@ -738,7 +698,7 @@ class DirectFrameExtractor(FrameExtractor):
                         case_sensitive:bool=False, fuzzy_match:bool=True, fuzzy_buffer_size:float=0.2, fuzzy_score_cutoff:float=0.8,
                         allow_overlap_entities:bool=False, return_messages_log:bool=False) -> List[LLMInformationExtractionFrame]:
         """
-        This method inputs a text and outputs a list of LLMInformationExtractionFrame
+        This method inputs a document text and outputs a list of LLMInformationExtractionFrame
         It use the extract() method and post-process outputs into frames.
 
         Parameters:
@@ -791,18 +751,21 @@ class DirectFrameExtractor(FrameExtractor):
                                                 verbose=verbose,
                                                 return_messages_log=return_messages_log)
             
-        llm_output_results, messages_log = extraction_results if return_messages_log else (extraction_results, None)
+        units, messages_log = extraction_results if return_messages_log else (extraction_results, None)
         
         frame_list = []
-        for res in llm_output_results:
+        for unit in units:
             entity_json = []
-            for entity in self._extract_json(gen_text=res.gen_text):
+            if unit.status != "success":
+                warnings.warn(f"Skipping failed unit ({unit.start}, {unit.end}): {unit.text}", RuntimeWarning)
+                continue
+            for entity in self._extract_json(gen_text=unit.gen_text):
                 if ENTITY_KEY in entity:
                     entity_json.append(entity)
                 else:
                     warnings.warn(f'Extractor output "{entity}" does not have entity_key ("{ENTITY_KEY}"). This frame will be dropped.', RuntimeWarning)
 
-            spans = self._find_entity_spans(text=res.text, 
+            spans = self._find_entity_spans(text=unit.text, 
                                             entities=[e[ENTITY_KEY] for e in entity_json], 
                                             case_sensitive=case_sensitive,
                                             fuzzy_match=fuzzy_match,
@@ -812,9 +775,9 @@ class DirectFrameExtractor(FrameExtractor):
             for ent, span in zip(entity_json, spans):
                 if span is not None:
                     start, end = span
-                    entity_text = res.text[start:end]
-                    start += res.start
-                    end += res.start
+                    entity_text = unit.text[start:end]
+                    start += unit.start
+                    end += unit.start
                     attr = {}
                     if "attr" in ent and ent["attr"] is not None:
                         attr = ent["attr"]
@@ -831,6 +794,208 @@ class DirectFrameExtractor(FrameExtractor):
         return frame_list
         
 
+    async def extract_frames_from_documents(self, text_contents:List[Union[str,Dict[str, any]]], document_key:str="text",
+            cpu_concurrency:int=4, llm_concurrency:int=32, case_sensitive:bool=False, 
+            fuzzy_match:bool=True, fuzzy_buffer_size:float=0.2, fuzzy_score_cutoff:float=0.8,
+            allow_overlap_entities:bool=False, return_messages_log:bool=False) -> AsyncGenerator[Dict[str, any], None]:
+        """
+        This method inputs a list of documents and yields the results for each document as soon as it is complete.
+
+        Parameters:
+        -----------
+        text_contents : List[Union[str,Dict[str, any]]]
+            a list of input text contents to put in prompt template. 
+            If str, the prompt template must has only 1 placeholder {{<placeholder name>}}, regardless of placeholder name.
+            If dict, all the keys must be included in the prompt template placeholder {{<placeholder name>}}.
+        document_key: str, optional
+            The key in the `text_contents` dictionaries that holds the document text.
+        cpu_concurrency: int, optional
+            The number of parallel threads to use for CPU-bound tasks like chunking.
+        llm_concurrency: int, optional
+            The number of concurrent requests to make to the LLM.
+        case_sensitive : bool, Optional
+            if True, entity text matching will be case-sensitive.
+        fuzzy_match : bool, Optional
+            if True, fuzzy matching will be applied to find entity text.
+        fuzzy_buffer_size : float, Optional
+            the buffer size for fuzzy matching. Default is 20% of entity text length.
+        fuzzy_score_cutoff : float, Optional
+            the Jaccard score cutoff for fuzzy matching. 
+            Matched entity text must have a score higher than this value or a None will be returned.
+        allow_overlap_entities : bool, Optional
+            if True, entities can overlap in the text.
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
+
+        Yields:
+        -------
+        AsyncGenerator[Dict[str, any], None]
+            A dictionary for each completed document, containing its 'idx' and extracted 'frames'.
+        """
+        # Validate text_contents must be a list of str or dict, and not both
+        if not isinstance(text_contents, list):
+            raise ValueError("text_contents must be a list of strings or dictionaries.")
+        if all(isinstance(doc, str) for doc in text_contents):
+            pass  
+        elif all(isinstance(doc, dict) for doc in text_contents):
+            pass
+        # Set CPU executor and queues
+        cpu_executor = ThreadPoolExecutor(max_workers=cpu_concurrency)
+        tasks_queue = asyncio.Queue(maxsize=llm_concurrency * 2)
+        # Store to track units and pending counts
+        results_store = {
+            idx: {'pending': 0, 'units': [], 'text': doc if isinstance(doc, str) else doc.get(document_key, "")}
+            for idx, doc in enumerate(text_contents)
+        }
+
+        output_queue = asyncio.Queue()
+        messages_logger = MessagesLogger() if return_messages_log else None
+
+        async def producer():
+            try:
+                for idx, text_content in enumerate(text_contents):
+                    text = text_content if isinstance(text_content, str) else text_content.get(document_key, "")
+                    if not text:
+                        warnings.warn(f"Document at index {idx} is empty or missing the document key '{document_key}'.")
+                        # signal that this document is done
+                        await output_queue.put({'idx': idx, 'frames': []})
+                        continue
+
+                    units = await self.unit_chunker.chunk_async(text, cpu_executor)
+                    await self.context_chunker.fit_async(text, units, cpu_executor)
+                    results_store[idx]['pending'] = len(units)
+
+                    # Handle cases where a document yields no units
+                    if not units: 
+                        # signal that this document is done
+                        await output_queue.put({'idx': idx, 'frames': []})
+                        continue
+
+                    # Iterate through units
+                    for unit in units:
+                        context = await self.context_chunker.chunk_async(unit, cpu_executor)
+                        messages = []
+                        if self.system_prompt:
+                            messages.append({'role': 'system', 'content': self.system_prompt})
+
+                        if not context:
+                            if isinstance(text_content, str):
+                                messages.append({'role': 'user', 'content': self._get_user_prompt(unit.text)})
+                            else:
+                                unit_content = text_content.copy()
+                                unit_content[document_key] = unit.text
+                                messages.append({'role': 'user', 'content': self._get_user_prompt(unit_content)})
+                        else:
+                            # insert context to user prompt
+                            if isinstance(text_content, str):
+                                messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                            else:
+                                context_content = text_content.copy()
+                                context_content[document_key] = context
+                                messages.append({'role': 'user', 'content': self._get_user_prompt(context_content)})
+                            # simulate conversation where assistant confirms
+                            messages.append({'role': 'assistant', 'content': 'Sure, please provide the unit text (e.g., sentence, line, chunk) of interest.'})
+                            # place unit of interest
+                            messages.append({'role': 'user', 'content': unit.text})
+
+                        await tasks_queue.put({'idx': idx, 'unit': unit, 'messages': messages})
+            finally:
+                for _ in range(llm_concurrency):
+                    await tasks_queue.put(None)
+
+        async def worker():
+            while True:
+                task_item = await tasks_queue.get()
+                if task_item is None:
+                    tasks_queue.task_done()
+                    break
+
+                idx = task_item['idx']
+                unit = task_item['unit']
+                doc_results = results_store[idx]
+
+                try:
+                    gen_text = await self.inference_engine.chat_async(
+                        messages=task_item['messages'], messages_logger=messages_logger
+                    )
+                    unit.set_generated_text(gen_text["response"])
+                    unit.set_status("success")
+                    doc_results['units'].append(unit)
+                except Exception as e:
+                    warnings.warn(f"Error processing unit for doc idx {idx}: {e}")
+                finally:
+                    doc_results['pending'] -= 1
+                    if doc_results['pending'] <= 0:
+                        final_frames = self._post_process_and_create_frames(doc_results, case_sensitive, fuzzy_match, fuzzy_buffer_size, fuzzy_score_cutoff, allow_overlap_entities)
+                        output_payload = {'idx': idx, 'frames': final_frames}
+                        if return_messages_log:
+                            output_payload['messages_log'] = messages_logger.get_messages_log()
+                        await output_queue.put(output_payload)
+
+                    tasks_queue.task_done()
+
+        # Start producer and workers
+        producer_task = asyncio.create_task(producer())
+        worker_tasks = [asyncio.create_task(worker()) for _ in range(llm_concurrency)]
+
+        # Main loop to gather results
+        docs_completed = 0
+        while docs_completed < len(text_contents):
+            result = await output_queue.get()
+            yield result
+            docs_completed += 1
+
+        # Final cleanup
+        await producer_task 
+        await tasks_queue.join()
+
+        # Cancel any lingering worker tasks
+        for task in worker_tasks:
+            task.cancel()
+        await asyncio.gather(*worker_tasks, return_exceptions=True)
+
+        cpu_executor.shutdown(wait=False)
+
+
+    def _post_process_and_create_frames(self, doc_results, case_sensitive, fuzzy_match, fuzzy_buffer_size, fuzzy_score_cutoff, allow_overlap_entities):
+        """Helper function to run post-processing logic for a completed document."""
+        ENTITY_KEY = "entity_text"
+        frame_list = []
+        for res in sorted(doc_results['units'], key=lambda r: r.start):
+            entity_json = []
+            for entity in self._extract_json(gen_text=res.gen_text):
+                if ENTITY_KEY in entity:
+                    entity_json.append(entity)
+                else:
+                    warnings.warn(f'Extractor output "{entity}" does not have entity_key ("{ENTITY_KEY}"). This frame will be dropped.', RuntimeWarning)
+
+            spans = self._find_entity_spans(
+                text=res.text,
+                entities=[e[ENTITY_KEY] for e in entity_json],
+                case_sensitive=case_sensitive,
+                fuzzy_match=fuzzy_match,
+                fuzzy_buffer_size=fuzzy_buffer_size,
+                fuzzy_score_cutoff=fuzzy_score_cutoff,
+                allow_overlap_entities=allow_overlap_entities
+            )
+            for ent, span in zip(entity_json, spans):
+                if span is not None:
+                    start, end = span
+                    entity_text = res.text[start:end]
+                    start += res.start
+                    end += res.start
+                    attr = ent.get("attr", {}) or {}
+                    frame = LLMInformationExtractionFrame(
+                        frame_id=f"{len(frame_list)}",
+                        start=start,
+                        end=end,
+                        entity_text=entity_text,
+                        attr=attr
+                    )
+                    frame_list.append(frame)
+        return frame_list
+
+
 class ReviewFrameExtractor(DirectFrameExtractor):
     def __init__(self, unit_chunker:UnitChunker, context_chunker:ContextChunker, inference_engine:InferenceEngine, 
                  prompt_template:str, review_mode:str, review_prompt:str=None, system_prompt:str=None):
@@ -902,7 +1067,7 @@ class ReviewFrameExtractor(DirectFrameExtractor):
             raise ValueError(f"Cannot find review prompt for {self.__class__.__name__} in the package. Please provide a review_prompt.")
 
     def extract(self, text_content:Union[str, Dict[str,str]], document_key:str=None, 
-                verbose:bool=False, return_messages_log:bool=False) -> List[FrameExtractionUnitResult]:
+                verbose:bool=False, return_messages_log:bool=False) -> List[FrameExtractionUnit]:
         """
         This method inputs a text and outputs a list of outputs per unit.
 
@@ -923,8 +1088,6 @@ class ReviewFrameExtractor(DirectFrameExtractor):
         Return : List[FrameExtractionUnitResult]
             the output from LLM for each unit. Contains the start, end, text, and generated text.
         """
-        # define output
-        output = []
         # unit chunking
         if isinstance(text_content, str):
             doc_text = text_content
@@ -937,9 +1100,9 @@ class ReviewFrameExtractor(DirectFrameExtractor):
         units = self.unit_chunker.chunk(doc_text)
         # context chunker init
         self.context_chunker.fit(doc_text, units)
-        # messages log
-        if return_messages_log:
-            messages_log = []
+
+        # messages logger init
+        messages_logger = MessagesLogger() if return_messages_log else None
 
         # generate unit by unit
         for i, unit in enumerate(units):
@@ -973,7 +1136,7 @@ class ReviewFrameExtractor(DirectFrameExtractor):
                 messages.append({'role': 'user', 'content': unit.text})
 
             if verbose:
-                print(f"\n\n{Fore.GREEN}Unit {i}:{Style.RESET_ALL}\n{unit.text}\n")
+                print(f"\n\n{Fore.GREEN}Unit {i + 1}/{len(units)}:{Style.RESET_ALL}\n{unit.text}\n")
                 if context != "":
                     print(f"{Fore.YELLOW}Context:{Style.RESET_ALL}\n{context}\n")
                 
@@ -983,7 +1146,8 @@ class ReviewFrameExtractor(DirectFrameExtractor):
             initial = self.inference_engine.chat(
                             messages=messages, 
                             verbose=verbose,
-                            stream=False
+                            stream=False,
+                            messages_logger=messages_logger
                         )
 
             # <--- Review step --->
@@ -996,7 +1160,8 @@ class ReviewFrameExtractor(DirectFrameExtractor):
             review = self.inference_engine.chat(
                             messages=messages, 
                             verbose=verbose,
-                            stream=False
+                            stream=False,
+                            messages_logger=messages_logger
                         )
 
             # Output
@@ -1005,28 +1170,14 @@ class ReviewFrameExtractor(DirectFrameExtractor):
             elif self.review_mode == "addition":
                 gen_text = initial["response"] + '\n' + review["response"]
 
-            if return_messages_log:
-                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
-            result = FrameExtractionUnitResult(
-                            start=unit.start,
-                            end=unit.end,
-                            text=unit.text,
-                            gen_text=gen_text)
-            output.append(result)
+            # add generated text to unit
+            unit.set_generated_text(gen_text)
+            unit.set_status("success")
             
         if return_messages_log:
-            return output, messages_log
+            return units, messages_logger.get_messages_log()
         
-        return output
+        return units
 
 
     def stream(self, text_content:Union[str, Dict[str,str]], document_key:str=None) -> Generator[str, None, None]:
@@ -1122,7 +1273,7 @@ class ReviewFrameExtractor(DirectFrameExtractor):
                 yield chunk
 
     async def extract_async(self, text_content:Union[str, Dict[str,str]], document_key:str=None,
-                            concurrent_batch_size:int=32, return_messages_log:bool=False, **kwrs) -> List[FrameExtractionUnitResult]:
+                            concurrent_batch_size:int=32, return_messages_log:bool=False, **kwrs) -> List[FrameExtractionUnit]:
         """
         This is the asynchronous version of the extract() method with the review step.
 
@@ -1154,11 +1305,15 @@ class ReviewFrameExtractor(DirectFrameExtractor):
         else:
             raise TypeError("text_content must be a string or a dictionary.")
 
+        # unit chunking
         units = self.unit_chunker.chunk(doc_text)
 
         # context chunker init
         self.context_chunker.fit(doc_text, units)
 
+        # messages logger init
+        messages_logger = MessagesLogger() if return_messages_log else None
+
         # <--- Initial generation step --->
         initial_tasks_input = []
         for i, unit in enumerate(units):
@@ -1202,7 +1357,8 @@ class ReviewFrameExtractor(DirectFrameExtractor):
 
             async with semaphore:
                 gen_text = await self.inference_engine.chat_async(
-                    messages=messages
+                    messages=messages,
+                    messages_logger=messages_logger
                 )
             # Return initial generation result along with the messages used and the unit
             out = {"original_index": original_index, "unit": unit, "initial_gen_text": gen_text["response"], "initial_messages": messages}
@@ -1253,16 +1409,11 @@ class ReviewFrameExtractor(DirectFrameExtractor):
 
             async with semaphore:
                 review_gen_text = await self.inference_engine.chat_async(
-                    messages=messages
+                    messages=messages,
+                    messages_logger=messages_logger
                 )
             # Combine initial and review results
             task_data["review_gen_text"] = review_gen_text["response"]
-            if return_messages_log:
-                # Log for the review call itself
-                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
@@ -1279,9 +1430,6 @@ class ReviewFrameExtractor(DirectFrameExtractor):
         final_results_raw.sort(key=lambda x: x["original_index"])
 
         # <--- Process final results --->
-        output: List[FrameExtractionUnitResult] = []
-        messages_log: Optional[List[List[Dict[str, str]]]] = [] if return_messages_log else None
-
         for result_data in final_results_raw:
             unit = result_data["unit"]
             initial_gen = result_data["initial_gen_text"]
@@ -1296,23 +1444,13 @@ class ReviewFrameExtractor(DirectFrameExtractor):
                 final_gen_text = review_gen # Default to revision if mode is somehow invalid
 
             # Create final result object
-            result = FrameExtractionUnitResult(
-                start=unit.start,
-                end=unit.end,
-                text=unit.text,
-                gen_text=final_gen_text # Use the combined/reviewed text
-            )
-            output.append(result)
-
-            # Append full conversation log if requested
-            if return_messages_log:
-                full_log_for_unit = result_data["full_review_log"]
-                messages_log.append(full_log_for_unit)
+            unit.set_generated_text(final_gen_text)
+            unit.set_status("success")
 
         if return_messages_log:
-            return output, messages_log
+            return units, messages_logger.get_messages_log()
         else:
-            return output
+            return units
 
 
 class BasicFrameExtractor(DirectFrameExtractor):
@@ -1549,6 +1687,9 @@ class AttributeExtractor(Extractor):
             a dictionary of attributes extracted from the frame.
             If return_messages_log is True, a list of messages will be returned as well.
         """
+        # messages logger init
+        messages_logger = MessagesLogger() if return_messages_log else None
+
         # construct chat messages
         messages = []
         if self.system_prompt:
@@ -1567,19 +1708,15 @@ class AttributeExtractor(Extractor):
         gen_text = self.inference_engine.chat(
                             messages=messages,
                             verbose=verbose,
-                            stream=False
+                            stream=False,
+                            messages_logger=messages_logger
                         )
-        if return_messages_log:
-            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["response"])
         if isinstance(attribute_list, list) and len(attribute_list) > 0:
             attributes = attribute_list[0]
             if return_messages_log:
-                return attributes, messages
+                return attributes, messages_logger.get_messages_log()
             return attributes
 
 
@@ -1620,7 +1757,7 @@ class AttributeExtractor(Extractor):
             if return_messages_log:
                 attr, messages = self._extract_from_frame(frame=frame, text=text, context_size=context_size,
                                                           verbose=verbose, return_messages_log=return_messages_log)
-                messages_log.append(messages)
+                messages_log.extend(messages)
             else: 
                 attr = self._extract_from_frame(frame=frame, text=text, context_size=context_size,
                                                 verbose=verbose, return_messages_log=return_messages_log)
@@ -1669,6 +1806,9 @@ class AttributeExtractor(Extractor):
         if not isinstance(text, str):
             raise TypeError(f"Expect text as str, received {type(text)} instead.")
 
+        # messages logger init
+        messages_logger = MessagesLogger() if return_messages_log else None
+
         # async helper
         semaphore = asyncio.Semaphore(concurrent_batch_size)
         
@@ -1681,14 +1821,7 @@ class AttributeExtractor(Extractor):
                 context = self._get_context(frame, text, context_size)
                 messages.append({'role': 'user', 'content': self._get_user_prompt({"context": context, "frame": str(frame.to_dict())})})
 
-                gen_text = await self.inference_engine.chat_async(messages=messages)
-                
-                if return_messages_log:
-                    message = {"role": "assistant", "content": gen_text["response"]}
-                    if "reasoning" in gen_text:
-                        message["reasoning"] = gen_text["reasoning"]
-                    messages.append(message)
-
+                gen_text = await self.inference_engine.chat_async(messages=messages, messages_logger=messages_logger)
                 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}
@@ -1699,12 +1832,8 @@ class AttributeExtractor(Extractor):
 
         # process results
         new_frames = []
-        messages_log = [] if return_messages_log else None
 
         for result in results:
-            if return_messages_log:
-                messages_log.append(result["messages"])
-
             if inplace:
                 result["frame"].attr.update(result["attributes"])
             else:
@@ -1714,9 +1843,9 @@ class AttributeExtractor(Extractor):
 
         # output
         if inplace:
-            return messages_log if return_messages_log else None
+            return messages_logger.get_messages_log() if return_messages_log else None
         else:
-            return (new_frames, messages_log) if return_messages_log else new_frames
+            return (new_frames, messages_logger.get_messages_log()) if return_messages_log else new_frames
 
     def extract_attributes(self, frames:List[LLMInformationExtractionFrame], text:str, context_size:int=256, 
                            concurrent:bool=False, concurrent_batch_size:int=32, verbose:bool=False, 
@@ -1839,7 +1968,7 @@ class RelationExtractor(Extractor):
                  return_messages_log: bool = False) -> Union[List[Dict], Tuple[List[Dict], List]]:
         pairs = itertools.combinations(doc.frames, 2)
         relations = []
-        messages_log = [] if return_messages_log else None
+        messages_logger = MessagesLogger() if return_messages_log else None
 
         for frame_1, frame_2 in pairs:
             task_payload = self._get_task_if_possible(frame_1, frame_2, doc.text, buffer_size)
@@ -1851,20 +1980,14 @@ class RelationExtractor(Extractor):
 
                 gen_text = self.inference_engine.chat(
                     messages=task_payload['messages'],
-                    verbose=verbose
+                    verbose=verbose,
+                    messages_logger=messages_logger
                 )
                 relation = self._post_process_result(gen_text["response"], task_payload)
                 if relation:
                     relations.append(relation)
 
-                if return_messages_log:
-                    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
+        return (relations, messages_logger.get_messages_log()) if return_messages_log else relations
 
     async def _extract_async(self, doc: LLMInformationExtractionDocument, buffer_size: int = 128, concurrent_batch_size: int = 32, return_messages_log: bool = False) -> Union[List[Dict], Tuple[List[Dict], List]]:
         pairs = list(itertools.combinations(doc.frames, 2))
@@ -1873,12 +1996,12 @@ class RelationExtractor(Extractor):
         tasks_input = [task for task in tasks_input if task is not None] 
 
         relations = []
-        messages_log = [] if return_messages_log else None
+        messages_logger = MessagesLogger() if return_messages_log else None
         semaphore = asyncio.Semaphore(concurrent_batch_size)
 
         async def semaphore_helper(task_payload: Dict):
             async with semaphore:
-                gen_text = await self.inference_engine.chat_async(messages=task_payload['messages'])
+                gen_text = await self.inference_engine.chat_async(messages=task_payload['messages'], messages_logger=messages_logger)
                 return gen_text, task_payload
 
         tasks = [asyncio.create_task(semaphore_helper(payload)) for payload in tasks_input]
@@ -1889,14 +2012,7 @@ class RelationExtractor(Extractor):
             if relation:
                 relations.append(relation)
 
-            if return_messages_log:
-                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
+        return (relations, messages_logger.get_messages_log()) if return_messages_log else relations
 
     def extract_relations(self, doc: LLMInformationExtractionDocument, buffer_size: int = 128, concurrent: bool = False, concurrent_batch_size: int = 32, verbose: bool = False, return_messages_log: bool = False) -> List[Dict]:
         if not doc.has_frame():