跳转至

Retrieval Module

erniebot_agent.retrieval

BaizhongSearch

A class for interacting with the Baizhong Search API.

Attributes:

Name Type Description
base_url str

The base URL for the AIStudio service.
access_token str

The access token for authentication.
knowledge_base_id int

The ID of the knowledge base being used (if applicable).
Source code in erniebot-agent/src/erniebot_agent/retrieval/baizhong_search.py 
 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
class BaizhongSearch:
    """
    A class for interacting with the Baizhong Search API.


    Attributes:
        base_url (str): The base URL for the AIStudio service.
        access_token (str): The access token for authentication.
        knowledge_base_id (int): The ID of the knowledge base being used (if applicable).
    """

    _AISTUDIO_BASE_URL: ClassVar[str] = "https://aistudio.baidu.com"

    def __init__(
        self,
        access_token: Optional[str] = None,
        knowledge_base_name: Optional[str] = None,
        knowledge_base_id: Optional[int] = None,
    ) -> None:
        """
        Initialize a BaizhongSearch object.

        Args:
            access_token (str): The access token for authentication.
            knowledge_base_name (Optional[str]): The name of the knowledge base to use (optional).
            knowledge_base_id (Optional[int]): The ID of an existing knowledge base to use (optional).

        Raises:
            BaizhongError: If neither knowledge_base_name nor knowledge_base_id is provided.

        """
        self._base_url = os.getenv("AISTUDIO_BASE_URL", self._AISTUDIO_BASE_URL)
        self.access_token: Optional[str] = None
        if access_token is not None:
            self.access_token = access_token
        elif C.get_global_access_token() is not None:
            self.access_token = C.get_global_access_token()
        else:
            raise BaizhongError(
                "Please ensure that either an access_token is provided or "
                "the EB_AGENT_ACCESS_TOKEN is set up as an environment variable."
            )
        if knowledge_base_id is not None:
            _logger.info(f"Loading existing project with `knowledge_base_id={knowledge_base_id}`")
            self.knowledge_base_id = knowledge_base_id
        elif knowledge_base_name is not None:
            self.knowledge_base_id = self.create_knowledge_base(knowledge_base_name)
        else:
            raise BaizhongError("You must provide either a `knowledge_base_name` or a `knowledge_base_id`.")

    def create_knowledge_base(self, knowledge_base_name: str):
        """
        Create a JSON payload with the provided knowledge base name.

        Args:
            knowledge_base_name (str): The knowledge base name.

        Returns:
            Dict[str, Any]: A dictionary containing knowledge base results.

        Raises:
            BaizhongError: If the API request fails, this exception is raised with details about the error.
        """
        json_data = {"knowledgeBaseName": knowledge_base_name}
        res = requests.post(
            f"{self._base_url}/llm/knowledge/create",
            json=json_data,
            headers=self._get_authorization_headers(access_token=self.access_token),
        )
        if res.status_code == 200:
            result = res.json()
            if result["errorCode"] != 0:
                raise BaizhongError(message=result["errorMsg"], error_code=result["errorCode"])
            return result["result"]["knowledgeBaseId"]
        else:
            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)

    def _get_authorization_headers(self, access_token: Optional[str]) -> dict:
        """
        Initialize a dictionary for HTTP headers with Content-Type set to application/json.

        Args:
            access_token (str): The AIStudio access_token.

        Returns:
            Dict[str, Any]: A dictionary containing HTTP headers information.
        """
        headers = {"Content-Type": "application/json"}
        if access_token is None:
            _logger.warning("access_token is NOT provided, this may cause 403 HTTP error..")
        else:
            headers["Authorization"] = f"token {access_token}"
        return headers

    def search(self, query: str, top_k: int = 10, filters: Optional[Dict[str, Any]] = None):
        """
        Perform a search using the Baizhong common search API.

        Args:
            query (str): The search query.
            top_k (int, optional): The number of top results to retrieve (default is 10).
            filters (Optional[Dict[str, Any]], optional): Additional filters to apply to the search query
            (default is None).

        Returns:
            List[Dict[str, Any]]: A list of dictionaries containing search results.

        Raises:
            BaizhongError: If the API request fails, this exception is raised with details about the error.
        """
        json_data = {"knowledgeBaseId": self.knowledge_base_id, "query": query, "topK": top_k}
        if filters is not None:
            filter_terms = [{"term": item} for item in filters]
            filterConditions = {"filterConditions": {"bool": {"filter": filter_terms}}}
            json_data.update(filterConditions)
        res = requests.post(
            f"{self._base_url}/llm/knowledge/search",
            json=json_data,
            headers=self._get_authorization_headers(access_token=self.access_token),
        )
        if res.status_code == 200:
            result = res.json()
            if result["errorCode"] != 0:
                raise BaizhongError(message=result["errorMsg"], error_code=result["errorCode"])
            list_data = []
            for item in result["result"]:
                doc = json.loads(item["source"]["doc"])
                list_data.append(
                    {
                        "id": item["fileId"],
                        "content": doc["content_se"],
                        "title": item["source"]["title"],
                        "score": item["score"],
                    }
                )
            return list_data
        else:
            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)

    def add_documents(self, documents: List[Document], batch_size: int = 1, thread_count: int = 1):
        """
        Add a batch of documents to the Baizhong system using multi-threading.
        Args:
            documents (List[Document]): A list of Document objects to be added.
            batch_size (int, optional): The size of each batch of documents (defaults to 1).
            thread_count (int, optional): The number of threads to use for concurrent document addition
            (defaults to 1).
        Returns:
            List[Union[None, Exception]]: A list of results from the document addition process.
        Note:
            This function uses multi-threading to improve the efficiency of adding a large number of
            documents.
        """
        raise NotImplementedError

