diff --git a/src/compare_agent.py b/src/compare_agent.py index 5a0aa0a..1552cc0 100644 --- a/src/compare_agent.py +++ b/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 - import json, re + # Handle both string and LLMResult + if isinstance(response, str): + text = response + else: + text = getattr(response, "content", "") try: - data = json.loads(response.content.strip()) - except Exception as e: - # fallback: use regex to find list - m = re.search(r"\[.*?\]", response.content, re.S) + data = json.loads(text.strip()) + except Exception: + m = re.search(r"\[.*?\]", text, 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]: