fix: better json load and adding repair to loading prompts

Author: qfeuilla

src/minillmlib/core/chat_node.py

@@ -13,18 +13,21 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import httpx
+import json_repair
 import requests
 from anthropic import Anthropic, AsyncAnthropic
 from anthropic.types.message import Message
 from mistralai import Mistral
 from openai import AsyncOpenAI, OpenAI
 from openai.types.chat import ChatCompletion
 
-from ..models.generator_info import (HUGGINGFACE_ACTIVATED, GeneratorCompletionParameters, GeneratorInfo,
-                                     pretty_messages, torch)
+from ..models.generator_info import (HUGGINGFACE_ACTIVATED,
+                                     GeneratorCompletionParameters,
+                                     GeneratorInfo, pretty_messages, torch)
 from ..utils.json_utils import extract_json_from_completion, to_dict
 from ..utils.logging_utils import get_logger
-from ..utils.message_utils import (AudioData, ImageData, VideoData, NodeCompletionParameters,
+from ..utils.message_utils import (AudioData, ImageData,
+                                   NodeCompletionParameters, VideoData,
                                    base64_to_wav, format_prompt, get_payload,
                                    hf_process_messages,
                                    merge_contiguous_messages,
@@ -1345,6 +1348,7 @@ async def collapse_thread(self,
     def from_thread(cls,
         path: str | List[str] | None = None,
         messages: List[Dict[str, str]] | None = None,
+        fix_json: bool = False,
     ) -> ChatNode:
         """Load a thread from a JSON file or multiple JSON files, or directly from messages.
 
@@ -1401,7 +1405,7 @@ def from_thread(cls,
             return parent
 
         # Single path logic
-        data = json.load(open(path, "r", encoding="utf-8"))
+        data = (json_repair if fix_json else json).load(open(path, "r", encoding="utf-8"))
         if "required_kwargs" not in data:
             data["required_kwargs"] = {}
 

src/minillmlib/utils/json_utils.py

@@ -27,4 +27,4 @@ def to_dict(item: Any) -> Dict[Any, Any]:
 
 def extract_json_from_completion(completion: str) -> str:
     """Extract and parse JSON from completion string."""
-    return json.dumps(json.loads(json_repair.repair_json(completion)))
+    return json.dumps(json_repair.loads(completion))