Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

VectorStore for GenAI integrations #2528

Merged
merged 36 commits into from
Apr 30, 2024
Merged
Show file tree
Hide file tree
Changes from 34 commits
Commits
Show all changes
36 commits
Select commit Hold shift + click to select a range
f30f1ad
ElasticsearchStore
maxjakob Apr 17, 2024
e03a17f
Update elasticsearch/store/_utilities.py
maxjakob Apr 17, 2024
8ff1c7c
rename; depend on client; async only
maxjakob Apr 18, 2024
9be44fd
generate _sync files
maxjakob Apr 18, 2024
7ee3846
add cleanup step for _sync generation
maxjakob Apr 18, 2024
2fd89bd
fix formatting
maxjakob Apr 18, 2024
9387b74
more linting fixes
maxjakob Apr 18, 2024
b18d63d
batch embedding call; infer num_dimensions
maxjakob Apr 18, 2024
9f83408
revert accidental changes
maxjakob Apr 22, 2024
9803414
keep field names only in store; apply metadata mappings in store
maxjakob Apr 22, 2024
7647961
fix typos in file names
maxjakob Apr 22, 2024
d397982
use `elasticsearch_url` fixture; create conftest.py
maxjakob Apr 22, 2024
2f1fcb0
export relevant classes
maxjakob Apr 23, 2024
b19de27
remove Semantic strategy
maxjakob Apr 23, 2024
274911a
es_query is sync
maxjakob Apr 23, 2024
8cec9cc
async strategies
maxjakob Apr 23, 2024
bbf2be9
cleanup old file
maxjakob Apr 24, 2024
299cd94
add docker-compose service with model deployment
maxjakob Apr 24, 2024
5f0d98d
optional dependencies for MMR
maxjakob Apr 24, 2024
58c8b7d
only test sync parts
maxjakob Apr 24, 2024
994b412
cleanup unasync script
maxjakob Apr 24, 2024
5073af1
nox: install optional deps
maxjakob Apr 24, 2024
9c50c6d
fix tests with requests remembering Transport
maxjakob Apr 24, 2024
a99a4f4
fix numpy typing
maxjakob Apr 24, 2024
d3c2e62
add user agent default argument
maxjakob Apr 24, 2024
11c8825
move to `elasticsearch.helpers.vectorstore`
maxjakob Apr 24, 2024
0d94881
use Protocol over ABC
maxjakob Apr 24, 2024
6aa6d73
revert Protocol change because Python 3.7
maxjakob Apr 24, 2024
71ca330
address PR feedback:
maxjakob Apr 25, 2024
a5dea84
improve docstring
maxjakob Apr 25, 2024
6f81af9
fix metadata mappings issue
maxjakob Apr 26, 2024
881d56c
address PR feedback
maxjakob Apr 29, 2024
f32ceb2
add error tests for strategies
maxjakob Apr 29, 2024
9b1778e
canonical names, keyword args only
maxjakob Apr 30, 2024
a8d80f2
fix sparse vector strategy bug (duplicate `size`)
maxjakob Apr 30, 2024
d27f9f8
all wildcard deletes in compose ES
maxjakob Apr 30, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion dev-requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,13 @@ twine
build
nox

numpy
pandas
orjson

# mmr for vectorstore
numpy
simsimd

# Testing the 'search_mvt' API response
mapbox-vector-tile
# Python 3.7 gets an old version of mapbox-vector-tile, requiring an
Expand Down
62 changes: 62 additions & 0 deletions elasticsearch/helpers/vectorstore/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
# Licensed to Elasticsearch B.V. under one or more contributor
# license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright
# ownership. Elasticsearch B.V. licenses this file to you under
# the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.

from elasticsearch.helpers.vectorstore._async.embedding_service import (
AsyncElasticsearchEmbeddings,
AsyncEmbeddingService,
)
from elasticsearch.helpers.vectorstore._async.strategies import (
AsyncBM25Strategy,
AsyncDenseVectorScriptScoreStrategy,
AsyncDenseVectorStrategy,
AsyncRetrievalStrategy,
AsyncSparseVectorStrategy,
)
from elasticsearch.helpers.vectorstore._async.vectorstore import AsyncVectorStore
from elasticsearch.helpers.vectorstore._sync.embedding_service import (
maxjakob marked this conversation as resolved.
Show resolved Hide resolved
ElasticsearchEmbeddings,
EmbeddingService,
)
from elasticsearch.helpers.vectorstore._sync.strategies import (
BM25Strategy,
DenseVectorScriptScoreStrategy,
DenseVectorStrategy,
RetrievalStrategy,
SparseVectorStrategy,
)
from elasticsearch.helpers.vectorstore._sync.vectorstore import VectorStore
from elasticsearch.helpers.vectorstore._utils import DistanceMetric

