swarms 7.6.4__tar.gz → 7.6.5__tar.gz

{swarms-7.6.4 → swarms-7.6.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: swarms
-Version: 7.6.4
+Version: 7.6.5
 Summary: Swarms - TGSC
 Home-page: https://github.com/kyegomez/swarms
 License: MIT

{swarms-7.6.4 → swarms-7.6.5}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "7.6.4"
+version = "7.6.5"
 description = "Swarms - TGSC"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]

{swarms-7.6.4 → swarms-7.6.5}/swarms/structs/__init__.py

@@ -2,7 +2,7 @@ from swarms.structs.agent import Agent
 from swarms.structs.agent_builder import AgentsBuilder
 from swarms.structs.agents_available import showcase_available_agents
 from swarms.structs.async_workflow import AsyncWorkflow
-from swarms.structs.auto_swarm import AutoSwarm, AutoSwarmRouter
+from experimental.auto_swarm import AutoSwarm, AutoSwarmRouter
 from swarms.structs.base_structure import BaseStructure
 from swarms.structs.base_swarm import BaseSwarm
 from swarms.structs.base_workflow import BaseWorkflow
@@ -85,8 +85,6 @@ from swarms.structs.swarming_architectures import (
 __all__ = [
     "Agent",
     "AsyncWorkflow",
-    "AutoSwarm",
-    "AutoSwarmRouter",
     "BaseStructure",
     "BaseSwarm",
     "BaseWorkflow",

{swarms-7.6.4 → swarms-7.6.5}/swarms/structs/agent.py

@@ -58,6 +58,12 @@ from swarms.utils.litellm_tokenizer import count_tokens
 from swarms.utils.pdf_to_text import pdf_to_text
 from swarms.utils.str_to_dict import str_to_dict
 
+from swarms.tools.mcp_integration import (
+    batch_mcp_flow,
+    mcp_flow_get_tool_schema,
+    MCPServerSseParams,
+)
+
 
 # Utils
 # Custom stopping condition
@@ -352,6 +358,7 @@ class Agent:
         role: agent_roles = "worker",
         no_print: bool = False,
         tools_list_dictionary: Optional[List[Dict[str, Any]]] = None,
+        mcp_servers: List[MCPServerSseParams] = [],
         *args,
         **kwargs,
     ):
@@ -471,6 +478,7 @@ class Agent:
         self.role = role
         self.no_print = no_print
         self.tools_list_dictionary = tools_list_dictionary
+        self.mcp_servers = mcp_servers
 
         if (
             self.agent_name is not None
@@ -584,6 +592,12 @@ class Agent:
         if self.llm is None:
             self.llm = self.llm_handling()
 
+        if (
+            self.tools_list_dictionary is None
+            and self.mcp_servers is not None
+        ):
+            self.tools_list_dictionary = self.mcp_tool_handling()
+
     def llm_handling(self):
         from swarms.utils.litellm_wrapper import LiteLLM
 
@@ -631,6 +645,69 @@ class Agent:
             logger.error(f"Error in llm_handling: {e}")
             return None
 
+    def mcp_execution_flow(self, response: any):
+        """
+        Executes the MCP (Model Context Protocol) flow based on the provided response.
+
+        This method takes a response, converts it from a string to a dictionary format,
+        and checks for the presence of a tool name or a name in the response. If either
+        is found, it retrieves the tool name and proceeds to call the batch_mcp_flow
+        function to execute the corresponding tool actions.
+
+        Args:
+            response (any): The response to be processed, which can be in string format
+            that represents a dictionary.
+
+        Returns:
+            The output from the batch_mcp_flow function, which contains the results of
+            the tool execution. If an error occurs during processing, it logs the error
+            and returns None.
+
+        Raises:
+            Exception: Logs any exceptions that occur during the execution flow.
+        """
+        try:
+            response = str_to_dict(response)
+
+            tool_output = batch_mcp_flow(
+                self.mcp_servers,
+                function_call=response,
+            )
+
+            return tool_output
+        except Exception as e:
+            logger.error(f"Error in mcp_execution_flow: {e}")
+            return None
+
+    def mcp_tool_handling(self):
+        """
+        Handles the retrieval of tool schemas from the MCP servers.
+
+        This method iterates over the list of MCP servers, retrieves the tool schema
+        for each server using the mcp_flow_get_tool_schema function, and compiles
+        these schemas into a list. The resulting list is stored in the
+        tools_list_dictionary attribute.
+
+        Returns:
+            list: A list of tool schemas retrieved from the MCP servers. If an error
+            occurs during the retrieval process, it logs the error and returns None.
+
+        Raises:
+            Exception: Logs any exceptions that occur during the tool handling process.
+        """
+        try:
+            self.tools_list_dictionary = []
+
+            for mcp_server in self.mcp_servers:
+                tool_schema = mcp_flow_get_tool_schema(mcp_server)
+                self.tools_list_dictionary.append(tool_schema)
+
+            print(self.tools_list_dictionary)
+            return self.tools_list_dictionary
+        except Exception as e:
+            logger.error(f"Error in mcp_tool_handling: {e}")
+            return None
+
     def setup_config(self):
         # The max_loops will be set dynamically if the dynamic_loop
         if self.dynamic_loops is True:

swarms-7.6.5/swarms/tools/mcp_integration.py

@@ -0,0 +1,392 @@
+from __future__ import annotations
+
+from typing import Any, List
+
+
+from loguru import logger
+
+import abc
+import asyncio
+from contextlib import AbstractAsyncContextManager, AsyncExitStack
+from pathlib import Path
+from typing import Literal
+
+from anyio.streams.memory import (
+    MemoryObjectReceiveStream,
+    MemoryObjectSendStream,
+)
+from mcp import (
+    ClientSession,
+    StdioServerParameters,
+    Tool as MCPTool,
+    stdio_client,
+)
+from mcp.client.sse import sse_client
+from mcp.types import CallToolResult, JSONRPCMessage
+from typing_extensions import NotRequired, TypedDict
+
+from swarms.utils.any_to_str import any_to_str
+
+
+class MCPServer(abc.ABC):
+    """Base class for Model Context Protocol servers."""
+
+    @abc.abstractmethod
+    async def connect(self):
+        """Connect to the server. For example, this might mean spawning a subprocess or
+        opening a network connection. The server is expected to remain connected until
+        `cleanup()` is called.
+        """
+        pass
+
+    @property
+    @abc.abstractmethod
+    def name(self) -> str:
+        """A readable name for the server."""
+        pass
+
+    @abc.abstractmethod
+    async def cleanup(self):
+        """Cleanup the server. For example, this might mean closing a subprocess or
+        closing a network connection.
+        """
+        pass
+
+    @abc.abstractmethod
+    async def list_tools(self) -> list[MCPTool]:
+        """List the tools available on the server."""
+        pass
+
+    @abc.abstractmethod
+    async def call_tool(
+        self, tool_name: str, arguments: dict[str, Any] | None
+    ) -> CallToolResult:
+        """Invoke a tool on the server."""
+        pass
+
+
+class _MCPServerWithClientSession(MCPServer, abc.ABC):
+    """Base class for MCP servers that use a `ClientSession` to communicate with the server."""
+
+    def __init__(self, cache_tools_list: bool):
+        """
+        Args:
+            cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
+            cached and only fetched from the server once. If `False`, the tools list will be
+            fetched from the server on each call to `list_tools()`. The cache can be invalidated
+            by calling `invalidate_tools_cache()`. You should set this to `True` if you know the
+            server will not change its tools list, because it can drastically improve latency
+            (by avoiding a round-trip to the server every time).
+        """
+        self.session: ClientSession | None = None
+        self.exit_stack: AsyncExitStack = AsyncExitStack()
+        self._cleanup_lock: asyncio.Lock = asyncio.Lock()
+        self.cache_tools_list = cache_tools_list
+
+        # The cache is always dirty at startup, so that we fetch tools at least once
+        self._cache_dirty = True
+        self._tools_list: list[MCPTool] | None = None
+
+    @abc.abstractmethod
+    def create_streams(
+        self,
+    ) -> AbstractAsyncContextManager[
+        tuple[
+            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+            MemoryObjectSendStream[JSONRPCMessage],
+        ]
+    ]:
+        """Create the streams for the server."""
+        pass
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.cleanup()
+
+    def invalidate_tools_cache(self):
+        """Invalidate the tools cache."""
+        self._cache_dirty = True
+
+    async def connect(self):
+        """Connect to the server."""
+        try:
+            transport = await self.exit_stack.enter_async_context(
+                self.create_streams()
+            )
+            read, write = transport
+            session = await self.exit_stack.enter_async_context(
+                ClientSession(read, write)
+            )
+            await session.initialize()
+            self.session = session
+        except Exception as e:
+            logger.error(f"Error initializing MCP server: {e}")
+            await self.cleanup()
+            raise
+
+    async def list_tools(self) -> list[MCPTool]:
+        """List the tools available on the server."""
+        if not self.session:
+            raise Exception(
+                "Server not initialized. Make sure you call `connect()` first."
+            )
+
+        # Return from cache if caching is enabled, we have tools, and the cache is not dirty
+        if (
+            self.cache_tools_list
+            and not self._cache_dirty
+            and self._tools_list
+        ):
+            return self._tools_list
+
+        # Reset the cache dirty to False
+        self._cache_dirty = False
+
+        # Fetch the tools from the server
+        self._tools_list = (await self.session.list_tools()).tools
+        return self._tools_list
+
+    async def call_tool(
+        self, arguments: dict[str, Any] | None
+    ) -> CallToolResult:
+        """Invoke a tool on the server."""
+        tool_name = arguments.get("tool_name") or arguments.get(
+            "name"
+        )
+
+        if not tool_name:
+            raise Exception("No tool name found in arguments")
+
+        if not self.session:
+            raise Exception(
+                "Server not initialized. Make sure you call `connect()` first."
+            )
+
+        return await self.session.call_tool(tool_name, arguments)
+
+    async def cleanup(self):
+        """Cleanup the server."""
+        async with self._cleanup_lock:
+            try:
+                await self.exit_stack.aclose()
+                self.session = None
+            except Exception as e:
+                logger.error(f"Error cleaning up server: {e}")
+
+
+class MCPServerStdioParams(TypedDict):
+    """Mirrors `mcp.client.stdio.StdioServerParameters`, but lets you pass params without another
+    import.
+    """
+
+    command: str
+    """The executable to run to start the server. For example, `python` or `node`."""
+
+    args: NotRequired[list[str]]
+    """Command line args to pass to the `command` executable. For example, `['foo.py']` or
+    `['server.js', '--port', '8080']`."""
+
+    env: NotRequired[dict[str, str]]
+    """The environment variables to set for the server. ."""
+
+    cwd: NotRequired[str | Path]
+    """The working directory to use when spawning the process."""
+
+    encoding: NotRequired[str]
+    """The text encoding used when sending/receiving messages to the server. Defaults to `utf-8`."""
+
+    encoding_error_handler: NotRequired[
+        Literal["strict", "ignore", "replace"]
+    ]
+    """The text encoding error handler. Defaults to `strict`.
+
+    See https://docs.python.org/3/library/codecs.html#codec-base-classes for
+    explanations of possible values.
+    """
+
+
+class MCPServerStdio(_MCPServerWithClientSession):
+    """MCP server implementation that uses the stdio transport. See the [spec]
+    (https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio) for
+    details.
+    """
+
+    def __init__(
+        self,
+        params: MCPServerStdioParams,
+        cache_tools_list: bool = False,
+        name: str | None = None,
+    ):
+        """Create a new MCP server based on the stdio transport.
+
+        Args:
+            params: The params that configure the server. This includes the command to run to
+                start the server, the args to pass to the command, the environment variables to
+                set for the server, the working directory to use when spawning the process, and
+                the text encoding used when sending/receiving messages to the server.
+            cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
+                cached and only fetched from the server once. If `False`, the tools list will be
+                fetched from the server on each call to `list_tools()`. The cache can be
+                invalidated by calling `invalidate_tools_cache()`. You should set this to `True`
+                if you know the server will not change its tools list, because it can drastically
+                improve latency (by avoiding a round-trip to the server every time).
+            name: A readable name for the server. If not provided, we'll create one from the
+                command.
+        """
+        super().__init__(cache_tools_list)
+
+        self.params = StdioServerParameters(
+            command=params["command"],
+            args=params.get("args", []),
+            env=params.get("env"),
+            cwd=params.get("cwd"),
+            encoding=params.get("encoding", "utf-8"),
+            encoding_error_handler=params.get(
+                "encoding_error_handler", "strict"
+            ),
+        )
+
+        self._name = name or f"stdio: {self.params.command}"
+
+    def create_streams(
+        self,
+    ) -> AbstractAsyncContextManager[
+        tuple[
+            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+            MemoryObjectSendStream[JSONRPCMessage],
+        ]
+    ]:
+        """Create the streams for the server."""
+        return stdio_client(self.params)
+
+    @property
+    def name(self) -> str:
+        """A readable name for the server."""
+        return self._name
+
+
+class MCPServerSseParams(TypedDict):
+    """Mirrors the params in`mcp.client.sse.sse_client`."""
+
+    url: str
+    """The URL of the server."""
+
+    headers: NotRequired[dict[str, str]]
+    """The headers to send to the server."""
+
+    timeout: NotRequired[float]
+    """The timeout for the HTTP request. Defaults to 5 seconds."""
+
+    sse_read_timeout: NotRequired[float]
+    """The timeout for the SSE connection, in seconds. Defaults to 5 minutes."""
+
+
+class MCPServerSse(_MCPServerWithClientSession):
+    """MCP server implementation that uses the HTTP with SSE transport. See the [spec]
+    (https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse)
+    for details.
+    """
+
+    def __init__(
+        self,
+        params: MCPServerSseParams,
+        cache_tools_list: bool = False,
+        name: str | None = None,
+    ):
+        """Create a new MCP server based on the HTTP with SSE transport.
+
+        Args:
+            params: The params that configure the server. This includes the URL of the server,
+                the headers to send to the server, the timeout for the HTTP request, and the
+                timeout for the SSE connection.
+
+            cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
+                cached and only fetched from the server once. If `False`, the tools list will be
+                fetched from the server on each call to `list_tools()`. The cache can be
+                invalidated by calling `invalidate_tools_cache()`. You should set this to `True`
+                if you know the server will not change its tools list, because it can drastically
+                improve latency (by avoiding a round-trip to the server every time).
+
+            name: A readable name for the server. If not provided, we'll create one from the
+                URL.
+        """
+        super().__init__(cache_tools_list)
+
+        self.params = params
+        self._name = name or f"sse: {self.params['url']}"
+
+    def create_streams(
+        self,
+    ) -> AbstractAsyncContextManager[
+        tuple[
+            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+            MemoryObjectSendStream[JSONRPCMessage],
+        ]
+    ]:
+        """Create the streams for the server."""
+        return sse_client(
+            url=self.params["url"],
+            headers=self.params.get("headers", None),
+            timeout=self.params.get("timeout", 5),
+            sse_read_timeout=self.params.get(
+                "sse_read_timeout", 60 * 5
+            ),
+        )
+
+    @property
+    def name(self) -> str:
+        """A readable name for the server."""
+        return self._name
+
+
+def mcp_flow_get_tool_schema(
+    params: MCPServerSseParams,
+) -> MCPServer:
+    server = MCPServerSse(params, cache_tools_list=True)
+
+    # Connect the server
+    asyncio.run(server.connect())
+
+    # Return the server
+    output = asyncio.run(server.list_tools())
+
+    # Cleanup the server
+    asyncio.run(server.cleanup())
+
+    return output.model_dump()
+
+
+def mcp_flow(
+    params: MCPServerSseParams,
+    function_call: dict[str, Any],
+) -> MCPServer:
+    server = MCPServerSse(params, cache_tools_list=True)
+
+    # Connect the server
+    asyncio.run(server.connect())
+
+    # Return the server
+    output = asyncio.run(server.call_tool(function_call))
+
+    output = output.model_dump()
+
+    # Cleanup the server
+    asyncio.run(server.cleanup())
+
+    return any_to_str(output)
+
+
+def batch_mcp_flow(
+    params: List[MCPServerSseParams],
+    function_call: List[dict[str, Any]] = [],
+) -> MCPServer:
+    output_list = []
+
+    for param in params:
+        output = mcp_flow(param, function_call)
+        output_list.append(output)
+
+    return output_list

swarms-7.6.5/swarms/utils/vllm_wrapper.py

@@ -0,0 +1,146 @@
+from typing import List, Optional, Dict, Any
+from loguru import logger
+
+try:
+    from vllm import LLM, SamplingParams
+except ImportError:
+    import subprocess
+    import sys
+
+    print("Installing vllm")
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "-U", "vllm"]
+    )
+    print("vllm installed")
+    from vllm import LLM, SamplingParams
+
+
+class VLLMWrapper:
+    """
+    A wrapper class for vLLM that provides a similar interface to LiteLLM.
+    This class handles model initialization and inference using vLLM.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "meta-llama/Llama-2-7b-chat-hf",
+        system_prompt: Optional[str] = None,
+        stream: bool = False,
+        temperature: float = 0.5,
+        max_tokens: int = 4000,
+        max_completion_tokens: int = 4000,
+        tools_list_dictionary: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: str = "auto",
+        parallel_tool_calls: bool = False,
+        *args,
+        **kwargs,
+    ):
+        """
+        Initialize the vLLM wrapper with the given parameters.
+
+        Args:
+            model_name (str): The name of the model to use. Defaults to "meta-llama/Llama-2-7b-chat-hf".
+            system_prompt (str, optional): The system prompt to use. Defaults to None.
+            stream (bool): Whether to stream the output. Defaults to False.
+            temperature (float): The temperature for sampling. Defaults to 0.5.
+            max_tokens (int): The maximum number of tokens to generate. Defaults to 4000.
+            max_completion_tokens (int): The maximum number of completion tokens. Defaults to 4000.
+            tools_list_dictionary (List[Dict[str, Any]], optional): List of available tools. Defaults to None.
+            tool_choice (str): How to choose tools. Defaults to "auto".
+            parallel_tool_calls (bool): Whether to allow parallel tool calls. Defaults to False.
+        """
+        self.model_name = model_name
+        self.system_prompt = system_prompt
+        self.stream = stream
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
+        self.tools_list_dictionary = tools_list_dictionary
+        self.tool_choice = tool_choice
+        self.parallel_tool_calls = parallel_tool_calls
+
+        # Initialize vLLM
+        self.llm = LLM(model=model_name, **kwargs)
+        self.sampling_params = SamplingParams(
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+
+    def _prepare_prompt(self, task: str) -> str:
+        """
+        Prepare the prompt for the given task.
+
+        Args:
+            task (str): The task to prepare the prompt for.
+
+        Returns:
+            str: The prepared prompt.
+        """
+        if self.system_prompt:
+            return f"{self.system_prompt}\n\nUser: {task}\nAssistant:"
+        return f"User: {task}\nAssistant:"
+
+    def run(self, task: str, *args, **kwargs) -> str:
+        """
+        Run the model for the given task.
+
+        Args:
+            task (str): The task to run the model for.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The model's response.
+        """
+        try:
+            prompt = self._prepare_prompt(task)
+
+            outputs = self.llm.generate(prompt, self.sampling_params)
+            response = outputs[0].outputs[0].text.strip()
+
+            return response
+
+        except Exception as error:
+            logger.error(f"Error in VLLMWrapper: {error}")
+            raise error
+
+    def __call__(self, task: str, *args, **kwargs) -> str:
+        """
+        Call the model for the given task.
+
+        Args:
+            task (str): The task to run the model for.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The model's response.
+        """
+        return self.run(task, *args, **kwargs)
+
+    def batched_run(
+        self, tasks: List[str], batch_size: int = 10
+    ) -> List[str]:
+        """
+        Run the model for multiple tasks in batches.
+
+        Args:
+            tasks (List[str]): List of tasks to run.
+            batch_size (int): Size of each batch. Defaults to 10.
+
+        Returns:
+            List[str]: List of model responses.
+        """
+        logger.info(
+            f"Running tasks in batches of size {batch_size}. Total tasks: {len(tasks)}"
+        )
+        results = []
+
+        for i in range(0, len(tasks), batch_size):
+            batch = tasks[i : i + batch_size]
+            for task in batch:
+                logger.info(f"Running task: {task}")
+                results.append(self.run(task))
+
+        logger.info("Completed all tasks.")
+        return results