File size: 5,552 Bytes
ed4d993
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
import asyncio
from functools import partial
from typing import Any, Dict, List, Optional, Type

from langchain_core.callbacks.manager import (
    AsyncCallbackManagerForToolRun,
    CallbackManagerForToolRun,
)
from langchain_core.pydantic_v1 import BaseModel, Field, create_model, root_validator
from langchain_core.tools import BaseTool

from langchain_community.tools.connery.models import Action, Parameter


class ConneryAction(BaseTool):
    """Connery Action tool."""

    name: str
    description: str
    args_schema: Type[BaseModel]

    action: Action
    connery_service: Any

    def _run(
        self,
        run_manager: Optional[CallbackManagerForToolRun] = None,
        **kwargs: Dict[str, str],
    ) -> Dict[str, str]:
        """
        Runs the Connery Action with the provided input.
        Parameters:
            kwargs (Dict[str, str]): The input dictionary expected by the action.
        Returns:
            Dict[str, str]: The output of the action.
        """

        return self.connery_service.run_action(self.action.id, kwargs)

    async def _arun(
        self,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
        **kwargs: Dict[str, str],
    ) -> Dict[str, str]:
        """
        Runs the Connery Action asynchronously with the provided input.
        Parameters:
            kwargs (Dict[str, str]): The input dictionary expected by the action.
        Returns:
            Dict[str, str]: The output of the action.
        """

        func = partial(self._run, **kwargs)
        return await asyncio.get_event_loop().run_in_executor(None, func)

    def get_schema_json(self) -> str:
        """
        Returns the JSON representation of the Connery Action Tool schema.
        This is useful for debugging.
        Returns:
            str: The JSON representation of the Connery Action Tool schema.
        """

        return self.args_schema.schema_json(indent=2)

    @root_validator()
    def validate_attributes(cls, values: dict) -> dict:
        """
        Validate the attributes of the ConneryAction class.
        Parameters:
            values (dict): The arguments to validate.
        Returns:
            dict: The validated arguments.
        """

        # Import ConneryService here and check if it is an instance
        # of ConneryService to avoid circular imports
        from .service import ConneryService

        if not isinstance(values.get("connery_service"), ConneryService):
            raise ValueError(
                "The attribute 'connery_service' must be an instance of ConneryService."
            )

        if not values.get("name"):
            raise ValueError("The attribute 'name' must be set.")
        if not values.get("description"):
            raise ValueError("The attribute 'description' must be set.")
        if not values.get("args_schema"):
            raise ValueError("The attribute 'args_schema' must be set.")
        if not values.get("action"):
            raise ValueError("The attribute 'action' must be set.")
        if not values.get("connery_service"):
            raise ValueError("The attribute 'connery_service' must be set.")

        return values

    @classmethod
    def create_instance(cls, action: Action, connery_service: Any) -> "ConneryAction":
        """
        Creates a Connery Action Tool from a Connery Action.
        Parameters:
            action (Action): The Connery Action to wrap in a Connery Action Tool.
            connery_service (ConneryService): The Connery Service
            to run the Connery Action. We use Any here to avoid circular imports.
        Returns:
            ConneryAction: The Connery Action Tool.
        """

        # Import ConneryService here and check if it is an instance
        # of ConneryService to avoid circular imports
        from .service import ConneryService

        if not isinstance(connery_service, ConneryService):
            raise ValueError(
                "The connery_service must be an instance of ConneryService."
            )

        input_schema = cls._create_input_schema(action.inputParameters)
        description = action.title + (
            ": " + action.description if action.description else ""
        )

        instance = cls(
            name=action.id,
            description=description,
            args_schema=input_schema,
            action=action,
            connery_service=connery_service,
        )

        return instance

    @classmethod
    def _create_input_schema(cls, inputParameters: List[Parameter]) -> Type[BaseModel]:
        """
        Creates an input schema for a Connery Action Tool
        based on the input parameters of the Connery Action.
        Parameters:
            inputParameters: List of input parameters of the Connery Action.
        Returns:
            Type[BaseModel]: The input schema for the Connery Action Tool.
        """

        dynamic_input_fields: Dict[str, Any] = {}

        for param in inputParameters:
            default = ... if param.validation and param.validation.required else None
            title = param.title
            description = param.title + (
                ": " + param.description if param.description else ""
            )
            type = param.type

            dynamic_input_fields[param.key] = (
                str,
                Field(default, title=title, description=description, type=type),
            )

        InputModel = create_model("InputSchema", **dynamic_input_fields)
        return InputModel