llama-index

1
"""ChatGPT Plugin."""
2

3
import os
4
from typing import Any, List, Optional
5

6
import requests
7
from requests.adapters import HTTPAdapter, Retry
8

9
from llama_index.legacy.readers.base import BaseReader
10
from llama_index.legacy.schema import Document
11

12

13
class ChatGPTRetrievalPluginReader(BaseReader):
14
    """ChatGPT Retrieval Plugin reader."""
15

16
    def __init__(
17
        self,
18
        endpoint_url: str,
19
        bearer_token: Optional[str] = None,
20
        retries: Optional[Retry] = None,
21
        batch_size: int = 100,
22
    ) -> None:
23
        """Chatgpt Retrieval Plugin."""
24
        self._endpoint_url = endpoint_url
25
        self._bearer_token = bearer_token or os.getenv("BEARER_TOKEN")
26
        self._retries = retries
27
        self._batch_size = batch_size
28

29
        self._s = requests.Session()
30
        self._s.mount("http://", HTTPAdapter(max_retries=self._retries))
31

32
    def load_data(
33
        self,
34
        query: str,
35
        top_k: int = 10,
36
        separate_documents: bool = True,
37
        **kwargs: Any,
38
    ) -> List[Document]:
39
        """Load data from ChatGPT Retrieval Plugin."""
40
        headers = {"Authorization": f"Bearer {self._bearer_token}"}
41
        queries = [{"query": query, "top_k": top_k}]
42
        res = requests.post(
43
            f"{self._endpoint_url}/query", headers=headers, json={"queries": queries}
44
        )
45
        documents: List[Document] = []
46
        for query_result in res.json()["results"]:
47
            for result in query_result["results"]:
48
                result_id = result["id"]
49
                result_txt = result["text"]
50
                result_embedding = result["embedding"]
51
                document = Document(
52
                    text=result_txt,
53
                    id_=result_id,
54
                    embedding=result_embedding,
55
                )
56
                documents.append(document)
57

58
            # NOTE: there should only be one query
59
            break
60

61
        if not separate_documents:
62
            text_list = [doc.get_content() for doc in documents]
63
            text = "\n\n".join(text_list)
64
            documents = [Document(text=text)]
65

66
        return documents
67