llama-index

1
from typing import Any, Dict
2

3
import requests
4

5
from llama_index.legacy.core.base_query_engine import BaseQueryEngine
6
from llama_index.legacy.core.response.schema import Response
7
from llama_index.legacy.schema import QueryBundle
8

9

10
class CogniswitchQueryEngine(BaseQueryEngine):
11
    def __init__(self, cs_token: str, OAI_token: str, apiKey: str) -> None:
12
        """The required fields.
13

14
        Args:
15
            cs_token (str): Cogniswitch token.
16
            OAI_token (str): OpenAI token.
17
            apiKey (str): Oauth token.
18
        """
19
        self.cs_token = cs_token
20
        self.OAI_token = OAI_token
21
        self.apiKey = apiKey
22
        self.knowledge_request_endpoint = (
23
            "https://api.cogniswitch.ai:8243/cs-api/0.0.1/cs/knowledgeRequest"
24
        )
25
        self.headers = {
26
            "apiKey": self.apiKey,
27
            "platformToken": self.cs_token,
28
            "openAIToken": self.OAI_token,
29
        }
30

31
    def query_knowledge(self, query: str) -> Response:
32
        """
33
        Send a query to the Cogniswitch service and retrieve the response.
34

35
        Args:
36
            query (str): Query to be answered.
37

38
        Returns:
39
            dict: Response JSON from the Cogniswitch service.
40
        """
41
        data = {"query": query}
42
        response = requests.post(
43
            self.knowledge_request_endpoint,
44
            headers=self.headers,
45
            verify=False,
46
            data=data,
47
        )
48
        if response.status_code == 200:
49
            resp = response.json()
50
            answer = resp["data"]["answer"]
51

52
            return Response(response=answer)
53
        else:
54
            error_message = response.json()["message"]
55
            return Response(response=error_message)
56

57
    def _query(self, query_bundle: QueryBundle) -> Response:
58
        return self.query_knowledge(query_bundle.query_str)
59

60
    async def _aquery(self, query_bundle: QueryBundle) -> Response:
61
        return self.query_knowledge(query_bundle.query_str)
62

63
    def _get_prompt_modules(self) -> Dict[str, Any]:
64
        """Get prompts."""
65
        return {}
66