File size: 13,896 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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
from __future__ import annotations

from typing import Any, Dict, Optional

import requests
from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_core.tools import BaseTool


class CogniswitchKnowledgeRequest(BaseTool):
    """Tool that uses the Cogniswitch service to answer questions.

    name: str = "cogniswitch_knowledge_request"
    description: str = (
        "A wrapper around cogniswitch service to answer the question
        from the knowledge base."
        "Input should be a search query."
    )
    """

    name: str = "cogniswitch_knowledge_request"
    description: str = """A wrapper around cogniswitch service to 
    answer the question from the knowledge base."""
    cs_token: str
    OAI_token: str
    apiKey: str
    api_url = "https://api.cogniswitch.ai:8243/cs-api/0.0.1/cs/knowledgeRequest"

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> Dict[str, Any]:
        """
        Use the tool to answer a query.

        Args:
            query (str): Natural language query,
              that you would like to ask to your knowledge graph.
            run_manager (Optional[CallbackManagerForChainRun]):
            Manager for chain run callbacks.

        Returns:
            Dict[str, Any]: Output dictionary containing
            the 'response' from the service.
        """
        response = self.answer_cs(self.cs_token, self.OAI_token, query, self.apiKey)
        return response

    def answer_cs(self, cs_token: str, OAI_token: str, query: str, apiKey: str) -> dict:
        """
        Send a query to the Cogniswitch service and retrieve the response.

        Args:
            cs_token (str): Cogniswitch token.
            OAI_token (str): OpenAI token.
            apiKey (str): OAuth token.
            query (str): Query to be answered.

        Returns:
            dict: Response JSON from the Cogniswitch service.
        """
        if not cs_token:
            raise ValueError("Missing cs_token")
        if not OAI_token:
            raise ValueError("Missing OpenAI token")
        if not apiKey:
            raise ValueError("Missing cogniswitch OAuth token")
        if not query:
            raise ValueError("Missing input query")

        headers = {
            "apiKey": apiKey,
            "platformToken": cs_token,
            "openAIToken": OAI_token,
        }

        data = {"query": query}
        response = requests.post(self.api_url, headers=headers, verify=False, data=data)
        return response.json()


