Fix: Enhance robustness and clarity in kg-solver components (OpenSPG#548)

This commit addresses several issues identified during a code analysis, focusing on improving the robustness, error handling, and clarity of various components within the KAG solver.

The following changes have been made:

1.  **KAGIterativePlanner**:
    *   Corrected `is_static()` method to return `False`, aligning with its iterative behavior.

2.  **KAGRetrievedResponse**:
    *   Removed a misleading note about an f-string formatting error from the `to_string()` method's docstring, as the error was not present in the code.

3.  **KAGStaticPlanner**:
    *   Improved `finish_judger` error handling: If the LLM call to judge the answer fails, it now logs a warning and returns `False` (treating the answer as potentially bad) instead of defaulting to `True`.

4.  **ChunkRetrievedExecutor**:
    *   Clarified schema name: Changed the `name` field in its schema dictionary from "Retriever" to "ChunkRetriever" to better differentiate it from other retriever executors like `KagHybridExecutor`.

5.  **PyBasedMathExecutor**:
    *   Added a configurable timeout (defaulting to 5 seconds) to the `subprocess.run()` call within the `run_py_code` function. This prevents indefinite hangs from long-running or stuck Python scripts generated by the LLM. Includes handling for `subprocess.TimeoutExpired`.

6.  **DefaultStaticPlanningPrompt**:
    *   Enhanced `parse_response` method: Implemented more robust JSON decoding and structural validation for the LLM-generated DAG plan. It now raises more descriptive `ValueError` exceptions, including details of the malformed data, when `KeyError` or `TypeError` occurs during task creation from the DAG, aiding in debugging.

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: thesteganos

kag/solver/executor/math/py_based_math_executor.py

@@ -29,21 +29,36 @@
 
 
 def run_py_code(python_code: str, **kwargs):
+    # Default timeout in seconds
+    default_timeout = 5 
+    # Allow timeout to be passed via kwargs if needed for more flexibility
+    timeout_duration = kwargs.get("timeout", default_timeout)
+
     with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as temp_file:
         temp_file.write(python_code.encode("utf-8"))
         temp_file_path = temp_file.name
 
+    stdout_value = None
+    stderr_value = None
+
     try:
         python_executable = sys.executable
         result = subprocess.run(
-            [python_executable, temp_file_path], capture_output=True, text=True
+            [python_executable, temp_file_path],
+            capture_output=True,
+            text=True,
+            timeout=timeout_duration # Added timeout
         )
+        stdout_value = result.stdout
+        stderr_value = result.stderr
+    except subprocess.TimeoutExpired as e:
+        stderr_value = f"Code execution timed out after {timeout_duration} seconds: {e}"
+    except Exception as e: # Catch other potential errors during subprocess.run
+        stderr_value = f"An unexpected error occurred during code execution: {e}"
     finally:
         os.remove(temp_file_path)
 
-    stdout_value = result.stdout
-    stderr_value = result.stderr
-    if len(stderr_value) > 0:
+    if stderr_value: # If there's any error (timeout or other execution error)
         return None, stderr_value, python_code
     return stdout_value, None, python_code
 

kag/solver/executor/retriever/local_knowledge_base/chunk_retrieved_executor.py

@@ -87,7 +87,7 @@ def schema(self) -> dict:
             dict: Schema definition in OpenAI Function format
         """
         return {
-            "name": "Retriever",
+            "name": "ChunkRetriever", # Changed from "Retriever"
             "description": "Retrieve relevant knowledge from the local knowledge base.",
             "parameters": {
                 "query": {

kag/solver/executor/retriever/local_knowledge_base/kag_retriever/kag_hybrid_executor.py

@@ -133,10 +133,6 @@ def to_string(self) -> str:
 
         Returns:
             str: Formatted string containing task description and sub-question results
-
-        Note:
-            Contains formatting error: "task: f{self.retrieved_task}"
-            should be corrected to "task: {self.retrieved_task}"
         """
         refer_docs = self.to_reference_list()
         for doc in refer_docs:

kag/solver/planner/kag_iterative_planner.py

@@ -80,6 +80,9 @@ def invoke(self, query, **kwargs) -> List[Task]:
             **kwargs,
         )
 
+    def is_static(self):
+        return False
+
     async def ainvoke(self, query, **kwargs) -> List[Task]:
         """Asynchronously generates task plan using LLM.
 

kag/solver/planner/kag_static_planner.py

@@ -9,6 +9,7 @@
 # Unless required by applicable law or agreed to in writing, software distributed under the License
 # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 # or implied.
+import logging
 import re
 from typing import List
 
@@ -86,11 +87,10 @@ async def finish_judger(self, query: str, answer: str):
                 return False
             return True
         except Exception as e:
-            print(f"Failed to run finish_judger, info: {e}")
-            import traceback
-
-            traceback.print_exc()
-            return True
+            # import logging # Make sure logging is imported if not already at the top of the file
+            logger = logging.getLogger(__name__) # Get a logger instance
+            logger.warning(f"LLM call failed in finish_judger for query '{query}'. Error: {e}", exc_info=True)
+            return False # Treat as potentially bad answer
 
     async def query_rewrite(self, task: Task, **kwargs):
         """Performs asynchronous query rewriting using LLM and context.

kag/solver/prompt/static_planning_prompt.py

@@ -144,9 +144,38 @@ def template_variables(self) -> List[str]:
 
     def parse_response(self, response: str, **kwargs):
         if isinstance(response, str):
-            response = json.loads(response)
-        if not isinstance(response, dict):
-            raise ValueError(f"response should be a dict, but got {type(response)}")
-        if "output" in response:
-            response = response["output"]
-        return Task.create_tasks_from_dag(response)
+            try:
+                response_json = json.loads(response)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Failed to decode LLM response as JSON: {e}. Response: {response}")
+        elif isinstance(response, dict):
+            response_json = response # If it's already a dict (e.g. from direct LLM client parsing)
+        else:
+            raise ValueError(f"LLM response is not a JSON string or a dictionary. Got type: {type(response)}. Response: {response}")
+
+        if not isinstance(response_json, dict):
+            # This case might be redundant if json.loads already ensures a dict or list, 
+            # but good for safety if the initial response could be a non-dict JSON type.
+            raise ValueError(f"Parsed LLM response should be a dict, but got {type(response_json)}. Response: {response_json}")
+
+        # Handle if the LLM wraps the DAG in an "output" key, as per original logic
+        actual_dag_data = response_json.get("output", response_json)
+        
+        if not isinstance(actual_dag_data, dict):
+            raise ValueError(f"The core plan data (after handling potential 'output' key) is not a dictionary. Got type: {type(actual_dag_data)}. Data: {actual_dag_data}")
+
+        try:
+            return Task.create_tasks_from_dag(actual_dag_data)
+        except (KeyError, TypeError) as e:
+            error_message = (
+                f"LLM response for static planning was malformed. Error: {e}. "
+                f"Each task in the DAG dictionary must define 'executor', 'arguments', and 'dependent_task_ids'. "
+                f"Problematic DAG data: {actual_dag_data}"
+            )
+            raise ValueError(error_message)
+        except Exception as e: # Catch any other unexpected errors from create_tasks_from_dag
+            error_message = (
+                f"An unexpected error occurred while creating tasks from DAG. Error: {e}. "
+                f"Problematic DAG data: {actual_dag_data}"
+            )
+            raise ValueError(error_message)