__init__(access_token=None, knowledge_base_name=None, knowledge_base_id=None)

Initialize a BaizhongSearch object.

Parameters:

Name Type Description Default
access_token str

The access token for authentication.

 None
knowledge_base_name Optional[str]

The name of the knowledge base to use (optional).

 None
knowledge_base_id Optional[int]

The ID of an existing knowledge base to use (optional).

 None

Raises:

Type Description
BaizhongError

If neither knowledge_base_name nor knowledge_base_id is provided.
Source code in erniebot-agent/src/erniebot_agent/retrieval/baizhong_search.py 
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
def __init__(
    self,
    access_token: Optional[str] = None,
    knowledge_base_name: Optional[str] = None,
    knowledge_base_id: Optional[int] = None,
) -> None:
    """
    Initialize a BaizhongSearch object.

    Args:
        access_token (str): The access token for authentication.
        knowledge_base_name (Optional[str]): The name of the knowledge base to use (optional).
        knowledge_base_id (Optional[int]): The ID of an existing knowledge base to use (optional).

    Raises:
        BaizhongError: If neither knowledge_base_name nor knowledge_base_id is provided.

    """
    self._base_url = os.getenv("AISTUDIO_BASE_URL", self._AISTUDIO_BASE_URL)
    self.access_token: Optional[str] = None
    if access_token is not None:
        self.access_token = access_token
    elif C.get_global_access_token() is not None:
        self.access_token = C.get_global_access_token()
    else:
        raise BaizhongError(
            "Please ensure that either an access_token is provided or "
            "the EB_AGENT_ACCESS_TOKEN is set up as an environment variable."
        )
    if knowledge_base_id is not None:
        _logger.info(f"Loading existing project with `knowledge_base_id={knowledge_base_id}`")
        self.knowledge_base_id = knowledge_base_id
    elif knowledge_base_name is not None:
        self.knowledge_base_id = self.create_knowledge_base(knowledge_base_name)
    else:
        raise BaizhongError("You must provide either a `knowledge_base_name` or a `knowledge_base_id`.")

add_documents(documents, batch_size=1, thread_count=1)

Add a batch of documents to the Baizhong system using multi-threading. Args: documents (List[Document]): A list of Document objects to be added. batch_size (int, optional): The size of each batch of documents (defaults to 1). thread_count (int, optional): The number of threads to use for concurrent document addition (defaults to 1). Returns: List[Union[None, Exception]]: A list of results from the document addition process. Note: This function uses multi-threading to improve the efficiency of adding a large number of documents.

