agno 1.7.6__py3-none-any.whl → 1.7.7__py3-none-any.whl

agno/tools/mcp.py

@@ -1,3 +1,5 @@
+import asyncio
+import weakref
 from contextlib import AsyncExitStack
 from dataclasses import asdict, dataclass
 from datetime import timedelta
@@ -18,6 +20,28 @@ except (ImportError, ModuleNotFoundError):
     raise ImportError("`mcp` not installed. Please install using `pip install mcp`")
 
 
+def _prepare_command(command: str) -> list[str]:
+    """Sanitize a command and split it into parts before using it to run a MCP server."""
+    from shlex import split
+
+    # Block dangerous characters
+    if any(char in command for char in ["&", "|", ";", "`", "$", "(", ")"]):
+        raise ValueError("MCP command can't contain shell metacharacters")
+
+    parts = split(command)
+    if not parts:
+        raise ValueError("MCP command can't be empty")
+
+    # Only allow specific executables
+    ALLOWED_COMMANDS = {"python", "python3", "node", "npm", "npx"}
+
+    executable = parts[0].split("/")[-1]
+    if executable not in ALLOWED_COMMANDS:
+        raise ValueError(f"MCP command needs to use one of the following executables: {ALLOWED_COMMANDS}")
+
+    return parts
+
+
 @dataclass
 class SSEClientParams:
     """Parameters for SSE client connection."""
@@ -138,11 +162,7 @@ class MCPTools(Toolkit):
             env = get_default_environment()
 
         if command is not None and transport not in ["sse", "streamable-http"]:
-            from shlex import split
-
-            parts = split(command)
-            if not parts:
-                raise ValueError("Empty command string")
+            parts = _prepare_command(command)
             cmd = parts[0]
             arguments = parts[1:] if len(parts) > 1 else []
             self.server_params = StdioServerParameters(command=cmd, args=arguments, env=env)
@@ -151,23 +171,49 @@ class MCPTools(Toolkit):
         self._context = None
         self._session_context = None
         self._initialized = False
+        self._connection_task = None
 
-    async def __aenter__(self) -> "MCPTools":
-        """Enter the async context manager."""
+        def cleanup():
+            """Cancel active connections"""
+            if self._connection_task and not self._connection_task.done():
+                self._connection_task.cancel()
+
+        # Setup cleanup logic before the instance is garbage collected
+        self._cleanup_finalizer = weakref.finalize(self, cleanup)
+
+    async def connect(self):
+        """Initialize a MCPTools instance and connect to the contextual MCP server"""
+        if self._initialized:
+            return
+
+        await self._connect()
+
+    def _start_connection(self):
+        """Ensure there are no active connections and setup a new one"""
+        if self._connection_task is None or self._connection_task.done():
+            self._connection_task = asyncio.create_task(self._connect())  # type: ignore
+
+    async def _connect(self) -> None:
+        """Connects to the MCP server and initializes the tools"""
+        if self._initialized:
+            return
 
         if self.session is not None:
-            # Already has a session, just initialize
-            if not self._initialized:
-                await self.initialize()
-            return self
+            await self.initialize()
+            return
+
+        if not hasattr(self, "_active_contexts"):
+            self._active_contexts: list[Any] = []
 
-        # Create a new session using stdio_client, sse_client or streamablehttp_client based on transport
+        # Create a new studio session
         if self.transport == "sse":
             sse_params = asdict(self.server_params) if self.server_params is not None else {}  # type: ignore
             if "url" not in sse_params:
                 sse_params["url"] = self.url
             self._context = sse_client(**sse_params)  # type: ignore
             client_timeout = min(self.timeout_seconds, sse_params.get("timeout", self.timeout_seconds))
+
+        # Create a new streamable HTTP session
         elif self.transport == "streamable-http":
             streamable_http_params = asdict(self.server_params) if self.server_params is not None else {}  # type: ignore
             if "url" not in streamable_http_params:
@@ -177,6 +223,7 @@ class MCPTools(Toolkit):
             if isinstance(params_timeout, timedelta):
                 params_timeout = int(params_timeout.total_seconds())
             client_timeout = min(self.timeout_seconds, params_timeout)
+
         else:
             if self.server_params is None:
                 raise ValueError("server_params must be provided when using stdio transport.")
@@ -184,24 +231,42 @@ class MCPTools(Toolkit):
             client_timeout = self.timeout_seconds
 
         session_params = await self._context.__aenter__()  # type: ignore
