perf: defer litellm and mcp out of import swarms (2.2s -> 0.37s)

swarms/__init__.py

@@ -14,3 +14,18 @@
 from swarms.telemetry import *  # noqa: E402, F403
 from swarms.tools import *  # noqa: E402, F403
 from swarms.utils import *  # noqa: E402, F403
+
+
+# Reachable as `swarms.MCPManager` as before, but resolved on access so the mcp
+# package is not imported until something actually uses it.
+_LAZY_TOOLS = frozenset({"MCPManager", "MCPFileTokenStorage"})
+
+
+def __getattr__(name: str):
+    if name in _LAZY_TOOLS:
+        from swarms import tools
+
+        return getattr(tools, name)
+    raise AttributeError(
+        f"module {__name__!r} has no attribute {name!r}"
+    )

swarms/agents/context_compressor.py

@@ -1,15 +1,5 @@
-"""
-Context window compression for autonomous (``max_loops="auto"``) agent runs.
-
-When an agent's conversation history approaches its context window limit,
-this module summarizes the accumulated history and replaces it with a dense
-summary, preserving the system prompt and keeping the agent within its token
-budget for an unbounded run.
-"""
-
 from typing import Any, Optional
 
-from litellm import completion
 from loguru import logger
 
 from swarms.utils.litellm_tokenizer import count_tokens
@@ -91,6 +81,8 @@
         return self.usage_ratio(agent) >= self.threshold
 
     def _summarize(self, agent: Any, history: str) -> str:
+        from litellm import completion
+
         model = self.summarizer_model or getattr(
             agent, "model_name", None
         )

swarms/agents/llm_manager.py

@@ -31,16 +31,6 @@
 import traceback
 from typing import Any, Callable, List, Optional, Union
 
-from litellm.exceptions import (
-    AuthenticationError,
-    BadRequestError,
-    InternalServerError,
-)
-from litellm.utils import (
-    supports_function_calling,
-    supports_parallel_function_calling,
-    supports_vision,
-)
 from loguru import logger
 
 from swarms.schemas.agent_errors import AgentLLMInitializationError
@@ -323,6 +313,12 @@
         Args:
             img (str, optional): Image input to check vision support for.
         """
+        from litellm.utils import (
+            supports_function_calling,
+            supports_parallel_function_calling,
+            supports_vision,
+        )
+
         agent = self.agent
 
         # Only check vision support if an image is provided
@@ -517,6 +513,12 @@
             AgentLLMError, BadRequestError, InternalServerError,
             AuthenticationError, Exception: re-raised for upstream handling.
         """
+        from litellm.exceptions import (
+            AuthenticationError,
+            BadRequestError,
+            InternalServerError,
+        )
+
         agent = self.agent
 
         # Filter out is_last from kwargs if present

swarms/structs/agent.py

@@ -18,17 +18,6 @@
 
 import toml
 import yaml
-from litellm import model_list
-from litellm.exceptions import (
-    AuthenticationError,
-    BadRequestError,
-    InternalServerError,
-)
-from litellm.utils import (
-    get_max_tokens,
-    get_model_info,
-    supports_function_calling,
-)
 from loguru import logger
 from pydantic import BaseModel
 
@@ -93,15 +82,14 @@
 from swarms.tools.base_tool import BaseTool
 from swarms.tools.handoffs_tool import handoff_task
 from swarms.tools.handoffs_tool_schema import get_handoff_tool_schema
-from swarms.tools.mcp_manager import MCPManager
 from swarms.tools.py_func_to_openai_func_str import (
     convert_multiple_functions_to_openai_function_schema,
 )
 from swarms.utils.file_processing import create_file_in_folder
 from swarms.utils.formatter import formatter
 from swarms.utils.generate_id import generate_id
 from swarms.utils.generate_keys import generate_api_key
-from swarms.utils.get_reasoning_efforts import get_reasoning_efforts
+from swarms.utils.get_reasoning_efforts import REASONING_EFFORTS
 from swarms.utils.history_output_formatter import (
     history_output_formatter,
 )
@@ -391,7 +379,7 @@
         reasoning_prompt_on: bool = True,
         dynamic_context_window: bool = True,
         show_tool_execution_output: bool = True,
-        reasoning_effort: Literal[get_reasoning_efforts()] = None,
+        reasoning_effort: Literal[REASONING_EFFORTS] = None,
         thinking_tokens: int = 1024,
         think_tool: bool = False,
         dynamic_tools: bool = True,
@@ -517,6 +505,10 @@
         self.publish_to_marketplace = publish_to_marketplace
         self.marketplace_prompt_id = marketplace_prompt_id
 
+        # Imported here rather than at module scope: mcp_manager pulls in the
+        # mcp package, which is a few hundred ms of `import swarms`.
+        from swarms.tools.mcp_manager import MCPManager
+
         self.mcp_manager = MCPManager(
             mcp_url=self.mcp_url,
             mcp_urls=self.mcp_urls,
@@ -1341,6 +1333,12 @@
             ...     streaming_callback=on_token
             ... )
         """
+        from litellm.exceptions import (
+            AuthenticationError,
+            BadRequestError,
+            InternalServerError,
+        )
+
         try:
             self.check_if_no_prompt_then_autogenerate(task)
 
@@ -2460,6 +2458,8 @@
               which may not correspond to available context for input (e.g., 32768 for gpt-4.1 output, but
               over a million for certain input windows).
         """