Source code in erniebot-agent/src/erniebot_agent/retrieval/baizhong_search.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def add_documents(self, documents: List[Document], batch_size: int = 1, thread_count: int = 1):
    """
    Add a batch of documents to the Baizhong system using multi-threading.
    Args:
        documents (List[Document]): A list of Document objects to be added.
        batch_size (int, optional): The size of each batch of documents (defaults to 1).
        thread_count (int, optional): The number of threads to use for concurrent document addition
        (defaults to 1).
    Returns:
        List[Union[None, Exception]]: A list of results from the document addition process.
    Note:
        This function uses multi-threading to improve the efficiency of adding a large number of
        documents.
    """
    raise NotImplementedError

create_knowledge_base(knowledge_base_name)

Create a JSON payload with the provided knowledge base name.

Parameters:

Name Type Description Default
knowledge_base_name str

The knowledge base name.

 required

Returns:

Type Description

Dict[str, Any]: A dictionary containing knowledge base results.

Raises:

Type Description
BaizhongError

If the API request fails, this exception is raised with details about the error.
Source code in erniebot-agent/src/erniebot_agent/retrieval/baizhong_search.py 
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
def create_knowledge_base(self, knowledge_base_name: str):
    """
    Create a JSON payload with the provided knowledge base name.

    Args:
        knowledge_base_name (str): The knowledge base name.

    Returns:
        Dict[str, Any]: A dictionary containing knowledge base results.

    Raises:
        BaizhongError: If the API request fails, this exception is raised with details about the error.
    """
    json_data = {"knowledgeBaseName": knowledge_base_name}
    res = requests.post(
        f"{self._base_url}/llm/knowledge/create",
        json=json_data,
        headers=self._get_authorization_headers(access_token=self.access_token),
    )
    if res.status_code == 200:
        result = res.json()
        if result["errorCode"] != 0:
            raise BaizhongError(message=result["errorMsg"], error_code=result["errorCode"])
        return result["result"]["knowledgeBaseId"]
    else:
        raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)

search(query, top_k=10, filters=None)

Perform a search using the Baizhong common search API.

Parameters:

Name Type Description Default
query str

The search query.

 required
top_k int

The number of top results to retrieve (default is 10).

 10
filters Optional[Dict[str, Any]]

Additional filters to apply to the search query

 None

Returns:

Type Description

List[Dict[str, Any]]: A list of dictionaries containing search results.

Raises:

Type Description
BaizhongError

If the API request fails, this exception is raised with details about the error.
Source code in erniebot-agent/src/erniebot_agent/retrieval/baizhong_search.py 
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
def search(self, query: str, top_k: int = 10, filters: Optional[Dict[str, Any]] = None):
    """
    Perform a search using the Baizhong common search API.

    Args:
        query (str): The search query.
        top_k (int, optional): The number of top results to retrieve (default is 10).
        filters (Optional[Dict[str, Any]], optional): Additional filters to apply to the search query
        (default is None).

    Returns:
        List[Dict[str, Any]]: A list of dictionaries containing search results.

    Raises:
        BaizhongError: If the API request fails, this exception is raised with details about the error.
    """
    json_data = {"knowledgeBaseId": self.knowledge_base_id, "query": query, "topK": top_k}
    if filters is not None:
        filter_terms = [{"term": item} for item in filters]
        filterConditions = {"filterConditions": {"bool": {"filter": filter_terms}}}
        json_data.update(filterConditions)
    res = requests.post(
        f"{self._base_url}/llm/knowledge/search",
        json=json_data,
        headers=self._get_authorization_headers(access_token=self.access_token),
    )
    if res.status_code == 200:
        result = res.json()
        if result["errorCode"] != 0:
            raise BaizhongError(message=result["errorMsg"], error_code=result["errorCode"])
        list_data = []
        for item in result["result"]:
            doc = json.loads(item["source"]["doc"])
            list_data.append(
                {
                    "id": item["fileId"],
                    "content": doc["content_se"],
                    "title": item["source"]["title"],
                    "score": item["score"],
                }
            )
        return list_data
    else:
        raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)