+        self._active_contexts.append(self._context)
         read, write = session_params[0:2]
 
         self._session_context = ClientSession(read, write, read_timeout_seconds=timedelta(seconds=client_timeout))  # type: ignore
         self.session = await self._session_context.__aenter__()  # type: ignore
+        self._active_contexts.append(self._session_context)
 
         # Initialize with the new session
         await self.initialize()
+
+    async def close(self) -> None:
+        """Close the MCP connection and clean up resources"""
+        if self._session_context is not None:
+            await self._session_context.__aexit__(None, None, None)
+            self.session = None
+            self._session_context = None
+
+        if self._context is not None:
+            await self._context.__aexit__(None, None, None)
+            self._context = None
+
+        self._initialized = False
+
+    async def __aenter__(self) -> "MCPTools":
+        await self._connect()
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, _exc_type, _exc_val, _exc_tb):
         """Exit the async context manager."""
         if self._session_context is not None:
-            await self._session_context.__aexit__(exc_type, exc_val, exc_tb)
+            await self._session_context.__aexit__(_exc_type, _exc_val, _exc_tb)
             self.session = None
             self._session_context = None
 
         if self._context is not None:
-            await self._context.__aexit__(exc_type, exc_val, exc_tb)
+            await self._context.__aexit__(_exc_type, _exc_val, _exc_tb)
             self._context = None
 
         self._initialized = False
@@ -213,7 +278,7 @@ class MCPTools(Toolkit):
 
         try:
             if self.session is None:
-                raise ValueError("Session is not available. Use as context manager or provide a session.")
+                raise ValueError("Failed to establish session connection")
 
             # Initialize the session if not already initialized
             await self.session.initialize()
@@ -343,12 +408,8 @@ class MultiMCPTools(Toolkit):
             env = get_default_environment()
 
         if commands is not None:
-            from shlex import split
-
             for command in commands:
-                parts = split(command)
-                if not parts:
-                    raise ValueError("Empty command string")
+                parts = _prepare_command(command)
                 cmd = parts[0]
                 arguments = parts[1:] if len(parts) > 1 else []
                 self.server_params_list.append(StdioServerParameters(command=cmd, args=arguments, env=env))
@@ -365,28 +426,92 @@ class MultiMCPTools(Toolkit):
                     self.server_params_list.append(StreamableHTTPClientParams(url=url))
 
         self._async_exit_stack = AsyncExitStack()
+        self._initialized = False
+        self._connection_task = None
+        self._active_contexts: list[Any] = []
+        self._used_as_context_manager = False
 
         self._client = client
 
-    async def __aenter__(self) -> "MultiMCPTools":
-        """Enter the async context manager."""
+        def cleanup():
+            """Cancel active connections"""
+            if self._connection_task and not self._connection_task.done():
+                self._connection_task.cancel()
+
+        # Setup cleanup logic before the instance is garbage collected
+        self._cleanup_finalizer = weakref.finalize(self, cleanup)
+
+    async def connect(self):
+        """Initialize a MultiMCPTools instance and connect to the MCP servers"""
+        if self._initialized:
+            return
+
+        await self._connect()
+
+    @classmethod
+    async def create_and_connect(
+        cls,
+        commands: Optional[List[str]] = None,
+        urls: Optional[List[str]] = None,
+        urls_transports: Optional[List[Literal["sse", "streamable-http"]]] = None,
+        *,
+        env: Optional[dict[str, str]] = None,
+        server_params_list: Optional[
+            List[Union[SSEClientParams, StdioServerParameters, StreamableHTTPClientParams]]
+        ] = None,
+        timeout_seconds: int = 5,
+        client=None,
+        include_tools: Optional[list[str]] = None,
+        exclude_tools: Optional[list[str]] = None,
+        **kwargs,
+    ) -> "MultiMCPTools":
+        """Initialize a MultiMCPTools instance and connect to the MCP servers"""
+        instance = cls(
+            commands=commands,
+            urls=urls,
+            urls_transports=urls_transports,
+            env=env,
+            server_params_list=server_params_list,
+            timeout_seconds=timeout_seconds,
+            client=client,
+            include_tools=include_tools,
+            exclude_tools=exclude_tools,
+            **kwargs,
+        )
+
+        await instance._connect()
+        return instance
+
+    def _start_connection(self):
+        """Ensure there are no active connections and setup a new one"""
+        if self._connection_task is None or self._connection_task.done():
+            self._connection_task = asyncio.create_task(self._connect())  # type: ignore
+
+    async def _connect(self) -> None:
+        """Connects to the MCP servers and initializes the tools"""
+        if self._initialized:
+            return
 
         for server_params in self.server_params_list:
             # Handle stdio connections
             if isinstance(server_params, StdioServerParameters):
                 stdio_transport = await self._async_exit_stack.enter_async_context(stdio_client(server_params))
