llm-ie 1.2.3__tar.gz → 1.2.4__tar.gz

{llm_ie-1.2.3 → llm_ie-1.2.4}/PKG-INFO

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "llm-ie"
-version = "1.2.3"
+version = "1.2.4"
 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.3 → llm_ie-1.2.4}/src/llm_ie/__init__.py

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

llm_ie-1.2.4/src/llm_ie/asset/default_prompts/LLMUnitChunker_user_prompt.txt

@@ -0,0 +1,129 @@
+### Task description
+You are a helpful assistant that breaks down a text document into semantic units or chunks. Each chunk should represent a coherent section of the text, such as a paragraph, subsection, or topic.
+
+### Schema definition
+You will output a JSON array of objects. Each object should have the following fields:
+    - "title": Generate a brief title summarizing the content of the chunk.
+    - "anchor_text": the first line of text in the chunk used to locate it in the original document. Must be an exact match.
+        - if there is a title or heading for the chunk, use that as the anchor_text.
+        - if there is no title or heading, use the first sentence of the chunk as the anchor_text.
+
+```JSON
+[
+    {
+        "title": "<your title here>",
+        "anchor_text": "<the anchor text of the chunk here>"
+    },
+    {
+        "title": "<your title here>",
+        "anchor_text": "<the anchor text of the chunk here>"
+    }
+]
+```
+
+### Examples
+
+**Input**:
+"# Clinical Note
+
+**Patient Name**: Michael Green  
+**Medical Record Number**: 1122334455  
+**Date of Visit**: January 5, 2025  
+**Provider**: Dr. Emily Carter, MD  
+
+## Reason for Visit  
+Follow-up for poorly controlled type 2 diabetes and complaints of occasional dizziness and blurred vision.
+
+## Summary of Visit  
+Michael Green, a 62-year-old male with a known history of type 2 diabetes, hypertension, and obesity, presents for follow-up regarding his glycemic control. Despite recent adjustments to his treatment plan, his glucose readings have remained elevated, averaging 180-220 mg/dL. He reports occasional episodes of dizziness and blurred vision, particularly in the morning before meals. He denies chest pain, palpitations, or recent falls. He reports compliance with his medication regimen but admits to difficulty following a consistent low-carbohydrate diet.
+
+Michael has been using a glucose meter to monitor his blood sugar levels and logs them daily. His last hemoglobin A1c, performed three months ago, was 9.2%. He reports no recent hospitalizations, infections, or significant stressors.
+
+## Notable History  
+- **Chronic Conditions**:  
+  - Type 2 diabetes mellitus, diagnosed 10 years ago.  
+  - Hypertension, well-controlled on medication.  
+  - Hyperlipidemia, on statin therapy.  
+- **Past Surgical History**:  
+  - Knee arthroscopy for a meniscal tear, age 50.  
+- **Family History**:  
+  - Mother: Deceased at 75, complications from diabetes.  
+  - Father: Deceased at 70, myocardial infarction."
+
+**Output**:
+```JSON
+[
+    {
+        "title": "Patient Information",
+        "anchor_text": "# Clinical Note"
+    },
+    {
+        "title": "Reason for Visit",
+        "anchor_text": "## Reason for Visit"
+    },
+    {
+        "title": "Summary of Visit",
+        "anchor_text": "## Summary of Visit"
+    },
+    {
+        "title": "Notable History",
+        "anchor_text": "## Notable History"
+    }
+]
+```
+
+**Input**:
+"In the [**Hospital1 18**] ED, 35.3 102 133/58 100%AC 500x20, 5, 1.0 with an
+ABG 7.16/66/162. He had a CTH which was unremarkable. He then
+had a CTA chest, afterwhich he went into PEA arrest.
+Rescucitation last approximately 10-15 minutes with multiple
+rounds of epi and bicarb, with ROSC. He was then admitted to the
+MICU for further management.
+.
+Currently, the patient is intubated, sedated, and parlyzed.
+
+Past Medical History:
+Asthma
+Dilated cardiomyopathy
+Multiple admissions for dyspnea this winter ([**1-26**]).
+Anxiety/depression
+CKD
+HLD
+Obesity
+HTN
+
+Social History:
+Unknown
+
+Family History:
+Unknown"
+
+**Output**:
+```JSON
+[
+    {
+        "title": "Patient Presentation and Initial Management",
+        "anchor_text": "In the [**Hospital1 18**] ED, 35.3 102 133/58 100%AC 500x20, 5, 1.0 with an"
+    },
+    {
+        "title": "Current Status of the Patient",
+        "anchor_text": "Currently, the patient is intubated, sedated, and parlyzed."
+    },
+    {
+        "title": "Past Medical History",
+        "anchor_text": "Past Medical History:"
+    },
+    {
+        "title": "Social History",
+        "anchor_text": "Social History:"
+    },
+    {
+        "title": "Family History",
+        "anchor_text": "Family History:"
+    }
+]
+```
+
+### Document text
+
+"{{document_text}}"

