Fix high severity issues across monitoring, agents, memory, and MoE modules

Author: harshdesaii

multimind/agents/agent.py

@@ -2,6 +2,7 @@
 Base Agent class for Multimind SDK.
 """
 
+import re
 from typing import List, Dict, Any, Optional
 from multimind.models.base import BaseLLM
 from multimind.agents.memory import AgentMemory
@@ -38,8 +39,11 @@ async def run(self, task: str, **kwargs) -> Dict[str, Any]:
     async def _process_task(self, task: str, **kwargs) -> Dict[str, Any]:
         """Process a task using available tools and the model."""
         # Try to match a tool by name
+        task_lower = task.lower()
         for tool in self.tools:
-            if tool.name.lower() in task.lower():
+            tool_name = tool.name.lower().strip()
+            # Match on whole tokens only (e.g., `calc` won't match `calculate`).
+            if re.search(rf"\b{re.escape(tool_name)}\b", task_lower):
                 try:
                     # Extract parameters for the tool from kwargs
                     params = {k: v for k, v in kwargs.items() if k in tool.get_parameters().get("required", [])}

multimind/agents/agent_loader.py

@@ -33,8 +33,22 @@ def load_agent(
     ) -> Agent:
         """Load an agent from a configuration file."""
         # Load config
-        with open(config_path, 'r') as f:
-            config = json.load(f)
+        try:
+            with open(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
+        except json.JSONDecodeError as e:
+            raise ValueError(
+                f"Invalid JSON in agent config file: {config_path}. {e}"
+            ) from e
+        except OSError as e:
+            raise RuntimeError(
+                f"Failed to read agent config file: {config_path}. {e}"
+            ) from e
+
+        if not isinstance(config, dict):
+            raise ValueError(f"Agent config must be a JSON object: {config_path}")
 
         # Validate config
         required_keys = {"model", "system_prompt"}

multimind/agents/agent_registry.py

@@ -22,7 +22,39 @@ def set_fallback(self, name: str, fallback_name: str):
     def get_agent(self, name: str) -> Optional[Callable]:
         return self.agents.get(name)
 
-    def run_agent(self, name: str, *args, session_id: Optional[str] = None, **kwargs):
+    def run_agent(
+        self,
+        name: str,
+        *args,
+        session_id: Optional[str] = None,
+        _visited: Optional[set] = None,
+        _depth: int = 0,
+        _max_depth: int = 10,
+        **kwargs,
+    ):
+        """
+        Run an agent by name with fallback support.
+
+        Cycle-protected: if fallbacks point back to an already tried agent (A->B->A),
+        we stop to avoid infinite recursion.
+        """
+        if _visited is None:
+            _visited = set()
+
+        if name in _visited:
+            self.logger.error(
+                f"Fallback recursion detected for agent '{name}'. Aborting. Visited={_visited}"
+            )
+            return None
+
+        if _depth >= _max_depth:
+            self.logger.error(
+                f"Max fallback depth reached while running agent '{name}'. Aborting."
+            )
+            return None
+
+        _visited.add(name)
+
         agent = self.get_agent(name)
         if not agent:
             self.logger.warning(f"Agent {name} not found.")
@@ -33,14 +65,29 @@ def run_agent(self, name: str, *args, session_id: Optional[str] = None, **kwargs
             result = agent(*args, state=state, **kwargs)
             # Optionally update state
             if session_id is not None:
-                self.state_memory[session_id] = result.get("state", state) if isinstance(result, dict) else state
+                self.state_memory[session_id] = (
+                    result.get("state", state) if isinstance(result, dict) else state
+                )
             return result
         except Exception as e:
             self.logger.error(f"Agent {name} failed: {e}")
             fallback = self.fallbacks.get(name)
             if fallback:
+                if fallback in _visited:
+                    self.logger.error(
+                        f"Fallback cycle detected: '{name}' -> '{fallback}'. Aborting."
+                    )
+                    return None
                 self.logger.info(f"Retrying with fallback agent: {fallback}")
-                return self.run_agent(fallback, *args, session_id=session_id, **kwargs)
+                return self.run_agent(
+                    fallback,
+                    *args,
+                    session_id=session_id,
+                    _visited=_visited,
+                    _depth=_depth + 1,
+                    _max_depth=_max_depth,
+                    **kwargs,
+                )
             return None
 
     def get_state(self, session_id: str):

multimind/agents/memory.py

@@ -11,21 +11,29 @@ class AgentMemory:
     def __init__(self, max_history: int = 100):
         self.max_history = max_history
         self.tasks: List[str] = []
+        self.task_timestamps: List[datetime] = []
         self.responses: List[Dict[str, Any]] = []
+        self.response_timestamps: List[datetime] = []
         self.state: Dict[str, Any] = {}
         self.created_at = datetime.now()
 
     def add_task(self, task: str) -> None:
         """Add a task to memory."""
         self.tasks.append(task)
+        self.task_timestamps.append(datetime.now())
         if len(self.tasks) > self.max_history:
             self.tasks.pop(0)
+            if self.task_timestamps:
+                self.task_timestamps.pop(0)
 
     def add_response(self, response: Dict[str, Any]) -> None:
         """Add a response to memory."""
         self.responses.append(response)
+        self.response_timestamps.append(datetime.now())
         if len(self.responses) > self.max_history:
             self.responses.pop(0)
+            if self.response_timestamps:
+                self.response_timestamps.pop(0)
 
     def update_state(self, key: str, value: Any) -> None:
         """Update agent state."""
@@ -41,16 +49,31 @@ def get_history(self, n: Optional[int] = None) -> List[Dict[str, Any]]:
             n = self.max_history
 
         history = []
-        for task, response in zip(self.tasks[-n:], self.responses[-n:]):
+        # Use timestamps recorded at insertion time (task/response) instead of `datetime.now()`
+        recent_tasks = self.tasks[-n:]
+        recent_task_timestamps = self.task_timestamps[-n:]
+        recent_responses = self.responses[-n:]
+        recent_response_timestamps = self.response_timestamps[-n:]
+
+        for task, response, task_ts, resp_ts in zip(
+            recent_tasks,
+            recent_responses,
+            recent_task_timestamps,
+            recent_response_timestamps,
+        ):
             history.append({
                 "task": task,
                 "response": response,
-                "timestamp": datetime.now().isoformat()
+                # Prefer response timestamp because it reflects when the completion arrived.
+                "timestamp": resp_ts.isoformat() if isinstance(resp_ts, datetime) else datetime.now().isoformat(),
+                "task_timestamp": task_ts.isoformat() if isinstance(task_ts, datetime) else None,
             })
         return history
 
     def clear(self) -> None:
         """Clear all memory."""
         self.tasks.clear()
+        self.task_timestamps.clear()
         self.responses.clear()
+        self.response_timestamps.clear()
         self.state.clear()

multimind/api/multi_model_api.py

@@ -2,14 +2,59 @@
 FastAPI-based API interface for the MultiModelWrapper.
 """
 