+        from litellm.utils import get_model_info
+
         try:
             return (
                 get_model_info(self.model_name).get(
@@ -2481,6 +2481,8 @@
         Returns:
             int: The maximum number of output tokens for the model. Returns 16000 if undetermined.
         """
+        from litellm.utils import get_model_info
+
         # get_model_info raises for unmapped ids, which would otherwise take down __init__ for custom models.
         try:
             return (
@@ -2493,6 +2495,11 @@
             return 16000
 
     def reliability_check(self):
+        from litellm import model_list
+        from litellm.utils import (
+            get_max_tokens,
+            supports_function_calling,
+        )
 
         if self.system_prompt is None:
             logger.warning(
@@ -3274,6 +3281,11 @@
             ...     img_base64 = base64.b64encode(f.read()).decode("utf-8")
             >>> agent.run("Describe this image", img=img_base64)
         """
+        from litellm.exceptions import (
+            AuthenticationError,
+            BadRequestError,
+            InternalServerError,
+        )
 
         # If no task is provided, prompt for one only in interactive mode.
         # Outside interactive mode, fail fast instead of blocking on stdin.

swarms/structs/agent_loader.py

@@ -14,7 +14,6 @@
 )
 
 import yaml
-from litellm import model_list
 from pydantic import BaseModel
 from tqdm import tqdm
 
@@ -99,6 +98,8 @@
         config: Union[BaseModel, Dict[str, Any]],
     ) -> AgentConfigDict:
         """Validate a config supplied as a pydantic model or a plain dict."""
+        from litellm import model_list
+
         try:
             if isinstance(config, BaseModel):
                 config = config.model_dump()

swarms/structs/agent_router.py

@@ -1,7 +1,6 @@
 import math
 from typing import Any, Callable, List, Optional, Union
 
-from litellm import embedding
 from tenacity import retry, stop_after_attempt, wait_exponential
 
 from swarms.structs.ma_blocks import find_agent_by_name
@@ -56,6 +55,8 @@
         Returns:
             List[float]: The embedding vector as a list of floats.
         """
+        from litellm import embedding
+
         try:
             # Prepare parameters for the embedding call
             params = {"model": self.embedding_model, "input": [text]}

swarms/structs/tree_swarm.py

@@ -3,7 +3,6 @@
 from datetime import datetime, timezone
 from typing import Any, List, Optional
 
-from litellm import embedding
 from pydantic import BaseModel, Field
 
 from swarms.structs.execution_utils import batched_run
@@ -194,14 +193,15 @@
         Returns:
             List[float]: Embedding vector
         """
+        from litellm import embedding
+
         try:
             response = embedding(
                 model=self.embedding_model_name, input=[text]
             )
             if self.verbose:
                 logger.info(f"Embedding type: {type(response)}")
-            # print(response)
-            # Handle different response structures from litellm
+
             if hasattr(response, "data") and response.data:
                 if hasattr(response.data[0], "embedding"):
                     return response.data[0].embedding
@@ -225,8 +225,6 @@
         except Exception as e:
             if self.verbose:
                 logger.error(f"Error getting embedding: {e}")
-            # Return a zero vector as fallback
-            return [0.0] * 1536  # Default OpenAI embedding dimension
 
     def calculate_distance(self, other_agent: "TreeAgent") -> float:
         """

swarms/tools/__init__.py

@@ -1,8 +1,4 @@
 from swarms.tools.base_tool import BaseTool
-from swarms.tools.mcp_manager import (
-    MCPFileTokenStorage,
-    MCPManager,
-)
 from swarms.tools.py_func_to_openai_func_str import (
     Function,
     ToolFunction,
@@ -39,6 +35,24 @@
     "BaseTool",
     "ToolStorage",
     "tool_registry",
-    "MCPManager",
-    "MCPFileTokenStorage",
 ]
+
+
+# Lazy-load MCP-related classes to avoid unnecessary import overhead.
+_MCP_EXPORTS = frozenset({"MCPManager", "MCPFileTokenStorage"})
+
+
+def __getattr__(name: str):
+    if name in _MCP_EXPORTS:
+        from swarms.tools import mcp_manager
+
+        return getattr(mcp_manager, name)
+    raise AttributeError(
+        f"module {__name__!r} has no attribute {name!r}"
+    )
+
+
+def __dir__():
+    # Kept out of __all__ so a star-import does not resolve them, but still
+    # advertised here for dir() and autocomplete.
+    return sorted(set(__all__) | _MCP_EXPORTS)

swarms/utils/litellm_tokenizer.py

@@ -1,4 +1,3 @@
-from litellm import encode, model_list
 from loguru import logger
 from typing import Optional
 from functools import lru_cache
@@ -33,6 +32,8 @@
     # Set fallback encoder
     fallback_model = default_encoder or DEFAULT_MODEL
 
+    from litellm import encode
+
     # First attempt with the requested model
     try:
         tokens = encode(model=model, text=text)
@@ -74,6 +75,8 @@
 def get_supported_models() -> list:
     """Get list of supported models from litellm."""
     try:
+        from litellm import model_list
+
         return model_list
     except Exception as e:
         logger.warning(f"Could not retrieve model list: {e}")