{llm_ie-1.2.3 → llm_ie-1.2.4}/src/llm_ie/chunkers.py

@@ -1,8 +1,11 @@
 import abc
-from typing import Set, List, Dict, Tuple, Union, Callable
+from typing import List
 import asyncio
 import uuid
+import importlib.resources
+from llm_ie.utils import extract_json, apply_prompt_template
 from llm_ie.data_types import FrameExtractionUnit
+from llm_ie.engines import InferenceEngine
 
 
 class UnitChunker(abc.ABC):
@@ -74,13 +77,14 @@ class SeparatorUnitChunker(UnitChunker):
         text : str
             The document text.
         """
+        doc_id = doc_id if doc_id is not None else str(uuid.uuid4())
         paragraphs = text.split(self.sep)
         paragraph_units = []
         start = 0
         for paragraph in paragraphs:
             end = start + len(paragraph)
             paragraph_units.append(FrameExtractionUnit(
-                doc_id=doc_id if doc_id is not None else str(uuid.uuid4()),
+                doc_id=doc_id,
                 start=start,
                 end=end,
                 text=paragraph
@@ -104,10 +108,11 @@ class SentenceUnitChunker(UnitChunker):
         text : str
             The document text.
         """
+        doc_id = doc_id if doc_id is not None else str(uuid.uuid4())
         sentences = []
         for start, end in self.PunktSentenceTokenizer().span_tokenize(text):
             sentences.append(FrameExtractionUnit(
-                doc_id=doc_id if doc_id is not None else str(uuid.uuid4()),
+                doc_id=doc_id,
                 start=start,
                 end=end,
                 text=text[start:end]
@@ -129,13 +134,14 @@ class TextLineUnitChunker(UnitChunker):
         text : str
             The document text.
         """
+        doc_id = doc_id if doc_id is not None else str(uuid.uuid4())
         lines = text.split('\n')
         line_units = []
         start = 0
         for line in lines:
             end = start + len(line)
             line_units.append(FrameExtractionUnit(
-                doc_id=doc_id if doc_id is not None else str(uuid.uuid4()),
+                doc_id=doc_id,
                 start=start,
                 end=end,
                 text=line
@@ -143,6 +149,100 @@ class TextLineUnitChunker(UnitChunker):
             start = end + 1 
         return line_units
     
+class LLMUnitChunker(UnitChunker):
+    def __init__(self, inference_engine:InferenceEngine, prompt_template:str=None, system_prompt:str=None):
+        """
+        This class prompt an LLM for document segmentation (e.g., sections, paragraphs).
+
+        Parameters:
+        ----------
+        inference_engine : InferenceEngine
+            the LLM inferencing engine object.
+        prompt_template : str
+            the prompt template that defines how to chunk the document. Must define a JSON schema with 
+            ```json
+            [
+                {
+                    "title": "<your title here>",
+                    "anchor_text": "<the anchor text of the chunk here>"
+                },
+                {
+                    "title": "<your title here>",
+                    "anchor_text": "<the anchor text of the chunk here>"
+                }
+            ]
+            ```
+        system_prompt : str, optional
+            The system prompt.
+        """
+        self.inference_engine = inference_engine
+
+        if prompt_template is None:
+            file_path = importlib.resources.files('llm_ie.asset.default_prompts').joinpath("LLMUnitChunker_user_prompt.txt")
+            with open(file_path, 'r', encoding="utf-8") as f:
+                self.prompt_template = f.read()
+        else:
+            self.prompt_template = prompt_template
+        
+        self.system_prompt = system_prompt
+
+    def chunk(self, text, doc_id=None) -> List[FrameExtractionUnit]:
+        """
+        Parameters:
+        -----------
+        text : str
+            the document text.
+        doc_id : str, optional
+            the document id.
+        """
+        doc_id = doc_id if doc_id is not None else str(uuid.uuid4())
+        user_prompt = apply_prompt_template(prompt_template=self.prompt_template, text_content=text)
+        messages = []
+        if self.system_prompt is not None:
+            messages.append({'role': 'system', 'content': self.system_prompt})
+        messages.append({'role': 'user', 'content': user_prompt})
+
+        gen_text = self.inference_engine.chat(messages=messages)
+
+        header_list = extract_json(gen_text=gen_text["response"])
+        units = []
+        start = 0
+        prev_end = 0
+        for header in header_list:
+            if "anchor_text" not in header:
+                Warning.warn(f"Missing anchor_text in header: {header}. Skipping this header.")
+                continue
+            if not isinstance(header["anchor_text"], str):
+                Warning.warn(f"Invalid anchor_text: {header['anchor_text']}. Skipping this header.")
+                continue
+
+            start = prev_end
+            # find the first instandce of the leading sentence in the rest of the text
+            end = text.find(header["anchor_text"], start)
+            # if not found, skip this header
+            if end == -1:
+                continue
+            # if start == end (empty text), skip this header
+            if start == end:
+                continue
+            # create a frame extraction unit
+            units.append(FrameExtractionUnit(
+                doc_id=doc_id,
+                start=start,
+                end=end,
+                text=text[start:end]
+            ))
+            prev_end = end
+        # add the last section
+        if prev_end < len(text):
+            units.append(FrameExtractionUnit(
+                doc_id=doc_id,
+                start=prev_end,
+                end=len(text),
+                text=text[prev_end:]
+            ))
+        return units
+
 
 class ContextChunker(abc.ABC):
     def __init__(self):

{llm_ie-1.2.3 → llm_ie-1.2.4}/src/llm_ie/engines.py

@@ -1060,6 +1060,50 @@ class VLLMInferenceEngine(OpenAICompatibleInferenceEngine):
         return {"reasoning": getattr(response.choices[0].message, "reasoning_content", ""),
                 "response": getattr(response.choices[0].message, "content", "")}
         
+class SGLangInferenceEngine(OpenAICompatibleInferenceEngine):
+    def __init__(self, model:str, api_key:str="", base_url:str="http://localhost:30000/v1", config:LLMConfig=None, **kwrs):
+        """
+        SGLang OpenAI compatible API inference engine.
+        https://docs.sglang.ai/basic_usage/openai_api.html
+
+        Parameters:
+        ----------
+        model_name : str
+            model name as shown in the vLLM server
+        api_key : str, Optional
+            the API key for the vLLM server.
+        base_url : str, Optional
+            the base url for the vLLM server. 
+        config : LLMConfig
+            the LLM configuration.
+        """
+        super().__init__(model, api_key, base_url, config, **kwrs)
+
+
+    def _format_response(self, response: Any) -> Dict[str, str]:
+        """
+        This method format the response from OpenAI API to a dict with keys "type" and "data".
+
+        Parameters:
+        ----------
+        response : Any
+            the response from OpenAI-compatible API. Could be a dict, generator, or object.
+        """
+        if isinstance(response, self.ChatCompletionChunk):
+            if hasattr(response.choices[0].delta, "reasoning_content") and getattr(response.choices[0].delta, "reasoning_content") is not None:
+                chunk_text = getattr(response.choices[0].delta, "reasoning_content", "")
+                if chunk_text is None:
+                    chunk_text = ""
+                return {"type": "reasoning", "data": chunk_text}
+            else:
+                chunk_text = getattr(response.choices[0].delta, "content", "")
+                if chunk_text is None:
+                    chunk_text = ""
+                return {"type": "response", "data": chunk_text}
+
+        return {"reasoning": getattr(response.choices[0].message, "reasoning_content", ""),
+                "response": getattr(response.choices[0].message, "content", "")}
+    
 
 class OpenRouterInferenceEngine(OpenAICompatibleInferenceEngine):
     def __init__(self, model:str, api_key:str=None, base_url:str="https://openrouter.ai/api/v1", config:LLMConfig=None, **kwrs):

{llm_ie-1.2.3 → llm_ie-1.2.4}/src/llm_ie/extractors.py

@@ -1,7 +1,5 @@
 import abc
 import re
-import json
-import json_repair
 import inspect
 import importlib.resources
 import warnings
@@ -10,6 +8,7 @@ import asyncio
 import nest_asyncio
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Set, List, Dict, Tuple, Union, Callable, Generator, Optional, AsyncGenerator
+from llm_ie.utils import extract_json, apply_prompt_template
 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
@@ -96,79 +95,8 @@ class Extractor:
         Returns : str
             a user prompt.
         """
-        pattern = re.compile(r'{{(.*?)}}')
-        if isinstance(text_content, str):
-            matches = pattern.findall(self.prompt_template)
-            if len(matches) != 1:
-                raise ValueError("When text_content is str, the prompt template must has exactly 1 placeholder {{<placeholder name>}}.")
-            text = re.sub(r'\\', r'\\\\', text_content)
-            prompt = pattern.sub(text, self.prompt_template)
+        return apply_prompt_template(self.prompt_template, text_content)
 
-        elif isinstance(text_content, dict):
-            # Check if all values are str
-            if not all([isinstance(v, str) for v in text_content.values()]):
-                raise ValueError("All values in text_content must be str.")
-            # Check if all keys are in the prompt template
-            placeholders = pattern.findall(self.prompt_template)
-            if len(placeholders) != len(text_content):
-                raise ValueError(f"Expect text_content ({len(text_content)}) and prompt template placeholder ({len(placeholders)}) to have equal size.")
-            if not all([k in placeholders for k, _ in text_content.items()]):
-                raise ValueError(f"All keys in text_content ({text_content.keys()}) must match placeholders in prompt template ({placeholders}).")
-
-            prompt = pattern.sub(lambda match: re.sub(r'\\', r'\\\\', text_content[match.group(1)]), self.prompt_template)
-
-        return prompt
-    
-    def _find_dict_strings(self, text: str) -> List[str]:
-        """
-        Extracts balanced JSON-like dictionaries from a string, even if nested.
-
-        Parameters:
-        -----------
-        text : str
-            the input text containing JSON-like structures.
-
-        Returns : List[str]
-            A list of valid JSON-like strings representing dictionaries.
-        """
-        open_brace = 0
-        start = -1
-        json_objects = []
-
-        for i, char in enumerate(text):
-            if char == '{':
-                if open_brace == 0:
-                    # start of a new JSON object
-                    start = i 
-                open_brace += 1
-            elif char == '}':
-                open_brace -= 1
-                if open_brace == 0 and start != -1:
-                    json_objects.append(text[start:i + 1])
-                    start = -1
-
-        return json_objects
-    
-    
-    def _extract_json(self, gen_text:str) -> List[Dict[str, str]]:
-        """ 
-        This method inputs a generated text and output a JSON of information tuples
-        """
-        out = []
-        dict_str_list = self._find_dict_strings(gen_text)
-        for dict_str in dict_str_list:
-            try:
-                dict_obj = json.loads(dict_str)
-                out.append(dict_obj)
-            except json.JSONDecodeError:
-                dict_obj = json_repair.repair_json(dict_str, skip_json_loads=True, return_objects=True)
-                if dict_obj:
-                    warnings.warn(f'JSONDecodeError detected, fixed with repair_json:\n{dict_str}', RuntimeWarning)
-                    out.append(dict_obj)
-                else:
-                    warnings.warn(f'JSONDecodeError could not be fixed:\n{dict_str}', RuntimeWarning)
-        return out
-    
 
 class FrameExtractor(Extractor):
     from nltk.tokenize import RegexpTokenizer
@@ -759,7 +687,7 @@ class DirectFrameExtractor(FrameExtractor):
             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):
+            for entity in extract_json(gen_text=unit.gen_text):
                 if ENTITY_KEY in entity:
                     entity_json.append(entity)
                 else:
@@ -963,7 +891,7 @@ class DirectFrameExtractor(FrameExtractor):
         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):