+import logging
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, Field
 from typing import List, Dict, Optional, Union
 import asyncio
+import json
+from typing import Tuple, Any
+from functools import lru_cache
 from ..models.factory import ModelFactory
 from ..models.multi_model import MultiModelWrapper
 
 app = FastAPI(title="Multi-Model API")
+logger = logging.getLogger(__name__)
+
+# Reuse a single factory across requests to avoid re-loading env / re-allocating caches.
+_MODEL_FACTORY = ModelFactory()
+
+# Cache MultiModelWrapper instances by request parameters.
+# Note: wrapper init can be expensive because it initializes provider model instances.
+_WRAPPER_CACHE: Dict[Tuple[str, Tuple[str, ...], str], MultiModelWrapper] = {}
+_WRAPPER_LOCKS: Dict[Tuple[str, Tuple[str, ...], str], asyncio.Lock] = {}
+
+
+def _weights_key(model_weights: Optional[Dict[str, float]]) -> str:
+    # Stable string key for dict weights (used for caching).
+    return json.dumps(model_weights or {}, sort_keys=True, default=str)
+
+
+async def _get_multi_model(
+    *,
+    primary_model: str,
+    fallback_models: List[str],
+    model_weights: Optional[Dict[str, float]],
+) -> MultiModelWrapper:
+    fallback_tuple = tuple(fallback_models or [])
+    key = (primary_model, fallback_tuple, _weights_key(model_weights))
+
+    if key in _WRAPPER_CACHE:
+        return _WRAPPER_CACHE[key]
+
+    lock = _WRAPPER_LOCKS.setdefault(key, asyncio.Lock())
+    async with lock:
+        if key in _WRAPPER_CACHE:
+            return _WRAPPER_CACHE[key]
+
+        wrapper = MultiModelWrapper(
+            model_factory=_MODEL_FACTORY,
+            primary_model=primary_model,
+            fallback_models=list(fallback_tuple),
+            model_weights=model_weights,
+        )
+        _WRAPPER_CACHE[key] = wrapper
+        return wrapper
 
 class GenerateRequest(BaseModel):
     prompt: str
