From bbc858bbfad84b3d6338287156aad5e422d00c60 Mon Sep 17 00:00:00 2001
From: Muhammad Talha Imran <92talhaimran@gmail.com>
Date: Mon, 9 Jun 2025 19:51:50 +0200
Subject: [PATCH] Handle extraneous text when parsing JSON

---
 Agents/AI Service/agency.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/Agents/AI Service/agency.py b/Agents/AI Service/agency.py
index bae0896..8c338ee 100644
--- a/Agents/AI Service/agency.py	
+++ b/Agents/AI Service/agency.py	
@@ -157,6 +157,16 @@ def __init__(self, role: str, prompt: ChatPromptTemplate):
         logger.info(f"Loading {role} model: {model_name}")
         self.llm = ChatOllama(model=model_name)
 
+    @staticmethod
+    def _parse_json(text: str) -> Any:
+        """Extract and parse the first JSON object found in text."""
+        start = text.find("{")
+        if start == -1:
+            raise json.JSONDecodeError("No JSON object found", text, 0)
+        decoder = json.JSONDecoder()
+        obj, _ = decoder.raw_decode(text[start:])
+        return obj
+
     async def run(self, project_json: str, context: Dict[str, Any], timeout: Optional[float] = 60.0) -> Any:
         """
         Invoke the LLM with formatted messages and return parsed JSON.
@@ -171,11 +181,12 @@ async def run(self, project_json: str, context: Dict[str, Any], timeout: Optiona
         ).to_messages()
 
         try:
-            raw: str = await asyncio.wait_for(
-                asyncio.to_thread(lambda: self.llm.invoke(messages)),
+            message = await asyncio.wait_for(
+                asyncio.to_thread(self.llm.invoke, messages),
                 timeout,
             )
-            parsed = json.loads(raw)
+            raw = getattr(message, "content", message)
+            parsed = self._parse_json(raw)
             logger.info(f"{self.role} responded successfully")
             return parsed
         except asyncio.TimeoutError: