add better error handling for json decoding

landing-ai · Mar 7, 2024 · 27d0a5a · 27d0a5a
1 parent b059ea1
commit 27d0a5a
Showing 1 changed file with 21 additions and 3 deletions.
diff --git a/vision_agent/lmm/lmm.py b/vision_agent/lmm/lmm.py
@@ -110,7 +110,13 @@ def generate_classifier(self, prompt: str) -> ImageTool:
  {"role": "user", "content": prompt},
  ],
  )
- prompt = json.loads(response.choices[0].message.content)["prompt"]
+
+ try:
+ prompt = json.loads(cast(str, response.choices[0].message.content))["prompt"]
+ except json.JSONDecodeError:
+ _LOGGER.error(f"Failed to decode response: {response.choices[0].message.content}")
+ raise ValueError("Failed to decode response")
+
  return CLIP(prompt)
 
  def generate_detector(self, prompt: str) -> ImageTool:
@@ -123,7 +129,13 @@ def generate_detector(self, prompt: str) -> ImageTool:
  {"role": "user", "content": prompt},
  ],
  )
- prompt = json.loads(response.choices[0].message.content)["prompt"]
+
+ try:
+ prompt = json.loads(cast(str, response.choices[0].message.content))["prompt"]
+ except json.JSONDecodeError:
+ _LOGGER.error(f"Failed to decode response: {response.choices[0].message.content}")
+ raise ValueError("Failed to decode response")
+
  return GroundingDINO(prompt)
 
  def generate_segmentor(self, prompt: str) -> ImageTool:
@@ -136,7 +148,13 @@ def generate_segmentor(self, prompt: str) -> ImageTool:
  {"role": "user", "content": prompt},
  ],
  )
- prompt = json.loads(response.choices[0].message.content)["prompt"]
+
+ try:
+ prompt = json.loads(cast(str, response.choices[0].message.content))["prompt"]
+ except json.JSONDecodeError:
+ _LOGGER.error(f"Failed to decode response: {response.choices[0].message.content}")
+ raise ValueError("Failed to decode response")
+
  return GroundingSAM(prompt)