Source code for aidputils.agents.tools.base_tool

from abc import ABC, abstractmethod
from typing import final
from aidputils.agents.toolkit import chat_context

import oci
import logging

logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )



class BaseTool(ABC):
    tool_class_registry = {}
    tool_function_registry = {}



    @classmethod
    @final
    def register(cls, subclass):
        """
        Decorator for automatic tool registration.

        Usage:
            @BaseTool.register
            class HttpTool(BaseTool):
                ...

        This will automatically add HttpTool to the global tool registry,
        allowing it to be discovered and invoked without explicit manual registration.

        Implementors should use this decorator on any new tool class they want
        to make available for agent usage.
        """
        cls.tool_class_registry[subclass.__name__] = subclass
        return subclass

    


    @classmethod
    def logger(cls, tool_name=None):
        """
        Returns a logger for the given tool name (or class).
        
        Args:
            tool_name (str): Optional. If provided, used as logger name and for caching.
                             If not provided, uses the class name.
        """
        # Ensure _loggers cache exists
        if not hasattr(cls, "_loggers"):
            cls._loggers = {}

        name = tool_name if tool_name else cls.__name__

        # If a logger for this name doesn't exist, create and cache it.
        if name not in cls._loggers:
            new_logger = logging.getLogger(name)
            cls._loggers[name] = new_logger

        return cls._loggers[name]

    


    @classmethod
    @final
    def invoke(cls, conf, runtime_params, **context_vars) -> dict:
        token = chat_context.tool_context_var.set(cls.__name__)
        try:
            response = cls._invoke_tool(conf, runtime_params, **context_vars)
            mcp_format_response = cls._format_output_as_mcp(response)
            return mcp_format_response
        finally:
            if token:
                chat_context.tool_context_var.reset(token)




    @classmethod
    def make_security_token_signer(cls, oci_config):
        pk = oci.signer.load_private_key_from_file(oci_config.get("key_file"), None)
        with open(oci_config.get("security_token_file")) as f:
            st_string = f.read()
        return oci.auth.signers.SecurityTokenSigner(st_string, pk)


    @classmethod
    @abstractmethod
    def _invoke_tool(cls, conf, runtime_params, **context_vars) -> dict:
        raise NotImplementedError("Each tool must implement its own _invoke_tool method.")

    @classmethod
    @abstractmethod
    def _format_output_as_mcp(cls, response):
        raise NotImplementedError("Each tool must implement its own format_output_as_mcp method.")

    # ========== MCP reference methods (for future MCP compliance) ==========



    @classmethod
    @final
    def register_tool_function(cls, conf):
        """
        Reference-only: For future MCP compliance, not for immediate agent or tool use.
        """
        cls.tool_function_registry[conf["name"]] = conf




    @classmethod
    @final
    def list_tools(cls):
        """
        Reference-only: Returns all MCP tool schemas in the local registry.
        """
        tools_schema = [cls._generate_mcp_schema(conf) for tool_name, conf in cls.tool_class_registry.items()]
        return {"tools": tools_schema}


    @classmethod
    @final
    def _tool_schema(cls, tool_name):
        """
        Reference-only: Returns MCP tool schema for a given registered tool.
        """
        return cls._generate_mcp_schema(cls.tool_class_registry[tool_name])

    @staticmethod
    def _generate_mcp_schema(conf):
        """
        Reference-only: Returns an MCP-compliant schema using the param specs found in conf["params"].
        Expects conf = {..., "params": {param_name: {"name": str, "description": str, "type": str}}}
        """
        params = conf.get("params", {})
        schema_properties = {}
        required = []

        for param_key, param_info in params.items():
            name = param_info.get("name", param_key)
            schema_properties[name] = {
                "type": param_info.get("type", "string"),
                "description": param_info.get("description", "")
            }
            required.append(name)

        input_schema = {
            "type": "object",
            "properties": schema_properties,
            "required": required
        }

        return {
            "description": conf.get("description", ""),
            "name": conf.get("name", ""),
            "inputSchema": input_schema
        }

    # ========== End MCP reference methods ==========