Improve parsing of JSON for tool calls

app.py

@@ -30,10 +30,10 @@ print(f"Using embedding checkpoints from {embedding_ckpt_dir}")
 
 # Download and extract data if data directory is not present
 if not os.path.isdir(db_dir):
-    print("Downloading data ... ", end = "")
+    print("Downloading data ... ", end="")
     download_data()
     print("done!")
-    print("Extracting data ... ", end = "")
+    print("Extracting data ... ", end="")
     extract_data()
     print("done!")
 
@@ -44,6 +44,7 @@ search_type = "hybrid"
 # https://www.gradio.app/guides/state-in-blocks
 graph_instances = {"local": {}, "remote": {}}
 
+
 def cleanup_graph(request: gr.Request):
     if request.session_hash in graph_instances["local"]:
         del graph_instances["local"][request.session_hash]

data.py

@@ -4,6 +4,7 @@ import shutil
 import boto3
 import os
 
+
 def download_file_from_bucket(bucket_name, s3_key, output_file):
     """Download file from S3 bucket"""
 
@@ -17,6 +18,7 @@ def download_file_from_bucket(bucket_name, s3_key, output_file):
     bucket = s3_resource.Bucket(bucket_name)
     bucket.download_file(Key=s3_key, Filename=output_file)
 
+
 def download_dropbox_file(shared_url, output_file):
     """Download file from Dropbox"""
 
@@ -35,9 +37,8 @@ def download_dropbox_file(shared_url, output_file):
                 file.write(chunk)
         print(f"File downloaded successfully as '{output_file}'")
     else:
-        print(
-            f"Failed to download file. HTTP Status Code: {response.status_code}"
-        )
+        print(f"Failed to download file. HTTP Status Code: {response.status_code}")
+
 
 def download_data():
     """Download the email database"""
@@ -50,6 +51,7 @@ def download_data():
         # output_filename = "db.zip"
         # download_dropbox_file(shared_link, output_filename)
 
+
 def extract_data():
     """Extract the db.zip file"""
 

main.py

@@ -41,7 +41,7 @@ model_id = os.getenv("MODEL_ID")
 if model_id is None:
     # model_id = "HuggingFaceTB/SmolLM3-3B"
     # model_id = "google/gemma-3-12b-it"
-    model_id = "Qwen/Qwen3-8B"
+    model_id = "Qwen/Qwen3-14B"
 
 # Suppress these messages:
 # INFO:httpx:HTTP Request: POST https://api.openai.com/v1/embeddings "HTTP/1.1 200 OK"

mods/tool_calling_llm.py

@@ -181,16 +181,20 @@ class ToolCallingLLM(BaseChatModel, ABC):
         # print("post_think")
         # print(post_think)
 
+        # Remove trailing comma (if there is one)
+        post_think = post_think.rstrip(",")
         # Parse output for JSON (support multiple objects separated by commas)
         try:
-            # Works for one or more JSON objects not enclosed in "[]"
-            parsed_json_results = json.loads(f"[{post_think}]")
+            # Works for one JSON object, or multiple JSON objects enclosed in "[]"
+            parsed_json_results = json.loads(f"{post_think}")
+            if not isinstance(parsed_json_results, list):
+                parsed_json_results = [parsed_json_results]
         except:
             try:
-                # Works for one or more JSON objects already enclosed in "[]"
-                parsed_json_results = json.loads(f"{post_think}")
+                # Works for multiple JSON objects not enclosed in "[]"
+                parsed_json_results = json.loads(f"[{post_think}]")
             except json.JSONDecodeError:
-                # Return entire response if JSON wasn't parsed (or is missing)
+                # Return entire response if JSON wasn't parsed or is missing
                 return AIMessage(content=response_message.content)
 
         # print("parsed_json_results")