+                self._active_contexts.append(stdio_transport)
                 read, write = stdio_transport
                 session = await self._async_exit_stack.enter_async_context(
                     ClientSession(read, write, read_timeout_seconds=timedelta(seconds=self.timeout_seconds))
                 )
+                self._active_contexts.append(session)
                 await self.initialize(session)
             # Handle SSE connections
             elif isinstance(server_params, SSEClientParams):
                 client_connection = await self._async_exit_stack.enter_async_context(
                     sse_client(**asdict(server_params))
                 )
+                self._active_contexts.append(client_connection)
                 read, write = client_connection
                 session = await self._async_exit_stack.enter_async_context(ClientSession(read, write))
+                self._active_contexts.append(session)
                 await self.initialize(session)
 
             # Handle Streamable HTTP connections
@@ -394,10 +519,22 @@ class MultiMCPTools(Toolkit):
                 client_connection = await self._async_exit_stack.enter_async_context(
                     streamablehttp_client(**asdict(server_params))
                 )
+                self._active_contexts.append(client_connection)
                 read, write = client_connection[0:2]
                 session = await self._async_exit_stack.enter_async_context(ClientSession(read, write))
+                self._active_contexts.append(session)
                 await self.initialize(session)
 
+        self._initialized = True
+
+    async def close(self) -> None:
+        """Close the MCP connections and clean up resources"""
+        await self._async_exit_stack.aclose()
+        self._initialized = False
+
+    async def __aenter__(self) -> "MultiMCPTools":
+        """Enter the async context manager."""
+        await self._connect()
         return self
 
     async def __aexit__(

agno/tools/models/morph.py

@@ -0,0 +1,186 @@
+import os
+from os import getenv
+from textwrap import dedent
+from typing import Optional
+
+from agno.tools import Toolkit
+from agno.utils.log import log_debug, log_error
+
+try:
+    from openai import OpenAI
+except ImportError:
+    raise ImportError("`openai` not installed. Please install using `pip install openai`")
+
+
+class MorphTools(Toolkit):
+    """Tools for interacting with Morph's Fast Apply API for code editing"""
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: str = "https://api.morphllm.com/v1",
+        instructions: Optional[str] = None,
+        add_instructions: bool = True,
+        model: str = "morph-v3-large",
+        **kwargs,
+    ):
+        """Initialize Morph Fast Apply tools.
+
+        Args:
+            api_key: Morph API key. If not provided, will look for MORPH_API_KEY environment variable.
+            base_url: The base URL for the Morph API.
+            model: The Morph model to use. Options:
+                  - "morph-v3-fast" (4500+ tok/sec, 96% accuracy)
+                  - "morph-v3-large" (2500+ tok/sec, 98% accuracy)
+                  - "auto" (automatic selection)
+            **kwargs: Additional arguments to pass to Toolkit.
+        """
+        # Set up instructions
+        if instructions is None:
+            self.instructions = self.DEFAULT_INSTRUCTIONS
+        else:
+            self.instructions = instructions
+
+        super().__init__(
+            name="morph_tools",
+            tools=[self.edit_file],
+            instructions=self.instructions,
+            add_instructions=add_instructions,
+            **kwargs,
+        )
+
+        self.api_key = api_key or getenv("MORPH_API_KEY")
+        if not self.api_key:
+            raise ValueError("MORPH_API_KEY not set. Please set the MORPH_API_KEY environment variable.")
+
+        self.base_url = base_url
+        self.model = model
+        self._morph_client: Optional[OpenAI] = None
+
+    def _get_client(self):
+        """Get or create the Morph OpenAI client."""
+        if self._morph_client is None:
+            self._morph_client = OpenAI(
+                api_key=self.api_key,
+                base_url=self.base_url,
+            )
+        return self._morph_client
+
+    def edit_file(
+        self,
+        target_file: str,
+        instructions: str,
+        code_edit: str,
+        original_code: Optional[str] = None,
+    ) -> str:
+        """
+        Apply code edits to a target file using Morph's Fast Apply API.
+
+        This function reads the specified file, sends its content along with
+        editing instructions and code edits to Morph's API, and writes the
+        resulting code back to the file. A backup of the original file is
+        created before writing changes.
+
+        Args:
+            target_file (str): Path to the file to be edited.
+            instructions (str): High-level instructions describing the intended change.
+            code_edit (str): Specific code edit or change to apply.
+            original_code (Optional[str], optional): Original content of the file.
+                If not provided, the function reads from target_file.
+
+        Returns:
+            str: Result message indicating success or failure, and details about
+                the backup and any errors encountered.
+        """
+        try:
+            # Always read the actual file content for backup purposes
+            actual_file_content = None
+            if os.path.exists(target_file):
+                try:
+                    with open(target_file, "r", encoding="utf-8") as f:
+                        actual_file_content = f.read()
+                except Exception as e:
+                    return f"Error reading {target_file} for backup: {e}"
+            else:
+                return f"Error: File {target_file} does not exist."
+
+            # Use provided original_code or fall back to file content
+            code_to_process = original_code if original_code is not None else actual_file_content
+
+            # Format the message for Morph's Fast Apply API
+            content = f"<instruction>{instructions}</instruction>\n<code>{code_to_process}</code>\n<update>{code_edit}</update>"
+
+            log_debug(f"Input to Morph: {content}")
+
+            client = self._get_client()
+
+            response = client.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": content,
+                    }
+                ],
+            )
+
+            if response.choices and response.choices[0].message.content:
+                final_code = response.choices[0].message.content
+
+                try:
+                    backup_file = f"{target_file}.backup"
+                    with open(backup_file, "w", encoding="utf-8") as f:
+                        f.write(actual_file_content)
+
+                    # Write the new code
+                    with open(target_file, "w", encoding="utf-8") as f:
+                        f.write(final_code)
+                    return f"Successfully applied edit to {target_file} using Morph Fast Apply! Original content backed up as {backup_file}"
+
+                except Exception as e:
+                    return f"Successfully applied edit but failed to write back to {target_file}: {e}"
+
+            else:
+                return f"Failed to apply edit to {target_file}: No response from Morph API"
+
+        except Exception as e:
+            log_error(f"Failed to apply edit using Morph Fast Apply: {e}")
+            return f"Failed to apply edit to {target_file}: {e}"
+
+    DEFAULT_INSTRUCTIONS = dedent("""\
+            You have access to Morph Fast Apply for ultra-fast code editing with 98% accuracy at 2500+ tokens/second.
+
+            ## How to use the edit_file tool:
+
+            **Critical Requirements:**
+            1. **Instructions Parameter**: Generate clear first-person instructions describing what you're doing
+            - Example: "I am adding type hints to all functions and methods"
+            - Example: "I am refactoring the error handling to use try-catch blocks"
+
+            2. **Code Edit Parameter**: Specify ONLY the lines you want to change
+            - Use `# ... existing code ...` (or `// ... existing code ...` for JS/Java) to represent unchanged sections
+            - NEVER write out unchanged code in the code_edit parameter
+            - Include sufficient context around changes to resolve ambiguity
+
+            3. **Single Edit Call**: Make ALL edits to a file in a single edit_file call. The apply model can handle many distinct edits at once.
+
+            **Example Format:**
+            ```
+            # ... existing code ...
+            def add(a: int, b: int) -> int:
+                \"\"\"Add two numbers together.\"\"\"
+                return a + b
+            # ... existing code ...
+            def multiply(x: int, y: int) -> int:
+                \"\"\"Multiply two numbers.\"\"\"
+                return x * y
+            # ... existing code ...
+            ```
+
+            **Important Guidelines:**
+            - Bias towards repeating as few lines as possible while conveying the change clearly
+            - Each edit should contain sufficient context of unchanged lines around the code you're editing
+            - DO NOT omit spans of pre-existing code without using the `# ... existing code ...` comment
+            - If deleting a section, provide context before and after to clearly indicate the deletion
+            - The tool automatically creates backup files before applying changes\
+        """)

