codekingpro/portable-devtools

1import asyncio2from functools import partial3from typing import Any, Dict, List, Optional, Type4 5from langchain_core.callbacks.manager import (6    AsyncCallbackManagerForToolRun,7    CallbackManagerForToolRun,8)9from langchain_core.tools import BaseTool10from pydantic import BaseModel, Field, create_model, model_validator11 12from langchain_community.tools.connery.models import Action, Parameter13 14 15class ConneryAction(BaseTool):16    """Connery Action tool."""17 18    name: str19    description: str20    args_schema: Type[BaseModel]21 22    action: Action23    connery_service: Any24 25    def _run(26        self,27        run_manager: Optional[CallbackManagerForToolRun] = None,28        **kwargs: Any,29    ) -> Dict[str, str]:30        """31        Runs the Connery Action with the provided input.32        Parameters:33            kwargs (Dict[str, str]): The input dictionary expected by the action.34        Returns:35            Dict[str, str]: The output of the action.36        """37 38        return self.connery_service.run_action(self.action.id, kwargs)39 40    async def _arun(41        self,42        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,43        **kwargs: Any,44    ) -> Dict[str, str]:45        """46        Runs the Connery Action asynchronously with the provided input.47        Parameters:48            kwargs (Dict[str, str]): The input dictionary expected by the action.49        Returns:50            Dict[str, str]: The output of the action.51        """52 53        func = partial(self._run, **kwargs)54        return await asyncio.get_event_loop().run_in_executor(None, func)55 56    def get_schema_json(self) -> str:57        """58        Returns the JSON representation of the Connery Action Tool schema.59        This is useful for debugging.60        Returns:61            str: The JSON representation of the Connery Action Tool schema.62        """63 64        return self.args_schema.schema_json(indent=2)65 66    @model_validator(mode="before")67    @classmethod68    def validate_attributes(cls, values: dict) -> Any:69        """70        Validate the attributes of the ConneryAction class.71        Parameters:72            values (dict): The arguments to validate.73        Returns:74            dict: The validated arguments.75        """76 77        # Import ConneryService here and check if it is an instance78        # of ConneryService to avoid circular imports79        from .service import ConneryService80 81        if not isinstance(values.get("connery_service"), ConneryService):82            raise ValueError(83                "The attribute 'connery_service' must be an instance of ConneryService."84            )85 86        if not values.get("name"):87            raise ValueError("The attribute 'name' must be set.")88        if not values.get("description"):89            raise ValueError("The attribute 'description' must be set.")90        if not values.get("args_schema"):91            raise ValueError("The attribute 'args_schema' must be set.")92        if not values.get("action"):93            raise ValueError("The attribute 'action' must be set.")94        if not values.get("connery_service"):95            raise ValueError("The attribute 'connery_service' must be set.")96 97        return values98 99    @classmethod100    def create_instance(cls, action: Action, connery_service: Any) -> "ConneryAction":101        """102        Creates a Connery Action Tool from a Connery Action.103        Parameters:104            action (Action): The Connery Action to wrap in a Connery Action Tool.105            connery_service (ConneryService): The Connery Service106            to run the Connery Action. We use Any here to avoid circular imports.107        Returns:108            ConneryAction: The Connery Action Tool.109        """110 111        # Import ConneryService here and check if it is an instance112        # of ConneryService to avoid circular imports113        from .service import ConneryService114 115        if not isinstance(connery_service, ConneryService):116            raise ValueError(117                "The connery_service must be an instance of ConneryService."118            )119 120        input_schema = cls._create_input_schema(action.inputParameters)121        description = action.title + (122            ": " + action.description if action.description else ""123        )124 125        instance = cls(126            name=action.id,127            description=description,128            args_schema=input_schema,129            action=action,130            connery_service=connery_service,131        )132 133        return instance134 135    @classmethod136    def _create_input_schema(cls, inputParameters: List[Parameter]) -> Type[BaseModel]:137        """138        Creates an input schema for a Connery Action Tool139        based on the input parameters of the Connery Action.140        Parameters:141            inputParameters: List of input parameters of the Connery Action.142        Returns:143            Type[BaseModel]: The input schema for the Connery Action Tool.144        """145 146        dynamic_input_fields: Dict[str, Any] = {}147 148        for param in inputParameters:149            default = ... if param.validation and param.validation.required else None150            title = param.title151            description = param.title + (152                ": " + param.description if param.description else ""153            )154            type = param.type155 156            dynamic_input_fields[param.key] = (157                type,158                Field(default, title=title, description=description),159            )160 161        InputModel = create_model("InputSchema", **dynamic_input_fields)162        return InputModel163