agno 0.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/utils/log.py

@@ -1,19 +1,92 @@
 import logging
+from functools import lru_cache
 from os import getenv
+from typing import Any, Literal, Optional
 
 from rich.logging import RichHandler
+from rich.text import Text
 
 LOGGER_NAME = "agno"
+TEAM_LOGGER_NAME = f"{LOGGER_NAME}-team"
+WORKFLOW_LOGGER_NAME = f"{LOGGER_NAME}-workflow"
 
+# Define custom styles for different log sources
+LOG_STYLES = {
+    "agent": {
+        "debug": "green",
+        "info": "blue",
+    },
+    "team": {
+        "debug": "magenta",
+        "info": "steel_blue1",
+    },
+    "workflow": {
+        "debug": "sandy_brown",
+        "info": "orange3",
+    },
+}
+
+
+class ColoredRichHandler(RichHandler):
+    def __init__(self, *args, source_type: Optional[str] = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.source_type = source_type
+
+    def get_level_text(self, record: logging.LogRecord) -> Text:
+        # Return empty Text if message is empty
+        if not record.msg:
+            return Text("")
+
+        level_name = record.levelname.lower()
+        if self.source_type and self.source_type in LOG_STYLES:
+            if level_name in LOG_STYLES[self.source_type]:
+                color = LOG_STYLES[self.source_type][level_name]
+                return Text(record.levelname, style=color)
+        else:
+            if level_name in LOG_STYLES["agent"]:
+                color = LOG_STYLES["agent"][level_name]
+                return Text(record.levelname, style=color)
+        return super().get_level_text(record)
+
+
+class AgnoLogger(logging.Logger):
+    def __init__(self, name: str, level: int = logging.NOTSET):
+        super().__init__(name, level)
+
+    def debug(self, msg: str, center: bool = False, symbol: str = "*", *args, **kwargs):  # type: ignore
+        if center:
+            msg = center_header(str(msg), symbol)
+        super().debug(msg, *args, **kwargs)
+
+    def info(self, msg: str, center: bool = False, symbol: str = "*", *args, **kwargs):  # type: ignore
+        if center:
+            msg = center_header(str(msg), symbol)
+        super().info(msg, *args, **kwargs)
+
+
+def build_logger(logger_name: str, source_type: Optional[str] = None) -> Any:
+    # If a logger with the name "agno.{source_type}" is already set, we want to use that one
+    _logger = logging.getLogger(f"agno.{logger_name}")
+    if _logger.handlers or _logger.level != logging.NOTSET:
+        return _logger
+
+    # Set the custom logger class as the default for this logger
+    logging.setLoggerClass(AgnoLogger)
+
+    # Create logger with custom class
+    _logger = logging.getLogger(logger_name)
+
+    # Reset logger class to default to avoid affecting other loggers
+    logging.setLoggerClass(logging.Logger)
 
-def get_logger(logger_name: str) -> logging.Logger:
     # https://rich.readthedocs.io/en/latest/reference/logging.html#rich.logging.RichHandler
     # https://rich.readthedocs.io/en/latest/logging.html#handle-exceptions
-    rich_handler = RichHandler(
+    rich_handler = ColoredRichHandler(
         show_time=False,
         rich_tracebacks=False,
         show_path=True if getenv("AGNO_API_RUNTIME") == "dev" else False,
         tracebacks_show_locals=False,
+        source_type=source_type or "agent",
     )
     rich_handler.setFormatter(
         logging.Formatter(
@@ -22,21 +95,161 @@ def get_logger(logger_name: str) -> logging.Logger:
         )
     )
 
-    _logger = logging.getLogger(logger_name)
     _logger.addHandler(rich_handler)
     _logger.setLevel(logging.INFO)
     _logger.propagate = False
     return _logger
 
 
-logger: logging.Logger = get_logger(LOGGER_NAME)
+agent_logger: AgnoLogger = build_logger(LOGGER_NAME, source_type="agent")
+team_logger: AgnoLogger = build_logger(TEAM_LOGGER_NAME, source_type="team")
+workflow_logger: AgnoLogger = build_logger(WORKFLOW_LOGGER_NAME, source_type="workflow")
 
+# Set the default logger to the agent logger
+logger: AgnoLogger = agent_logger
 
-def set_log_level_to_debug():
-    _logger = logging.getLogger(LOGGER_NAME)
+
+debug_on: bool = False
+debug_level: Literal[1, 2] = 1
+
+
+def set_log_level_to_debug(source_type: Optional[str] = None, level: Literal[1, 2] = 1):
+    if source_type is None:
+        use_agent_logger()
+
+    _logger = logging.getLogger(LOGGER_NAME if source_type is None else f"{LOGGER_NAME}-{source_type}")
     _logger.setLevel(logging.DEBUG)
 
+    global debug_on
+    debug_on = True
 
-def set_log_level_to_info():
-    _logger = logging.getLogger(LOGGER_NAME)
+    global debug_level
+    debug_level = level
+
+
+def set_log_level_to_info(source_type: Optional[str] = None):
+    _logger = logging.getLogger(LOGGER_NAME if source_type is None else f"{LOGGER_NAME}-{source_type}")
     _logger.setLevel(logging.INFO)
+
+    global debug_on
+    debug_on = False
+
+
+def set_log_level_to_warning(source_type: Optional[str] = None):
+    _logger = logging.getLogger(LOGGER_NAME if source_type is None else f"{LOGGER_NAME}-{source_type}")
+    _logger.setLevel(logging.WARNING)
+
+    global debug_on
+    debug_on = False
+
+
+def set_log_level_to_error(source_type: Optional[str] = None):
+    _logger = logging.getLogger(LOGGER_NAME if source_type is None else f"{LOGGER_NAME}-{source_type}")
+    _logger.setLevel(logging.ERROR)
+
+    global debug_on
+    debug_on = False
+
+
+def center_header(message: str, symbol: str = "*") -> str:
+    try:
+        import shutil
+
+        terminal_width = shutil.get_terminal_size().columns
+    except Exception:
+        terminal_width = 80  # fallback width
+
+    header = f" {message} "
+    return f"{header.center(terminal_width - 20, symbol)}"
+
+
+def use_team_logger():
+    """Switch the default logger to use team_logger"""
+    global logger
+    logger = team_logger
+
+
+def use_agent_logger():
+    """Switch the default logger to use the default agent logger"""
+    global logger
+    logger = agent_logger
+
+
+def use_workflow_logger():
+    """Switch the default logger to use workflow_logger"""
+    global logger
+    logger = workflow_logger
+
+
+@lru_cache(maxsize=128)
+def _using_default_logger(logger_instance: Any) -> bool:
+    """Return True if the currently active logger is our default AgnoLogger"""
+    return isinstance(logger_instance, AgnoLogger)
+
+
+def log_debug(msg, center: bool = False, symbol: str = "*", log_level: Literal[1, 2] = 1, *args, **kwargs):
+    global logger
+    global debug_on
+    global debug_level
+
+    if debug_on:
+        if debug_level >= log_level:
+            if _using_default_logger(logger):
+                logger.debug(msg, center, symbol, *args, **kwargs)
+            else:
+                logger.debug(msg, *args, **kwargs)
+
+
+def log_info(msg, center: bool = False, symbol: str = "*", *args, **kwargs):
+    global logger
+    if _using_default_logger(logger):
+        logger.info(msg, center, symbol, *args, **kwargs)
+    else:
+        logger.info(msg, *args, **kwargs)
+
+
+def log_warning(msg, *args, **kwargs):
+    global logger
+    logger.warning(msg, *args, **kwargs)
+
+
+def log_error(msg, *args, **kwargs):
+    global logger
+    logger.error(msg, *args, **kwargs)
+
+
+def log_exception(msg, *args, **kwargs):
+    global logger
+    logger.exception(msg, *args, **kwargs)
+
+
+def configure_agno_logging(
+    custom_default_logger: Optional[Any] = None,
+    custom_agent_logger: Optional[Any] = None,
+    custom_team_logger: Optional[Any] = None,
+    custom_workflow_logger: Optional[Any] = None,
+) -> None:
+    """
+    Util to set custom loggers. These will be used everywhere across the Agno library.
+
+    Args:
+        custom_default_logger: Default logger to use (overrides agent_logger for default)
+        custom_agent_logger: Custom logger for agent operations
+        custom_team_logger: Custom logger for team operations
+        custom_workflow_logger: Custom logger for workflow operations
+    """
+    if custom_default_logger is not None:
+        global logger
+        logger = custom_default_logger
+
+    if custom_agent_logger is not None:
+        global agent_logger
+        agent_logger = custom_agent_logger
+
+    if custom_team_logger is not None:
+        global team_logger
+        team_logger = custom_team_logger
+
+    if custom_workflow_logger is not None:
+        global workflow_logger
+        workflow_logger = custom_workflow_logger

agno/utils/mcp.py

@@ -0,0 +1,214 @@
+import json
+from functools import partial
+from uuid import uuid4
+
+from agno.utils.log import log_debug, log_exception
+
+try:
+    from mcp import ClientSession
+    from mcp.types import CallToolResult, EmbeddedResource, ImageContent, TextContent
+    from mcp.types import Tool as MCPTool
+except (ImportError, ModuleNotFoundError):
+    raise ImportError("`mcp` not installed. Please install using `pip install mcp`")
+
+
+from agno.media import Image
+from agno.tools.function import ToolResult
+
+
+def get_entrypoint_for_tool(tool: MCPTool, session: ClientSession):
+    """
+    Return an entrypoint for an MCP tool.
+
+    Args:
+        tool: The MCP tool to create an entrypoint for
+        session: The session to use
+
+    Returns:
+        Callable: The entrypoint function for the tool
+    """
+
+    async def call_tool(tool_name: str, **kwargs) -> ToolResult:
+        try:
+            await session.send_ping()
+        except Exception as e:
+            log_exception(e)
+
+        try:
+            log_debug(f"Calling MCP Tool '{tool_name}' with args: {kwargs}")
+            result: CallToolResult = await session.call_tool(tool_name, kwargs)  # type: ignore
+
+            # Return an error if the tool call failed
+            if result.isError:
+                return ToolResult(content=f"Error from MCP tool '{tool_name}': {result.content}")
+
+            # Process the result content
+            response_str = ""
+            images = []
+
+            for content_item in result.content:
+                if isinstance(content_item, TextContent):
+                    text_content = content_item.text
+
+                    # Parse as JSON to check for custom image format
+                    try:
+                        parsed_json = json.loads(text_content)
+                        if (
+                            isinstance(parsed_json, dict)
+                            and parsed_json.get("type") == "image"
+                            and "data" in parsed_json
+                        ):
+                            log_debug("Found custom JSON image format in TextContent")
+
+                            # Extract image data
+                            image_data = parsed_json.get("data")
+                            mime_type = parsed_json.get("mimeType", "image/png")
+
+                            if image_data and isinstance(image_data, str):
+                                import base64
+
+                                try:
+                                    image_bytes = base64.b64decode(image_data)
+                                except Exception as e:
+                                    log_debug(f"Failed to decode base64 image data: {e}")
+                                    image_bytes = None
+
+                                if image_bytes:
+                                    img_artifact = Image(
+                                        id=str(uuid4()),
+                                        url=None,
+                                        content=image_bytes,
+                                        mime_type=mime_type,
+                                    )
+                                    images.append(img_artifact)
+                                    response_str += "Image has been generated and added to the response.\n"
+                                    continue
+
+                    except (json.JSONDecodeError, TypeError):
+                        pass
+
+                    response_str += text_content + "\n"
+
+                elif isinstance(content_item, ImageContent):
+                    # Handle standard MCP ImageContent
+                    image_data = getattr(content_item, "data", None)
+
+                    if image_data and isinstance(image_data, str):
+                        import base64
+
+                        try:
+                            image_data = base64.b64decode(image_data)
+                        except Exception as e:
+                            log_debug(f"Failed to decode base64 image data: {e}")
+                            image_data = None
+
+                    img_artifact = Image(
+                        id=str(uuid4()),
+                        url=getattr(content_item, "url", None),
+                        content=image_data,
+                        mime_type=getattr(content_item, "mimeType", "image/png"),
+                    )
+                    images.append(img_artifact)
+                    response_str += "Image has been generated and added to the response.\n"
+                elif isinstance(content_item, EmbeddedResource):
+                    # Handle embedded resources
+                    response_str += f"[Embedded resource: {content_item.resource.model_dump_json()}]\n"
+                else:
+                    # Handle other content types
+                    response_str += f"[Unsupported content type: {content_item.type}]\n"
+
+            return ToolResult(
+                content=response_str.strip(),
+                images=images if images else None,
+            )
+        except Exception as e:
+            log_exception(f"Failed to call MCP tool '{tool_name}': {e}")
+            return ToolResult(content=f"Error: {e}")
+
+    return partial(call_tool, tool_name=tool.name)
+
+
+def prepare_command(command: str) -> list[str]:
+    """Sanitize a command and split it into parts before using it to run a MCP server."""
+    import os
+    import shutil
+    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
+        "python",
+        "python3",
+        "uv",
+        "uvx",
+        "pipx",
+        # Node
+        "node",
+        "npm",
+        "npx",
+        "yarn",
+        "pnpm",
+        "bun",
+        # Other runtimes
+        "deno",
+        "java",
+        "ruby",
+        "docker",
+    }
+
+    executable = parts[0].split("/")[-1]
+
+    # Check if it's a relative path starting with ./ or ../
+    if executable.startswith("./") or executable.startswith("../"):
+        # Allow relative paths to binaries
+        return parts
+
+    # Check if it's an absolute path to a binary
+    if executable.startswith("/") and os.path.isfile(executable):
+        # Allow absolute paths to existing files
+        return parts
+
+    # Check if it's a binary in current directory without ./
+    if "/" not in executable and os.path.isfile(executable):
+        # Allow binaries in current directory
+        return parts
+
+    # Check if it's a binary in PATH
+    if shutil.which(executable):
+        return parts
+
+    if executable not in ALLOWED_COMMANDS:
+        raise ValueError(f"MCP command needs to use one of the following executables: {ALLOWED_COMMANDS}")
+
+    first_part = parts[0]
+    executable = first_part.split("/")[-1]
+
+    # Allow known commands
+    if executable in ALLOWED_COMMANDS:
+        return parts
+
+    # Allow relative paths to custom binaries
+    if first_part.startswith(("./", "../")):
+        return parts
+
+    # Allow absolute paths to existing files
+    if first_part.startswith("/") and os.path.isfile(first_part):
+        return parts
+
+    # Allow binaries in current directory without ./
+    if "/" not in first_part and os.path.isfile(first_part):
+        return parts
+
+    # Allow binaries in PATH
+    if shutil.which(first_part):
+        return parts
+
+    raise ValueError(f"MCP command needs to use one of the following executables: {ALLOWED_COMMANDS}")