agno/tools/postgres.py

@@ -2,14 +2,12 @@ import csv
 from typing import Any, Dict, List, Optional
 
 try:
-    import psycopg2
-    from psycopg2 import sql
-    from psycopg2.extensions import connection as PgConnection
-    from psycopg2.extras import DictCursor
+    import psycopg
+    from psycopg import sql
+    from psycopg.connection import Connection as PgConnection
+    from psycopg.rows import DictRow, dict_row
 except ImportError:
-    raise ImportError(
-        "`psycopg2` not installed. Please install using `pip install psycopg2`. If you face issues, try `pip install psycopg2-binary`."
-    )
+    raise ImportError("`psycopg` not installed. Please install using `pip install 'psycopg-binary'`.")
 
 from agno.tools import Toolkit
 from agno.utils.log import log_debug, log_error
@@ -18,7 +16,7 @@ from agno.utils.log import log_debug, log_error
 class PostgresTools(Toolkit):
     def __init__(
         self,
-        connection: Optional[PgConnection] = None,
+        connection: Optional[PgConnection[DictRow]] = None,
         db_name: Optional[str] = None,
         user: Optional[str] = None,
         password: Optional[str] = None,
@@ -31,7 +29,7 @@ class PostgresTools(Toolkit):
         table_schema: str = "public",
         **kwargs,
     ):
