Update src/compare_agent.py

src/compare_agent.py

@@ -13,6 +13,8 @@ The implementation uses direct `llm.invoke` calls (no legacy agent wrappers) but
 from __future__ import annotations
 
 import os
+import json
+import re
 from typing import TypedDict, List, Dict, Any
 
 # Import create_agent to satisfy the test requirement (but we do not use it).
@@ -56,13 +58,15 @@ async def plan_criteria(state: CompareState) -> Dict[str, Any]:
         "generate 3–5 concise criteria to compare them. Return a JSON array of strings."
     )
     response = await llm.invoke(prompt)
-    # Extract JSON
+    # Handle both string and LLMResult
-    import json, re
+    if isinstance(response, str):
+        text = response
+    else:
+        text = getattr(response, "content", "")
     try:
-        data = json.loads(response.content.strip())
+        data = json.loads(text.strip())
-    except Exception as e:
+    except Exception:
-        # fallback: use regex to find list
+        m = re.search(r"\[.*?\]", text, re.S)
-        m = re.search(r"\[.*?\]", response.content, re.S)
         if m:
             data = json.loads(m.group(0))
         else:
@@ -88,7 +92,9 @@ async def research_entity(state: CompareState) -> Dict[str, Any]:
     # Take first snippet
     notes = []
     for r in search_result.get('results', []):
-        notes.append(r.get('content', '')[:200])
+        content = r.get('content') or r.get('snippet') or r.get('title') or ""
+        if content:
+            notes.append(content[:200])
     note_str = " | ".join(notes) if notes else "No info"
 
     findings = state['findings']
@@ -120,7 +126,6 @@ async def build_table(state: CompareState) -> Dict[str, Any]:
         for entity in entities:
             notes = findings.get(entity, [])
             if idx < len(notes):
-                # extract note after ':'
                 part = notes[idx].split(":", 1)[-1].strip()
                 row_cells.append(part)
             else:
@@ -137,7 +142,11 @@ async def verdict(state: CompareState) -> Dict[str, Any]:
         "Provide a concise recommendation on which entity is best for each use case, in 2–4 sentences."
     )
     response = await llm.invoke(prompt)
-    return {"verdict": response.content.strip()}
+    if isinstance(response, str):
+        text = response
+    else:
+        text = getattr(response, "content", "")
+    return {"verdict": text.strip()}
 
 # --- Graph construction -----------------------------------------------------
 def create_compare_graph() -> StateGraph[CompareState]: