Skip to content

Elasticsearch

ElasticsearchDocumentStore

Bases: BaseDocumentStore

Simple memory document store that store document in a dictionary

Source code in libs/kotaemon/kotaemon/storages/docstores/elasticsearch.py 
 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
class ElasticsearchDocumentStore(BaseDocumentStore):
    """Simple memory document store that store document in a dictionary"""

    def __init__(
        self,
        collection_name: str = "docstore",
        elasticsearch_url: str = "http://localhost:9200",
        k1: float = 2.0,
        b: float = 0.75,
        **kwargs,
    ):
        try:
            from elasticsearch import Elasticsearch
            from elasticsearch.helpers import bulk
        except ImportError:
            raise ImportError(
                "To use ElaticsearchDocstore please install `pip install elasticsearch`"
            )

        self.elasticsearch_url = elasticsearch_url
        self.index_name = collection_name
        self.k1 = k1
        self.b = b

        # Create an Elasticsearch client instance
        self.client = Elasticsearch(elasticsearch_url, **kwargs)
        self.es_bulk = bulk
        # Define the index settings and mappings
        settings = {
            "analysis": {"analyzer": {"default": {"type": "standard"}}},
            "similarity": {
                "custom_bm25": {
                    "type": "BM25",
                    "k1": k1,
                    "b": b,
                }
            },
        }
        mappings = {
            "properties": {
                "content": {
                    "type": "text",
                    "similarity": "custom_bm25",  # Use the custom BM25 similarity
                }
            }
        }

        # Create the index with the specified settings and mappings
        if not self.client.indices.exists(index=self.index_name):
            self.client.indices.create(
                index=self.index_name, mappings=mappings, settings=settings
            )

    def add(
        self,
        docs: Union[Document, List[Document]],
        ids: Optional[Union[List[str], str]] = None,
        refresh_indices: bool = True,
        **kwargs,
    ):
        """Add document into document store

        Args:
            docs: list of documents to add
            ids: specify the ids of documents to add or use existing doc.doc_id
            refresh_indices: request Elasticsearch to update its index (default to True)
        """
        if ids and not isinstance(ids, list):
            ids = [ids]
        if not isinstance(docs, list):
            docs = [docs]
        doc_ids = ids if ids else [doc.doc_id for doc in docs]

        requests = []
        for doc_id, doc in zip(doc_ids, docs):
            text = doc.text
            metadata = doc.metadata
            request = {
                "_op_type": "index",
                "_index": self.index_name,
                "content": text,
                "metadata": metadata,
                "_id": doc_id,
            }
            requests.append(request)

        success, failed = self.es_bulk(self.client, requests)
        print("Added/Updated documents to index", success)
        print("Failed documents to index", failed)

        if refresh_indices:
            self.client.indices.refresh(index=self.index_name)

    def query_raw(self, query: dict) -> List[Document]:
        """Query Elasticsearch store using query format of ES client

        Args:
            query (dict): Elasticsearch query format

        Returns:
            List[Document]: List of result documents
        """
        res = self.client.search(index=self.index_name, body=query)
        docs = []
        for r in res["hits"]["hits"]:
            docs.append(
                Document(
                    id_=r["_id"],
                    text=r["_source"]["content"],
                    metadata=r["_source"]["metadata"],
                )
            )
        return docs

    def query(
        self, query: str, top_k: int = 10, doc_ids: Optional[list] = None
    ) -> List[Document]:
        """Search Elasticsearch docstore using search query (BM25)

        Args:
            query (str): query text
            top_k (int, optional): number of
                top documents to return. Defaults to 10.

        Returns:
            List[Document]: List of result documents
        """
        query_dict: dict = {"match": {"content": query}}
        if doc_ids is not None:
            query_dict = {"bool": {"must": [query_dict, {"terms": {"_id": doc_ids}}]}}
        query_dict = {"query": query_dict, "size": top_k}
        return self.query_raw(query_dict)

    def get(self, ids: Union[List[str], str]) -> List[Document]:
        """Get document by id"""
        if not isinstance(ids, list):
            ids = [ids]
        query_dict = {"query": {"terms": {"_id": ids}}, "size": 10000}
        return self.query_raw(query_dict)

    def count(self) -> int:
        """Count number of documents"""
        count = int(
            self.client.cat.count(index=self.index_name, format="json")[0]["count"]
        )
        return count

    def get_all(self) -> List[Document]:
        """Get all documents"""
        query_dict = {"query": {"match_all": {}}, "size": MAX_DOCS_TO_GET}
        return self.query_raw(query_dict)

    def delete(self, ids: Union[List[str], str]):
        """Delete document by id"""
        if not isinstance(ids, list):
            ids = [ids]

        query = {"query": {"terms": {"_id": ids}}}
        self.client.delete_by_query(index=self.index_name, body=query)
        self.client.indices.refresh(index=self.index_name)

    def drop(self):
        """Drop the document store"""
        self.client.indices.delete(index=self.index_name)
        self.client.indices.refresh(index=self.index_name)

    def __persist_flow__(self):
        return {
            "index_name": self.index_name,
            "elasticsearch_url": self.elasticsearch_url,
            "k1": self.k1,
            "b": self.b,
        }

add 