-        self._connection: Optional[PgConnection] = connection
+        self._connection: Optional[PgConnection[DictRow]] = connection
         self.db_name: Optional[str] = db_name
         self.user: Optional[str] = user
         self.password: Optional[str] = password
@@ -55,16 +53,16 @@ class PostgresTools(Toolkit):
         super().__init__(name="postgres_tools", tools=tools, **kwargs)
 
     @property
-    def connection(self) -> PgConnection:
+    def connection(self) -> PgConnection[DictRow]:
         """
-        Returns the Postgres psycopg2 connection.
-        :return psycopg2.extensions.connection: psycopg2 connection
+        Returns the Postgres psycopg connection.
+        :return psycopg.connection.Connection: psycopg connection
         """
         if self._connection is None or self._connection.closed:
             log_debug("Establishing new PostgreSQL connection.")
-            connection_kwargs: Dict[str, Any] = {"cursor_factory": DictCursor}
+            connection_kwargs: Dict[str, Any] = {"row_factory": dict_row}
             if self.db_name:
-                connection_kwargs["database"] = self.db_name
+                connection_kwargs["dbname"] = self.db_name
             if self.user:
                 connection_kwargs["user"] = self.user
             if self.password:
@@ -76,8 +74,8 @@ class PostgresTools(Toolkit):
 
             connection_kwargs["options"] = f"-c search_path={self.table_schema}"
 
-            self._connection = psycopg2.connect(**connection_kwargs)
-            self._connection.set_session(readonly=True)
+            self._connection = psycopg.connect(**connection_kwargs)
+            self._connection.read_only = True
 
         return self._connection
 
@@ -110,10 +108,10 @@ class PostgresTools(Toolkit):
                     return f"Query returned no results.\nColumns: {', '.join(columns)}"
 
                 header = ",".join(columns)
-                data_rows = [",".join(map(str, row)) for row in rows]
+                data_rows = [",".join(map(str, row.values())) for row in rows]
                 return f"{header}\n" + "\n".join(data_rows)
 
-        except psycopg2.Error as e:
+        except psycopg.Error as e:
             log_error(f"Database error: {e}")
             if self.connection and not self.connection.closed:
                 self.connection.rollback()
@@ -203,15 +201,18 @@ class PostgresTools(Toolkit):
                         cursor.execute(query)
                         stats = cursor.fetchone()
                         summary_parts.append(f"\n--- Column: {col_name} (Type: {data_type}) ---")
-                        for key, value in stats.items():
-                            val_str = (
-                                f"{value:.2f}" if isinstance(value, (float, int)) and value is not None else str(value)
-                            )
-                            summary_parts.append(f"  {key}: {val_str}")
+                        if stats is not None:
+                            for key, value in stats.items():
+                                val_str = (
+                                    f"{value:.2f}" if isinstance(value, float) and value is not None else str(value)
+                                )
+                                summary_parts.append(f"  {key}: {val_str}")
+                        else:
+                            summary_parts.append("  No statistics available")
 
                 return "\n".join(summary_parts)
 
-        except psycopg2.Error as e:
+        except psycopg.Error as e:
             return f"Error summarizing table: {e}"
 
     def inspect_query(self, query: str) -> str:
@@ -239,15 +240,19 @@ class PostgresTools(Toolkit):
         try:
             with self.connection.cursor() as cursor:
                 cursor.execute(stmt)
+
+                if cursor.description is None:
+                    return f"Error: Query returned no description for table '{table}'."
+
                 columns = [desc[0] for desc in cursor.description]
 
                 with open(path, "w", newline="", encoding="utf-8") as f:
                     writer = csv.writer(f)
                     writer.writerow(columns)
-                    writer.writerows(cursor)
+                    writer.writerows(row.values() for row in cursor)
 
             return f"Successfully exported table '{table}' to '{path}'."
-        except (psycopg2.Error, IOError) as e:
+        except (psycopg.Error, IOError) as e:
             return f"Error exporting table: {e}"
 
     def run_query(self, query: str) -> str: