Merge pull request #80 from multimindlab/medium_issues_solved

Fix routing reliability, kwargs safety, async session reuse, and secu…

Author: niralbhalodia

multimind/agents/agent_loader.py

@@ -16,6 +16,19 @@ class AgentLoader:
     def __init__(self, model_registry: Optional[Dict[str, BaseLLM]] = None):
         self.model_registry = model_registry or {}
         self.tool_registry: Dict[str, BaseTool] = {}
+        self._base_dir = Path.cwd().resolve()
+
+    def _resolve_safe_path(self, path: str) -> Path:
+        """Resolve and validate path stays under the loader base directory."""
+        candidate = Path(path).expanduser()
+        resolved = candidate.resolve()
+        try:
+            resolved.relative_to(self._base_dir)
+        except ValueError as e:
+            raise ValueError(
+                f"Path traversal detected. '{path}' escapes base directory '{self._base_dir}'."
+            ) from e
+        return resolved
 
     def register_model(self, name: str, model: BaseLLM) -> None:
         """Register a model for use in agents."""
@@ -32,23 +45,27 @@ def load_agent(
         tools: Optional[List[BaseTool]] = None
     ) -> Agent:
         """Load an agent from a configuration file."""
+        safe_config_path = self._resolve_safe_path(config_path)
+        if safe_config_path.suffix.lower() != ".json":
+            raise ValueError(f"Agent config must be a JSON file: {config_path}")
+
         # Load config
         try:
-            with open(config_path, "r", encoding="utf-8") as f:
+            with open(safe_config_path, "r", encoding="utf-8") as f:
                 config = json.load(f)
         except FileNotFoundError as e:
-            raise FileNotFoundError(f"Agent config file not found: {config_path}") from e
+            raise FileNotFoundError(f"Agent config file not found: {safe_config_path}") from e
         except json.JSONDecodeError as e:
             raise ValueError(
-                f"Invalid JSON in agent config file: {config_path}. {e}"
+                f"Invalid JSON in agent config file: {safe_config_path}. {e}"
             ) from e
         except OSError as e:
             raise RuntimeError(
-                f"Failed to read agent config file: {config_path}. {e}"
+                f"Failed to read agent config file: {safe_config_path}. {e}"
             ) from e
 
         if not isinstance(config, dict):
-            raise ValueError(f"Agent config must be a JSON object: {config_path}")
+            raise ValueError(f"Agent config must be a JSON object: {safe_config_path}")
 
         # Validate config
         required_keys = {"model", "system_prompt"}
@@ -93,7 +110,9 @@ def load_agents_from_dir(
     ) -> Dict[str, Agent]:
         """Load multiple agents from a directory of config files."""
         agents = {}
-        config_dir = Path(dir_path)
+        config_dir = self._resolve_safe_path(dir_path)
+        if not config_dir.is_dir():
+            raise FileNotFoundError(f"Agent config directory not found: {config_dir}")
 
         for config_file in config_dir.glob("*.json"):
             agent_name = config_file.stem

multimind/agents/memory.py

@@ -65,7 +65,7 @@ def get_history(self, n: Optional[int] = None) -> List[Dict[str, Any]]:
                 "task": task,
                 "response": response,
                 # Prefer response timestamp because it reflects when the completion arrived.
-                "timestamp": resp_ts.isoformat() if isinstance(resp_ts, datetime) else datetime.now().isoformat(),
+                "timestamp": resp_ts.isoformat() if isinstance(resp_ts, datetime) else None,
                 "task_timestamp": task_ts.isoformat() if isinstance(task_ts, datetime) else None,
             })
         return history

multimind/agents/prompt_correction.py

@@ -30,15 +30,17 @@ def monitor(self, prompt: str, output: str, trace: Dict = None) -> str:
                 self.logger.warning(f"Detected issue in output: {output}")
                 for hook in self.error_hooks:
                     hook(prompt, Exception("Detected hallucination"), trace)
-                # Apply correction hooks
+                # Apply correction hooks to the *output* and return corrected output.
+                # (Correction hooks are expected to take a string and trace, and return a string.)
+                corrected_output = output
                 for hook in self.correction_hooks:
-                    prompt = hook(prompt, trace)
-                self.logger.info(f"Corrected prompt: {prompt}")
-                return prompt
-            return prompt
+                    corrected_output = hook(corrected_output, trace)
+                self.logger.info(f"Corrected output: {corrected_output}")
+                return corrected_output
+            return output
         except Exception as e:
             self.logger.error(f"Prompt correction failed: {e}")