class CogniswitchKnowledgeStatus(BaseTool):
    """Tool that uses the Cogniswitch services to get the
     status of the document or url uploaded.

    name: str = "cogniswitch_knowledge_status"
    description: str = (
        "A wrapper around cogniswitch services to know the status of
         the document uploaded from a url or a file. "
        "Input should be a file name or the url link"
    )
    """

    name: str = "cogniswitch_knowledge_status"
    description: str = """A wrapper around cogniswitch services to know 
    the status of the document uploaded from a url or a file."""
    cs_token: str
    OAI_token: str
    apiKey: str
    knowledge_status_url = (
        "https://api.cogniswitch.ai:8243/cs-api/0.0.1/cs/knowledgeSource/status"
    )

    def _run(
        self,
        document_name: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> Dict[str, Any]:
        """
        Use the tool to know the status of the document uploaded.

        Args:
            document_name (str): name of the document or
            the url uploaded
            run_manager (Optional[CallbackManagerForChainRun]):
            Manager for chain run callbacks.

        Returns:
            Dict[str, Any]: Output dictionary containing
            the 'response' from the service.
        """
        response = self.knowledge_status(document_name)
        return response

    def knowledge_status(self, document_name: str) -> dict:
        """
        Use this function to know the status of the document or the URL uploaded
        Args:
            document_name (str): The document name or the url that is uploaded.

        Returns:
            dict: Response JSON from the Cogniswitch service.
        """

        params = {"docName": document_name, "platformToken": self.cs_token}
        headers = {
            "apiKey": self.apiKey,
            "openAIToken": self.OAI_token,
            "platformToken": self.cs_token,
        }
        response = requests.get(
            self.knowledge_status_url,
            headers=headers,
            params=params,
            verify=False,
        )
        if response.status_code == 200:
            source_info = response.json()
            source_data = dict(source_info[-1])
            status = source_data.get("status")
            if status == 0:
                source_data["status"] = "SUCCESS"
            elif status == 1:
                source_data["status"] = "PROCESSING"
            elif status == 2:
                source_data["status"] = "UPLOADED"
            elif status == 3:
                source_data["status"] = "FAILURE"
            elif status == 4:
                source_data["status"] = "UPLOAD_FAILURE"
            elif status == 5:
                source_data["status"] = "REJECTED"

            if "filePath" in source_data.keys():
                source_data.pop("filePath")
            if "savedFileName" in source_data.keys():
                source_data.pop("savedFileName")
            if "integrationConfigId" in source_data.keys():
                source_data.pop("integrationConfigId")
            if "metaData" in source_data.keys():
                source_data.pop("metaData")
            if "docEntryId" in source_data.keys():
                source_data.pop("docEntryId")
            return source_data
        else:
            # error_message = response.json()["message"]
            return {
                "message": response.status_code,
            }


class CogniswitchKnowledgeSourceFile(BaseTool):
    """Tool that uses the Cogniswitch services to store data from file.

    name: str = "cogniswitch_knowledge_source_file"
    description: str = (
        "This calls the CogniSwitch services to analyze & store data from a file.
        If the input looks like a file path, assign that string value to file key.
        Assign document name & description only if provided in input."
    )
    """

    name: str = "cogniswitch_knowledge_source_file"
    description: str = """
        This calls the CogniSwitch services to analyze & store data from a file. 
        If the input looks like a file path, assign that string value to file key. 
        Assign document name & description only if provided in input.
        """
    cs_token: str
    OAI_token: str
    apiKey: str
    knowledgesource_file = (
        "https://api.cogniswitch.ai:8243/cs-api/0.0.1/cs/knowledgeSource/file"
    )

    def _run(
        self,
        file: Optional[str] = None,
        document_name: Optional[str] = None,
        document_description: Optional[str] = None,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> Dict[str, Any]:
        """
        Execute the tool to store the data given from a file.
        This calls the CogniSwitch services to analyze & store data from a file.
        If the input looks like a file path, assign that string value to file key.
        Assign document name & description only if provided in input.

        Args:
            file Optional[str]: The file path of your knowledge
            document_name Optional[str]: Name of your knowledge document
            document_description Optional[str]: Description of your knowledge document
            run_manager (Optional[CallbackManagerForChainRun]):
            Manager for chain run callbacks.

        Returns:
            Dict[str, Any]: Output dictionary containing
            the 'response' from the service.
        """
        if not file:
            return {
                "message": "No input provided",
            }
        else:
            response = self.store_data(
                file=file,
                document_name=document_name,
                document_description=document_description,
            )
            return response

    def store_data(
        self,
        file: Optional[str],
        document_name: Optional[str],
        document_description: Optional[str],
    ) -> dict:
        """
        Store data using the Cogniswitch service.
        This calls the CogniSwitch services to analyze & store data from a file.
        If the input looks like a file path, assign that string value to file key.
        Assign document name & description only if provided in input.

        Args:
            file (Optional[str]): file path of your file.
            the current files supported by the files are
            .txt, .pdf, .docx, .doc, .html
            document_name (Optional[str]): Name of the document you are uploading.
            document_description (Optional[str]): Description of the document.

        Returns:
            dict: Response JSON from the Cogniswitch service.
        """
        headers = {
            "apiKey": self.apiKey,
            "openAIToken": self.OAI_token,
            "platformToken": self.cs_token,
        }
        data: Dict[str, Any]
        if not document_name:
            document_name = ""
        if not document_description:
            document_description = ""

        if file is not None:
            files = {"file": open(file, "rb")}

        data = {
            "documentName": document_name,
            "documentDescription": document_description,
        }
        response = requests.post(
            self.knowledgesource_file,
            headers=headers,
            verify=False,
            data=data,
            files=files,
        )
        if response.status_code == 200:
            return response.json()
        else:
            return {"message": "Bad Request"}


class CogniswitchKnowledgeSourceURL(BaseTool):
    """Tool that uses the Cogniswitch services to store data from a URL.

    name: str = "cogniswitch_knowledge_source_url"
    description: str = (
        "This calls the CogniSwitch services to analyze & store data from a url.
        the URL is provided in input, assign that value to the url key.
        Assign document name & description only if provided in input"
    )
    """

    name: str = "cogniswitch_knowledge_source_url"
    description: str = """
    This calls the CogniSwitch services to analyze & store data from a url. 
        the URL is provided in input, assign that value to the url key. 
        Assign document name & description only if provided in input"""
    cs_token: str
    OAI_token: str
    apiKey: str
    knowledgesource_url = (
        "https://api.cogniswitch.ai:8243/cs-api/0.0.1/cs/knowledgeSource/url"
    )

    def _run(
        self,
        url: Optional[str] = None,
        document_name: Optional[str] = None,
        document_description: Optional[str] = None,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> Dict[str, Any]:
        """
        Execute the tool to store the data given from a url.
        This calls the CogniSwitch services to analyze & store data from a url.
        the URL is provided in input, assign that value to the url key.
        Assign document name & description only if provided in input.

        Args:
            url Optional[str]: The website/url link of your knowledge
            document_name Optional[str]: Name of your knowledge document
            document_description Optional[str]: Description of your knowledge document
            run_manager (Optional[CallbackManagerForChainRun]):
            Manager for chain run callbacks.

        Returns:
            Dict[str, Any]: Output dictionary containing
            the 'response' from the service.
        """
        if not url:
            return {
                "message": "No input provided",
            }
        response = self.store_data(
            url=url,
            document_name=document_name,
            document_description=document_description,
        )
        return response

    def store_data(
        self,
        url: Optional[str],
        document_name: Optional[str],
        document_description: Optional[str],
    ) -> dict:
        """
        Store data using the Cogniswitch service.
        This calls the CogniSwitch services to analyze & store data from a url.
        the URL is provided in input, assign that value to the url key.
        Assign document name & description only if provided in input.

        Args:
            url (Optional[str]): URL link.
            document_name (Optional[str]): Name of the document you are uploading.
            document_description (Optional[str]): Description of the document.

        Returns:
            dict: Response JSON from the Cogniswitch service.
        """
        headers = {
            "apiKey": self.apiKey,
            "openAIToken": self.OAI_token,
            "platformToken": self.cs_token,
        }
        data: Dict[str, Any]
        if not document_name:
            document_name = ""
        if not document_description:
            document_description = ""
        if not url:
            return {
                "message": "No input provided",
            }
        else:
            data = {"url": url}
            response = requests.post(
                self.knowledgesource_url,
                headers=headers,
                verify=False,
                data=data,
            )
        if response.status_code == 200:
            return response.json()
        else:
            return {"message": "Bad Request"}