Update kig_core/planner.py

kig_core/planner.py

@@ -156,37 +156,45 @@ def generate_structured_issues(state: PlannerState) -> Dict[str, Any]:
 
     chain = prompt | issue_llm | output_parser
 
+try:
+    structured_issues_obj = invoke_llm(chain, {
+        "user_query": user_query,
+        "context": full_context
+    })
+    print(f"structured_issues_obj => type : {type(structured_issues_obj)}, value : {structured_issues_obj}")
+
+    # If the output is a dict with a key 'key_issues', extract it
+    if isinstance(structured_issues_obj, dict) and 'key_issues' in structured_issues_obj:
+        issues_data = structured_issues_obj['key_issues']
+    else:
+        issues_data = structured_issues_obj  # Assume it's already a list of dicts
+
+    # Always convert to KeyIssueInvoke objects
+    key_issues_list = [KeyIssueInvoke(**issue_dict) for issue_dict in issues_data]
+
+    # Ensure IDs are sequential if the LLM didn't assign them correctly
+    for i, issue in enumerate(key_issues_list):
+        issue.id = i + 1
+
+    logger.info(f"Successfully generated {len(key_issues_list)} structured key issues.")
+    final_message = f"Generated {len(key_issues_list)} Key Issues based on the query '{user_query}'."
+    return {
+        "key_issues": key_issues_list,
+        "messages": [AIMessage(content=final_message)],
+        "error": None
+    }
+except Exception as e:
+    logger.error(f"Failed to generate or parse structured key issues: {e}", exc_info=True)
+    # Attempt to get raw output for debugging if possible
+    raw_output = "Could not retrieve raw output."
     try:
-        structured_issues_obj = invoke_llm(chain,{
-            "user_query": user_query,
-            "context": full_context
-        })
-        print(f"structured_issues_obj => type : {type(structured_issues_obj)}, value : {structured_issues_obj}")
-        structured_issues = structured_issues_obj #.key_issues
-        key_issues_list = [KeyIssueInvoke(**issue_dict) for issue_dict in structured_issues_obj]
-        # Ensure IDs are sequential if the LLM didn't assign them correctly
-        for i, issue in enumerate(structured_issues):
-            issue.id = i + 1
-
-        logger.info(f"Successfully generated {len(structured_issues)} structured key issues.")
-        final_message = f"Generated {len(structured_issues)} Key Issues based on the query '{user_query}'."
-        return {
-            "key_issues": structured_issues,
-            "messages": [AIMessage(content=final_message)],
-            "error": None
-        }
-    except Exception as e:
-        logger.error(f"Failed to generate or parse structured key issues: {e}", exc_info=True)
-        # Attempt to get raw output for debugging if possible
-        raw_output = "Could not retrieve raw output."
-        try:
-             raw_chain = prompt | issue_llm | StrOutputParser()
-             raw_output = invoke_llm(raw_chain,{"user_query": user_query, "context": full_context})
-             logger.debug(f"Raw output from failed JSON parsing:\n{raw_output}")
-        except Exception as raw_e:
-             logger.error(f"Could not even get raw output: {raw_e}")
-
-        return {"error": f"Failed to generate structured key issues: {e}. Raw output hint: {raw_output[:500]}..."}
+         raw_chain = prompt | issue_llm | StrOutputParser()
+         raw_output = invoke_llm(raw_chain, {"user_query": user_query, "context": full_context})
+         logger.debug(f"Raw output from failed JSON parsing:\n{raw_output}")
+    except Exception as raw_e:
+         logger.error(f"Could not even get raw output: {raw_e}")
+
+    return {"error": f"Failed to generate structured key issues: {e}. Raw output hint: {raw_output[:500]}..."}
 
 
 # --- Conditional Edges ---