-            return prompt
+            return output
 
     def update_adapter(self, adapter_key: str, new_adapter_path: str):
         """
@@ -61,6 +63,6 @@ def adapter_updater(adapter_key, new_path):
     pcl.add_correction_hook(simple_correction)
     pcl.add_adapter_update_hook(adapter_updater)
     # Simulate monitoring
-    new_prompt = pcl.monitor("What is the capital of France?", "[error] hallucination detected", {"step": 1})
-    print("New prompt after correction:", new_prompt)
+    corrected_output = pcl.monitor("What is the capital of France?", "[error] hallucination detected", {"step": 1})
+    print("Corrected output after correction:", corrected_output)
     pcl.update_adapter("user123", "lora_adapter_v2") 

multimind/agents/react_toolchain.py

@@ -1,5 +1,10 @@
 from typing import List, Callable, Any, Dict
 
+
+class ReasoningChainExecutionError(RuntimeError):
+    """Raised when a reasoning chain step fails."""
+
+
 class ReasoningStep:
     """
     Represents a single step in a reasoning/toolchain. Can be a model call, tool call, or custom function.
@@ -21,13 +26,29 @@ def __init__(self, steps: List[ReasoningStep]):
         self.hooks = []  # List of callables: hook(step, input, output)
     def add_hook(self, hook: Callable[[ReasoningStep, Any, Any], None]):
         self.hooks.append(hook)
+
     def run(self, input_data: Any, context: Dict = None):
         context = context or {}
         data = input_data
         for step in self.steps:
-            output = step(data, context=context)
+            try:
+                output = step(data, context=context)
+            except Exception as e:
+                context["last_failed_step"] = step.name
+                context.setdefault("errors", []).append(
+                    {"step": step.name, "error": str(e)}
+                )
+                raise ReasoningChainExecutionError(
+                    f"Reasoning chain failed at step '{step.name}': {e}"
+                ) from e
+
             for hook in self.hooks:
-                hook(step, data, output)
+                try:
+                    hook(step, data, output)
+                except Exception as e:
+                    context.setdefault("hook_errors", []).append(
+                        {"step": step.name, "hook": getattr(hook, "__name__", str(hook)), "error": str(e)}
+                    )
             data = output
         return data
 

multimind/agents/tools/calculator.py

@@ -11,6 +11,10 @@
 class CalculatorTool(BaseTool):
     """A tool for performing mathematical calculations."""
 
+    MAX_EXPRESSION_LENGTH = 512
+    MAX_AST_NODES = 128
+    MAX_AST_DEPTH = 24
+
     def __init__(self):
         super().__init__(
             name="calculator",
@@ -53,12 +57,25 @@ def get_parameters(self) -> Dict[str, Any]:
 
     def _evaluate(self, expression: str) -> Union[int, float]:
         """Safely evaluate a mathematical expression."""
+        if len(expression) > self.MAX_EXPRESSION_LENGTH:
+            raise ValueError(
+                f"Expression too long (max {self.MAX_EXPRESSION_LENGTH} characters)"
+            )
+
+        def _ast_depth(node: ast.AST) -> int:
+            children = list(ast.iter_child_nodes(node))
+            if not children:
+                return 1
+            return 1 + max(_ast_depth(child) for child in children)
+
         def _eval(node) -> Union[int, float]:
-            if isinstance(node, ast.Num):
-                val = node.n
+            if isinstance(node, ast.Constant):
+                val = node.value
                 if isinstance(val, complex):
                     raise ValueError("Complex numbers are not supported")
-                return float(val) if isinstance(val, Real) else val
+                if not isinstance(val, Real):
+                    raise TypeError(f"Unsupported constant type: {type(val)}")
+                return float(val)
             elif isinstance(node, ast.BinOp):
                 return self.operators[type(node.op)](
                     _eval(node.left),
@@ -70,6 +87,12 @@ def _eval(node) -> Union[int, float]:
                 raise TypeError(f"Unsupported operation: {type(node)}")
 
         tree = ast.parse(expression, mode='eval')
+        node_count = sum(1 for _ in ast.walk(tree))
+        if node_count > self.MAX_AST_NODES:
+            raise ValueError(f"Expression too complex (max {self.MAX_AST_NODES} AST nodes)")
+        if _ast_depth(tree) > self.MAX_AST_DEPTH:
+            raise ValueError(f"Expression too deeply nested (max depth {self.MAX_AST_DEPTH})")
+
         result = _eval(tree.body)
         if isinstance(result, complex):
             raise ValueError("Complex numbers are not supported")