__all__ = [
"AsyncBM25Strategy",
"AsyncDenseVectorScriptScoreStrategy",
"AsyncDenseVectorStrategy",
"AsyncElasticsearchEmbeddings",
"AsyncEmbeddingService",
"AsyncRetrievalStrategy",
"AsyncSparseVectorStrategy",
"AsyncVectorStore",
"BM25Strategy",
"DenseVectorScriptScoreStrategy",
"DenseVectorStrategy",
"DistanceMetric",
"ElasticsearchEmbeddings",
"EmbeddingService",
"RetrievalStrategy",
"SparseVectorStrategy",
"VectorStore",
]
16 changes: 16 additions & 0 deletions elasticsearch/helpers/vectorstore/_async/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
# Licensed to Elasticsearch B.V. under one or more contributor
# license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright
# ownership. Elasticsearch B.V. licenses this file to you under
# the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
39 changes: 39 additions & 0 deletions elasticsearch/helpers/vectorstore/_async/_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
# Licensed to Elasticsearch B.V. under one or more contributor
# license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright
# ownership. Elasticsearch B.V. licenses this file to you under
# the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.

from elasticsearch import AsyncElasticsearch, BadRequestError, NotFoundError


async def model_must_be_deployed(client: AsyncElasticsearch, model_id: str) -> None:
"""
:raises [NotFoundError]: if the model is neither downloaded nor deployed.
:raises [ConflictError]: if the model is downloaded but not yet deployed.
"""
doc = {"text_field": f"test if the model '{model_id}' is deployed"}
try:
await client.ml.infer_trained_model(model_id=model_id, docs=[doc])
except BadRequestError:
# The model is deployed but expects a different input field name.
pass


async def model_is_deployed(client: AsyncElasticsearch, model_id: str) -> bool:
try:
await model_must_be_deployed(client, model_id)
return True
except NotFoundError:
return False
89 changes: 89 additions & 0 deletions elasticsearch/helpers/vectorstore/_async/embedding_service.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
# Licensed to Elasticsearch B.V. under one or more contributor
# license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright
# ownership. Elasticsearch B.V. licenses this file to you under
# the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.

from abc import ABC, abstractmethod
from typing import List

from elasticsearch import AsyncElasticsearch
from elasticsearch._version import __versionstr__ as lib_version


class AsyncEmbeddingService(ABC):
@abstractmethod
async def embed_documents(self, texts: List[str]) -> List[List[float]]:
"""Generate embeddings for a list of documents.

:param texts: A list of document strings to generate embeddings for.

:return: A list of embeddings, one for each document in the input.
"""

@abstractmethod
async def embed_query(self, query: str) -> List[float]:
"""Generate an embedding for a single query text.

:param text: The query text to generate an embedding for.

:return: The embedding for the input query text.
"""


class AsyncElasticsearchEmbeddings(AsyncEmbeddingService):
"""Elasticsearch as a service for embedding model inference.

You need to have an embedding model downloaded and deployed in Elasticsearch:
- https://www.elastic.co/guide/en/elasticsearch/reference/current/infer-trained-model.html
- https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-deploy-models.html
""" # noqa: E501

def __init__(
self,
*,
client: AsyncElasticsearch,
model_id: str,
input_field: str = "text_field",
user_agent: str = f"elasticsearch-py-es/{lib_version}",
):
"""
:param agent_header: user agent header specific to the 3rd party integration.
Used for usage tracking in Elastic Cloud.
:param model_id: The model_id of the model deployed in the Elasticsearch cluster.
:param input_field: The name of the key for the input text field in the
document. Defaults to 'text_field'.
:param client: Elasticsearch client connection. Alternatively specify the
Elasticsearch connection with the other es_* parameters.
"""
# Add integration-specific usage header for tracking usage in Elastic Cloud.
# client.options preserves existing (non-user-agent) headers.
client = client.options(headers={"User-Agent": user_agent})

self.client = client
self.model_id = model_id
self.input_field = input_field

async def embed_documents(self, texts: List[str]) -> List[List[float]]:
return await self._embedding_func(texts)

async def embed_query(self, text: str) -> List[float]:
result = await self._embedding_func([text])
return result[0]

async def _embedding_func(self, texts: List[str]) -> List[List[float]]:
response = await self.client.ml.infer_trained_model(
model_id=self.model_id, docs=[{self.input_field: text} for text in texts]
)
return [doc["predicted_value"] for doc in response["inference_results"]]
Loading