Update kig_core/planner.py

kig_core/planner.py

@@ -13,7 +13,7 @@ from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
 from .config import settings
 from .schemas import PlannerState, KeyIssue, GraphConfig # Import schemas
 from .prompts import get_initial_planner_prompt, KEY_ISSUE_STRUCTURING_PROMPT
-from .llm_interface import get_llm
+from .llm_interface import get_llm, invoke_llm
 from .graph_operations import (
     generate_cypher_auto, generate_cypher_guided,
     retrieve_documents, evaluate_documents
@@ -36,7 +36,7 @@ def start_planning(state: PlannerState) -> Dict[str, Any]:
     chain = initial_prompt | llm | StrOutputParser()
 
     try:
-        plan_text = chain.invoke({}) # Prompt already includes query
+        plan_text = invoke_llm(chain,{}) # Prompt already includes query
         logger.debug(f"Raw plan text: {plan_text}")
 
         # Extract plan steps (simple regex, might need refinement)
@@ -151,7 +151,7 @@ def generate_structured_issues(state: PlannerState) -> Dict[str, Any]:
     chain = prompt | issue_llm | output_parser
 
     try:
-        structured_issues_obj = chain.invoke({
+        structured_issues_obj = invoke_llm(chain,{
             "user_query": user_query,
             "context": full_context
         })
@@ -169,12 +169,11 @@ def generate_structured_issues(state: PlannerState) -> Dict[str, Any]:
         }
     except Exception as e:
         logger.error(f"Failed to generate or parse structured key issues: {e}", exc_info=True)
-        time.sleep(60)
         # 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 = raw_chain.invoke({"user_query": user_query, "context": full_context})
+             raw_output = llm_invoke(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}")