1
add(docs, ids=None, refresh_indices=True, **kwargs)

Add document into document store

Parameters:

Name Type Description Default
docs Union[Document, List[Document]]

list of documents to add

 required
ids Optional[Union[List[str], str]]

specify the ids of documents to add or use existing doc.doc_id

 None
refresh_indices bool

request Elasticsearch to update its index (default to True)

 True
Source code in libs/kotaemon/kotaemon/storages/docstores/elasticsearch.py 
 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
def add(
    self,
    docs: Union[Document, List[Document]],
    ids: Optional[Union[List[str], str]] = None,
    refresh_indices: bool = True,
    **kwargs,
):
    """Add document into document store

    Args:
        docs: list of documents to add
        ids: specify the ids of documents to add or use existing doc.doc_id
        refresh_indices: request Elasticsearch to update its index (default to True)
    """
    if ids and not isinstance(ids, list):
        ids = [ids]
    if not isinstance(docs, list):
        docs = [docs]
    doc_ids = ids if ids else [doc.doc_id for doc in docs]

    requests = []
    for doc_id, doc in zip(doc_ids, docs):
        text = doc.text
        metadata = doc.metadata
        request = {
            "_op_type": "index",
            "_index": self.index_name,
            "content": text,
            "metadata": metadata,
            "_id": doc_id,
        }
        requests.append(request)

    success, failed = self.es_bulk(self.client, requests)
    print("Added/Updated documents to index", success)
    print("Failed documents to index", failed)

    if refresh_indices:
        self.client.indices.refresh(index=self.index_name)

query_raw 

1
query_raw(query)

Query Elasticsearch store using query format of ES client

Parameters:

Name Type Description Default
query dict

Elasticsearch query format

 required

Returns:

Type Description
List[Document]

List[Document]: List of result documents
Source code in libs/kotaemon/kotaemon/storages/docstores/elasticsearch.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def query_raw(self, query: dict) -> List[Document]:
    """Query Elasticsearch store using query format of ES client

    Args:
        query (dict): Elasticsearch query format

    Returns:
        List[Document]: List of result documents
    """
    res = self.client.search(index=self.index_name, body=query)
    docs = []
    for r in res["hits"]["hits"]:
        docs.append(
            Document(
                id_=r["_id"],
                text=r["_source"]["content"],
                metadata=r["_source"]["metadata"],
            )
        )
    return docs

query 

1
query(query, top_k=10, doc_ids=None)

Search Elasticsearch docstore using search query (BM25)

Parameters:

Name Type Description Default
query str

query text

 required
top_k int

number of top documents to return. Defaults to 10.

 10

Returns:

Type Description
List[Document]

List[Document]: List of result documents
Source code in libs/kotaemon/kotaemon/storages/docstores/elasticsearch.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def query(
    self, query: str, top_k: int = 10, doc_ids: Optional[list] = None
) -> List[Document]:
    """Search Elasticsearch docstore using search query (BM25)

    Args:
        query (str): query text
        top_k (int, optional): number of
            top documents to return. Defaults to 10.

    Returns:
        List[Document]: List of result documents
    """
    query_dict: dict = {"match": {"content": query}}
    if doc_ids is not None:
        query_dict = {"bool": {"must": [query_dict, {"terms": {"_id": doc_ids}}]}}
    query_dict = {"query": query_dict, "size": top_k}
    return self.query_raw(query_dict)

get 

1
get(ids)

Get document by id

Source code in libs/kotaemon/kotaemon/storages/docstores/elasticsearch.py 
143
144
145
146
147
148
def get(self, ids: Union[List[str], str]) -> List[Document]:
    """Get document by id"""
    if not isinstance(ids, list):
        ids = [ids]
    query_dict = {"query": {"terms": {"_id": ids}}, "size": 10000}
    return self.query_raw(query_dict)

count 

1
count()

Count number of documents

Source code in libs/kotaemon/kotaemon/storages/docstores/elasticsearch.py 
150
151
152
153
154
155
def count(self) -> int:
    """Count number of documents"""
    count = int(
        self.client.cat.count(index=self.index_name, format="json")[0]["count"]
    )
    return count

get_all 

1
get_all()

Get all documents

Source code in libs/kotaemon/kotaemon/storages/docstores/elasticsearch.py 
157
158
159
160
def get_all(self) -> List[Document]:
    """Get all documents"""
    query_dict = {"query": {"match_all": {}}, "size": MAX_DOCS_TO_GET}
    return self.query_raw(query_dict)

delete 

1
delete(ids)

Delete document by id

Source code in libs/kotaemon/kotaemon/storages/docstores/elasticsearch.py 
162
163
164
165
166
167
168
169
def delete(self, ids: Union[List[str], str]):
    """Delete document by id"""
    if not isinstance(ids, list):
        ids = [ids]

    query = {"query": {"terms": {"_id": ids}}}
    self.client.delete_by_query(index=self.index_name, body=query)
    self.client.indices.refresh(index=self.index_name)

drop 

1
drop()

Drop the document store

Source code in libs/kotaemon/kotaemon/storages/docstores/elasticsearch.py 
171
172
173
174
def drop(self):
    """Drop the document store"""
    self.client.indices.delete(index=self.index_name)
    self.client.indices.refresh(index=self.index_name)