+            for entity in extract_json(gen_text=res.gen_text):
                 if ENTITY_KEY in entity:
                     entity_json.append(entity)
                 else:
@@ -1712,7 +1640,7 @@ class AttributeExtractor(Extractor):
                             messages_logger=messages_logger
                         )
 
-        attribute_list = self._extract_json(gen_text=gen_text["response"])
+        attribute_list = 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:
@@ -1822,7 +1750,7 @@ class AttributeExtractor(Extractor):
                 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, messages_logger=messages_logger)
-                attribute_list = self._extract_json(gen_text=gen_text["response"])
+                attribute_list = 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}
 
@@ -2075,7 +2003,7 @@ class BinaryRelationExtractor(RelationExtractor):
         return None
 
     def _post_process_result(self, gen_text: str, pair_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
-        rel_json = self._extract_json(gen_text)
+        rel_json = extract_json(gen_text)
         if len(rel_json) > 0 and "Relation" in rel_json[0]:
             rel = rel_json[0]["Relation"]
             if (isinstance(rel, bool) and rel) or (isinstance(rel, str) and rel.lower() == 'true'):
@@ -2141,7 +2069,7 @@ class MultiClassRelationExtractor(RelationExtractor):
         return None
 
     def _post_process_result(self, gen_text: str, pair_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
-        rel_json = self._extract_json(gen_text)
+        rel_json = extract_json(gen_text)
         pos_rel_types = pair_data['pos_rel_types']
         if len(rel_json) > 0 and "RelationType" in rel_json[0]:
             rel_type = rel_json[0]["RelationType"]

{llm_ie-1.2.3 → llm_ie-1.2.4}/src/llm_ie/prompt_editor.py

@@ -2,6 +2,7 @@ import sys
 import warnings
 from typing import List, Dict, Generator
 import importlib.resources
+from llm_ie.utils import apply_prompt_template
 from llm_ie.engines import InferenceEngine
 from llm_ie.extractors import FrameExtractor
 import re
@@ -45,30 +46,6 @@ class PromptEditor:
 
         # internal memory (history messages) for the `chat` method
         self.messages = []
-
-    def _apply_prompt_template(self, text_content:Dict[str,str], prompt_template:str) -> str:
-        """
-        This method applies text_content to prompt_template and returns a prompt.
-
-        Parameters
-        ----------
-        text_content : Dict[str,str]
-            the input text content to put in prompt template. 
-            all the keys must be included in the prompt template placeholder {{<placeholder name>}}.
-
-        Returns : str
-            a prompt.
-        """
-        pattern = re.compile(r'{{(.*?)}}')
-        placeholders = pattern.findall(prompt_template)
-        if len(placeholders) != len(text_content):
-            raise ValueError(f"Expect text_content ({len(text_content)}) and prompt template placeholder ({len(placeholders)}) to have equal size.")
-        if not all([k in placeholders for k, _ in text_content.items()]):
-            raise ValueError(f"All keys in text_content ({text_content.keys()}) must match placeholders in prompt template ({placeholders}).")
-
-        prompt = pattern.sub(lambda match: re.sub(r'\\', r'\\\\', text_content[match.group(1)]), prompt_template)
-
-        return prompt
     
 
     def rewrite(self, draft:str) -> str:
@@ -80,8 +57,8 @@ class PromptEditor:
         with open(file_path, 'r') as f:
             rewrite_prompt_template = f.read()
 
-        prompt = self._apply_prompt_template(text_content={"draft": draft, "prompt_guideline": self.prompt_guide}, 
-                                             prompt_template=rewrite_prompt_template)
+        prompt = apply_prompt_template(prompt_template=rewrite_prompt_template,
+                                       text_content={"draft": draft, "prompt_guideline": self.prompt_guide})
         messages = [{"role": "system", "content": self.system_prompt},
                     {"role": "user", "content": prompt}]
         res = self.inference_engine.chat(messages, verbose=True)
@@ -96,8 +73,8 @@ class PromptEditor:
         with open(file_path, 'r') as f:
             comment_prompt_template = f.read()
 
-        prompt = self._apply_prompt_template(text_content={"draft": draft, "prompt_guideline": self.prompt_guide}, 
-                                             prompt_template=comment_prompt_template)
+        prompt = apply_prompt_template(prompt_template=comment_prompt_template,
+                                       text_content={"draft": draft, "prompt_guideline": self.prompt_guide})
         messages = [{"role": "system", "content": self.system_prompt},
                     {"role": "user", "content": prompt}]
         res = self.inference_engine.chat(messages, verbose=True)
@@ -254,8 +231,8 @@ class PromptEditor:
             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)
+            guideline = apply_prompt_template(prompt_template=chat_prompt_template, 
+                                              text_content={"prompt_guideline": self.prompt_guide})
 
             self.messages = [{"role": "system", "content": self.system_prompt + guideline}]
 
@@ -288,8 +265,8 @@ class PromptEditor:
         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)
