"""Agent tool utilities"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/04_pipelines.agent.tools.ipynb.

# %% auto 0
__all__ = ['DEFAULT_TOOLS', 'FunctionTool', 'fromdocs', 'createtool', 'VectorStoreTool']

# %% ../../../nbs/04_pipelines.agent.tools.ipynb 3
from transformers.utils import chat_template_utils, TypeHintParsingException

from smolagents import Tool
from smolagents import PythonInterpreterTool, Tool, tool as CreateTool, VisitWebpageTool, WebSearchTool
import inspect
from types import FunctionType, MethodType


DEFAULT_TOOLS = {
    "python": PythonInterpreterTool(), 
    "websearch": WebSearchTool(), 
    "webview": VisitWebpageTool()
}


class FunctionTool(Tool):
    """
    Creates a FunctionTool. A FunctionTool takes descriptive configuration and injects it along with a target function
    into an LLM prompt.
    """

    # pylint: disable=W0231
    def __init__(self, config):
        """
        Creates a FunctionTool.

        Args:
            config: `name`, `description`, `inputs`, `output` and `target` configuration
        """

        # Tool parameters
        self.name = config["name"]
        self.description = config["description"]
        self.inputs = config["inputs"]
        self.output_type = config.get("output", config.get("output_type", "any"))
        self.target = config["target"]

        # Skip forward signature validation
        self.skip_forward_signature_validation = True

        # Validate parameters and initialize tool
        super().__init__()

    def forward(self, *args, **kwargs):
        """
        Runs target function.

        Args:
            args: positional args
            kwargs: keyword args

        Returns:
            result
        """

        return self.target(*args, **kwargs)

    
def fromdocs(target, config):
    """
    Creates a tool from method documentation.

    Args:
        target: target object or function
        config: tool configuration

    Returns:
        Tool
    """

    # Get function name and target - use target if it's a function or method, else use target.__call__
    name = target.__name__ if isinstance(target, (FunctionType, MethodType)) or not hasattr(target, "__call__") else target.__class__.__name__
    target = target if isinstance(target, (FunctionType, MethodType)) or not hasattr(target, "__call__") else target.__call__

    # Extract target documentation
    doc = inspect.getdoc(target)
    description, parameters, _ = chat_template_utils.parse_google_format_docstring(doc.strip()) if doc else (None, {}, None)

    # Get list of required parameters
    signature = inspect.signature(target)
    inputs = {}
    for pname, param in signature.parameters.items():
        if param.default == inspect.Parameter.empty and pname in parameters:
            inputs[pname] = {"type": "any", "description": parameters[pname]}

    # Create function tool
    return FunctionTool(
        {
            "name": config.get("name", name.lower()),
            "description": config.get("description", description),
            "inputs": config.get("inputs", inputs),
            "target": config.get("target", target),
        }
    )

# %% ../../../nbs/04_pipelines.agent.tools.ipynb 4
def createtool(target, config=None):
    """
    Creates a new Tool.

    Args:
        target: target object or function
        config: optional tool configuration

    Returns:
        Tool
    """

    try:
        # Try to create using CreateTool function - this fails when no annotations are available
        return CreateTool(target)
    except (TypeHintParsingException, TypeError):
        return ToolFactory.fromdocs(target, config if config else {})

# %% ../../../nbs/04_pipelines.agent.tools.ipynb 5
from ...ingest.stores.base import VectorStore

class VectorStoreTool(Tool):
    """
    Tool to execute an VectorStore search.
    """

    def __init__(self, name:str, description:str, store:VectorStore):
        """
        Creates a new VectorStoreTool.

        Args:
            name: Name for vectorstore
            description: Description of what the vectostore contains.
            store: A VectorStore instance
        """

        # Tool parameters
        self.name = name
        self.description = f"""{description}. Results are returned as a list of dict elements. Each result has keys 'id', 'text', 'score'."""

        # Input and output descriptions
        self.inputs = {"query": {"type": "string", "description": "The search query to perform."}}
        self.output_type = "any"

        # Load embeddings instance
        self.store = store

        # Validate parameters and initialize tool
        super().__init__()

    def forward(self, query):
        """
        Runs a search.

        Args:
            query: input query

        Returns:
            search results
        """
        # Handle case where LLM passes input schema instead of string value
        if isinstance(query, dict) and 'type' in query and 'description' in query:
            query = query['description']
        results = self.store.semantic_search(query, 5)
        return [{'id': r.id, 'score': r.metadata['score'], 'text': r.page_content} for r in results]
