Parsing LLM response with JSON parsing Library

.pre-commit-config.yaml

@@ -30,7 +30,7 @@ repos:
     hooks:
       - id: bandit
         args: ["-c", "pyproject.toml"]
-        additional_dependencies: ["bandit[toml]"]
+        additional_dependencies: ["bandit[toml]", "pbr"]
 
   # Markdown linting
   - repo: https://github.com/igorshubovych/markdownlint-cli

pyproject.toml

@@ -20,7 +20,9 @@ dependencies = [
     "loguru>=0.7.3,<0.8.0",
     "mcp>=1.0.0",
     "google-generativeai>=0.8.0,<1.0.0",
-    "python-dotenv>=1.0.0,<2.0.0"
+    "python-dotenv>=1.0.0,<2.0.0",
+    "dirtyjson (>=1.0.8,<2.0.0)",
+    "pre-commit (>=4.3.0,<5.0.0)"
 ]
 
 [project.scripts]
@@ -132,3 +134,10 @@ inherit = false
 match = "(?!test_).*\\.py"
 match-dir = "(?!tests|migrations|__pycache__).*"
 ignore = "D100,D104"  # Ignore missing docstrings in modules and packages for now
+
+[dependency-groups]
+dev = [
+    "pre-commit (>=4.3.0,<5.0.0)",
+    "bandit (>=1.8.6,<2.0.0)",
+    "pbr (>=7.0.1,<8.0.0)"
+]

src/shardguard/core/planning.py

@@ -5,6 +5,8 @@
 import re
 from typing import Protocol
 
+import dirtyjson
+
 from .llm_providers import LLMProviderFactory
 from .mcp_integration import MCPClient
 
@@ -71,21 +73,23 @@ async def get_available_tools_description(self) -> str:
 
     def _extract_json_from_response(self, response: str) -> str:
         """Extract JSON from LLM response that might contain extra text."""
-        # Try to find JSON block enclosed in curly braces
-        matches = re.findall(r"\{.*\}", response, re.DOTALL)
-
-        if matches:
-            # Return the longest JSON-like match
-            json_candidate = max(matches, key=len)
-            # Validate that it's actually valid JSON
-            try:
-                json.loads(json_candidate)
-                return json_candidate
-            except json.JSONDecodeError:
-                pass
-
-        # If no valid JSON found, return the original response
-        return response
+        try:
+            obj = dirtyjson.loads(response, search_for_first_object=True)
+            return json.dumps(obj)
+        except Exception:
+            # Try to find a JSON block enclosed in curly braces as a fallback
+            matches = re.findall(r"\{.*\}", response, re.DOTALL)
+            if matches:
+                # Return the longest JSON-like match
+                json_candidate = max(matches, key=len)
+                # Validate that it's actually valid JSON
+                try:
+                    json.loads(json_candidate)
+                    return json_candidate
+                except json.JSONDecodeError:
+                    pass
+            # If no valid JSON is found, return the original response
+            return response
 
     def _create_fallback_response(self, prompt: str, error: str) -> str:
         """Create a fallback response when plan generation fails."""