@@ -37,12 +82,10 @@ class EmbeddingsRequest(BaseModel):
 async def generate(request: GenerateRequest):
     """Generate text using the multi-model wrapper."""
     try:
-        factory = ModelFactory()
-        multi_model = MultiModelWrapper(
-            model_factory=factory,
+        multi_model = await _get_multi_model(
             primary_model=request.primary_model,
             fallback_models=request.fallback_models,
-            model_weights=request.model_weights
+            model_weights=request.model_weights,
         )
 
         response = await multi_model.generate(
@@ -52,18 +95,17 @@ async def generate(request: GenerateRequest):
         )
         return {"response": response}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.exception("Unhandled error in /generate")
+        raise HTTPException(status_code=500, detail="Internal server error")
 
 @app.post("/chat")
 async def chat(request: ChatRequest):
     """Generate chat completion using the multi-model wrapper."""
     try:
-        factory = ModelFactory()
-        multi_model = MultiModelWrapper(
-            model_factory=factory,
+        multi_model = await _get_multi_model(
             primary_model=request.primary_model,
             fallback_models=request.fallback_models,
-            model_weights=request.model_weights
+            model_weights=request.model_weights,
         )
 
         response = await multi_model.chat(
@@ -73,24 +115,24 @@ async def chat(request: ChatRequest):
         )
         return {"response": response}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.exception("Unhandled error in /chat")
+        raise HTTPException(status_code=500, detail="Internal server error")
 
 @app.post("/embeddings")
 async def embeddings(request: EmbeddingsRequest):
     """Generate embeddings using the multi-model wrapper."""
     try:
-        factory = ModelFactory()
-        multi_model = MultiModelWrapper(
-            model_factory=factory,
+        multi_model = await _get_multi_model(
             primary_model=request.primary_model,
             fallback_models=request.fallback_models,
-            model_weights=request.model_weights
+            model_weights=request.model_weights,
         )
 
         embeddings = await multi_model.embeddings(request.text)
         return {"embeddings": embeddings}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.exception("Unhandled error in /embeddings")
+        raise HTTPException(status_code=500, detail="Internal server error")
 
 @app.get("/health")
 async def health_check():

multimind/api/unified_api.py

@@ -6,7 +6,7 @@
 from pydantic import BaseModel, Field
 from typing import Dict, List, Any, Optional, Union
 import asyncio
-import traceback
+import logging
 import os
 import base64
 import io
@@ -15,8 +15,13 @@
 from ..models.moe import Expert
 from ..types import UnifiedRequest, UnifiedResponse, ModalityInput
 
+logger = logging.getLogger(__name__)
+
 app = FastAPI(title="Unified Multi-Modal API")
 
+# Reuse a single factory across requests to avoid re-creating model caches.
+_MODEL_FACTORY = ModelFactory()
+
 
 class _TextExpertAdapter(Expert):
     """Expert wrapper around a model instance for text."""
@@ -113,7 +118,6 @@ async def process(self, input_data: Any) -> Any:
 def _build_experts(modalities: List[str], router: Any) -> Dict[str, Expert]:
     """Build available experts for modality MoE."""
     experts: Dict[str, Expert] = {}
-    factory = ModelFactory()
 
     for modality in modalities:
         model = None
@@ -122,9 +126,9 @@ def _build_experts(modalities: List[str], router: Any) -> Dict[str, Expert]:
             model = next(iter(model_map.values()), None)
 
         if model is None and modality == "text":
-            for provider in factory.available_models():
+            for provider in _MODEL_FACTORY.available_models():
                 try:
-                    model = factory.get_model(provider)
+                    model = _MODEL_FACTORY.get_model(provider)
                     break
                 except Exception:
                     continue
@@ -280,10 +284,8 @@ async def process_request(request: UnifiedRequest):
         # Preserve intended HTTP status codes (e.g., 400 for invalid input).
         raise
     except Exception as e:
-        raise HTTPException(
-            status_code=500,
-            detail=f"Error processing request: {str(e)}\nTraceback:\n{traceback.format_exc()}"
-        )
+        logger.exception("Error processing request")
+        raise HTTPException(status_code=500, detail="Internal server error")
 
 @app.get("/v1/models")
 async def list_models():