heymenn commited on
Commit
ead7603
·
verified ·
1 Parent(s): 7e33637

Update kig_core/planner.py

Browse files
Files changed (1) hide show
  1. kig_core/planner.py +4 -5
kig_core/planner.py CHANGED
@@ -13,7 +13,7 @@ from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
13
  from .config import settings
14
  from .schemas import PlannerState, KeyIssue, GraphConfig # Import schemas
15
  from .prompts import get_initial_planner_prompt, KEY_ISSUE_STRUCTURING_PROMPT
16
- from .llm_interface import get_llm
17
  from .graph_operations import (
18
  generate_cypher_auto, generate_cypher_guided,
19
  retrieve_documents, evaluate_documents
@@ -36,7 +36,7 @@ def start_planning(state: PlannerState) -> Dict[str, Any]:
36
  chain = initial_prompt | llm | StrOutputParser()
37
 
38
  try:
39
- plan_text = chain.invoke({}) # Prompt already includes query
40
  logger.debug(f"Raw plan text: {plan_text}")
41
 
42
  # Extract plan steps (simple regex, might need refinement)
@@ -151,7 +151,7 @@ def generate_structured_issues(state: PlannerState) -> Dict[str, Any]:
151
  chain = prompt | issue_llm | output_parser
152
 
153
  try:
154
- structured_issues_obj = chain.invoke({
155
  "user_query": user_query,
156
  "context": full_context
157
  })
@@ -169,12 +169,11 @@ def generate_structured_issues(state: PlannerState) -> Dict[str, Any]:
169
  }
170
  except Exception as e:
171
  logger.error(f"Failed to generate or parse structured key issues: {e}", exc_info=True)
172
- time.sleep(60)
173
  # Attempt to get raw output for debugging if possible
174
  raw_output = "Could not retrieve raw output."
175
  try:
176
  raw_chain = prompt | issue_llm | StrOutputParser()
177
- raw_output = raw_chain.invoke({"user_query": user_query, "context": full_context})
178
  logger.debug(f"Raw output from failed JSON parsing:\n{raw_output}")
179
  except Exception as raw_e:
180
  logger.error(f"Could not even get raw output: {raw_e}")
 
13
  from .config import settings
14
  from .schemas import PlannerState, KeyIssue, GraphConfig # Import schemas
15
  from .prompts import get_initial_planner_prompt, KEY_ISSUE_STRUCTURING_PROMPT
16
+ from .llm_interface import get_llm, invoke_llm
17
  from .graph_operations import (
18
  generate_cypher_auto, generate_cypher_guided,
19
  retrieve_documents, evaluate_documents
 
36
  chain = initial_prompt | llm | StrOutputParser()
37
 
38
  try:
39
+ plan_text = invoke_llm(chain,{}) # Prompt already includes query
40
  logger.debug(f"Raw plan text: {plan_text}")
41
 
42
  # Extract plan steps (simple regex, might need refinement)
 
151
  chain = prompt | issue_llm | output_parser
152
 
153
  try:
154
+ structured_issues_obj = invoke_llm(chain,{
155
  "user_query": user_query,
156
  "context": full_context
157
  })
 
169
  }
170
  except Exception as e:
171
  logger.error(f"Failed to generate or parse structured key issues: {e}", exc_info=True)
 
172
  # Attempt to get raw output for debugging if possible
173
  raw_output = "Could not retrieve raw output."
174
  try:
175
  raw_chain = prompt | issue_llm | StrOutputParser()
176
+ raw_output = llm_invoke(raw_chain,{"user_query": user_query, "context": full_context})
177
  logger.debug(f"Raw output from failed JSON parsing:\n{raw_output}")
178
  except Exception as raw_e:
179
  logger.error(f"Could not even get raw output: {raw_e}")