Update kig_core/planner.py

kig_core/planner.py

@@ -4,6 +4,8 @@ from typing import List, Dict, Any
 from langgraph.graph import StateGraph, END
 from langgraph.checkpoint.memory import MemorySaver # Or SqliteSaver etc.
 
+from pydantic import BaseModel, Field
+
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
 
@@ -107,6 +109,14 @@ def execute_plan_step(state: PlannerState) -> Dict[str, Any]:
         "step_outputs": current_step_outputs
     }
 
+class KeyIssue(BaseModel):
+    # define your fields here
+    id: int
+    description: str
+
+class KeyIssueList(BaseModel):
+    key_issues: List[KeyIssue] = Field(description="List of key issues")
+
 
 def generate_structured_issues(state: PlannerState) -> Dict[str, Any]:
     """Generates the final structured Key Issues based on all gathered context."""
@@ -130,8 +140,9 @@ def generate_structured_issues(state: PlannerState) -> Dict[str, Any]:
     # --- Call LLM for Structured Output ---
     issue_llm = get_llm(settings.main_llm_model)
     # Use PydanticOutputParser for robust parsing
-    output_parser = JsonOutputParser(pydantic_object=List[KeyIssue])
+    output_parser = JsonOutputParser(pydantic_object=KeyIssueList)
 
+    
     prompt = KEY_ISSUE_STRUCTURING_PROMPT.partial(
         schema=output_parser.get_format_instructions(), # Inject schema instructions if needed by prompt
     )
@@ -139,11 +150,11 @@ def generate_structured_issues(state: PlannerState) -> Dict[str, Any]:
     chain = prompt | issue_llm | output_parser
 
     try:
-        structured_issues = chain.invoke({
+        structured_issues_obj = chain.invoke({
             "user_query": user_query,
             "context": full_context
         })
-
+        structured_issues = structured_issues_obj.key_issues
         # Ensure IDs are sequential if the LLM didn't assign them correctly
         for i, issue in enumerate(structured_issues):
             issue.id = i + 1
@@ -152,8 +163,8 @@ def generate_structured_issues(state: PlannerState) -> Dict[str, Any]:
         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)], # Final summary message
-            "error": None # Clear any previous errors
+            "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)