+        guideline = apply_prompt_template(prompt_template=chat_prompt_template, 
+                                          text_content={"prompt_guideline": self.prompt_guide})
 
         messages = [{"role": "system", "content": self.system_prompt + guideline}] + messages
 

llm_ie-1.2.4/src/llm_ie/utils.py

@@ -0,0 +1,95 @@
+from typing import List, Dict, Union
+import re
+import json
+import warnings
+import json_repair
+
+def _find_dict_strings(text: str) -> List[str]:
+    """
+    Extracts balanced JSON-like dictionaries from a string, even if nested.
+
+    Parameters:
+    -----------
+    text : str
+        the input text containing JSON-like structures.
+
+    Returns : List[str]
+        A list of valid JSON-like strings representing dictionaries.
+    """
+    open_brace = 0
+    start = -1
+    json_objects = []
+
+    for i, char in enumerate(text):
+        if char == '{':
+            if open_brace == 0:
+                # start of a new JSON object
+                start = i 
+            open_brace += 1
+        elif char == '}':
+            open_brace -= 1
+            if open_brace == 0 and start != -1:
+                json_objects.append(text[start:i + 1])
+                start = -1
+
+    return json_objects
+    
+    
+def extract_json(gen_text:str) -> List[Dict[str, str]]:
+    """ 
+    This method inputs a generated text and output a JSON of information tuples
+    """
+    out = []
+    dict_str_list = _find_dict_strings(gen_text)
+    for dict_str in dict_str_list:
+        try:
+            dict_obj = json.loads(dict_str)
+            out.append(dict_obj)
+        except json.JSONDecodeError:
+            dict_obj = json_repair.repair_json(dict_str, skip_json_loads=True, return_objects=True)
+            if dict_obj:
+                warnings.warn(f'JSONDecodeError detected, fixed with repair_json:\n{dict_str}', RuntimeWarning)
+                out.append(dict_obj)
+            else:
+                warnings.warn(f'JSONDecodeError could not be fixed:\n{dict_str}', RuntimeWarning)
+    return out
+
+
+def apply_prompt_template(prompt_template:str, text_content:Union[str, Dict[str,str]]) -> str:
+    """
+    This method applies text_content to prompt_template and returns a prompt.
+
+    Parameters:
+    ----------
+    prompt_template : str
+        the prompt template with placeholders {{<placeholder name>}}.
+    text_content : Union[str, Dict[str,str]]
+        the input text content 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>}}. All values must be str.
+
+    Returns : str
+        a user prompt.
+    """
+    pattern = re.compile(r'{{(.*?)}}')
+    if isinstance(text_content, str):
+        matches = pattern.findall(prompt_template)
+        if len(matches) != 1:
+            raise ValueError("When text_content is str, the prompt template must has exactly 1 placeholder {{<placeholder name>}}.")
+        text = re.sub(r'\\', r'\\\\', text_content)
+        prompt = pattern.sub(text, prompt_template)
+
+    elif isinstance(text_content, dict):
+        # Check if all values are str
+        if not all([isinstance(v, str) for v in text_content.values()]):
+            raise ValueError("All values in text_content must be str.")
+        # Check if all keys are in the prompt template
+        placeholders = pattern.findall(prompt_template)
+        if len(placeholders) != len(text_content):
+            raise ValueError(f"Expect text_content ({len(text_content)}) and prompt template placeholder ({len(placeholders)}) to have equal size.")
+        if not all([k in placeholders for k, _ in text_content.items()]):
+            raise ValueError(f"All keys in text_content ({text_content.keys()}) must match placeholders in prompt template ({placeholders}).")
+
+        prompt = pattern.sub(lambda match: re.sub(r'\\', r'\\\\', text_content[match.group(1)]), prompt_template)
+
+    return prompt