Spaces:
Running
Running
Update kig_core/planner.py
Browse files- 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
|
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
|
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
|
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}")
|