diff --git a/.buildkite/Dockerfile b/.buildkite/Dockerfile index 3c429a37b..7b0eb2e8e 100644 --- a/.buildkite/Dockerfile +++ b/.buildkite/Dockerfile @@ -1,4 +1,4 @@ -ARG PYTHON_VERSION=3.8 +ARG PYTHON_VERSION=3.12 FROM python:${PYTHON_VERSION} # Default UID/GID to 1000 diff --git a/dev-requirements.txt b/dev-requirements.txt index c42af4eab..330cb2701 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -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 diff --git a/docs/examples/00fea15cbca83be9d5f1a024ff2ec708.asciidoc b/docs/examples/00fea15cbca83be9d5f1a024ff2ec708.asciidoc index 27278ffd3..4e63a4e90 100644 --- a/docs/examples/00fea15cbca83be9d5f1a024ff2ec708.asciidoc +++ b/docs/examples/00fea15cbca83be9d5f1a024ff2ec708.asciidoc @@ -1,4 +1,4 @@ -// inference/put-inference.asciidoc:275 +// inference/put-inference.asciidoc:381 [source, python] ---- diff --git a/docs/examples/10c3fe2265bb34964bd1005f9da66773.asciidoc b/docs/examples/10c3fe2265bb34964bd1005f9da66773.asciidoc index d4aa29269..08c365b91 100644 --- a/docs/examples/10c3fe2265bb34964bd1005f9da66773.asciidoc +++ b/docs/examples/10c3fe2265bb34964bd1005f9da66773.asciidoc @@ -1,4 +1,4 @@ -// inference/put-inference.asciidoc:369 +// inference/put-inference.asciidoc:488 [source, python] ---- diff --git a/docs/examples/13ecdf99114098c76b050397d9c3d4e6.asciidoc b/docs/examples/13ecdf99114098c76b050397d9c3d4e6.asciidoc index 5d5454084..f3d348b3e 100644 --- a/docs/examples/13ecdf99114098c76b050397d9c3d4e6.asciidoc +++ b/docs/examples/13ecdf99114098c76b050397d9c3d4e6.asciidoc @@ -1,4 +1,4 @@ -// inference/post-inference.asciidoc:73 +// inference/post-inference.asciidoc:197 [source, python] ---- diff --git a/docs/examples/1aa91d3d48140d6367b6cabca8737b8f.asciidoc b/docs/examples/1aa91d3d48140d6367b6cabca8737b8f.asciidoc index c4553e246..7bc426db3 100644 --- a/docs/examples/1aa91d3d48140d6367b6cabca8737b8f.asciidoc +++ b/docs/examples/1aa91d3d48140d6367b6cabca8737b8f.asciidoc @@ -1,4 +1,4 @@ -// docs/bulk.asciidoc:634 +// docs/bulk.asciidoc:632 [source, python] ---- diff --git a/docs/examples/316cd43feb3b86396483903af1a048b1.asciidoc b/docs/examples/316cd43feb3b86396483903af1a048b1.asciidoc index 6d040bd6d..0641456c4 100644 --- a/docs/examples/316cd43feb3b86396483903af1a048b1.asciidoc +++ b/docs/examples/316cd43feb3b86396483903af1a048b1.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:781 +// aggregations/bucket/datehistogram-aggregation.asciidoc:782 [source, python] ---- diff --git a/docs/examples/3d1ff6097e2359f927c88c2ccdb36252.asciidoc b/docs/examples/3d1ff6097e2359f927c88c2ccdb36252.asciidoc deleted file mode 100644 index 4a576819a..000000000 --- a/docs/examples/3d1ff6097e2359f927c88c2ccdb36252.asciidoc +++ /dev/null @@ -1,7 +0,0 @@ -// tab-widgets/api-call.asciidoc:13 - -[source, python] ----- -resp = client.info() -print(resp) ----- \ No newline at end of file diff --git a/docs/examples/4e3414fc712b16311f9e433dd366f49d.asciidoc b/docs/examples/4e3414fc712b16311f9e433dd366f49d.asciidoc index c64de0814..eebe7e709 100644 --- a/docs/examples/4e3414fc712b16311f9e433dd366f49d.asciidoc +++ b/docs/examples/4e3414fc712b16311f9e433dd366f49d.asciidoc @@ -1,4 +1,4 @@ -// inference/delete-inference.asciidoc:53 +// inference/delete-inference.asciidoc:55 [source, python] ---- diff --git a/docs/examples/51b40610ae05730b4c6afd25647d7ae0.asciidoc b/docs/examples/51b40610ae05730b4c6afd25647d7ae0.asciidoc index c662e2685..f5af5e526 100644 --- a/docs/examples/51b40610ae05730b4c6afd25647d7ae0.asciidoc +++ b/docs/examples/51b40610ae05730b4c6afd25647d7ae0.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:488 +// aggregations/bucket/datehistogram-aggregation.asciidoc:489 [source, python] ---- diff --git a/docs/examples/5203560189ccab7122c03500147701ef.asciidoc b/docs/examples/5203560189ccab7122c03500147701ef.asciidoc index 3e0285303..412284e48 100644 --- a/docs/examples/5203560189ccab7122c03500147701ef.asciidoc +++ b/docs/examples/5203560189ccab7122c03500147701ef.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:568 +// aggregations/bucket/datehistogram-aggregation.asciidoc:569 [source, python] ---- diff --git a/docs/examples/59d736a4d064ed2013c7ead8e32e0998.asciidoc b/docs/examples/59d736a4d064ed2013c7ead8e32e0998.asciidoc new file mode 100644 index 000000000..25a961d96 --- /dev/null +++ b/docs/examples/59d736a4d064ed2013c7ead8e32e0998.asciidoc @@ -0,0 +1,17 @@ +// inference/put-inference.asciidoc:529 + +[source, python] +---- +resp = client.inference.put_model( + task_type="completion", + inference_id="openai-completion", + body={ + "service": "openai", + "service_settings": { + "api_key": "", + "model_id": "gpt-3.5-turbo", + }, + }, +) +print(resp) +---- \ No newline at end of file diff --git a/docs/examples/77b90f6787195767b6da60d8532714b4.asciidoc b/docs/examples/77b90f6787195767b6da60d8532714b4.asciidoc new file mode 100644 index 000000000..4393f9458 --- /dev/null +++ b/docs/examples/77b90f6787195767b6da60d8532714b4.asciidoc @@ -0,0 +1,19 @@ +// inference/put-inference.asciidoc:552 + +[source, python] +---- +resp = client.inference.put_model( + task_type="text_embedding", + inference_id="azure_openai_embeddings", + body={ + "service": "azureopenai", + "service_settings": { + "api_key": "", + "resource_name": "", + "deployment_id": "", + "api_version": "2024-02-01", + }, + }, +) +print(resp) +---- \ No newline at end of file diff --git a/docs/examples/8619bd17bbfe33490b1f277007f654db.asciidoc b/docs/examples/8619bd17bbfe33490b1f277007f654db.asciidoc new file mode 100644 index 000000000..32ff8da77 --- /dev/null +++ b/docs/examples/8619bd17bbfe33490b1f277007f654db.asciidoc @@ -0,0 +1,18 @@ +// inference/put-inference.asciidoc:353 + +[source, python] +---- +resp = client.inference.put_model( + task_type="rerank", + inference_id="cohere-rerank", + body={ + "service": "cohere", + "service_settings": { + "api_key": "", + "model_id": "rerank-english-v3.0", + }, + "task_settings": {"top_n": 10, "return_documents": True}, + }, +) +print(resp) +---- \ No newline at end of file diff --git a/docs/examples/8cd00a3aba7c3c158277bc032aac2830.asciidoc b/docs/examples/8cd00a3aba7c3c158277bc032aac2830.asciidoc index 639a269cc..452631e77 100644 --- a/docs/examples/8cd00a3aba7c3c158277bc032aac2830.asciidoc +++ b/docs/examples/8cd00a3aba7c3c158277bc032aac2830.asciidoc @@ -1,4 +1,4 @@ -// docs/bulk.asciidoc:612 +// docs/bulk.asciidoc:610 [source, python] ---- diff --git a/docs/examples/9a02bd47c000a3d9a8911233c37c890f.asciidoc b/docs/examples/9a02bd47c000a3d9a8911233c37c890f.asciidoc index fd12e4bd9..d6793b366 100644 --- a/docs/examples/9a02bd47c000a3d9a8911233c37c890f.asciidoc +++ b/docs/examples/9a02bd47c000a3d9a8911233c37c890f.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:366 +// aggregations/bucket/datehistogram-aggregation.asciidoc:367 [source, python] ---- diff --git a/docs/examples/9a203aae3e1412d919546276fb52a5ca.asciidoc b/docs/examples/9a203aae3e1412d919546276fb52a5ca.asciidoc index 6b6a5531b..8f030467d 100644 --- a/docs/examples/9a203aae3e1412d919546276fb52a5ca.asciidoc +++ b/docs/examples/9a203aae3e1412d919546276fb52a5ca.asciidoc @@ -1,4 +1,4 @@ -// inference/put-inference.asciidoc:253 +// inference/put-inference.asciidoc:335 [source, python] ---- diff --git a/docs/examples/6dda348069f89ffb2379b3a281a24cc9.asciidoc b/docs/examples/9f16fca9813304e398ee052aa857dbcd.asciidoc similarity index 77% rename from docs/examples/6dda348069f89ffb2379b3a281a24cc9.asciidoc rename to docs/examples/9f16fca9813304e398ee052aa857dbcd.asciidoc index fd9eda50d..1d854a9c9 100644 --- a/docs/examples/6dda348069f89ffb2379b3a281a24cc9.asciidoc +++ b/docs/examples/9f16fca9813304e398ee052aa857dbcd.asciidoc @@ -1,10 +1,10 @@ -// inference/put-inference.asciidoc:394 +// inference/put-inference.asciidoc:513 [source, python] ---- resp = client.inference.put_model( task_type="text_embedding", - inference_id="openai_embeddings", + inference_id="openai-embeddings", body={ "service": "openai", "service_settings": { diff --git a/docs/examples/a4a3c3cd09efa75168dab90105afb2e9.asciidoc b/docs/examples/a4a3c3cd09efa75168dab90105afb2e9.asciidoc index 79b43f6c1..5f21ffc5e 100644 --- a/docs/examples/a4a3c3cd09efa75168dab90105afb2e9.asciidoc +++ b/docs/examples/a4a3c3cd09efa75168dab90105afb2e9.asciidoc @@ -1,4 +1,4 @@ -// inference/get-inference.asciidoc:69 +// inference/get-inference.asciidoc:73 [source, python] ---- diff --git a/docs/examples/a84ffbaa4ffa68b22f6fe42d3b4f8dd5.asciidoc b/docs/examples/a84ffbaa4ffa68b22f6fe42d3b4f8dd5.asciidoc index 2bcfac3ae..37789b6fd 100644 --- a/docs/examples/a84ffbaa4ffa68b22f6fe42d3b4f8dd5.asciidoc +++ b/docs/examples/a84ffbaa4ffa68b22f6fe42d3b4f8dd5.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:588 +// aggregations/bucket/datehistogram-aggregation.asciidoc:589 [source, python] ---- diff --git a/docs/examples/ae9ccfaa146731ab9176df90670db1c2.asciidoc b/docs/examples/ae9ccfaa146731ab9176df90670db1c2.asciidoc index 72d9604b9..668510e72 100644 --- a/docs/examples/ae9ccfaa146731ab9176df90670db1c2.asciidoc +++ b/docs/examples/ae9ccfaa146731ab9176df90670db1c2.asciidoc @@ -1,4 +1,4 @@ -// docs/bulk.asciidoc:501 +// docs/bulk.asciidoc:499 [source, python] ---- diff --git a/docs/examples/b45a8c6fc746e9c90fd181e69a605fad.asciidoc b/docs/examples/b45a8c6fc746e9c90fd181e69a605fad.asciidoc new file mode 100644 index 000000000..8a1fc3fc3 --- /dev/null +++ b/docs/examples/b45a8c6fc746e9c90fd181e69a605fad.asciidoc @@ -0,0 +1,11 @@ +// inference/post-inference.asciidoc:102 + +[source, python] +---- +resp = client.inference.inference( + task_type="completion", + inference_id="openai_chat_completions", + body={"input": "What is Elastic?"}, +) +print(resp) +---- \ No newline at end of file diff --git a/docs/examples/ba59a3b9a0a2694704b2bf9c6ad4a8cf.asciidoc b/docs/examples/ba59a3b9a0a2694704b2bf9c6ad4a8cf.asciidoc index 77b897af5..8622ad4b6 100644 --- a/docs/examples/ba59a3b9a0a2694704b2bf9c6ad4a8cf.asciidoc +++ b/docs/examples/ba59a3b9a0a2694704b2bf9c6ad4a8cf.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:610 +// aggregations/bucket/datehistogram-aggregation.asciidoc:611 [source, python] ---- diff --git a/docs/examples/bfdad8a928ea30d7cf60d0a0a6bc6e2e.asciidoc b/docs/examples/bfdad8a928ea30d7cf60d0a0a6bc6e2e.asciidoc index c1d524d1c..7d843b759 100644 --- a/docs/examples/bfdad8a928ea30d7cf60d0a0a6bc6e2e.asciidoc +++ b/docs/examples/bfdad8a928ea30d7cf60d0a0a6bc6e2e.asciidoc @@ -1,4 +1,4 @@ -// docs/bulk.asciidoc:713 +// docs/bulk.asciidoc:711 [source, python] ---- diff --git a/docs/examples/c9c396b94bb88098477e2b08b55a12ee.asciidoc b/docs/examples/c9c396b94bb88098477e2b08b55a12ee.asciidoc index a0950395e..f205a86e4 100644 --- a/docs/examples/c9c396b94bb88098477e2b08b55a12ee.asciidoc +++ b/docs/examples/c9c396b94bb88098477e2b08b55a12ee.asciidoc @@ -1,4 +1,4 @@ -// docs/bulk.asciidoc:766 +// docs/bulk.asciidoc:764 [source, python] ---- diff --git a/docs/examples/cedb56a71cc743d80263ce352bb21720.asciidoc b/docs/examples/cedb56a71cc743d80263ce352bb21720.asciidoc index c9e868528..2be37beca 100644 --- a/docs/examples/cedb56a71cc743d80263ce352bb21720.asciidoc +++ b/docs/examples/cedb56a71cc743d80263ce352bb21720.asciidoc @@ -1,4 +1,4 @@ -// inference/put-inference.asciidoc:300 +// inference/put-inference.asciidoc:406 [source, python] ---- diff --git a/docs/examples/da3cecc36a7313385d32c7f52ccfb7e3.asciidoc b/docs/examples/da3cecc36a7313385d32c7f52ccfb7e3.asciidoc index 99cd507e8..5fe27c39e 100644 --- a/docs/examples/da3cecc36a7313385d32c7f52ccfb7e3.asciidoc +++ b/docs/examples/da3cecc36a7313385d32c7f52ccfb7e3.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:815 +// aggregations/bucket/datehistogram-aggregation.asciidoc:816 [source, python] ---- diff --git a/docs/examples/e6b972611c0ec8ab4c240f33f323d85b.asciidoc b/docs/examples/e6b972611c0ec8ab4c240f33f323d85b.asciidoc index 8dac1cb20..a90bb7bfc 100644 --- a/docs/examples/e6b972611c0ec8ab4c240f33f323d85b.asciidoc +++ b/docs/examples/e6b972611c0ec8ab4c240f33f323d85b.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:417 +// aggregations/bucket/datehistogram-aggregation.asciidoc:418 [source, python] ---- diff --git a/docs/examples/ecc57597f6b791d1151ad79d9f4ce67b.asciidoc b/docs/examples/ecc57597f6b791d1151ad79d9f4ce67b.asciidoc index 415bcd506..483bfbae0 100644 --- a/docs/examples/ecc57597f6b791d1151ad79d9f4ce67b.asciidoc +++ b/docs/examples/ecc57597f6b791d1151ad79d9f4ce67b.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:642 +// aggregations/bucket/datehistogram-aggregation.asciidoc:643 [source, python] ---- diff --git a/docs/examples/eec051555c8050d017d3fe38ea59e3a0.asciidoc b/docs/examples/eec051555c8050d017d3fe38ea59e3a0.asciidoc index 659204a42..12c64a679 100644 --- a/docs/examples/eec051555c8050d017d3fe38ea59e3a0.asciidoc +++ b/docs/examples/eec051555c8050d017d3fe38ea59e3a0.asciidoc @@ -1,4 +1,4 @@ -// search/search.asciidoc:956 +// search/search.asciidoc:907 [source, python] ---- diff --git a/docs/examples/eee6110831c08b9c1b3f56b24656e95b.asciidoc b/docs/examples/eee6110831c08b9c1b3f56b24656e95b.asciidoc index 3c44c5a60..286378a6b 100644 --- a/docs/examples/eee6110831c08b9c1b3f56b24656e95b.asciidoc +++ b/docs/examples/eee6110831c08b9c1b3f56b24656e95b.asciidoc @@ -1,4 +1,4 @@ -// inference/put-inference.asciidoc:339 +// inference/put-inference.asciidoc:445 [source, python] ---- diff --git a/docs/examples/ef779b87b3b0fb6e6bae9c8875e3a1cf.asciidoc b/docs/examples/ef779b87b3b0fb6e6bae9c8875e3a1cf.asciidoc index 27719ab27..ea871d67d 100644 --- a/docs/examples/ef779b87b3b0fb6e6bae9c8875e3a1cf.asciidoc +++ b/docs/examples/ef779b87b3b0fb6e6bae9c8875e3a1cf.asciidoc @@ -1,4 +1,4 @@ -// aggregations/bucket/datehistogram-aggregation.asciidoc:698 +// aggregations/bucket/datehistogram-aggregation.asciidoc:699 [source, python] ---- diff --git a/docs/examples/f1b24217b1d9ba6ea5e4fa6e6f412022.asciidoc b/docs/examples/f1b24217b1d9ba6ea5e4fa6e6f412022.asciidoc new file mode 100644 index 000000000..89cd690dc --- /dev/null +++ b/docs/examples/f1b24217b1d9ba6ea5e4fa6e6f412022.asciidoc @@ -0,0 +1,14 @@ +// inference/post-inference.asciidoc:133 + +[source, python] +---- +resp = client.inference.inference( + task_type="rerank", + inference_id="cohere_rerank", + body={ + "input": ["luke", "like", "leia", "chewy", "r2d2", "star", "wars"], + "query": "star wars main character", + }, +) +print(resp) +---- \ No newline at end of file diff --git a/docs/guide/release-notes.asciidoc b/docs/guide/release-notes.asciidoc index bcad88527..a198f096d 100644 --- a/docs/guide/release-notes.asciidoc +++ b/docs/guide/release-notes.asciidoc @@ -1,6 +1,7 @@ [[release-notes]] == Release notes +* <> * <> * <> * <> @@ -36,6 +37,14 @@ * <> * <> +[discrete] +[[rn-8-13-1]] +=== 8.13.1 (2024-05-03) + +- Added ``force_synthetic_source`` to the Get API +- Added ``wait_for_completion`` to the Create trained model API +- Added ``typed_keys`` to the Query API key information API + [discrete] [[rn-8-13-0]] === 8.13.0 (2024-03-22) diff --git a/elasticsearch/_async/client/cluster.py b/elasticsearch/_async/client/cluster.py index db807d710..990653d27 100644 --- a/elasticsearch/_async/client/cluster.py +++ b/elasticsearch/_async/client/cluster.py @@ -661,7 +661,7 @@ async def post_voting_config_exclusions( ) @_rewrite_parameters( - body_fields=("template", "allow_auto_create", "meta", "version"), + body_fields=("template", "deprecated", "meta", "version"), parameter_aliases={"_meta": "meta"}, ) async def put_component_template( @@ -669,8 +669,9 @@ async def put_component_template( *, name: str, template: t.Optional[t.Mapping[str, t.Any]] = None, - allow_auto_create: t.Optional[bool] = None, + cause: t.Optional[str] = None, create: t.Optional[bool] = None, + deprecated: t.Optional[bool] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, @@ -698,13 +699,12 @@ async def put_component_template( update settings API. :param template: The template to be applied which includes mappings, settings, or aliases configuration. - :param allow_auto_create: This setting overrides the value of the `action.auto_create_index` - cluster setting. If set to `true` in a template, then indices can be automatically - created using that template even if auto-creation of indices is disabled - via `actions.auto_create_index`. If set to `false` then data streams matching - the template must always be explicitly created. + :param cause: :param create: If `true`, this request cannot replace or update existing component templates. + :param deprecated: Marks this index template as deprecated. When creating or + updating a non-deprecated index template that uses deprecated components, + Elasticsearch will emit a deprecation warning. :param master_timeout: Period to wait for a connection to the master node. If no response is received before the timeout expires, the request fails and returns an error. @@ -724,6 +724,8 @@ async def put_component_template( __path = f'/_component_template/{__path_parts["name"]}' __query: t.Dict[str, t.Any] = {} __body: t.Dict[str, t.Any] = body if body is not None else {} + if cause is not None: + __query["cause"] = cause if create is not None: __query["create"] = create if error_trace is not None: @@ -739,8 +741,8 @@ async def put_component_template( if not __body: if template is not None: __body["template"] = template - if allow_auto_create is not None: - __body["allow_auto_create"] = allow_auto_create + if deprecated is not None: + __body["deprecated"] = deprecated if meta is not None: __body["_meta"] = meta if version is not None: diff --git a/elasticsearch/_async/client/connector.py b/elasticsearch/_async/client/connector.py index 8cdebdb19..c7fd33f67 100644 --- a/elasticsearch/_async/client/connector.py +++ b/elasticsearch/_async/client/connector.py @@ -70,6 +70,7 @@ async def delete( self, *, connector_id: str, + delete_sync_jobs: bool, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, @@ -81,12 +82,17 @@ async def delete( ``_ :param connector_id: The unique identifier of the connector to be deleted + :param delete_sync_jobs: Determines whether associated sync jobs are also deleted. """ if connector_id in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'connector_id'") + if delete_sync_jobs is None: + raise ValueError("Empty value passed for parameter 'delete_sync_jobs'") __path_parts: t.Dict[str, str] = {"connector_id": _quote(connector_id)} __path = f'/_connector/{__path_parts["connector_id"]}' __query: t.Dict[str, t.Any] = {} + if delete_sync_jobs is not None: + __query["delete_sync_jobs"] = delete_sync_jobs if error_trace is not None: __query["error_trace"] = error_trace if filter_path is not None: @@ -486,6 +492,302 @@ async def put( path_parts=__path_parts, ) + @_rewrite_parameters() + async def sync_job_cancel( + self, + *, + connector_sync_job_id: str, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Cancels a connector sync job. + + ``_ + + :param connector_sync_job_id: The unique identifier of the connector sync job + """ + if connector_sync_job_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_sync_job_id'") + __path_parts: t.Dict[str, str] = { + "connector_sync_job_id": _quote(connector_sync_job_id) + } + __path = ( + f'/_connector/_sync_job/{__path_parts["connector_sync_job_id"]}/_cancel' + ) + __query: t.Dict[str, t.Any] = {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + __headers = {"accept": "application/json"} + return await self.perform_request( # type: ignore[return-value] + "PUT", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.sync_job_cancel", + path_parts=__path_parts, + ) + + @_rewrite_parameters() + async def sync_job_delete( + self, + *, + connector_sync_job_id: str, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Deletes a connector sync job. + + ``_ + + :param connector_sync_job_id: The unique identifier of the connector sync job + to be deleted + """ + if connector_sync_job_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_sync_job_id'") + __path_parts: t.Dict[str, str] = { + "connector_sync_job_id": _quote(connector_sync_job_id) + } + __path = f'/_connector/_sync_job/{__path_parts["connector_sync_job_id"]}' + __query: t.Dict[str, t.Any] = {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + __headers = {"accept": "application/json"} + return await self.perform_request( # type: ignore[return-value] + "DELETE", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.sync_job_delete", + path_parts=__path_parts, + ) + + @_rewrite_parameters() + async def sync_job_get( + self, + *, + connector_sync_job_id: str, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Returns the details about a connector sync job. + + ``_ + + :param connector_sync_job_id: The unique identifier of the connector sync job + """ + if connector_sync_job_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_sync_job_id'") + __path_parts: t.Dict[str, str] = { + "connector_sync_job_id": _quote(connector_sync_job_id) + } + __path = f'/_connector/_sync_job/{__path_parts["connector_sync_job_id"]}' + __query: t.Dict[str, t.Any] = {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + __headers = {"accept": "application/json"} + return await self.perform_request( # type: ignore[return-value] + "GET", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.sync_job_get", + path_parts=__path_parts, + ) + + @_rewrite_parameters( + parameter_aliases={"from": "from_"}, + ) + async def sync_job_list( + self, + *, + connector_id: t.Optional[str] = None, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + from_: t.Optional[int] = None, + human: t.Optional[bool] = None, + job_type: t.Optional[ + t.Sequence[ + t.Union["t.Literal['access_control', 'full', 'incremental']", str] + ] + ] = None, + pretty: t.Optional[bool] = None, + size: t.Optional[int] = None, + status: t.Optional[ + t.Union[ + "t.Literal['canceled', 'canceling', 'completed', 'error', 'in_progress', 'pending', 'suspended']", + str, + ] + ] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Lists all connector sync jobs. + + ``_ + + :param connector_id: A connector id to fetch connector sync jobs for + :param from_: Starting offset (default: 0) + :param job_type: A comma-separated list of job types to fetch the sync jobs for + :param size: Specifies a max number of results to get + :param status: A sync job status to fetch connector sync jobs for + """ + __path_parts: t.Dict[str, str] = {} + __path = "/_connector/_sync_job" + __query: t.Dict[str, t.Any] = {} + if connector_id is not None: + __query["connector_id"] = connector_id + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if from_ is not None: + __query["from"] = from_ + if human is not None: + __query["human"] = human + if job_type is not None: + __query["job_type"] = job_type + if pretty is not None: + __query["pretty"] = pretty + if size is not None: + __query["size"] = size + if status is not None: + __query["status"] = status + __headers = {"accept": "application/json"} + return await self.perform_request( # type: ignore[return-value] + "GET", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.sync_job_list", + path_parts=__path_parts, + ) + + @_rewrite_parameters( + body_fields=("id", "job_type", "trigger_method"), + ) + async def sync_job_post( + self, + *, + id: t.Optional[str] = None, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + job_type: t.Optional[ + t.Union["t.Literal['access_control', 'full', 'incremental']", str] + ] = None, + pretty: t.Optional[bool] = None, + trigger_method: t.Optional[ + t.Union["t.Literal['on_demand', 'scheduled']", str] + ] = None, + body: t.Optional[t.Dict[str, t.Any]] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Creates a connector sync job. + + ``_ + + :param id: The id of the associated connector + :param job_type: + :param trigger_method: + """ + if id is None and body is None: + raise ValueError("Empty value passed for parameter 'id'") + __path_parts: t.Dict[str, str] = {} + __path = "/_connector/_sync_job" + __query: t.Dict[str, t.Any] = {} + __body: t.Dict[str, t.Any] = body if body is not None else {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + if not __body: + if id is not None: + __body["id"] = id + if job_type is not None: + __body["job_type"] = job_type + if trigger_method is not None: + __body["trigger_method"] = trigger_method + __headers = {"accept": "application/json", "content-type": "application/json"} + return await self.perform_request( # type: ignore[return-value] + "POST", + __path, + params=__query, + headers=__headers, + body=__body, + endpoint_id="connector.sync_job_post", + path_parts=__path_parts, + ) + + @_rewrite_parameters() + async def update_active_filtering( + self, + *, + connector_id: str, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Activates the draft filtering rules if they are in a validated state. + + ``_ + + :param connector_id: The unique identifier of the connector to be updated + """ + if connector_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_id'") + __path_parts: t.Dict[str, str] = {"connector_id": _quote(connector_id)} + __path = f'/_connector/{__path_parts["connector_id"]}/_filtering/_activate' + __query: t.Dict[str, t.Any] = {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + __headers = {"accept": "application/json"} + return await self.perform_request( # type: ignore[return-value] + "PUT", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.update_active_filtering", + path_parts=__path_parts, + ) + @_rewrite_parameters( body_fields=("api_key_id", "api_key_secret_id"), ) @@ -647,17 +949,19 @@ async def update_error( ) @_rewrite_parameters( - body_fields=("filtering",), + body_fields=("advanced_snippet", "filtering", "rules"), ) async def update_filtering( self, *, connector_id: str, - filtering: t.Optional[t.Sequence[t.Mapping[str, t.Any]]] = None, + advanced_snippet: t.Optional[t.Mapping[str, t.Any]] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + filtering: t.Optional[t.Sequence[t.Mapping[str, t.Any]]] = None, human: t.Optional[bool] = None, pretty: t.Optional[bool] = None, + rules: t.Optional[t.Sequence[t.Mapping[str, t.Any]]] = None, body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: """ @@ -666,12 +970,12 @@ async def update_filtering( ``_ :param connector_id: The unique identifier of the connector to be updated + :param advanced_snippet: :param filtering: + :param rules: """ if connector_id in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'connector_id'") - if filtering is None and body is None: - raise ValueError("Empty value passed for parameter 'filtering'") __path_parts: t.Dict[str, str] = {"connector_id": _quote(connector_id)} __path = f'/_connector/{__path_parts["connector_id"]}/_filtering' __query: t.Dict[str, t.Any] = {} @@ -685,8 +989,12 @@ async def update_filtering( if pretty is not None: __query["pretty"] = pretty if not __body: + if advanced_snippet is not None: + __body["advanced_snippet"] = advanced_snippet if filtering is not None: __body["filtering"] = filtering + if rules is not None: + __body["rules"] = rules __headers = {"accept": "application/json", "content-type": "application/json"} return await self.perform_request( # type: ignore[return-value] "PUT", @@ -698,6 +1006,58 @@ async def update_filtering( path_parts=__path_parts, ) + @_rewrite_parameters( + body_fields=("validation",), + ) + async def update_filtering_validation( + self, + *, + connector_id: str, + validation: t.Optional[t.Mapping[str, t.Any]] = None, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + body: t.Optional[t.Dict[str, t.Any]] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Updates the validation info of the draft filtering rules. + + ``_ + + :param connector_id: The unique identifier of the connector to be updated + :param validation: + """ + if connector_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_id'") + if validation is None and body is None: + raise ValueError("Empty value passed for parameter 'validation'") + __path_parts: t.Dict[str, str] = {"connector_id": _quote(connector_id)} + __path = f'/_connector/{__path_parts["connector_id"]}/_filtering/_validation' + __query: t.Dict[str, t.Any] = {} + __body: t.Dict[str, t.Any] = body if body is not None else {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + if not __body: + if validation is not None: + __body["validation"] = validation + __headers = {"accept": "application/json", "content-type": "application/json"} + return await self.perform_request( # type: ignore[return-value] + "PUT", + __path, + params=__query, + headers=__headers, + body=__body, + endpoint_id="connector.update_filtering_validation", + path_parts=__path_parts, + ) + @_rewrite_parameters( body_fields=("index_name",), ) diff --git a/elasticsearch/_async/client/esql.py b/elasticsearch/_async/client/esql.py index d39a86f28..072bedbc1 100644 --- a/elasticsearch/_async/client/esql.py +++ b/elasticsearch/_async/client/esql.py @@ -26,13 +26,14 @@ class EsqlClient(NamespacedClient): @_rewrite_parameters( - body_fields=("query", "columnar", "filter", "locale", "params"), + body_fields=("query", "version", "columnar", "filter", "locale", "params"), ignore_deprecated_options={"params"}, ) async def query( self, *, query: t.Optional[str] = None, + version: t.Optional[t.Union["t.Literal['2024.04.01']", str]] = None, columnar: t.Optional[bool] = None, delimiter: t.Optional[str] = None, error_trace: t.Optional[bool] = None, @@ -52,6 +53,8 @@ async def query( :param query: The ES|QL query API accepts an ES|QL query string in the query parameter, runs it, and returns the results. + :param version: The version of the ES|QL language in which the "query" field + was written. :param columnar: By default, ES|QL returns results as rows. For example, FROM returns each individual document as one row. For the JSON, YAML, CBOR and smile formats, ES|QL can return the results in a columnar fashion where one @@ -68,6 +71,8 @@ async def query( """ if query is None and body is None: raise ValueError("Empty value passed for parameter 'query'") + if version is None and body is None: + raise ValueError("Empty value passed for parameter 'version'") __path_parts: t.Dict[str, str] = {} __path = "/_query" __query: t.Dict[str, t.Any] = {} @@ -87,6 +92,8 @@ async def query( if not __body: if query is not None: __body["query"] = query + if version is not None: + __body["version"] = version if columnar is not None: __body["columnar"] = columnar if filter is not None: diff --git a/elasticsearch/_async/client/indices.py b/elasticsearch/_async/client/indices.py index 31ffc82f6..b73b03dff 100644 --- a/elasticsearch/_async/client/indices.py +++ b/elasticsearch/_async/client/indices.py @@ -2868,8 +2868,11 @@ async def put_data_lifecycle( @_rewrite_parameters( body_fields=( + "allow_auto_create", "composed_of", "data_stream", + "deprecated", + "ignore_missing_component_templates", "index_patterns", "meta", "priority", @@ -2882,13 +2885,20 @@ async def put_index_template( self, *, name: str, + allow_auto_create: t.Optional[bool] = None, + cause: t.Optional[str] = None, composed_of: t.Optional[t.Sequence[str]] = None, create: t.Optional[bool] = None, data_stream: t.Optional[t.Mapping[str, t.Any]] = None, + deprecated: t.Optional[bool] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, + ignore_missing_component_templates: t.Optional[t.Sequence[str]] = None, index_patterns: t.Optional[t.Union[str, t.Sequence[str]]] = None, + master_timeout: t.Optional[ + t.Union["t.Literal[-1]", "t.Literal[0]", str] + ] = None, meta: t.Optional[t.Mapping[str, t.Any]] = None, pretty: t.Optional[bool] = None, priority: t.Optional[int] = None, @@ -2902,6 +2912,13 @@ async def put_index_template( ``_ :param name: Index or template name + :param allow_auto_create: This setting overrides the value of the `action.auto_create_index` + cluster setting. If set to `true` in a template, then indices can be automatically + created using that template even if auto-creation of indices is disabled + via `actions.auto_create_index`. If set to `false`, then indices or data + streams matching the template must always be explicitly created, and may + never be automatically created. + :param cause: User defined reason for creating/updating the index template :param composed_of: An ordered list of component template names. Component templates are merged in the order specified, meaning that the last component template specified has the highest precedence. @@ -2910,7 +2927,16 @@ async def put_index_template( :param data_stream: If this object is included, the template is used to create data streams and their backing indices. Supports an empty object. Data streams require a matching index template with a `data_stream` object. + :param deprecated: Marks this index template as deprecated. When creating or + updating a non-deprecated index template that uses deprecated components, + Elasticsearch will emit a deprecation warning. + :param ignore_missing_component_templates: The configuration option ignore_missing_component_templates + can be used when an index template references a component template that might + not exist :param index_patterns: Name of the index template to create. + :param master_timeout: Period to wait for a connection to the master node. If + no response is received before the timeout expires, the request fails and + returns an error. :param meta: Optional user metadata about the index template. May have any contents. This map is not automatically generated by Elasticsearch. :param priority: Priority to determine index template precedence when a new data @@ -2929,6 +2955,8 @@ async def put_index_template( __path = f'/_index_template/{__path_parts["name"]}' __query: t.Dict[str, t.Any] = {} __body: t.Dict[str, t.Any] = body if body is not None else {} + if cause is not None: + __query["cause"] = cause if create is not None: __query["create"] = create if error_trace is not None: @@ -2937,13 +2965,23 @@ async def put_index_template( __query["filter_path"] = filter_path if human is not None: __query["human"] = human + if master_timeout is not None: + __query["master_timeout"] = master_timeout if pretty is not None: __query["pretty"] = pretty if not __body: + if allow_auto_create is not None: + __body["allow_auto_create"] = allow_auto_create if composed_of is not None: __body["composed_of"] = composed_of if data_stream is not None: __body["data_stream"] = data_stream + if deprecated is not None: + __body["deprecated"] = deprecated + if ignore_missing_component_templates is not None: + __body["ignore_missing_component_templates"] = ( + ignore_missing_component_templates + ) if index_patterns is not None: __body["index_patterns"] = index_patterns if meta is not None: @@ -3250,10 +3288,10 @@ async def put_template( *, name: str, aliases: t.Optional[t.Mapping[str, t.Mapping[str, t.Any]]] = None, + cause: t.Optional[str] = None, create: t.Optional[bool] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, - flat_settings: t.Optional[bool] = None, human: t.Optional[bool] = None, index_patterns: t.Optional[t.Union[str, t.Sequence[str]]] = None, mappings: t.Optional[t.Mapping[str, t.Any]] = None, @@ -3263,7 +3301,6 @@ async def put_template( order: t.Optional[int] = None, pretty: t.Optional[bool] = None, settings: t.Optional[t.Mapping[str, t.Any]] = None, - timeout: t.Optional[t.Union["t.Literal[-1]", "t.Literal[0]", str]] = None, version: t.Optional[int] = None, body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: @@ -3274,9 +3311,9 @@ async def put_template( :param name: The name of the template :param aliases: Aliases for the index. + :param cause: :param create: If true, this request cannot replace or update existing index templates. - :param flat_settings: If `true`, returns settings in flat format. :param index_patterns: Array of wildcard expressions used to match the names of indices during creation. :param mappings: Mapping for fields in the index. @@ -3288,8 +3325,6 @@ async def put_template( Templates with higher 'order' values are merged later, overriding templates with lower values. :param settings: Configuration options for the index. - :param timeout: Period to wait for a response. If no response is received before - the timeout expires, the request fails and returns an error. :param version: Version number used to manage index templates externally. This number is not automatically generated by Elasticsearch. """ @@ -3299,22 +3334,20 @@ async def put_template( __path = f'/_template/{__path_parts["name"]}' __query: t.Dict[str, t.Any] = {} __body: t.Dict[str, t.Any] = body if body is not None else {} + if cause is not None: + __query["cause"] = cause if create is not None: __query["create"] = create if error_trace is not None: __query["error_trace"] = error_trace if filter_path is not None: __query["filter_path"] = filter_path - if flat_settings is not None: - __query["flat_settings"] = flat_settings if human is not None: __query["human"] = human if master_timeout is not None: __query["master_timeout"] = master_timeout if pretty is not None: __query["pretty"] = pretty - if timeout is not None: - __query["timeout"] = timeout if not __body: if aliases is not None: __body["aliases"] = aliases @@ -4004,91 +4037,37 @@ async def shrink( path_parts=__path_parts, ) - @_rewrite_parameters( - body_fields=( - "allow_auto_create", - "composed_of", - "data_stream", - "index_patterns", - "meta", - "priority", - "template", - "version", - ), - parameter_aliases={"_meta": "meta"}, - ) + @_rewrite_parameters() async def simulate_index_template( self, *, name: str, - allow_auto_create: t.Optional[bool] = None, - composed_of: t.Optional[t.Sequence[str]] = None, - create: t.Optional[bool] = None, - data_stream: t.Optional[t.Mapping[str, t.Any]] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, include_defaults: t.Optional[bool] = None, - index_patterns: t.Optional[t.Union[str, t.Sequence[str]]] = None, master_timeout: t.Optional[ t.Union["t.Literal[-1]", "t.Literal[0]", str] ] = None, - meta: t.Optional[t.Mapping[str, t.Any]] = None, pretty: t.Optional[bool] = None, - priority: t.Optional[int] = None, - template: t.Optional[t.Mapping[str, t.Any]] = None, - version: t.Optional[int] = None, - body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: """ Simulate matching the given index name against the index templates in the system ``_ - :param name: Index or template name to simulate - :param allow_auto_create: This setting overrides the value of the `action.auto_create_index` - cluster setting. If set to `true` in a template, then indices can be automatically - created using that template even if auto-creation of indices is disabled - via `actions.auto_create_index`. If set to `false`, then indices or data - streams matching the template must always be explicitly created, and may - never be automatically created. - :param composed_of: An ordered list of component template names. Component templates - are merged in the order specified, meaning that the last component template - specified has the highest precedence. - :param create: If `true`, the template passed in the body is only used if no - existing templates match the same index patterns. If `false`, the simulation - uses the template with the highest priority. Note that the template is not - permanently added or updated in either case; it is only used for the simulation. - :param data_stream: If this object is included, the template is used to create - data streams and their backing indices. Supports an empty object. Data streams - require a matching index template with a `data_stream` object. + :param name: Name of the index to simulate :param include_defaults: If true, returns all relevant default configurations for the index template. - :param index_patterns: Array of wildcard (`*`) expressions used to match the - names of data streams and indices during creation. :param master_timeout: Period to wait for a connection to the master node. If no response is received before the timeout expires, the request fails and returns an error. - :param meta: Optional user metadata about the index template. May have any contents. - This map is not automatically generated by Elasticsearch. - :param priority: Priority to determine index template precedence when a new data - stream or index is created. The index template with the highest priority - is chosen. If no priority is specified the template is treated as though - it is of priority 0 (lowest priority). This number is not automatically generated - by Elasticsearch. - :param template: Template to be applied. It may optionally include an `aliases`, - `mappings`, or `settings` configuration. - :param version: Version number used to manage index templates externally. This - number is not automatically generated by Elasticsearch. """ if name in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'name'") __path_parts: t.Dict[str, str] = {"name": _quote(name)} __path = f'/_index_template/_simulate_index/{__path_parts["name"]}' __query: t.Dict[str, t.Any] = {} - __body: t.Dict[str, t.Any] = body if body is not None else {} - if create is not None: - __query["create"] = create if error_trace is not None: __query["error_trace"] = error_trace if filter_path is not None: @@ -4101,56 +4080,55 @@ async def simulate_index_template( __query["master_timeout"] = master_timeout if pretty is not None: __query["pretty"] = pretty - if not __body: - if allow_auto_create is not None: - __body["allow_auto_create"] = allow_auto_create - if composed_of is not None: - __body["composed_of"] = composed_of - if data_stream is not None: - __body["data_stream"] = data_stream - if index_patterns is not None: - __body["index_patterns"] = index_patterns - if meta is not None: - __body["_meta"] = meta - if priority is not None: - __body["priority"] = priority - if template is not None: - __body["template"] = template - if version is not None: - __body["version"] = version - if not __body: - __body = None # type: ignore[assignment] __headers = {"accept": "application/json"} - if __body is not None: - __headers["content-type"] = "application/json" return await self.perform_request( # type: ignore[return-value] "POST", __path, params=__query, headers=__headers, - body=__body, endpoint_id="indices.simulate_index_template", path_parts=__path_parts, ) @_rewrite_parameters( - body_name="template", + body_fields=( + "allow_auto_create", + "composed_of", + "data_stream", + "deprecated", + "ignore_missing_component_templates", + "index_patterns", + "meta", + "priority", + "template", + "version", + ), + parameter_aliases={"_meta": "meta"}, ) async def simulate_template( self, *, name: t.Optional[str] = None, + allow_auto_create: t.Optional[bool] = None, + composed_of: t.Optional[t.Sequence[str]] = None, create: t.Optional[bool] = None, + data_stream: t.Optional[t.Mapping[str, t.Any]] = None, + deprecated: t.Optional[bool] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, + ignore_missing_component_templates: t.Optional[t.Sequence[str]] = None, include_defaults: t.Optional[bool] = None, + index_patterns: t.Optional[t.Union[str, t.Sequence[str]]] = None, master_timeout: t.Optional[ t.Union["t.Literal[-1]", "t.Literal[0]", str] ] = None, + meta: t.Optional[t.Mapping[str, t.Any]] = None, pretty: t.Optional[bool] = None, + priority: t.Optional[int] = None, template: t.Optional[t.Mapping[str, t.Any]] = None, - body: t.Optional[t.Mapping[str, t.Any]] = None, + version: t.Optional[int] = None, + body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: """ Simulate resolving the given template name or body @@ -4160,23 +4138,47 @@ async def simulate_template( :param name: Name of the index template to simulate. To test a template configuration before you add it to the cluster, omit this parameter and specify the template configuration in the request body. + :param allow_auto_create: This setting overrides the value of the `action.auto_create_index` + cluster setting. If set to `true` in a template, then indices can be automatically + created using that template even if auto-creation of indices is disabled + via `actions.auto_create_index`. If set to `false`, then indices or data + streams matching the template must always be explicitly created, and may + never be automatically created. + :param composed_of: An ordered list of component template names. Component templates + are merged in the order specified, meaning that the last component template + specified has the highest precedence. :param create: If true, the template passed in the body is only used if no existing templates match the same index patterns. If false, the simulation uses the template with the highest priority. Note that the template is not permanently added or updated in either case; it is only used for the simulation. + :param data_stream: If this object is included, the template is used to create + data streams and their backing indices. Supports an empty object. Data streams + require a matching index template with a `data_stream` object. + :param deprecated: Marks this index template as deprecated. When creating or + updating a non-deprecated index template that uses deprecated components, + Elasticsearch will emit a deprecation warning. + :param ignore_missing_component_templates: The configuration option ignore_missing_component_templates + can be used when an index template references a component template that might + not exist :param include_defaults: If true, returns all relevant default configurations for the index template. + :param index_patterns: Array of wildcard (`*`) expressions used to match the + names of data streams and indices during creation. :param master_timeout: Period to wait for a connection to the master node. If no response is received before the timeout expires, the request fails and returns an error. - :param template: + :param meta: Optional user metadata about the index template. May have any contents. + This map is not automatically generated by Elasticsearch. + :param priority: Priority to determine index template precedence when a new data + stream or index is created. The index template with the highest priority + is chosen. If no priority is specified the template is treated as though + it is of priority 0 (lowest priority). This number is not automatically generated + by Elasticsearch. + :param template: Template to be applied. It may optionally include an `aliases`, + `mappings`, or `settings` configuration. + :param version: Version number used to manage index templates externally. This + number is not automatically generated by Elasticsearch. """ - if template is None and body is None: - raise ValueError( - "Empty value passed for parameters 'template' and 'body', one of them should be set." - ) - elif template is not None and body is not None: - raise ValueError("Cannot set both 'template' and 'body'") __path_parts: t.Dict[str, str] if name not in SKIP_IN_PATH: __path_parts = {"name": _quote(name)} @@ -4185,6 +4187,7 @@ async def simulate_template( __path_parts = {} __path = "/_index_template/_simulate" __query: t.Dict[str, t.Any] = {} + __body: t.Dict[str, t.Any] = body if body is not None else {} if create is not None: __query["create"] = create if error_trace is not None: @@ -4199,9 +4202,31 @@ async def simulate_template( __query["master_timeout"] = master_timeout if pretty is not None: __query["pretty"] = pretty - __body = template if template is not None else body if not __body: - __body = None + if allow_auto_create is not None: + __body["allow_auto_create"] = allow_auto_create + if composed_of is not None: + __body["composed_of"] = composed_of + if data_stream is not None: + __body["data_stream"] = data_stream + if deprecated is not None: + __body["deprecated"] = deprecated + if ignore_missing_component_templates is not None: + __body["ignore_missing_component_templates"] = ( + ignore_missing_component_templates + ) + if index_patterns is not None: + __body["index_patterns"] = index_patterns + if meta is not None: + __body["_meta"] = meta + if priority is not None: + __body["priority"] = priority + if template is not None: + __body["template"] = template + if version is not None: + __body["version"] = version + if not __body: + __body = None # type: ignore[assignment] __headers = {"accept": "application/json"} if __body is not None: __headers["content-type"] = "application/json" diff --git a/elasticsearch/_async/client/inference.py b/elasticsearch/_async/client/inference.py index 07e4953cf..c7669d1e6 100644 --- a/elasticsearch/_async/client/inference.py +++ b/elasticsearch/_async/client/inference.py @@ -26,7 +26,7 @@ class InferenceClient(NamespacedClient): @_rewrite_parameters() - async def delete_model( + async def delete( self, *, inference_id: str, @@ -42,7 +42,7 @@ async def delete_model( pretty: t.Optional[bool] = None, ) -> ObjectApiResponse[t.Any]: """ - Delete model in the Inference API + Delete an inference endpoint ``_ @@ -78,36 +78,34 @@ async def delete_model( __path, params=__query, headers=__headers, - endpoint_id="inference.delete_model", + endpoint_id="inference.delete", path_parts=__path_parts, ) @_rewrite_parameters() - async def get_model( + async def get( self, *, - inference_id: str, task_type: t.Optional[ t.Union[ "t.Literal['completion', 'rerank', 'sparse_embedding', 'text_embedding']", str, ] ] = None, + inference_id: t.Optional[str] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, pretty: t.Optional[bool] = None, ) -> ObjectApiResponse[t.Any]: """ - Get a model in the Inference API + Get an inference endpoint ``_ - :param inference_id: The inference Id :param task_type: The task type + :param inference_id: The inference Id """ - if inference_id in SKIP_IN_PATH: - raise ValueError("Empty value passed for parameter 'inference_id'") __path_parts: t.Dict[str, str] if task_type not in SKIP_IN_PATH and inference_id not in SKIP_IN_PATH: __path_parts = { @@ -119,7 +117,8 @@ async def get_model( __path_parts = {"inference_id": _quote(inference_id)} __path = f'/_inference/{__path_parts["inference_id"]}' else: - raise ValueError("Couldn't find a path for the given parameters") + __path_parts = {} + __path = "/_inference" __query: t.Dict[str, t.Any] = {} if error_trace is not None: __query["error_trace"] = error_trace @@ -135,7 +134,7 @@ async def get_model( __path, params=__query, headers=__headers, - endpoint_id="inference.get_model", + endpoint_id="inference.get", path_parts=__path_parts, ) @@ -159,18 +158,21 @@ async def inference( pretty: t.Optional[bool] = None, query: t.Optional[str] = None, task_settings: t.Optional[t.Any] = None, + timeout: t.Optional[t.Union["t.Literal[-1]", "t.Literal[0]", str]] = None, body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: """ - Perform inference on a model + Perform inference ``_ :param inference_id: The inference Id - :param input: Text input to the model. Either a string or an array of strings. + :param input: Inference input. Either a string or an array of strings. :param task_type: The task type :param query: Query input, required for rerank task. Not required for other tasks. :param task_settings: Optional task settings + :param timeout: Specifies the amount of time to wait for the inference request + to complete. """ if inference_id in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'inference_id'") @@ -198,6 +200,8 @@ async def inference( __query["human"] = human if pretty is not None: __query["pretty"] = pretty + if timeout is not None: + __query["timeout"] = timeout if not __body: if input is not None: __body["input"] = input @@ -221,13 +225,13 @@ async def inference( ) @_rewrite_parameters( - body_name="model_config", + body_name="inference_config", ) - async def put_model( + async def put( self, *, inference_id: str, - model_config: t.Optional[t.Mapping[str, t.Any]] = None, + inference_config: t.Optional[t.Mapping[str, t.Any]] = None, body: t.Optional[t.Mapping[str, t.Any]] = None, task_type: t.Optional[ t.Union[ @@ -241,22 +245,22 @@ async def put_model( pretty: t.Optional[bool] = None, ) -> ObjectApiResponse[t.Any]: """ - Configure a model for use in the Inference API + Configure an inference endpoint for use in the Inference API ``_ :param inference_id: The inference Id - :param model_config: + :param inference_config: :param task_type: The task type """ if inference_id in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'inference_id'") - if model_config is None and body is None: + if inference_config is None and body is None: raise ValueError( - "Empty value passed for parameters 'model_config' and 'body', one of them should be set." + "Empty value passed for parameters 'inference_config' and 'body', one of them should be set." ) - elif model_config is not None and body is not None: - raise ValueError("Cannot set both 'model_config' and 'body'") + elif inference_config is not None and body is not None: + raise ValueError("Cannot set both 'inference_config' and 'body'") __path_parts: t.Dict[str, str] if task_type not in SKIP_IN_PATH and inference_id not in SKIP_IN_PATH: __path_parts = { @@ -278,7 +282,7 @@ async def put_model( __query["human"] = human if pretty is not None: __query["pretty"] = pretty - __body = model_config if model_config is not None else body + __body = inference_config if inference_config is not None else body __headers = {"accept": "application/json", "content-type": "application/json"} return await self.perform_request( # type: ignore[return-value] "PUT", @@ -286,6 +290,6 @@ async def put_model( params=__query, headers=__headers, body=__body, - endpoint_id="inference.put_model", + endpoint_id="inference.put", path_parts=__path_parts, ) diff --git a/elasticsearch/_async/client/security.py b/elasticsearch/_async/client/security.py index 3171ddfc8..1fecf18e8 100644 --- a/elasticsearch/_async/client/security.py +++ b/elasticsearch/_async/client/security.py @@ -1761,7 +1761,7 @@ async def has_privileges( cluster: t.Optional[ t.Sequence[ t.Union[ - "t.Literal['all', 'cancel_task', 'create_snapshot', 'grant_api_key', 'manage', 'manage_api_key', 'manage_ccr', 'manage_enrich', 'manage_ilm', 'manage_index_templates', 'manage_ingest_pipelines', 'manage_logstash_pipelines', 'manage_ml', 'manage_oidc', 'manage_own_api_key', 'manage_pipeline', 'manage_rollup', 'manage_saml', 'manage_security', 'manage_service_account', 'manage_slm', 'manage_token', 'manage_transform', 'manage_user_profile', 'manage_watcher', 'monitor', 'monitor_ml', 'monitor_rollup', 'monitor_snapshot', 'monitor_text_structure', 'monitor_transform', 'monitor_watcher', 'read_ccr', 'read_ilm', 'read_pipeline', 'read_slm', 'transport_client']", + "t.Literal['all', 'cancel_task', 'create_snapshot', 'cross_cluster_replication', 'cross_cluster_search', 'delegate_pki', 'grant_api_key', 'manage', 'manage_api_key', 'manage_autoscaling', 'manage_behavioral_analytics', 'manage_ccr', 'manage_data_frame_transforms', 'manage_data_stream_global_retention', 'manage_enrich', 'manage_ilm', 'manage_index_templates', 'manage_inference', 'manage_ingest_pipelines', 'manage_logstash_pipelines', 'manage_ml', 'manage_oidc', 'manage_own_api_key', 'manage_pipeline', 'manage_rollup', 'manage_saml', 'manage_search_application', 'manage_search_query_rules', 'manage_search_synonyms', 'manage_security', 'manage_service_account', 'manage_slm', 'manage_token', 'manage_transform', 'manage_user_profile', 'manage_watcher', 'monitor', 'monitor_data_frame_transforms', 'monitor_data_stream_global_retention', 'monitor_enrich', 'monitor_inference', 'monitor_ml', 'monitor_rollup', 'monitor_snapshot', 'monitor_text_structure', 'monitor_transform', 'monitor_watcher', 'none', 'post_behavioral_analytics_event', 'read_ccr', 'read_connector_secrets', 'read_fleet_secrets', 'read_ilm', 'read_pipeline', 'read_security', 'read_slm', 'transport_client', 'write_connector_secrets', 'write_fleet_secrets']", str, ] ] @@ -2084,7 +2084,7 @@ async def put_role( cluster: t.Optional[ t.Sequence[ t.Union[ - "t.Literal['all', 'cancel_task', 'create_snapshot', 'grant_api_key', 'manage', 'manage_api_key', 'manage_ccr', 'manage_enrich', 'manage_ilm', 'manage_index_templates', 'manage_ingest_pipelines', 'manage_logstash_pipelines', 'manage_ml', 'manage_oidc', 'manage_own_api_key', 'manage_pipeline', 'manage_rollup', 'manage_saml', 'manage_security', 'manage_service_account', 'manage_slm', 'manage_token', 'manage_transform', 'manage_user_profile', 'manage_watcher', 'monitor', 'monitor_ml', 'monitor_rollup', 'monitor_snapshot', 'monitor_text_structure', 'monitor_transform', 'monitor_watcher', 'read_ccr', 'read_ilm', 'read_pipeline', 'read_slm', 'transport_client']", + "t.Literal['all', 'cancel_task', 'create_snapshot', 'cross_cluster_replication', 'cross_cluster_search', 'delegate_pki', 'grant_api_key', 'manage', 'manage_api_key', 'manage_autoscaling', 'manage_behavioral_analytics', 'manage_ccr', 'manage_data_frame_transforms', 'manage_data_stream_global_retention', 'manage_enrich', 'manage_ilm', 'manage_index_templates', 'manage_inference', 'manage_ingest_pipelines', 'manage_logstash_pipelines', 'manage_ml', 'manage_oidc', 'manage_own_api_key', 'manage_pipeline', 'manage_rollup', 'manage_saml', 'manage_search_application', 'manage_search_query_rules', 'manage_search_synonyms', 'manage_security', 'manage_service_account', 'manage_slm', 'manage_token', 'manage_transform', 'manage_user_profile', 'manage_watcher', 'monitor', 'monitor_data_frame_transforms', 'monitor_data_stream_global_retention', 'monitor_enrich', 'monitor_inference', 'monitor_ml', 'monitor_rollup', 'monitor_snapshot', 'monitor_text_structure', 'monitor_transform', 'monitor_watcher', 'none', 'post_behavioral_analytics_event', 'read_ccr', 'read_connector_secrets', 'read_fleet_secrets', 'read_ilm', 'read_pipeline', 'read_security', 'read_slm', 'transport_client', 'write_connector_secrets', 'write_fleet_secrets']", str, ] ] diff --git a/elasticsearch/_sync/client/cluster.py b/elasticsearch/_sync/client/cluster.py index 722e4eac7..6cc0a3309 100644 --- a/elasticsearch/_sync/client/cluster.py +++ b/elasticsearch/_sync/client/cluster.py @@ -661,7 +661,7 @@ def post_voting_config_exclusions( ) @_rewrite_parameters( - body_fields=("template", "allow_auto_create", "meta", "version"), + body_fields=("template", "deprecated", "meta", "version"), parameter_aliases={"_meta": "meta"}, ) def put_component_template( @@ -669,8 +669,9 @@ def put_component_template( *, name: str, template: t.Optional[t.Mapping[str, t.Any]] = None, - allow_auto_create: t.Optional[bool] = None, + cause: t.Optional[str] = None, create: t.Optional[bool] = None, + deprecated: t.Optional[bool] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, @@ -698,13 +699,12 @@ def put_component_template( update settings API. :param template: The template to be applied which includes mappings, settings, or aliases configuration. - :param allow_auto_create: This setting overrides the value of the `action.auto_create_index` - cluster setting. If set to `true` in a template, then indices can be automatically - created using that template even if auto-creation of indices is disabled - via `actions.auto_create_index`. If set to `false` then data streams matching - the template must always be explicitly created. + :param cause: :param create: If `true`, this request cannot replace or update existing component templates. + :param deprecated: Marks this index template as deprecated. When creating or + updating a non-deprecated index template that uses deprecated components, + Elasticsearch will emit a deprecation warning. :param master_timeout: Period to wait for a connection to the master node. If no response is received before the timeout expires, the request fails and returns an error. @@ -724,6 +724,8 @@ def put_component_template( __path = f'/_component_template/{__path_parts["name"]}' __query: t.Dict[str, t.Any] = {} __body: t.Dict[str, t.Any] = body if body is not None else {} + if cause is not None: + __query["cause"] = cause if create is not None: __query["create"] = create if error_trace is not None: @@ -739,8 +741,8 @@ def put_component_template( if not __body: if template is not None: __body["template"] = template - if allow_auto_create is not None: - __body["allow_auto_create"] = allow_auto_create + if deprecated is not None: + __body["deprecated"] = deprecated if meta is not None: __body["_meta"] = meta if version is not None: diff --git a/elasticsearch/_sync/client/connector.py b/elasticsearch/_sync/client/connector.py index 93e292e72..e343cf62c 100644 --- a/elasticsearch/_sync/client/connector.py +++ b/elasticsearch/_sync/client/connector.py @@ -70,6 +70,7 @@ def delete( self, *, connector_id: str, + delete_sync_jobs: bool, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, @@ -81,12 +82,17 @@ def delete( ``_ :param connector_id: The unique identifier of the connector to be deleted + :param delete_sync_jobs: Determines whether associated sync jobs are also deleted. """ if connector_id in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'connector_id'") + if delete_sync_jobs is None: + raise ValueError("Empty value passed for parameter 'delete_sync_jobs'") __path_parts: t.Dict[str, str] = {"connector_id": _quote(connector_id)} __path = f'/_connector/{__path_parts["connector_id"]}' __query: t.Dict[str, t.Any] = {} + if delete_sync_jobs is not None: + __query["delete_sync_jobs"] = delete_sync_jobs if error_trace is not None: __query["error_trace"] = error_trace if filter_path is not None: @@ -486,6 +492,302 @@ def put( path_parts=__path_parts, ) + @_rewrite_parameters() + def sync_job_cancel( + self, + *, + connector_sync_job_id: str, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Cancels a connector sync job. + + ``_ + + :param connector_sync_job_id: The unique identifier of the connector sync job + """ + if connector_sync_job_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_sync_job_id'") + __path_parts: t.Dict[str, str] = { + "connector_sync_job_id": _quote(connector_sync_job_id) + } + __path = ( + f'/_connector/_sync_job/{__path_parts["connector_sync_job_id"]}/_cancel' + ) + __query: t.Dict[str, t.Any] = {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + __headers = {"accept": "application/json"} + return self.perform_request( # type: ignore[return-value] + "PUT", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.sync_job_cancel", + path_parts=__path_parts, + ) + + @_rewrite_parameters() + def sync_job_delete( + self, + *, + connector_sync_job_id: str, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Deletes a connector sync job. + + ``_ + + :param connector_sync_job_id: The unique identifier of the connector sync job + to be deleted + """ + if connector_sync_job_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_sync_job_id'") + __path_parts: t.Dict[str, str] = { + "connector_sync_job_id": _quote(connector_sync_job_id) + } + __path = f'/_connector/_sync_job/{__path_parts["connector_sync_job_id"]}' + __query: t.Dict[str, t.Any] = {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + __headers = {"accept": "application/json"} + return self.perform_request( # type: ignore[return-value] + "DELETE", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.sync_job_delete", + path_parts=__path_parts, + ) + + @_rewrite_parameters() + def sync_job_get( + self, + *, + connector_sync_job_id: str, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Returns the details about a connector sync job. + + ``_ + + :param connector_sync_job_id: The unique identifier of the connector sync job + """ + if connector_sync_job_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_sync_job_id'") + __path_parts: t.Dict[str, str] = { + "connector_sync_job_id": _quote(connector_sync_job_id) + } + __path = f'/_connector/_sync_job/{__path_parts["connector_sync_job_id"]}' + __query: t.Dict[str, t.Any] = {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + __headers = {"accept": "application/json"} + return self.perform_request( # type: ignore[return-value] + "GET", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.sync_job_get", + path_parts=__path_parts, + ) + + @_rewrite_parameters( + parameter_aliases={"from": "from_"}, + ) + def sync_job_list( + self, + *, + connector_id: t.Optional[str] = None, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + from_: t.Optional[int] = None, + human: t.Optional[bool] = None, + job_type: t.Optional[ + t.Sequence[ + t.Union["t.Literal['access_control', 'full', 'incremental']", str] + ] + ] = None, + pretty: t.Optional[bool] = None, + size: t.Optional[int] = None, + status: t.Optional[ + t.Union[ + "t.Literal['canceled', 'canceling', 'completed', 'error', 'in_progress', 'pending', 'suspended']", + str, + ] + ] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Lists all connector sync jobs. + + ``_ + + :param connector_id: A connector id to fetch connector sync jobs for + :param from_: Starting offset (default: 0) + :param job_type: A comma-separated list of job types to fetch the sync jobs for + :param size: Specifies a max number of results to get + :param status: A sync job status to fetch connector sync jobs for + """ + __path_parts: t.Dict[str, str] = {} + __path = "/_connector/_sync_job" + __query: t.Dict[str, t.Any] = {} + if connector_id is not None: + __query["connector_id"] = connector_id + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if from_ is not None: + __query["from"] = from_ + if human is not None: + __query["human"] = human + if job_type is not None: + __query["job_type"] = job_type + if pretty is not None: + __query["pretty"] = pretty + if size is not None: + __query["size"] = size + if status is not None: + __query["status"] = status + __headers = {"accept": "application/json"} + return self.perform_request( # type: ignore[return-value] + "GET", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.sync_job_list", + path_parts=__path_parts, + ) + + @_rewrite_parameters( + body_fields=("id", "job_type", "trigger_method"), + ) + def sync_job_post( + self, + *, + id: t.Optional[str] = None, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + job_type: t.Optional[ + t.Union["t.Literal['access_control', 'full', 'incremental']", str] + ] = None, + pretty: t.Optional[bool] = None, + trigger_method: t.Optional[ + t.Union["t.Literal['on_demand', 'scheduled']", str] + ] = None, + body: t.Optional[t.Dict[str, t.Any]] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Creates a connector sync job. + + ``_ + + :param id: The id of the associated connector + :param job_type: + :param trigger_method: + """ + if id is None and body is None: + raise ValueError("Empty value passed for parameter 'id'") + __path_parts: t.Dict[str, str] = {} + __path = "/_connector/_sync_job" + __query: t.Dict[str, t.Any] = {} + __body: t.Dict[str, t.Any] = body if body is not None else {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + if not __body: + if id is not None: + __body["id"] = id + if job_type is not None: + __body["job_type"] = job_type + if trigger_method is not None: + __body["trigger_method"] = trigger_method + __headers = {"accept": "application/json", "content-type": "application/json"} + return self.perform_request( # type: ignore[return-value] + "POST", + __path, + params=__query, + headers=__headers, + body=__body, + endpoint_id="connector.sync_job_post", + path_parts=__path_parts, + ) + + @_rewrite_parameters() + def update_active_filtering( + self, + *, + connector_id: str, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Activates the draft filtering rules if they are in a validated state. + + ``_ + + :param connector_id: The unique identifier of the connector to be updated + """ + if connector_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_id'") + __path_parts: t.Dict[str, str] = {"connector_id": _quote(connector_id)} + __path = f'/_connector/{__path_parts["connector_id"]}/_filtering/_activate' + __query: t.Dict[str, t.Any] = {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + __headers = {"accept": "application/json"} + return self.perform_request( # type: ignore[return-value] + "PUT", + __path, + params=__query, + headers=__headers, + endpoint_id="connector.update_active_filtering", + path_parts=__path_parts, + ) + @_rewrite_parameters( body_fields=("api_key_id", "api_key_secret_id"), ) @@ -647,17 +949,19 @@ def update_error( ) @_rewrite_parameters( - body_fields=("filtering",), + body_fields=("advanced_snippet", "filtering", "rules"), ) def update_filtering( self, *, connector_id: str, - filtering: t.Optional[t.Sequence[t.Mapping[str, t.Any]]] = None, + advanced_snippet: t.Optional[t.Mapping[str, t.Any]] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + filtering: t.Optional[t.Sequence[t.Mapping[str, t.Any]]] = None, human: t.Optional[bool] = None, pretty: t.Optional[bool] = None, + rules: t.Optional[t.Sequence[t.Mapping[str, t.Any]]] = None, body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: """ @@ -666,12 +970,12 @@ def update_filtering( ``_ :param connector_id: The unique identifier of the connector to be updated + :param advanced_snippet: :param filtering: + :param rules: """ if connector_id in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'connector_id'") - if filtering is None and body is None: - raise ValueError("Empty value passed for parameter 'filtering'") __path_parts: t.Dict[str, str] = {"connector_id": _quote(connector_id)} __path = f'/_connector/{__path_parts["connector_id"]}/_filtering' __query: t.Dict[str, t.Any] = {} @@ -685,8 +989,12 @@ def update_filtering( if pretty is not None: __query["pretty"] = pretty if not __body: + if advanced_snippet is not None: + __body["advanced_snippet"] = advanced_snippet if filtering is not None: __body["filtering"] = filtering + if rules is not None: + __body["rules"] = rules __headers = {"accept": "application/json", "content-type": "application/json"} return self.perform_request( # type: ignore[return-value] "PUT", @@ -698,6 +1006,58 @@ def update_filtering( path_parts=__path_parts, ) + @_rewrite_parameters( + body_fields=("validation",), + ) + def update_filtering_validation( + self, + *, + connector_id: str, + validation: t.Optional[t.Mapping[str, t.Any]] = None, + error_trace: t.Optional[bool] = None, + filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, + human: t.Optional[bool] = None, + pretty: t.Optional[bool] = None, + body: t.Optional[t.Dict[str, t.Any]] = None, + ) -> ObjectApiResponse[t.Any]: + """ + Updates the validation info of the draft filtering rules. + + ``_ + + :param connector_id: The unique identifier of the connector to be updated + :param validation: + """ + if connector_id in SKIP_IN_PATH: + raise ValueError("Empty value passed for parameter 'connector_id'") + if validation is None and body is None: + raise ValueError("Empty value passed for parameter 'validation'") + __path_parts: t.Dict[str, str] = {"connector_id": _quote(connector_id)} + __path = f'/_connector/{__path_parts["connector_id"]}/_filtering/_validation' + __query: t.Dict[str, t.Any] = {} + __body: t.Dict[str, t.Any] = body if body is not None else {} + if error_trace is not None: + __query["error_trace"] = error_trace + if filter_path is not None: + __query["filter_path"] = filter_path + if human is not None: + __query["human"] = human + if pretty is not None: + __query["pretty"] = pretty + if not __body: + if validation is not None: + __body["validation"] = validation + __headers = {"accept": "application/json", "content-type": "application/json"} + return self.perform_request( # type: ignore[return-value] + "PUT", + __path, + params=__query, + headers=__headers, + body=__body, + endpoint_id="connector.update_filtering_validation", + path_parts=__path_parts, + ) + @_rewrite_parameters( body_fields=("index_name",), ) diff --git a/elasticsearch/_sync/client/esql.py b/elasticsearch/_sync/client/esql.py index 1dee2e934..ee83ec3d2 100644 --- a/elasticsearch/_sync/client/esql.py +++ b/elasticsearch/_sync/client/esql.py @@ -26,13 +26,14 @@ class EsqlClient(NamespacedClient): @_rewrite_parameters( - body_fields=("query", "columnar", "filter", "locale", "params"), + body_fields=("query", "version", "columnar", "filter", "locale", "params"), ignore_deprecated_options={"params"}, ) def query( self, *, query: t.Optional[str] = None, + version: t.Optional[t.Union["t.Literal['2024.04.01']", str]] = None, columnar: t.Optional[bool] = None, delimiter: t.Optional[str] = None, error_trace: t.Optional[bool] = None, @@ -52,6 +53,8 @@ def query( :param query: The ES|QL query API accepts an ES|QL query string in the query parameter, runs it, and returns the results. + :param version: The version of the ES|QL language in which the "query" field + was written. :param columnar: By default, ES|QL returns results as rows. For example, FROM returns each individual document as one row. For the JSON, YAML, CBOR and smile formats, ES|QL can return the results in a columnar fashion where one @@ -68,6 +71,8 @@ def query( """ if query is None and body is None: raise ValueError("Empty value passed for parameter 'query'") + if version is None and body is None: + raise ValueError("Empty value passed for parameter 'version'") __path_parts: t.Dict[str, str] = {} __path = "/_query" __query: t.Dict[str, t.Any] = {} @@ -87,6 +92,8 @@ def query( if not __body: if query is not None: __body["query"] = query + if version is not None: + __body["version"] = version if columnar is not None: __body["columnar"] = columnar if filter is not None: diff --git a/elasticsearch/_sync/client/indices.py b/elasticsearch/_sync/client/indices.py index 773e27322..aeca66804 100644 --- a/elasticsearch/_sync/client/indices.py +++ b/elasticsearch/_sync/client/indices.py @@ -2868,8 +2868,11 @@ def put_data_lifecycle( @_rewrite_parameters( body_fields=( + "allow_auto_create", "composed_of", "data_stream", + "deprecated", + "ignore_missing_component_templates", "index_patterns", "meta", "priority", @@ -2882,13 +2885,20 @@ def put_index_template( self, *, name: str, + allow_auto_create: t.Optional[bool] = None, + cause: t.Optional[str] = None, composed_of: t.Optional[t.Sequence[str]] = None, create: t.Optional[bool] = None, data_stream: t.Optional[t.Mapping[str, t.Any]] = None, + deprecated: t.Optional[bool] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, + ignore_missing_component_templates: t.Optional[t.Sequence[str]] = None, index_patterns: t.Optional[t.Union[str, t.Sequence[str]]] = None, + master_timeout: t.Optional[ + t.Union["t.Literal[-1]", "t.Literal[0]", str] + ] = None, meta: t.Optional[t.Mapping[str, t.Any]] = None, pretty: t.Optional[bool] = None, priority: t.Optional[int] = None, @@ -2902,6 +2912,13 @@ def put_index_template( ``_ :param name: Index or template name + :param allow_auto_create: This setting overrides the value of the `action.auto_create_index` + cluster setting. If set to `true` in a template, then indices can be automatically + created using that template even if auto-creation of indices is disabled + via `actions.auto_create_index`. If set to `false`, then indices or data + streams matching the template must always be explicitly created, and may + never be automatically created. + :param cause: User defined reason for creating/updating the index template :param composed_of: An ordered list of component template names. Component templates are merged in the order specified, meaning that the last component template specified has the highest precedence. @@ -2910,7 +2927,16 @@ def put_index_template( :param data_stream: If this object is included, the template is used to create data streams and their backing indices. Supports an empty object. Data streams require a matching index template with a `data_stream` object. + :param deprecated: Marks this index template as deprecated. When creating or + updating a non-deprecated index template that uses deprecated components, + Elasticsearch will emit a deprecation warning. + :param ignore_missing_component_templates: The configuration option ignore_missing_component_templates + can be used when an index template references a component template that might + not exist :param index_patterns: Name of the index template to create. + :param master_timeout: Period to wait for a connection to the master node. If + no response is received before the timeout expires, the request fails and + returns an error. :param meta: Optional user metadata about the index template. May have any contents. This map is not automatically generated by Elasticsearch. :param priority: Priority to determine index template precedence when a new data @@ -2929,6 +2955,8 @@ def put_index_template( __path = f'/_index_template/{__path_parts["name"]}' __query: t.Dict[str, t.Any] = {} __body: t.Dict[str, t.Any] = body if body is not None else {} + if cause is not None: + __query["cause"] = cause if create is not None: __query["create"] = create if error_trace is not None: @@ -2937,13 +2965,23 @@ def put_index_template( __query["filter_path"] = filter_path if human is not None: __query["human"] = human + if master_timeout is not None: + __query["master_timeout"] = master_timeout if pretty is not None: __query["pretty"] = pretty if not __body: + if allow_auto_create is not None: + __body["allow_auto_create"] = allow_auto_create if composed_of is not None: __body["composed_of"] = composed_of if data_stream is not None: __body["data_stream"] = data_stream + if deprecated is not None: + __body["deprecated"] = deprecated + if ignore_missing_component_templates is not None: + __body["ignore_missing_component_templates"] = ( + ignore_missing_component_templates + ) if index_patterns is not None: __body["index_patterns"] = index_patterns if meta is not None: @@ -3250,10 +3288,10 @@ def put_template( *, name: str, aliases: t.Optional[t.Mapping[str, t.Mapping[str, t.Any]]] = None, + cause: t.Optional[str] = None, create: t.Optional[bool] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, - flat_settings: t.Optional[bool] = None, human: t.Optional[bool] = None, index_patterns: t.Optional[t.Union[str, t.Sequence[str]]] = None, mappings: t.Optional[t.Mapping[str, t.Any]] = None, @@ -3263,7 +3301,6 @@ def put_template( order: t.Optional[int] = None, pretty: t.Optional[bool] = None, settings: t.Optional[t.Mapping[str, t.Any]] = None, - timeout: t.Optional[t.Union["t.Literal[-1]", "t.Literal[0]", str]] = None, version: t.Optional[int] = None, body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: @@ -3274,9 +3311,9 @@ def put_template( :param name: The name of the template :param aliases: Aliases for the index. + :param cause: :param create: If true, this request cannot replace or update existing index templates. - :param flat_settings: If `true`, returns settings in flat format. :param index_patterns: Array of wildcard expressions used to match the names of indices during creation. :param mappings: Mapping for fields in the index. @@ -3288,8 +3325,6 @@ def put_template( Templates with higher 'order' values are merged later, overriding templates with lower values. :param settings: Configuration options for the index. - :param timeout: Period to wait for a response. If no response is received before - the timeout expires, the request fails and returns an error. :param version: Version number used to manage index templates externally. This number is not automatically generated by Elasticsearch. """ @@ -3299,22 +3334,20 @@ def put_template( __path = f'/_template/{__path_parts["name"]}' __query: t.Dict[str, t.Any] = {} __body: t.Dict[str, t.Any] = body if body is not None else {} + if cause is not None: + __query["cause"] = cause if create is not None: __query["create"] = create if error_trace is not None: __query["error_trace"] = error_trace if filter_path is not None: __query["filter_path"] = filter_path - if flat_settings is not None: - __query["flat_settings"] = flat_settings if human is not None: __query["human"] = human if master_timeout is not None: __query["master_timeout"] = master_timeout if pretty is not None: __query["pretty"] = pretty - if timeout is not None: - __query["timeout"] = timeout if not __body: if aliases is not None: __body["aliases"] = aliases @@ -4004,91 +4037,37 @@ def shrink( path_parts=__path_parts, ) - @_rewrite_parameters( - body_fields=( - "allow_auto_create", - "composed_of", - "data_stream", - "index_patterns", - "meta", - "priority", - "template", - "version", - ), - parameter_aliases={"_meta": "meta"}, - ) + @_rewrite_parameters() def simulate_index_template( self, *, name: str, - allow_auto_create: t.Optional[bool] = None, - composed_of: t.Optional[t.Sequence[str]] = None, - create: t.Optional[bool] = None, - data_stream: t.Optional[t.Mapping[str, t.Any]] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, include_defaults: t.Optional[bool] = None, - index_patterns: t.Optional[t.Union[str, t.Sequence[str]]] = None, master_timeout: t.Optional[ t.Union["t.Literal[-1]", "t.Literal[0]", str] ] = None, - meta: t.Optional[t.Mapping[str, t.Any]] = None, pretty: t.Optional[bool] = None, - priority: t.Optional[int] = None, - template: t.Optional[t.Mapping[str, t.Any]] = None, - version: t.Optional[int] = None, - body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: """ Simulate matching the given index name against the index templates in the system ``_ - :param name: Index or template name to simulate - :param allow_auto_create: This setting overrides the value of the `action.auto_create_index` - cluster setting. If set to `true` in a template, then indices can be automatically - created using that template even if auto-creation of indices is disabled - via `actions.auto_create_index`. If set to `false`, then indices or data - streams matching the template must always be explicitly created, and may - never be automatically created. - :param composed_of: An ordered list of component template names. Component templates - are merged in the order specified, meaning that the last component template - specified has the highest precedence. - :param create: If `true`, the template passed in the body is only used if no - existing templates match the same index patterns. If `false`, the simulation - uses the template with the highest priority. Note that the template is not - permanently added or updated in either case; it is only used for the simulation. - :param data_stream: If this object is included, the template is used to create - data streams and their backing indices. Supports an empty object. Data streams - require a matching index template with a `data_stream` object. + :param name: Name of the index to simulate :param include_defaults: If true, returns all relevant default configurations for the index template. - :param index_patterns: Array of wildcard (`*`) expressions used to match the - names of data streams and indices during creation. :param master_timeout: Period to wait for a connection to the master node. If no response is received before the timeout expires, the request fails and returns an error. - :param meta: Optional user metadata about the index template. May have any contents. - This map is not automatically generated by Elasticsearch. - :param priority: Priority to determine index template precedence when a new data - stream or index is created. The index template with the highest priority - is chosen. If no priority is specified the template is treated as though - it is of priority 0 (lowest priority). This number is not automatically generated - by Elasticsearch. - :param template: Template to be applied. It may optionally include an `aliases`, - `mappings`, or `settings` configuration. - :param version: Version number used to manage index templates externally. This - number is not automatically generated by Elasticsearch. """ if name in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'name'") __path_parts: t.Dict[str, str] = {"name": _quote(name)} __path = f'/_index_template/_simulate_index/{__path_parts["name"]}' __query: t.Dict[str, t.Any] = {} - __body: t.Dict[str, t.Any] = body if body is not None else {} - if create is not None: - __query["create"] = create if error_trace is not None: __query["error_trace"] = error_trace if filter_path is not None: @@ -4101,56 +4080,55 @@ def simulate_index_template( __query["master_timeout"] = master_timeout if pretty is not None: __query["pretty"] = pretty - if not __body: - if allow_auto_create is not None: - __body["allow_auto_create"] = allow_auto_create - if composed_of is not None: - __body["composed_of"] = composed_of - if data_stream is not None: - __body["data_stream"] = data_stream - if index_patterns is not None: - __body["index_patterns"] = index_patterns - if meta is not None: - __body["_meta"] = meta - if priority is not None: - __body["priority"] = priority - if template is not None: - __body["template"] = template - if version is not None: - __body["version"] = version - if not __body: - __body = None # type: ignore[assignment] __headers = {"accept": "application/json"} - if __body is not None: - __headers["content-type"] = "application/json" return self.perform_request( # type: ignore[return-value] "POST", __path, params=__query, headers=__headers, - body=__body, endpoint_id="indices.simulate_index_template", path_parts=__path_parts, ) @_rewrite_parameters( - body_name="template", + body_fields=( + "allow_auto_create", + "composed_of", + "data_stream", + "deprecated", + "ignore_missing_component_templates", + "index_patterns", + "meta", + "priority", + "template", + "version", + ), + parameter_aliases={"_meta": "meta"}, ) def simulate_template( self, *, name: t.Optional[str] = None, + allow_auto_create: t.Optional[bool] = None, + composed_of: t.Optional[t.Sequence[str]] = None, create: t.Optional[bool] = None, + data_stream: t.Optional[t.Mapping[str, t.Any]] = None, + deprecated: t.Optional[bool] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, + ignore_missing_component_templates: t.Optional[t.Sequence[str]] = None, include_defaults: t.Optional[bool] = None, + index_patterns: t.Optional[t.Union[str, t.Sequence[str]]] = None, master_timeout: t.Optional[ t.Union["t.Literal[-1]", "t.Literal[0]", str] ] = None, + meta: t.Optional[t.Mapping[str, t.Any]] = None, pretty: t.Optional[bool] = None, + priority: t.Optional[int] = None, template: t.Optional[t.Mapping[str, t.Any]] = None, - body: t.Optional[t.Mapping[str, t.Any]] = None, + version: t.Optional[int] = None, + body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: """ Simulate resolving the given template name or body @@ -4160,23 +4138,47 @@ def simulate_template( :param name: Name of the index template to simulate. To test a template configuration before you add it to the cluster, omit this parameter and specify the template configuration in the request body. + :param allow_auto_create: This setting overrides the value of the `action.auto_create_index` + cluster setting. If set to `true` in a template, then indices can be automatically + created using that template even if auto-creation of indices is disabled + via `actions.auto_create_index`. If set to `false`, then indices or data + streams matching the template must always be explicitly created, and may + never be automatically created. + :param composed_of: An ordered list of component template names. Component templates + are merged in the order specified, meaning that the last component template + specified has the highest precedence. :param create: If true, the template passed in the body is only used if no existing templates match the same index patterns. If false, the simulation uses the template with the highest priority. Note that the template is not permanently added or updated in either case; it is only used for the simulation. + :param data_stream: If this object is included, the template is used to create + data streams and their backing indices. Supports an empty object. Data streams + require a matching index template with a `data_stream` object. + :param deprecated: Marks this index template as deprecated. When creating or + updating a non-deprecated index template that uses deprecated components, + Elasticsearch will emit a deprecation warning. + :param ignore_missing_component_templates: The configuration option ignore_missing_component_templates + can be used when an index template references a component template that might + not exist :param include_defaults: If true, returns all relevant default configurations for the index template. + :param index_patterns: Array of wildcard (`*`) expressions used to match the + names of data streams and indices during creation. :param master_timeout: Period to wait for a connection to the master node. If no response is received before the timeout expires, the request fails and returns an error. - :param template: + :param meta: Optional user metadata about the index template. May have any contents. + This map is not automatically generated by Elasticsearch. + :param priority: Priority to determine index template precedence when a new data + stream or index is created. The index template with the highest priority + is chosen. If no priority is specified the template is treated as though + it is of priority 0 (lowest priority). This number is not automatically generated + by Elasticsearch. + :param template: Template to be applied. It may optionally include an `aliases`, + `mappings`, or `settings` configuration. + :param version: Version number used to manage index templates externally. This + number is not automatically generated by Elasticsearch. """ - if template is None and body is None: - raise ValueError( - "Empty value passed for parameters 'template' and 'body', one of them should be set." - ) - elif template is not None and body is not None: - raise ValueError("Cannot set both 'template' and 'body'") __path_parts: t.Dict[str, str] if name not in SKIP_IN_PATH: __path_parts = {"name": _quote(name)} @@ -4185,6 +4187,7 @@ def simulate_template( __path_parts = {} __path = "/_index_template/_simulate" __query: t.Dict[str, t.Any] = {} + __body: t.Dict[str, t.Any] = body if body is not None else {} if create is not None: __query["create"] = create if error_trace is not None: @@ -4199,9 +4202,31 @@ def simulate_template( __query["master_timeout"] = master_timeout if pretty is not None: __query["pretty"] = pretty - __body = template if template is not None else body if not __body: - __body = None + if allow_auto_create is not None: + __body["allow_auto_create"] = allow_auto_create + if composed_of is not None: + __body["composed_of"] = composed_of + if data_stream is not None: + __body["data_stream"] = data_stream + if deprecated is not None: + __body["deprecated"] = deprecated + if ignore_missing_component_templates is not None: + __body["ignore_missing_component_templates"] = ( + ignore_missing_component_templates + ) + if index_patterns is not None: + __body["index_patterns"] = index_patterns + if meta is not None: + __body["_meta"] = meta + if priority is not None: + __body["priority"] = priority + if template is not None: + __body["template"] = template + if version is not None: + __body["version"] = version + if not __body: + __body = None # type: ignore[assignment] __headers = {"accept": "application/json"} if __body is not None: __headers["content-type"] = "application/json" diff --git a/elasticsearch/_sync/client/inference.py b/elasticsearch/_sync/client/inference.py index 7d4a207ce..7427833b0 100644 --- a/elasticsearch/_sync/client/inference.py +++ b/elasticsearch/_sync/client/inference.py @@ -26,7 +26,7 @@ class InferenceClient(NamespacedClient): @_rewrite_parameters() - def delete_model( + def delete( self, *, inference_id: str, @@ -42,7 +42,7 @@ def delete_model( pretty: t.Optional[bool] = None, ) -> ObjectApiResponse[t.Any]: """ - Delete model in the Inference API + Delete an inference endpoint ``_ @@ -78,36 +78,34 @@ def delete_model( __path, params=__query, headers=__headers, - endpoint_id="inference.delete_model", + endpoint_id="inference.delete", path_parts=__path_parts, ) @_rewrite_parameters() - def get_model( + def get( self, *, - inference_id: str, task_type: t.Optional[ t.Union[ "t.Literal['completion', 'rerank', 'sparse_embedding', 'text_embedding']", str, ] ] = None, + inference_id: t.Optional[str] = None, error_trace: t.Optional[bool] = None, filter_path: t.Optional[t.Union[str, t.Sequence[str]]] = None, human: t.Optional[bool] = None, pretty: t.Optional[bool] = None, ) -> ObjectApiResponse[t.Any]: """ - Get a model in the Inference API + Get an inference endpoint ``_ - :param inference_id: The inference Id :param task_type: The task type + :param inference_id: The inference Id """ - if inference_id in SKIP_IN_PATH: - raise ValueError("Empty value passed for parameter 'inference_id'") __path_parts: t.Dict[str, str] if task_type not in SKIP_IN_PATH and inference_id not in SKIP_IN_PATH: __path_parts = { @@ -119,7 +117,8 @@ def get_model( __path_parts = {"inference_id": _quote(inference_id)} __path = f'/_inference/{__path_parts["inference_id"]}' else: - raise ValueError("Couldn't find a path for the given parameters") + __path_parts = {} + __path = "/_inference" __query: t.Dict[str, t.Any] = {} if error_trace is not None: __query["error_trace"] = error_trace @@ -135,7 +134,7 @@ def get_model( __path, params=__query, headers=__headers, - endpoint_id="inference.get_model", + endpoint_id="inference.get", path_parts=__path_parts, ) @@ -159,18 +158,21 @@ def inference( pretty: t.Optional[bool] = None, query: t.Optional[str] = None, task_settings: t.Optional[t.Any] = None, + timeout: t.Optional[t.Union["t.Literal[-1]", "t.Literal[0]", str]] = None, body: t.Optional[t.Dict[str, t.Any]] = None, ) -> ObjectApiResponse[t.Any]: """ - Perform inference on a model + Perform inference ``_ :param inference_id: The inference Id - :param input: Text input to the model. Either a string or an array of strings. + :param input: Inference input. Either a string or an array of strings. :param task_type: The task type :param query: Query input, required for rerank task. Not required for other tasks. :param task_settings: Optional task settings + :param timeout: Specifies the amount of time to wait for the inference request + to complete. """ if inference_id in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'inference_id'") @@ -198,6 +200,8 @@ def inference( __query["human"] = human if pretty is not None: __query["pretty"] = pretty + if timeout is not None: + __query["timeout"] = timeout if not __body: if input is not None: __body["input"] = input @@ -221,13 +225,13 @@ def inference( ) @_rewrite_parameters( - body_name="model_config", + body_name="inference_config", ) - def put_model( + def put( self, *, inference_id: str, - model_config: t.Optional[t.Mapping[str, t.Any]] = None, + inference_config: t.Optional[t.Mapping[str, t.Any]] = None, body: t.Optional[t.Mapping[str, t.Any]] = None, task_type: t.Optional[ t.Union[ @@ -241,22 +245,22 @@ def put_model( pretty: t.Optional[bool] = None, ) -> ObjectApiResponse[t.Any]: """ - Configure a model for use in the Inference API + Configure an inference endpoint for use in the Inference API ``_ :param inference_id: The inference Id - :param model_config: + :param inference_config: :param task_type: The task type """ if inference_id in SKIP_IN_PATH: raise ValueError("Empty value passed for parameter 'inference_id'") - if model_config is None and body is None: + if inference_config is None and body is None: raise ValueError( - "Empty value passed for parameters 'model_config' and 'body', one of them should be set." + "Empty value passed for parameters 'inference_config' and 'body', one of them should be set." ) - elif model_config is not None and body is not None: - raise ValueError("Cannot set both 'model_config' and 'body'") + elif inference_config is not None and body is not None: + raise ValueError("Cannot set both 'inference_config' and 'body'") __path_parts: t.Dict[str, str] if task_type not in SKIP_IN_PATH and inference_id not in SKIP_IN_PATH: __path_parts = { @@ -278,7 +282,7 @@ def put_model( __query["human"] = human if pretty is not None: __query["pretty"] = pretty - __body = model_config if model_config is not None else body + __body = inference_config if inference_config is not None else body __headers = {"accept": "application/json", "content-type": "application/json"} return self.perform_request( # type: ignore[return-value] "PUT", @@ -286,6 +290,6 @@ def put_model( params=__query, headers=__headers, body=__body, - endpoint_id="inference.put_model", + endpoint_id="inference.put", path_parts=__path_parts, ) diff --git a/elasticsearch/_sync/client/security.py b/elasticsearch/_sync/client/security.py index 2ce01eadb..e4cf338c8 100644 --- a/elasticsearch/_sync/client/security.py +++ b/elasticsearch/_sync/client/security.py @@ -1761,7 +1761,7 @@ def has_privileges( cluster: t.Optional[ t.Sequence[ t.Union[ - "t.Literal['all', 'cancel_task', 'create_snapshot', 'grant_api_key', 'manage', 'manage_api_key', 'manage_ccr', 'manage_enrich', 'manage_ilm', 'manage_index_templates', 'manage_ingest_pipelines', 'manage_logstash_pipelines', 'manage_ml', 'manage_oidc', 'manage_own_api_key', 'manage_pipeline', 'manage_rollup', 'manage_saml', 'manage_security', 'manage_service_account', 'manage_slm', 'manage_token', 'manage_transform', 'manage_user_profile', 'manage_watcher', 'monitor', 'monitor_ml', 'monitor_rollup', 'monitor_snapshot', 'monitor_text_structure', 'monitor_transform', 'monitor_watcher', 'read_ccr', 'read_ilm', 'read_pipeline', 'read_slm', 'transport_client']", + "t.Literal['all', 'cancel_task', 'create_snapshot', 'cross_cluster_replication', 'cross_cluster_search', 'delegate_pki', 'grant_api_key', 'manage', 'manage_api_key', 'manage_autoscaling', 'manage_behavioral_analytics', 'manage_ccr', 'manage_data_frame_transforms', 'manage_data_stream_global_retention', 'manage_enrich', 'manage_ilm', 'manage_index_templates', 'manage_inference', 'manage_ingest_pipelines', 'manage_logstash_pipelines', 'manage_ml', 'manage_oidc', 'manage_own_api_key', 'manage_pipeline', 'manage_rollup', 'manage_saml', 'manage_search_application', 'manage_search_query_rules', 'manage_search_synonyms', 'manage_security', 'manage_service_account', 'manage_slm', 'manage_token', 'manage_transform', 'manage_user_profile', 'manage_watcher', 'monitor', 'monitor_data_frame_transforms', 'monitor_data_stream_global_retention', 'monitor_enrich', 'monitor_inference', 'monitor_ml', 'monitor_rollup', 'monitor_snapshot', 'monitor_text_structure', 'monitor_transform', 'monitor_watcher', 'none', 'post_behavioral_analytics_event', 'read_ccr', 'read_connector_secrets', 'read_fleet_secrets', 'read_ilm', 'read_pipeline', 'read_security', 'read_slm', 'transport_client', 'write_connector_secrets', 'write_fleet_secrets']", str, ] ] @@ -2084,7 +2084,7 @@ def put_role( cluster: t.Optional[ t.Sequence[ t.Union[ - "t.Literal['all', 'cancel_task', 'create_snapshot', 'grant_api_key', 'manage', 'manage_api_key', 'manage_ccr', 'manage_enrich', 'manage_ilm', 'manage_index_templates', 'manage_ingest_pipelines', 'manage_logstash_pipelines', 'manage_ml', 'manage_oidc', 'manage_own_api_key', 'manage_pipeline', 'manage_rollup', 'manage_saml', 'manage_security', 'manage_service_account', 'manage_slm', 'manage_token', 'manage_transform', 'manage_user_profile', 'manage_watcher', 'monitor', 'monitor_ml', 'monitor_rollup', 'monitor_snapshot', 'monitor_text_structure', 'monitor_transform', 'monitor_watcher', 'read_ccr', 'read_ilm', 'read_pipeline', 'read_slm', 'transport_client']", + "t.Literal['all', 'cancel_task', 'create_snapshot', 'cross_cluster_replication', 'cross_cluster_search', 'delegate_pki', 'grant_api_key', 'manage', 'manage_api_key', 'manage_autoscaling', 'manage_behavioral_analytics', 'manage_ccr', 'manage_data_frame_transforms', 'manage_data_stream_global_retention', 'manage_enrich', 'manage_ilm', 'manage_index_templates', 'manage_inference', 'manage_ingest_pipelines', 'manage_logstash_pipelines', 'manage_ml', 'manage_oidc', 'manage_own_api_key', 'manage_pipeline', 'manage_rollup', 'manage_saml', 'manage_search_application', 'manage_search_query_rules', 'manage_search_synonyms', 'manage_security', 'manage_service_account', 'manage_slm', 'manage_token', 'manage_transform', 'manage_user_profile', 'manage_watcher', 'monitor', 'monitor_data_frame_transforms', 'monitor_data_stream_global_retention', 'monitor_enrich', 'monitor_inference', 'monitor_ml', 'monitor_rollup', 'monitor_snapshot', 'monitor_text_structure', 'monitor_transform', 'monitor_watcher', 'none', 'post_behavioral_analytics_event', 'read_ccr', 'read_connector_secrets', 'read_fleet_secrets', 'read_ilm', 'read_pipeline', 'read_security', 'read_slm', 'transport_client', 'write_connector_secrets', 'write_fleet_secrets']", str, ] ] diff --git a/elasticsearch/_version.py b/elasticsearch/_version.py index 43e7f92a8..20c5eb394 100644 --- a/elasticsearch/_version.py +++ b/elasticsearch/_version.py @@ -15,4 +15,4 @@ # specific language governing permissions and limitations # under the License. -__versionstr__ = "8.13.0" +__versionstr__ = "8.13.1" diff --git a/elasticsearch/helpers/vectorstore/__init__.py b/elasticsearch/helpers/vectorstore/__init__.py new file mode 100644 index 000000000..30a4c3d6e --- /dev/null +++ b/elasticsearch/helpers/vectorstore/__init__.py @@ -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 ( + 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", +] diff --git a/elasticsearch/helpers/vectorstore/_async/__init__.py b/elasticsearch/helpers/vectorstore/_async/__init__.py new file mode 100644 index 000000000..2a87d183f --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_async/__init__.py @@ -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. diff --git a/elasticsearch/helpers/vectorstore/_async/_utils.py b/elasticsearch/helpers/vectorstore/_async/_utils.py new file mode 100644 index 000000000..67b6b6a27 --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_async/_utils.py @@ -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 diff --git a/elasticsearch/helpers/vectorstore/_async/embedding_service.py b/elasticsearch/helpers/vectorstore/_async/embedding_service.py new file mode 100644 index 000000000..20005b665 --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_async/embedding_service.py @@ -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"]] diff --git a/elasticsearch/helpers/vectorstore/_async/strategies.py b/elasticsearch/helpers/vectorstore/_async/strategies.py new file mode 100644 index 000000000..a7f813f43 --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_async/strategies.py @@ -0,0 +1,466 @@ +# 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 Any, Dict, List, Optional, Tuple, Union, cast + +from elasticsearch import AsyncElasticsearch +from elasticsearch.helpers.vectorstore._async._utils import model_must_be_deployed +from elasticsearch.helpers.vectorstore._utils import DistanceMetric + + +class AsyncRetrievalStrategy(ABC): + @abstractmethod + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + """ + Returns the Elasticsearch query body for the given parameters. + The store will execute the query. + + :param query: The text query. Can be None if query_vector is given. + :param k: The total number of results to retrieve. + :param num_candidates: The number of results to fetch initially in knn search. + :param filter: List of filter clauses to apply to the query. + :param query_vector: The query vector. Can be None if a query string is given. + + :return: The Elasticsearch query body. + """ + + @abstractmethod + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Create the required index and do necessary preliminary work, like + creating inference pipelines or checking if a required model was deployed. + + :param client: Elasticsearch client connection. + :param text_field: The field containing the text data in the index. + :param vector_field: The field containing the vector representations in the index. + :param num_dimensions: If vectors are indexed, how many dimensions do they have. + + :return: Dictionary with field and field type pairs that describe the schema. + """ + + async def before_index_creation( + self, *, client: AsyncElasticsearch, text_field: str, vector_field: str + ) -> None: + """ + Executes before the index is created. Used for setting up + any required Elasticsearch resources like a pipeline. + Defaults to a no-op. + + :param client: The Elasticsearch client. + :param text_field: The field containing the text data in the index. + :param vector_field: The field containing the vector representations in the index. + """ + pass + + def needs_inference(self) -> bool: + """ + Some retrieval strategies index embedding vectors and allow search by embedding + vector, for example the `DenseVectorStrategy` strategy. Mapping a user input query + string to an embedding vector is called inference. Inference can be applied + in Elasticsearch (using a `model_id`) or outside of Elasticsearch (using an + `EmbeddingService` defined on the `VectorStore`). In the latter case, + this method has to return True. + """ + return False + + +class AsyncSparseVectorStrategy(AsyncRetrievalStrategy): + """Sparse retrieval strategy using the `text_expansion` processor.""" + + def __init__(self, model_id: str = ".elser_model_2"): + self.model_id = model_id + self._tokens_field = "tokens" + self._pipeline_name = f"{self.model_id}_sparse_embedding" + + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + if query_vector: + raise ValueError( + "Cannot do sparse retrieval with a query_vector. " + "Inference is currently always applied in Elasticsearch." + ) + if query is None: + raise ValueError("please specify a query string") + + return { + "query": { + "bool": { + "must": [ + { + "text_expansion": { + f"{vector_field}.{self._tokens_field}": { + "model_id": self.model_id, + "model_text": query, + } + } + } + ], + "filter": filter, + } + } + } + + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + mappings: Dict[str, Any] = { + "properties": { + vector_field: { + "properties": {self._tokens_field: {"type": "rank_features"}} + } + } + } + settings = {"default_pipeline": self._pipeline_name} + + return mappings, settings + + async def before_index_creation( + self, *, client: AsyncElasticsearch, text_field: str, vector_field: str + ) -> None: + if self.model_id: + await model_must_be_deployed(client, self.model_id) + + # Create a pipeline for the model + await client.ingest.put_pipeline( + id=self._pipeline_name, + description="Embedding pipeline for Python VectorStore", + processors=[ + { + "inference": { + "model_id": self.model_id, + "target_field": vector_field, + "field_map": {text_field: "text_field"}, + "inference_config": { + "text_expansion": {"results_field": self._tokens_field} + }, + } + } + ], + ) + + +class AsyncDenseVectorStrategy(AsyncRetrievalStrategy): + """K-nearest-neighbors retrieval.""" + + def __init__( + self, + *, + distance: DistanceMetric = DistanceMetric.COSINE, + model_id: Optional[str] = None, + hybrid: bool = False, + rrf: Union[bool, Dict[str, Any]] = True, + text_field: Optional[str] = "text_field", + ): + if hybrid and not text_field: + raise ValueError( + "to enable hybrid you have to specify a text_field (for BM25Strategy matching)" + ) + + self.distance = distance + self.model_id = model_id + self.hybrid = hybrid + self.rrf = rrf + self.text_field = text_field + + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + knn = { + "filter": filter, + "field": vector_field, + "k": k, + "num_candidates": num_candidates, + } + + if query_vector is not None: + knn["query_vector"] = query_vector + else: + # Inference in Elasticsearch. When initializing we make sure to always have + # a model_id if don't have an embedding_service. + knn["query_vector_builder"] = { + "text_embedding": { + "model_id": self.model_id, + "model_text": query, + } + } + + if self.hybrid: + return self._hybrid(query=cast(str, query), knn=knn, filter=filter) + + return {"knn": knn} + + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + if self.distance is DistanceMetric.COSINE: + similarity = "cosine" + elif self.distance is DistanceMetric.EUCLIDEAN_DISTANCE: + similarity = "l2_norm" + elif self.distance is DistanceMetric.DOT_PRODUCT: + similarity = "dot_product" + elif self.distance is DistanceMetric.MAX_INNER_PRODUCT: + similarity = "max_inner_product" + else: + raise ValueError(f"Similarity {self.distance} not supported.") + + mappings: Dict[str, Any] = { + "properties": { + vector_field: { + "type": "dense_vector", + "dims": num_dimensions, + "index": True, + "similarity": similarity, + }, + } + } + + return mappings, {} + + async def before_index_creation( + self, *, client: AsyncElasticsearch, text_field: str, vector_field: str + ) -> None: + if self.model_id: + await model_must_be_deployed(client, self.model_id) + + def _hybrid( + self, query: str, knn: Dict[str, Any], filter: List[Dict[str, Any]] + ) -> Dict[str, Any]: + # Add a query to the knn query. + # RRF is used to even the score from the knn query and text query + # RRF has two optional parameters: {'rank_constant':int, 'window_size':int} + # https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html + query_body = { + "knn": knn, + "query": { + "bool": { + "must": [ + { + "match": { + self.text_field: { + "query": query, + } + } + } + ], + "filter": filter, + } + }, + } + + if isinstance(self.rrf, Dict): + query_body["rank"] = {"rrf": self.rrf} + elif isinstance(self.rrf, bool) and self.rrf is True: + query_body["rank"] = {"rrf": {}} + + return query_body + + def needs_inference(self) -> bool: + return not self.model_id + + +class AsyncDenseVectorScriptScoreStrategy(AsyncRetrievalStrategy): + """Exact nearest neighbors retrieval using the `script_score` query.""" + + def __init__(self, distance: DistanceMetric = DistanceMetric.COSINE) -> None: + self.distance = distance + + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + if not query_vector: + raise ValueError("specify a query_vector") + + if self.distance is DistanceMetric.COSINE: + similarity_algo = ( + f"cosineSimilarity(params.query_vector, '{vector_field}') + 1.0" + ) + elif self.distance is DistanceMetric.EUCLIDEAN_DISTANCE: + similarity_algo = f"1 / (1 + l2norm(params.query_vector, '{vector_field}'))" + elif self.distance is DistanceMetric.DOT_PRODUCT: + similarity_algo = f""" + double value = dotProduct(params.query_vector, '{vector_field}'); + return sigmoid(1, Math.E, -value); + """ + elif self.distance is DistanceMetric.MAX_INNER_PRODUCT: + similarity_algo = f""" + double value = dotProduct(params.query_vector, '{vector_field}'); + if (dotProduct < 0) {{ + return 1 / (1 + -1 * dotProduct); + }} + return dotProduct + 1; + """ + else: + raise ValueError(f"Similarity {self.distance} not supported.") + + query_bool: Dict[str, Any] = {"match_all": {}} + if filter: + query_bool = {"bool": {"filter": filter}} + + return { + "query": { + "script_score": { + "query": query_bool, + "script": { + "source": similarity_algo, + "params": {"query_vector": query_vector}, + }, + }, + } + } + + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + mappings = { + "properties": { + vector_field: { + "type": "dense_vector", + "dims": num_dimensions, + "index": False, + } + } + } + + return mappings, {} + + def needs_inference(self) -> bool: + return True + + +class AsyncBM25Strategy(AsyncRetrievalStrategy): + def __init__( + self, + k1: Optional[float] = None, + b: Optional[float] = None, + ): + self.k1 = k1 + self.b = b + + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + return { + "query": { + "bool": { + "must": [ + { + "match": { + text_field: { + "query": query, + } + }, + }, + ], + "filter": filter, + }, + }, + } + + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + similarity_name = "custom_bm25" + + mappings: Dict[str, Any] = { + "properties": { + text_field: { + "type": "text", + "similarity": similarity_name, + }, + }, + } + + bm25: Dict[str, Any] = { + "type": "BM25", + } + if self.k1 is not None: + bm25["k1"] = self.k1 + if self.b is not None: + bm25["b"] = self.b + settings = { + "similarity": { + similarity_name: bm25, + } + } + + return mappings, settings diff --git a/elasticsearch/helpers/vectorstore/_async/vectorstore.py b/elasticsearch/helpers/vectorstore/_async/vectorstore.py new file mode 100644 index 000000000..b79e2dcaf --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_async/vectorstore.py @@ -0,0 +1,391 @@ +# 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. + +import logging +import uuid +from typing import Any, Callable, Dict, List, Optional + +from elasticsearch import AsyncElasticsearch +from elasticsearch._version import __versionstr__ as lib_version +from elasticsearch.helpers import BulkIndexError, async_bulk +from elasticsearch.helpers.vectorstore import ( + AsyncEmbeddingService, + AsyncRetrievalStrategy, +) +from elasticsearch.helpers.vectorstore._utils import maximal_marginal_relevance + +logger = logging.getLogger(__name__) + + +class AsyncVectorStore: + """ + VectorStore is a higher-level abstraction of indexing and search. + Users can pick from available retrieval strategies. + + Documents have up to 3 fields: + - text_field: the text to be indexed and searched. + - metadata: additional information about the document, either schema-free + or defined by the supplied metadata_mappings. + - vector_field (usually not filled by the user): the embedding vector of the text. + + Depending on the strategy, vector embeddings are + - created by the user beforehand + - created by this AsyncVectorStore class in Python + - created in-stack by inference pipelines. + """ + + def __init__( + self, + client: AsyncElasticsearch, + *, + index: str, + retrieval_strategy: AsyncRetrievalStrategy, + embedding_service: Optional[AsyncEmbeddingService] = None, + num_dimensions: Optional[int] = None, + text_field: str = "text_field", + vector_field: str = "vector_field", + metadata_mappings: Optional[Dict[str, Any]] = None, + user_agent: str = f"elasticsearch-py-vs/{lib_version}", + ) -> None: + """ + :param user_header: user agent header specific to the 3rd party integration. + Used for usage tracking in Elastic Cloud. + :param index: The name of the index to query. + :param retrieval_strategy: how to index and search the data. See the strategies + module for availble strategies. + :param text_field: Name of the field with the textual data. + :param vector_field: For strategies that perform embedding inference in Python, + the embedding vector goes in this 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}) + + if hasattr(retrieval_strategy, "text_field"): + retrieval_strategy.text_field = text_field + if hasattr(retrieval_strategy, "vector_field"): + retrieval_strategy.vector_field = vector_field + + self.client = client + self.index = index + self.retrieval_strategy = retrieval_strategy + self.embedding_service = embedding_service + self.num_dimensions = num_dimensions + self.text_field = text_field + self.vector_field = vector_field + self.metadata_mappings = metadata_mappings + + async def close(self) -> None: + return await self.client.close() + + async def add_texts( + self, + texts: List[str], + *, + metadatas: Optional[List[Dict[str, Any]]] = None, + vectors: Optional[List[List[float]]] = None, + ids: Optional[List[str]] = None, + refresh_indices: bool = True, + create_index_if_not_exists: bool = True, + bulk_kwargs: Optional[Dict[str, Any]] = None, + ) -> List[str]: + """Add documents to the Elasticsearch index. + + :param texts: List of text documents. + :param metadata: Optional list of document metadata. Must be of same length as + texts. + :param vectors: Optional list of embedding vectors. Must be of same length as + texts. + :param ids: Optional list of ID strings. Must be of same length as texts. + :param refresh_indices: Whether to refresh the index after deleting documents. + Defaults to True. + :param create_index_if_not_exists: Whether to create the index if it does not + exist. Defaults to True. + :param bulk_kwargs: Arguments to pass to the bulk function when indexing + (for example chunk_size). + + :return: List of IDs of the created documents, either echoing the provided one + or returning newly created ones. + """ + bulk_kwargs = bulk_kwargs or {} + ids = ids or [str(uuid.uuid4()) for _ in texts] + requests = [] + + if create_index_if_not_exists: + await self._create_index_if_not_exists() + + if self.embedding_service and not vectors: + vectors = await self.embedding_service.embed_documents(texts) + + for i, text in enumerate(texts): + metadata = metadatas[i] if metadatas else {} + + request: Dict[str, Any] = { + "_op_type": "index", + "_index": self.index, + self.text_field: text, + "metadata": metadata, + "_id": ids[i], + } + + if vectors: + request[self.vector_field] = vectors[i] + + requests.append(request) + + if len(requests) > 0: + try: + success, failed = await async_bulk( + self.client, + requests, + stats_only=True, + refresh=refresh_indices, + **bulk_kwargs, + ) + logger.debug(f"added texts {ids} to index") + return ids + except BulkIndexError as e: + logger.error(f"Error adding texts: {e}") + firstError = e.errors[0].get("index", {}).get("error", {}) + logger.error(f"First error reason: {firstError.get('reason')}") + raise e + + else: + logger.debug("No texts to add to index") + return [] + + async def delete( # type: ignore[no-untyped-def] + self, + *, + ids: Optional[List[str]] = None, + query: Optional[Dict[str, Any]] = None, + refresh_indices: bool = True, + **delete_kwargs, + ) -> bool: + """Delete documents from the Elasticsearch index. + + :param ids: List of IDs of documents to delete. + :param refresh_indices: Whether to refresh the index after deleting documents. + Defaults to True. + + :return: True if deletion was successful. + """ + if ids is not None and query is not None: + raise ValueError("one of ids or query must be specified") + elif ids is None and query is None: + raise ValueError("either specify ids or query") + + try: + if ids: + body = [ + {"_op_type": "delete", "_index": self.index, "_id": _id} + for _id in ids + ] + await async_bulk( + self.client, + body, + refresh=refresh_indices, + ignore_status=404, + **delete_kwargs, + ) + logger.debug(f"Deleted {len(body)} texts from index") + + else: + await self.client.delete_by_query( + index=self.index, + query=query, + refresh=refresh_indices, + **delete_kwargs, + ) + + except BulkIndexError as e: + logger.error(f"Error deleting texts: {e}") + firstError = e.errors[0].get("index", {}).get("error", {}) + logger.error(f"First error reason: {firstError.get('reason')}") + raise e + + return True + + async def search( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]] = None, + k: int = 4, + num_candidates: int = 50, + fields: Optional[List[str]] = None, + filter: Optional[List[Dict[str, Any]]] = None, + custom_query: Optional[ + Callable[[Dict[str, Any], Optional[str]], Dict[str, Any]] + ] = None, + ) -> List[Dict[str, Any]]: + """ + :param query: Input query string. + :param query_vector: Input embedding vector. If given, input query string is + ignored. + :param k: Number of returned results. + :param num_candidates: Number of candidates to fetch from data nodes in knn. + :param fields: List of field names to return. + :param filter: Elasticsearch filters to apply. + :param custom_query: Function to modify the Elasticsearch query body before it is + sent to Elasticsearch. + + :return: List of document hits. Includes _index, _id, _score and _source. + """ + if fields is None: + fields = [] + if "metadata" not in fields: + fields.append("metadata") + if self.text_field not in fields: + fields.append(self.text_field) + + if self.embedding_service and not query_vector: + if not query: + raise ValueError("specify a query or a query_vector to search") + query_vector = await self.embedding_service.embed_query(query) + + query_body = self.retrieval_strategy.es_query( + query=query, + query_vector=query_vector, + text_field=self.text_field, + vector_field=self.vector_field, + k=k, + num_candidates=num_candidates, + filter=filter or [], + ) + + if custom_query is not None: + query_body = custom_query(query_body, query) + logger.debug(f"Calling custom_query, Query body now: {query_body}") + + response = await self.client.search( + index=self.index, + **query_body, + size=k, + source=True, + source_includes=fields, + ) + hits: List[Dict[str, Any]] = response["hits"]["hits"] + + return hits + + async def _create_index_if_not_exists(self) -> None: + exists = await self.client.indices.exists(index=self.index) + if exists.meta.status == 200: + logger.debug(f"Index {self.index} already exists. Skipping creation.") + return + + if self.retrieval_strategy.needs_inference(): + if not self.num_dimensions and not self.embedding_service: + raise ValueError( + "retrieval strategy requires embeddings; either embedding_service " + "or num_dimensions need to be specified" + ) + if not self.num_dimensions and self.embedding_service: + vector = await self.embedding_service.embed_query("get num dimensions") + self.num_dimensions = len(vector) + + mappings, settings = self.retrieval_strategy.es_mappings_settings( + text_field=self.text_field, + vector_field=self.vector_field, + num_dimensions=self.num_dimensions, + ) + if self.metadata_mappings: + metadata = mappings["properties"].get("metadata", {"properties": {}}) + for key in self.metadata_mappings.keys(): + if key in metadata: + raise ValueError(f"metadata key {key} already exists in mappings") + + metadata = dict(**metadata["properties"], **self.metadata_mappings) + mappings["properties"]["metadata"] = {"properties": metadata} + + await self.retrieval_strategy.before_index_creation( + client=self.client, + text_field=self.text_field, + vector_field=self.vector_field, + ) + await self.client.indices.create( + index=self.index, mappings=mappings, settings=settings + ) + + async def max_marginal_relevance_search( + self, + *, + embedding_service: AsyncEmbeddingService, + query: str, + vector_field: str, + k: int = 4, + num_candidates: int = 20, + lambda_mult: float = 0.5, + fields: Optional[List[str]] = None, + custom_query: Optional[ + Callable[[Dict[str, Any], Optional[str]], Dict[str, Any]] + ] = None, + ) -> List[Dict[str, Any]]: + """Return docs selected using the maximal marginal relevance. + + Maximal marginal relevance optimizes for similarity to query AND diversity + among selected documents. + + :param query (str): Text to look up documents similar to. + :param k (int): Number of Documents to return. Defaults to 4. + :param fetch_k (int): Number of Documents to fetch to pass to MMR algorithm. + :param lambda_mult (float): Number between 0 and 1 that determines the degree + of diversity among the results with 0 corresponding + to maximum diversity and 1 to minimum diversity. + Defaults to 0.5. + :param fields: Other fields to get from elasticsearch source. These fields + will be added to the document metadata. + + :return: A list of Documents selected by maximal marginal relevance. + """ + remove_vector_query_field_from_metadata = True + if fields is None: + fields = [vector_field] + elif vector_field not in fields: + fields.append(vector_field) + else: + remove_vector_query_field_from_metadata = False + + # Embed the query + query_embedding = await embedding_service.embed_query(query) + + # Fetch the initial documents + got_hits = await self.search( + query=None, + query_vector=query_embedding, + k=num_candidates, + fields=fields, + custom_query=custom_query, + ) + + # Get the embeddings for the fetched documents + got_embeddings = [hit["_source"][vector_field] for hit in got_hits] + + # Select documents using maximal marginal relevance + selected_indices = maximal_marginal_relevance( + query_embedding, got_embeddings, lambda_mult=lambda_mult, k=k + ) + selected_hits = [got_hits[i] for i in selected_indices] + + if remove_vector_query_field_from_metadata: + for hit in selected_hits: + del hit["_source"][vector_field] + + return selected_hits diff --git a/elasticsearch/helpers/vectorstore/_sync/__init__.py b/elasticsearch/helpers/vectorstore/_sync/__init__.py new file mode 100644 index 000000000..2a87d183f --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_sync/__init__.py @@ -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. diff --git a/elasticsearch/helpers/vectorstore/_sync/_utils.py b/elasticsearch/helpers/vectorstore/_sync/_utils.py new file mode 100644 index 000000000..496aec970 --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_sync/_utils.py @@ -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 BadRequestError, Elasticsearch, NotFoundError + + +def model_must_be_deployed(client: Elasticsearch, 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: + 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 + + +def model_is_deployed(client: Elasticsearch, model_id: str) -> bool: + try: + model_must_be_deployed(client, model_id) + return True + except NotFoundError: + return False diff --git a/elasticsearch/helpers/vectorstore/_sync/embedding_service.py b/elasticsearch/helpers/vectorstore/_sync/embedding_service.py new file mode 100644 index 000000000..5b0163d98 --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_sync/embedding_service.py @@ -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 Elasticsearch +from elasticsearch._version import __versionstr__ as lib_version + + +class EmbeddingService(ABC): + @abstractmethod + 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 + 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 ElasticsearchEmbeddings(EmbeddingService): + """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: Elasticsearch, + 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 + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + return self._embedding_func(texts) + + def embed_query(self, text: str) -> List[float]: + result = self._embedding_func([text]) + return result[0] + + def _embedding_func(self, texts: List[str]) -> List[List[float]]: + response = 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"]] diff --git a/elasticsearch/helpers/vectorstore/_sync/strategies.py b/elasticsearch/helpers/vectorstore/_sync/strategies.py new file mode 100644 index 000000000..928d34143 --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_sync/strategies.py @@ -0,0 +1,466 @@ +# 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 Any, Dict, List, Optional, Tuple, Union, cast + +from elasticsearch import Elasticsearch +from elasticsearch.helpers.vectorstore._sync._utils import model_must_be_deployed +from elasticsearch.helpers.vectorstore._utils import DistanceMetric + + +class RetrievalStrategy(ABC): + @abstractmethod + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + """ + Returns the Elasticsearch query body for the given parameters. + The store will execute the query. + + :param query: The text query. Can be None if query_vector is given. + :param k: The total number of results to retrieve. + :param num_candidates: The number of results to fetch initially in knn search. + :param filter: List of filter clauses to apply to the query. + :param query_vector: The query vector. Can be None if a query string is given. + + :return: The Elasticsearch query body. + """ + + @abstractmethod + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Create the required index and do necessary preliminary work, like + creating inference pipelines or checking if a required model was deployed. + + :param client: Elasticsearch client connection. + :param text_field: The field containing the text data in the index. + :param vector_field: The field containing the vector representations in the index. + :param num_dimensions: If vectors are indexed, how many dimensions do they have. + + :return: Dictionary with field and field type pairs that describe the schema. + """ + + def before_index_creation( + self, *, client: Elasticsearch, text_field: str, vector_field: str + ) -> None: + """ + Executes before the index is created. Used for setting up + any required Elasticsearch resources like a pipeline. + Defaults to a no-op. + + :param client: The Elasticsearch client. + :param text_field: The field containing the text data in the index. + :param vector_field: The field containing the vector representations in the index. + """ + pass + + def needs_inference(self) -> bool: + """ + Some retrieval strategies index embedding vectors and allow search by embedding + vector, for example the `DenseVectorStrategy` strategy. Mapping a user input query + string to an embedding vector is called inference. Inference can be applied + in Elasticsearch (using a `model_id`) or outside of Elasticsearch (using an + `EmbeddingService` defined on the `VectorStore`). In the latter case, + this method has to return True. + """ + return False + + +class SparseVectorStrategy(RetrievalStrategy): + """Sparse retrieval strategy using the `text_expansion` processor.""" + + def __init__(self, model_id: str = ".elser_model_2"): + self.model_id = model_id + self._tokens_field = "tokens" + self._pipeline_name = f"{self.model_id}_sparse_embedding" + + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + if query_vector: + raise ValueError( + "Cannot do sparse retrieval with a query_vector. " + "Inference is currently always applied in Elasticsearch." + ) + if query is None: + raise ValueError("please specify a query string") + + return { + "query": { + "bool": { + "must": [ + { + "text_expansion": { + f"{vector_field}.{self._tokens_field}": { + "model_id": self.model_id, + "model_text": query, + } + } + } + ], + "filter": filter, + } + } + } + + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + mappings: Dict[str, Any] = { + "properties": { + vector_field: { + "properties": {self._tokens_field: {"type": "rank_features"}} + } + } + } + settings = {"default_pipeline": self._pipeline_name} + + return mappings, settings + + def before_index_creation( + self, *, client: Elasticsearch, text_field: str, vector_field: str + ) -> None: + if self.model_id: + model_must_be_deployed(client, self.model_id) + + # Create a pipeline for the model + client.ingest.put_pipeline( + id=self._pipeline_name, + description="Embedding pipeline for Python VectorStore", + processors=[ + { + "inference": { + "model_id": self.model_id, + "target_field": vector_field, + "field_map": {text_field: "text_field"}, + "inference_config": { + "text_expansion": {"results_field": self._tokens_field} + }, + } + } + ], + ) + + +class DenseVectorStrategy(RetrievalStrategy): + """K-nearest-neighbors retrieval.""" + + def __init__( + self, + *, + distance: DistanceMetric = DistanceMetric.COSINE, + model_id: Optional[str] = None, + hybrid: bool = False, + rrf: Union[bool, Dict[str, Any]] = True, + text_field: Optional[str] = "text_field", + ): + if hybrid and not text_field: + raise ValueError( + "to enable hybrid you have to specify a text_field (for BM25Strategy matching)" + ) + + self.distance = distance + self.model_id = model_id + self.hybrid = hybrid + self.rrf = rrf + self.text_field = text_field + + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + knn = { + "filter": filter, + "field": vector_field, + "k": k, + "num_candidates": num_candidates, + } + + if query_vector is not None: + knn["query_vector"] = query_vector + else: + # Inference in Elasticsearch. When initializing we make sure to always have + # a model_id if don't have an embedding_service. + knn["query_vector_builder"] = { + "text_embedding": { + "model_id": self.model_id, + "model_text": query, + } + } + + if self.hybrid: + return self._hybrid(query=cast(str, query), knn=knn, filter=filter) + + return {"knn": knn} + + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + if self.distance is DistanceMetric.COSINE: + similarity = "cosine" + elif self.distance is DistanceMetric.EUCLIDEAN_DISTANCE: + similarity = "l2_norm" + elif self.distance is DistanceMetric.DOT_PRODUCT: + similarity = "dot_product" + elif self.distance is DistanceMetric.MAX_INNER_PRODUCT: + similarity = "max_inner_product" + else: + raise ValueError(f"Similarity {self.distance} not supported.") + + mappings: Dict[str, Any] = { + "properties": { + vector_field: { + "type": "dense_vector", + "dims": num_dimensions, + "index": True, + "similarity": similarity, + }, + } + } + + return mappings, {} + + def before_index_creation( + self, *, client: Elasticsearch, text_field: str, vector_field: str + ) -> None: + if self.model_id: + model_must_be_deployed(client, self.model_id) + + def _hybrid( + self, query: str, knn: Dict[str, Any], filter: List[Dict[str, Any]] + ) -> Dict[str, Any]: + # Add a query to the knn query. + # RRF is used to even the score from the knn query and text query + # RRF has two optional parameters: {'rank_constant':int, 'window_size':int} + # https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html + query_body = { + "knn": knn, + "query": { + "bool": { + "must": [ + { + "match": { + self.text_field: { + "query": query, + } + } + } + ], + "filter": filter, + } + }, + } + + if isinstance(self.rrf, Dict): + query_body["rank"] = {"rrf": self.rrf} + elif isinstance(self.rrf, bool) and self.rrf is True: + query_body["rank"] = {"rrf": {}} + + return query_body + + def needs_inference(self) -> bool: + return not self.model_id + + +class DenseVectorScriptScoreStrategy(RetrievalStrategy): + """Exact nearest neighbors retrieval using the `script_score` query.""" + + def __init__(self, distance: DistanceMetric = DistanceMetric.COSINE) -> None: + self.distance = distance + + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + if not query_vector: + raise ValueError("specify a query_vector") + + if self.distance is DistanceMetric.COSINE: + similarity_algo = ( + f"cosineSimilarity(params.query_vector, '{vector_field}') + 1.0" + ) + elif self.distance is DistanceMetric.EUCLIDEAN_DISTANCE: + similarity_algo = f"1 / (1 + l2norm(params.query_vector, '{vector_field}'))" + elif self.distance is DistanceMetric.DOT_PRODUCT: + similarity_algo = f""" + double value = dotProduct(params.query_vector, '{vector_field}'); + return sigmoid(1, Math.E, -value); + """ + elif self.distance is DistanceMetric.MAX_INNER_PRODUCT: + similarity_algo = f""" + double value = dotProduct(params.query_vector, '{vector_field}'); + if (dotProduct < 0) {{ + return 1 / (1 + -1 * dotProduct); + }} + return dotProduct + 1; + """ + else: + raise ValueError(f"Similarity {self.distance} not supported.") + + query_bool: Dict[str, Any] = {"match_all": {}} + if filter: + query_bool = {"bool": {"filter": filter}} + + return { + "query": { + "script_score": { + "query": query_bool, + "script": { + "source": similarity_algo, + "params": {"query_vector": query_vector}, + }, + }, + } + } + + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + mappings = { + "properties": { + vector_field: { + "type": "dense_vector", + "dims": num_dimensions, + "index": False, + } + } + } + + return mappings, {} + + def needs_inference(self) -> bool: + return True + + +class BM25Strategy(RetrievalStrategy): + def __init__( + self, + k1: Optional[float] = None, + b: Optional[float] = None, + ): + self.k1 = k1 + self.b = b + + def es_query( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]], + text_field: str, + vector_field: str, + k: int, + num_candidates: int, + filter: List[Dict[str, Any]] = [], + ) -> Dict[str, Any]: + return { + "query": { + "bool": { + "must": [ + { + "match": { + text_field: { + "query": query, + } + }, + }, + ], + "filter": filter, + }, + }, + } + + def es_mappings_settings( + self, + *, + text_field: str, + vector_field: str, + num_dimensions: Optional[int], + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + similarity_name = "custom_bm25" + + mappings: Dict[str, Any] = { + "properties": { + text_field: { + "type": "text", + "similarity": similarity_name, + }, + }, + } + + bm25: Dict[str, Any] = { + "type": "BM25", + } + if self.k1 is not None: + bm25["k1"] = self.k1 + if self.b is not None: + bm25["b"] = self.b + settings = { + "similarity": { + similarity_name: bm25, + } + } + + return mappings, settings diff --git a/elasticsearch/helpers/vectorstore/_sync/vectorstore.py b/elasticsearch/helpers/vectorstore/_sync/vectorstore.py new file mode 100644 index 000000000..2feb96ec4 --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_sync/vectorstore.py @@ -0,0 +1,388 @@ +# 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. + +import logging +import uuid +from typing import Any, Callable, Dict, List, Optional + +from elasticsearch import Elasticsearch +from elasticsearch._version import __versionstr__ as lib_version +from elasticsearch.helpers import BulkIndexError, bulk +from elasticsearch.helpers.vectorstore import EmbeddingService, RetrievalStrategy +from elasticsearch.helpers.vectorstore._utils import maximal_marginal_relevance + +logger = logging.getLogger(__name__) + + +class VectorStore: + """ + VectorStore is a higher-level abstraction of indexing and search. + Users can pick from available retrieval strategies. + + Documents have up to 3 fields: + - text_field: the text to be indexed and searched. + - metadata: additional information about the document, either schema-free + or defined by the supplied metadata_mappings. + - vector_field (usually not filled by the user): the embedding vector of the text. + + Depending on the strategy, vector embeddings are + - created by the user beforehand + - created by this AsyncVectorStore class in Python + - created in-stack by inference pipelines. + """ + + def __init__( + self, + client: Elasticsearch, + *, + index: str, + retrieval_strategy: RetrievalStrategy, + embedding_service: Optional[EmbeddingService] = None, + num_dimensions: Optional[int] = None, + text_field: str = "text_field", + vector_field: str = "vector_field", + metadata_mappings: Optional[Dict[str, Any]] = None, + user_agent: str = f"elasticsearch-py-vs/{lib_version}", + ) -> None: + """ + :param user_header: user agent header specific to the 3rd party integration. + Used for usage tracking in Elastic Cloud. + :param index: The name of the index to query. + :param retrieval_strategy: how to index and search the data. See the strategies + module for availble strategies. + :param text_field: Name of the field with the textual data. + :param vector_field: For strategies that perform embedding inference in Python, + the embedding vector goes in this 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}) + + if hasattr(retrieval_strategy, "text_field"): + retrieval_strategy.text_field = text_field + if hasattr(retrieval_strategy, "vector_field"): + retrieval_strategy.vector_field = vector_field + + self.client = client + self.index = index + self.retrieval_strategy = retrieval_strategy + self.embedding_service = embedding_service + self.num_dimensions = num_dimensions + self.text_field = text_field + self.vector_field = vector_field + self.metadata_mappings = metadata_mappings + + def close(self) -> None: + return self.client.close() + + def add_texts( + self, + texts: List[str], + *, + metadatas: Optional[List[Dict[str, Any]]] = None, + vectors: Optional[List[List[float]]] = None, + ids: Optional[List[str]] = None, + refresh_indices: bool = True, + create_index_if_not_exists: bool = True, + bulk_kwargs: Optional[Dict[str, Any]] = None, + ) -> List[str]: + """Add documents to the Elasticsearch index. + + :param texts: List of text documents. + :param metadata: Optional list of document metadata. Must be of same length as + texts. + :param vectors: Optional list of embedding vectors. Must be of same length as + texts. + :param ids: Optional list of ID strings. Must be of same length as texts. + :param refresh_indices: Whether to refresh the index after deleting documents. + Defaults to True. + :param create_index_if_not_exists: Whether to create the index if it does not + exist. Defaults to True. + :param bulk_kwargs: Arguments to pass to the bulk function when indexing + (for example chunk_size). + + :return: List of IDs of the created documents, either echoing the provided one + or returning newly created ones. + """ + bulk_kwargs = bulk_kwargs or {} + ids = ids or [str(uuid.uuid4()) for _ in texts] + requests = [] + + if create_index_if_not_exists: + self._create_index_if_not_exists() + + if self.embedding_service and not vectors: + vectors = self.embedding_service.embed_documents(texts) + + for i, text in enumerate(texts): + metadata = metadatas[i] if metadatas else {} + + request: Dict[str, Any] = { + "_op_type": "index", + "_index": self.index, + self.text_field: text, + "metadata": metadata, + "_id": ids[i], + } + + if vectors: + request[self.vector_field] = vectors[i] + + requests.append(request) + + if len(requests) > 0: + try: + success, failed = bulk( + self.client, + requests, + stats_only=True, + refresh=refresh_indices, + **bulk_kwargs, + ) + logger.debug(f"added texts {ids} to index") + return ids + except BulkIndexError as e: + logger.error(f"Error adding texts: {e}") + firstError = e.errors[0].get("index", {}).get("error", {}) + logger.error(f"First error reason: {firstError.get('reason')}") + raise e + + else: + logger.debug("No texts to add to index") + return [] + + def delete( # type: ignore[no-untyped-def] + self, + *, + ids: Optional[List[str]] = None, + query: Optional[Dict[str, Any]] = None, + refresh_indices: bool = True, + **delete_kwargs, + ) -> bool: + """Delete documents from the Elasticsearch index. + + :param ids: List of IDs of documents to delete. + :param refresh_indices: Whether to refresh the index after deleting documents. + Defaults to True. + + :return: True if deletion was successful. + """ + if ids is not None and query is not None: + raise ValueError("one of ids or query must be specified") + elif ids is None and query is None: + raise ValueError("either specify ids or query") + + try: + if ids: + body = [ + {"_op_type": "delete", "_index": self.index, "_id": _id} + for _id in ids + ] + bulk( + self.client, + body, + refresh=refresh_indices, + ignore_status=404, + **delete_kwargs, + ) + logger.debug(f"Deleted {len(body)} texts from index") + + else: + self.client.delete_by_query( + index=self.index, + query=query, + refresh=refresh_indices, + **delete_kwargs, + ) + + except BulkIndexError as e: + logger.error(f"Error deleting texts: {e}") + firstError = e.errors[0].get("index", {}).get("error", {}) + logger.error(f"First error reason: {firstError.get('reason')}") + raise e + + return True + + def search( + self, + *, + query: Optional[str], + query_vector: Optional[List[float]] = None, + k: int = 4, + num_candidates: int = 50, + fields: Optional[List[str]] = None, + filter: Optional[List[Dict[str, Any]]] = None, + custom_query: Optional[ + Callable[[Dict[str, Any], Optional[str]], Dict[str, Any]] + ] = None, + ) -> List[Dict[str, Any]]: + """ + :param query: Input query string. + :param query_vector: Input embedding vector. If given, input query string is + ignored. + :param k: Number of returned results. + :param num_candidates: Number of candidates to fetch from data nodes in knn. + :param fields: List of field names to return. + :param filter: Elasticsearch filters to apply. + :param custom_query: Function to modify the Elasticsearch query body before it is + sent to Elasticsearch. + + :return: List of document hits. Includes _index, _id, _score and _source. + """ + if fields is None: + fields = [] + if "metadata" not in fields: + fields.append("metadata") + if self.text_field not in fields: + fields.append(self.text_field) + + if self.embedding_service and not query_vector: + if not query: + raise ValueError("specify a query or a query_vector to search") + query_vector = self.embedding_service.embed_query(query) + + query_body = self.retrieval_strategy.es_query( + query=query, + query_vector=query_vector, + text_field=self.text_field, + vector_field=self.vector_field, + k=k, + num_candidates=num_candidates, + filter=filter or [], + ) + + if custom_query is not None: + query_body = custom_query(query_body, query) + logger.debug(f"Calling custom_query, Query body now: {query_body}") + + response = self.client.search( + index=self.index, + **query_body, + size=k, + source=True, + source_includes=fields, + ) + hits: List[Dict[str, Any]] = response["hits"]["hits"] + + return hits + + def _create_index_if_not_exists(self) -> None: + exists = self.client.indices.exists(index=self.index) + if exists.meta.status == 200: + logger.debug(f"Index {self.index} already exists. Skipping creation.") + return + + if self.retrieval_strategy.needs_inference(): + if not self.num_dimensions and not self.embedding_service: + raise ValueError( + "retrieval strategy requires embeddings; either embedding_service " + "or num_dimensions need to be specified" + ) + if not self.num_dimensions and self.embedding_service: + vector = self.embedding_service.embed_query("get num dimensions") + self.num_dimensions = len(vector) + + mappings, settings = self.retrieval_strategy.es_mappings_settings( + text_field=self.text_field, + vector_field=self.vector_field, + num_dimensions=self.num_dimensions, + ) + if self.metadata_mappings: + metadata = mappings["properties"].get("metadata", {"properties": {}}) + for key in self.metadata_mappings.keys(): + if key in metadata: + raise ValueError(f"metadata key {key} already exists in mappings") + + metadata = dict(**metadata["properties"], **self.metadata_mappings) + mappings["properties"]["metadata"] = {"properties": metadata} + + self.retrieval_strategy.before_index_creation( + client=self.client, + text_field=self.text_field, + vector_field=self.vector_field, + ) + self.client.indices.create( + index=self.index, mappings=mappings, settings=settings + ) + + def max_marginal_relevance_search( + self, + *, + embedding_service: EmbeddingService, + query: str, + vector_field: str, + k: int = 4, + num_candidates: int = 20, + lambda_mult: float = 0.5, + fields: Optional[List[str]] = None, + custom_query: Optional[ + Callable[[Dict[str, Any], Optional[str]], Dict[str, Any]] + ] = None, + ) -> List[Dict[str, Any]]: + """Return docs selected using the maximal marginal relevance. + + Maximal marginal relevance optimizes for similarity to query AND diversity + among selected documents. + + :param query (str): Text to look up documents similar to. + :param k (int): Number of Documents to return. Defaults to 4. + :param fetch_k (int): Number of Documents to fetch to pass to MMR algorithm. + :param lambda_mult (float): Number between 0 and 1 that determines the degree + of diversity among the results with 0 corresponding + to maximum diversity and 1 to minimum diversity. + Defaults to 0.5. + :param fields: Other fields to get from elasticsearch source. These fields + will be added to the document metadata. + + :return: A list of Documents selected by maximal marginal relevance. + """ + remove_vector_query_field_from_metadata = True + if fields is None: + fields = [vector_field] + elif vector_field not in fields: + fields.append(vector_field) + else: + remove_vector_query_field_from_metadata = False + + # Embed the query + query_embedding = embedding_service.embed_query(query) + + # Fetch the initial documents + got_hits = self.search( + query=None, + query_vector=query_embedding, + k=num_candidates, + fields=fields, + custom_query=custom_query, + ) + + # Get the embeddings for the fetched documents + got_embeddings = [hit["_source"][vector_field] for hit in got_hits] + + # Select documents using maximal marginal relevance + selected_indices = maximal_marginal_relevance( + query_embedding, got_embeddings, lambda_mult=lambda_mult, k=k + ) + selected_hits = [got_hits[i] for i in selected_indices] + + if remove_vector_query_field_from_metadata: + for hit in selected_hits: + del hit["_source"][vector_field] + + return selected_hits diff --git a/elasticsearch/helpers/vectorstore/_utils.py b/elasticsearch/helpers/vectorstore/_utils.py new file mode 100644 index 000000000..df91b5cc9 --- /dev/null +++ b/elasticsearch/helpers/vectorstore/_utils.py @@ -0,0 +1,116 @@ +# 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 enum import Enum +from typing import TYPE_CHECKING, List, Union + +if TYPE_CHECKING: + import numpy as np + import numpy.typing as npt + +Matrix = Union[ + List[List[float]], List["npt.NDArray[np.float64]"], "npt.NDArray[np.float64]" +] + + +class DistanceMetric(str, Enum): + """Enumerator of all Elasticsearch dense vector distance metrics.""" + + COSINE = "COSINE" + DOT_PRODUCT = "DOT_PRODUCT" + EUCLIDEAN_DISTANCE = "EUCLIDEAN_DISTANCE" + MAX_INNER_PRODUCT = "MAX_INNER_PRODUCT" + + +def maximal_marginal_relevance( + query_embedding: List[float], + embedding_list: List[List[float]], + lambda_mult: float = 0.5, + k: int = 4, +) -> List[int]: + """Calculate maximal marginal relevance.""" + + try: + import numpy as np + except ModuleNotFoundError as e: + _raise_missing_mmr_deps_error(e) + + query_embedding_arr = np.array(query_embedding) + + if min(k, len(embedding_list)) <= 0: + return [] + if query_embedding_arr.ndim == 1: + query_embedding_arr = np.expand_dims(query_embedding_arr, axis=0) + similarity_to_query = _cosine_similarity(query_embedding_arr, embedding_list)[0] + most_similar = int(np.argmax(similarity_to_query)) + idxs = [most_similar] + selected = np.array([embedding_list[most_similar]]) + while len(idxs) < min(k, len(embedding_list)): + best_score = -np.inf + idx_to_add = -1 + similarity_to_selected = _cosine_similarity(embedding_list, selected) + for i, query_score in enumerate(similarity_to_query): + if i in idxs: + continue + redundant_score = max(similarity_to_selected[i]) + equation_score = ( + lambda_mult * query_score - (1 - lambda_mult) * redundant_score + ) + if equation_score > best_score: + best_score = equation_score + idx_to_add = i + idxs.append(idx_to_add) + selected = np.append(selected, [embedding_list[idx_to_add]], axis=0) + return idxs + + +def _cosine_similarity(X: Matrix, Y: Matrix) -> "npt.NDArray[np.float64]": + """Row-wise cosine similarity between two equal-width matrices.""" + + try: + import numpy as np + import simsimd as simd + except ModuleNotFoundError as e: + _raise_missing_mmr_deps_error(e) + + if len(X) == 0 or len(Y) == 0: + return np.array([]) + + X = np.array(X) + Y = np.array(Y) + if X.shape[1] != Y.shape[1]: + raise ValueError( + f"Number of columns in X and Y must be the same. X has shape {X.shape} " + f"and Y has shape {Y.shape}." + ) + + X = np.array(X, dtype=np.float32) + Y = np.array(Y, dtype=np.float32) + Z = 1 - np.array(simd.cdist(X, Y, metric="cosine")) + if isinstance(Z, float): + return np.array([Z]) + return np.array(Z) + + +def _raise_missing_mmr_deps_error(parent_error: ModuleNotFoundError) -> None: + import sys + + raise ModuleNotFoundError( + f"Failed to compute maximal marginal relevance because the required " + f"module '{parent_error.name}' is missing. You can install it by running: " + f"'{sys.executable} -m pip install elasticsearch[vectorstore_mmr]'" + ) from parent_error diff --git a/noxfile.py b/noxfile.py index c303fe26c..12ad4f02e 100644 --- a/noxfile.py +++ b/noxfile.py @@ -48,7 +48,9 @@ def pytest_argv(): @nox.session(python=["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]) def test(session): - session.install(".[async,requests,orjson]", env=INSTALL_ENV, silent=False) + session.install( + ".[async,requests,orjson,vectorstore_mmr]", env=INSTALL_ENV, silent=False + ) session.install("-r", "dev-requirements.txt", silent=False) session.run(*pytest_argv()) @@ -71,7 +73,7 @@ def test_otel(session): @nox.session() def format(session): - session.install("black~=24.0", "isort", "flynt", "unasync", "setuptools") + session.install("black~=24.0", "isort", "flynt", "unasync>=0.6.0") session.run("python", "utils/run-unasync.py") session.run("isort", "--profile=black", *SOURCE_FILES) @@ -95,7 +97,7 @@ def lint(session): session.run("flake8", *SOURCE_FILES) session.run("python", "utils/license-headers.py", "check", *SOURCE_FILES) - session.install(".[async,requests,orjson]", env=INSTALL_ENV) + session.install(".[async,requests,orjson,vectorstore_mmr]", env=INSTALL_ENV) # Run mypy on the package and then the type examples separately for # the two different mypy use-cases, ourselves and our users. diff --git a/setup.py b/setup.py index dc592dcc4..e9ee3a377 100644 --- a/setup.py +++ b/setup.py @@ -92,5 +92,7 @@ "requests": ["requests>=2.4.0, <3.0.0"], "async": ["aiohttp>=3,<4"], "orjson": ["orjson>=3"], + # Maximal Marginal Relevance (MMR) for search results + "vectorstore_mmr": ["numpy>=1", "simsimd>=3"], }, ) diff --git a/test_elasticsearch/test_server/conftest.py b/test_elasticsearch/test_server/conftest.py index 558d0b013..7b87fd1d3 100644 --- a/test_elasticsearch/test_server/conftest.py +++ b/test_elasticsearch/test_server/conftest.py @@ -30,19 +30,33 @@ ELASTICSEARCH_REST_API_TESTS = [] +def _create(elasticsearch_url, transport=None, node_class=None): + # Configure the client with certificates + kw = {} + if elasticsearch_url.startswith("https://"): + kw["ca_certs"] = CA_CERTS + + # Optionally configure an HTTP conn class depending on + # 'PYTHON_CONNECTION_CLASS' env var + if "PYTHON_CONNECTION_CLASS" in os.environ: + kw["node_class"] = os.environ["PYTHON_CONNECTION_CLASS"] + + if node_class is not None and "node_class" not in kw: + kw["node_class"] = node_class + + if transport: + kw["transport_class"] = transport + + # We do this little dance with the URL to force + # Requests to respect 'headers: None' within rest API spec tests. + return elasticsearch.Elasticsearch(elasticsearch_url, **kw) + + @pytest.fixture(scope="session") def sync_client_factory(elasticsearch_url): client = None try: - # Configure the client with certificates and optionally - # an HTTP conn class depending on 'PYTHON_CONNECTION_CLASS' envvar - kw = {"ca_certs": CA_CERTS} - if "PYTHON_CONNECTION_CLASS" in os.environ: - kw["node_class"] = os.environ["PYTHON_CONNECTION_CLASS"] - - # We do this little dance with the URL to force - # Requests to respect 'headers: None' within rest API spec tests. - client = elasticsearch.Elasticsearch(elasticsearch_url, **kw) + client = _create(elasticsearch_url) # Wipe the cluster before we start testing just in case it wasn't wiped # cleanly from the previous run of pytest? diff --git a/test_elasticsearch/test_server/test_mapbox_vector_tile.py b/test_elasticsearch/test_server/test_mapbox_vector_tile.py index 988210984..332e8d144 100644 --- a/test_elasticsearch/test_server/test_mapbox_vector_tile.py +++ b/test_elasticsearch/test_server/test_mapbox_vector_tile.py @@ -17,7 +17,9 @@ import pytest -from elasticsearch import Elasticsearch, RequestError +from elasticsearch import RequestError + +from .conftest import _create @pytest.fixture(scope="function") @@ -73,7 +75,8 @@ def mvt_setup(sync_client): @pytest.mark.parametrize("node_class", ["urllib3", "requests"]) def test_mapbox_vector_tile_error(elasticsearch_url, mvt_setup, node_class, ca_certs): - client = Elasticsearch(elasticsearch_url, node_class=node_class, ca_certs=ca_certs) + client = _create(elasticsearch_url, node_class=node_class) + client.search_mvt( index="museums", zoom=13, @@ -121,7 +124,7 @@ def test_mapbox_vector_tile_response( except ImportError: return pytest.skip("Requires the 'mapbox-vector-tile' package") - client = Elasticsearch(elasticsearch_url, node_class=node_class, ca_certs=ca_certs) + client = _create(elasticsearch_url, node_class=node_class) resp = client.search_mvt( index="museums", diff --git a/test_elasticsearch/test_server/test_vectorstore/__init__.py b/test_elasticsearch/test_server/test_vectorstore/__init__.py new file mode 100644 index 000000000..87710976a --- /dev/null +++ b/test_elasticsearch/test_server/test_vectorstore/__init__.py @@ -0,0 +1,81 @@ +# 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 typing import List + +from elastic_transport import Transport + +from elasticsearch.helpers.vectorstore import EmbeddingService + + +class RequestSavingTransport(Transport): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.requests: list = [] + + def perform_request(self, *args, **kwargs): + self.requests.append(kwargs) + return super().perform_request(*args, **kwargs) + + +class FakeEmbeddings(EmbeddingService): + """Fake embeddings functionality for testing.""" + + def __init__(self, dimensionality: int = 10) -> None: + self.dimensionality = dimensionality + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + """Return simple embeddings. Embeddings encode each text as its index.""" + return [ + [float(1.0)] * (self.dimensionality - 1) + [float(i)] + for i in range(len(texts)) + ] + + def embed_query(self, text: str) -> List[float]: + """Return constant query embeddings. + Embeddings are identical to embed_documents(texts)[0]. + Distance to each text will be that text's index, + as it was passed to embed_documents. + """ + return [float(1.0)] * (self.dimensionality - 1) + [float(0.0)] + + +class ConsistentFakeEmbeddings(FakeEmbeddings): + """Fake embeddings which remember all the texts seen so far to return consistent + vectors for the same texts.""" + + def __init__(self, dimensionality: int = 10) -> None: + self.known_texts: List[str] = [] + self.dimensionality = dimensionality + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + """Return consistent embeddings for each text seen so far.""" + out_vectors = [] + for text in texts: + if text not in self.known_texts: + self.known_texts.append(text) + vector = [float(1.0)] * (self.dimensionality - 1) + [ + float(self.known_texts.index(text)) + ] + out_vectors.append(vector) + return out_vectors + + def embed_query(self, text: str) -> List[float]: + """Return consistent embeddings for the text, if seen before, or a constant + one if the text is unknown.""" + result = self.embed_documents([text]) + return result[0] diff --git a/test_elasticsearch/test_server/test_vectorstore/conftest.py b/test_elasticsearch/test_server/test_vectorstore/conftest.py new file mode 100644 index 000000000..a0886a9c4 --- /dev/null +++ b/test_elasticsearch/test_server/test_vectorstore/conftest.py @@ -0,0 +1,60 @@ +# 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. + +import uuid + +import pytest + +from ...utils import wipe_cluster +from ..conftest import _create +from . import RequestSavingTransport + + +@pytest.fixture(scope="function") +def index() -> str: + return f"test_{uuid.uuid4().hex}" + + +@pytest.fixture(scope="function") +def sync_client_request_saving_factory(elasticsearch_url): + client = None + + try: + client = _create(elasticsearch_url) + # Wipe the cluster before we start testing just in case it wasn't wiped + # cleanly from the previous run of pytest? + wipe_cluster(client) + finally: + client.close() + + try: + # Recreate client with a transport that saves requests. + client = _create(elasticsearch_url, RequestSavingTransport) + + yield client + finally: + if client: + client.close() + + +@pytest.fixture(scope="function") +def sync_client_request_saving(sync_client_request_saving_factory): + try: + yield sync_client_request_saving_factory + finally: + # Wipe the cluster clean after every test execution. + wipe_cluster(sync_client_request_saving_factory) diff --git a/test_elasticsearch/test_server/test_vectorstore/docker-compose.yml b/test_elasticsearch/test_server/test_vectorstore/docker-compose.yml new file mode 100644 index 000000000..009c427a8 --- /dev/null +++ b/test_elasticsearch/test_server/test_vectorstore/docker-compose.yml @@ -0,0 +1,34 @@ +version: "3" + +services: + elasticsearch: + image: elasticsearch:8.13.0 + environment: + - action.destructive_requires_name=false # allow wildcard index deletions + - discovery.type=single-node + - xpack.license.self_generated.type=trial + - xpack.security.enabled=false # disable password and TLS; never do this in production! + ports: + - "9200:9200" + healthcheck: + test: + [ + "CMD-SHELL", + "curl --silent --fail http://localhost:9200/_cluster/health || exit 1" + ] + interval: 10s + retries: 60 + + # Currently fails on Mac: https://github.com/elastic/elasticsearch/issues/106206 + elasticsearch-with-model: + image: docker.elastic.co/eland/eland + depends_on: + - elasticsearch + restart: no + command: sh -c " + sleep 10 && + eland_import_hub_model \ + --hub-model-id sentence-transformers/all-minilm-l6-v2 \ + --url http://elasticsearch:9200 \ + --start + " diff --git a/test_elasticsearch/test_server/test_vectorstore/test_embedding_service.py b/test_elasticsearch/test_server/test_vectorstore/test_embedding_service.py new file mode 100644 index 000000000..1f9b196a7 --- /dev/null +++ b/test_elasticsearch/test_server/test_vectorstore/test_embedding_service.py @@ -0,0 +1,91 @@ +# 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. + +import os +import re + +import pytest + +from elasticsearch import Elasticsearch +from elasticsearch.helpers.vectorstore import ElasticsearchEmbeddings +from elasticsearch.helpers.vectorstore._sync._utils import model_is_deployed + +# deployed with +# https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-text-emb-vector-search-example.html +MODEL_ID = os.getenv("MODEL_ID", "sentence-transformers__all-minilm-l6-v2") +NUM_DIMENSIONS = int(os.getenv("NUM_DIMENSIONS", "384")) + + +def test_elasticsearch_embedding_documents(sync_client: Elasticsearch) -> None: + """Test Elasticsearch embedding documents.""" + + if not model_is_deployed(sync_client, MODEL_ID): + pytest.skip(f"{MODEL_ID} model is not deployed in ML Node, skipping test") + + documents = ["foo bar", "bar foo", "foo"] + embedding = ElasticsearchEmbeddings( + client=sync_client, user_agent="test", model_id=MODEL_ID + ) + output = embedding.embed_documents(documents) + assert len(output) == 3 + assert len(output[0]) == NUM_DIMENSIONS + assert len(output[1]) == NUM_DIMENSIONS + assert len(output[2]) == NUM_DIMENSIONS + + +def test_elasticsearch_embedding_query(sync_client: Elasticsearch) -> None: + """Test Elasticsearch embedding query.""" + + if not model_is_deployed(sync_client, MODEL_ID): + pytest.skip(f"{MODEL_ID} model is not deployed in ML Node, skipping test") + + document = "foo bar" + embedding = ElasticsearchEmbeddings( + client=sync_client, user_agent="test", model_id=MODEL_ID + ) + output = embedding.embed_query(document) + assert len(output) == NUM_DIMENSIONS + + +def test_user_agent_default( + sync_client: Elasticsearch, sync_client_request_saving: Elasticsearch +) -> None: + """Test to make sure the user-agent is set correctly.""" + + if not model_is_deployed(sync_client, MODEL_ID): + pytest.skip(f"{MODEL_ID} model is not deployed in ML Node, skipping test") + + embeddings = ElasticsearchEmbeddings( + client=sync_client_request_saving, model_id=MODEL_ID + ) + + expected_pattern = r"^elasticsearch-py-es/\d+\.\d+\.\d+$" + + got_agent = embeddings.client._headers["User-Agent"] + assert ( + re.match(expected_pattern, got_agent) is not None + ), f"The user agent '{got_agent}' does not match the expected pattern." + + embeddings.embed_query("foo bar") + + requests = embeddings.client.transport.requests # type: ignore + assert len(requests) == 1 + + got_request_agent = requests[0]["headers"]["User-Agent"] + assert ( + re.match(expected_pattern, got_request_agent) is not None + ), f"The user agent '{got_request_agent}' does not match the expected pattern." diff --git a/test_elasticsearch/test_server/test_vectorstore/test_vectorstore.py b/test_elasticsearch/test_server/test_vectorstore/test_vectorstore.py new file mode 100644 index 000000000..bb15d3dc7 --- /dev/null +++ b/test_elasticsearch/test_server/test_vectorstore/test_vectorstore.py @@ -0,0 +1,909 @@ +# 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. + +import logging +import re +from functools import partial +from typing import Any, List, Optional, Union + +import pytest + +from elasticsearch import Elasticsearch, NotFoundError +from elasticsearch.helpers import BulkIndexError +from elasticsearch.helpers.vectorstore import ( + BM25Strategy, + DenseVectorScriptScoreStrategy, + DenseVectorStrategy, + DistanceMetric, + SparseVectorStrategy, + VectorStore, +) +from elasticsearch.helpers.vectorstore._sync._utils import model_is_deployed + +from . import ConsistentFakeEmbeddings, FakeEmbeddings + +logging.basicConfig(level=logging.DEBUG) + +""" +docker-compose up elasticsearch + +By default runs against local docker instance of Elasticsearch. +To run against Elastic Cloud, set the following environment variables: +- ES_CLOUD_ID +- ES_API_KEY + +Some of the tests require the following models to be deployed in the ML Node: +- elser (can be downloaded and deployed through Kibana and trained models UI) +- sentence-transformers__all-minilm-l6-v2 (can be deployed through the API, + loaded via eland) + +These tests that require the models to be deployed are skipped by default. +Enable them by adding the model name to the modelsDeployed list below. +""" + +ELSER_MODEL_ID = ".elser_model_2" +TRANSFORMER_MODEL_ID = "sentence-transformers__all-minilm-l6-v2" + + +class TestVectorStore: + def test_search_without_metadata( + self, sync_client: Elasticsearch, index: str + ) -> None: + """Test end to end construction and search without metadata.""" + + def assert_query(query_body: dict, query: Optional[str]) -> dict: + assert query_body == { + "knn": { + "field": "vector_field", + "filter": [], + "k": 1, + "num_candidates": 50, + "query_vector": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0], + } + } + return query_body + + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + + texts = ["foo", "bar", "baz"] + store.add_texts(texts) + + output = store.search(query="foo", k=1, custom_query=assert_query) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + + def test_search_without_metadata_async( + self, sync_client: Elasticsearch, index: str + ) -> None: + """Test end to end construction and search without metadata.""" + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + + texts = ["foo", "bar", "baz"] + store.add_texts(texts) + + output = store.search(query="foo", k=1) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + + def test_add_vectors(self, sync_client: Elasticsearch, index: str) -> None: + """ + Test adding pre-built embeddings instead of using inference for the texts. + This allows you to separate the embeddings text and the page_content + for better proximity between user's question and embedded text. + For example, your embedding text can be a question, whereas page_content + is the answer. + """ + embeddings = ConsistentFakeEmbeddings() + texts = ["foo1", "foo2", "foo3"] + metadatas = [{"page": i} for i in range(len(texts))] + + embedding_vectors = embeddings.embed_documents(texts) + + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(), + embedding_service=embeddings, + client=sync_client, + ) + + store.add_texts(texts=texts, vectors=embedding_vectors, metadatas=metadatas) + output = store.search(query="foo1", k=1) + assert [doc["_source"]["text_field"] for doc in output] == ["foo1"] + assert [doc["_source"]["metadata"]["page"] for doc in output] == [0] + + def test_search_with_metadata(self, sync_client: Elasticsearch, index: str) -> None: + """Test end to end construction and search with metadata.""" + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(), + embedding_service=ConsistentFakeEmbeddings(), + client=sync_client, + ) + + texts = ["foo", "bar", "baz"] + metadatas = [{"page": i} for i in range(len(texts))] + store.add_texts(texts=texts, metadatas=metadatas) + + output = store.search(query="foo", k=1) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + assert [doc["_source"]["metadata"]["page"] for doc in output] == [0] + + output = store.search(query="bar", k=1) + assert [doc["_source"]["text_field"] for doc in output] == ["bar"] + assert [doc["_source"]["metadata"]["page"] for doc in output] == [1] + + def test_search_with_filter(self, sync_client: Elasticsearch, index: str) -> None: + """Test end to end construction and search with metadata.""" + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + + texts = ["foo", "foo", "foo"] + metadatas = [{"page": i} for i in range(len(texts))] + store.add_texts(texts=texts, metadatas=metadatas) + + def assert_query(query_body: dict, query: Optional[str]) -> dict: + assert query_body == { + "knn": { + "field": "vector_field", + "filter": [{"term": {"metadata.page": "1"}}], + "k": 3, + "num_candidates": 50, + "query_vector": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0], + } + } + return query_body + + output = store.search( + query="foo", + k=3, + filter=[{"term": {"metadata.page": "1"}}], + custom_query=assert_query, + ) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + assert [doc["_source"]["metadata"]["page"] for doc in output] == [1] + + def test_search_script_score(self, sync_client: Elasticsearch, index: str) -> None: + """Test end to end construction and search with metadata.""" + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorScriptScoreStrategy(), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + + texts = ["foo", "bar", "baz"] + store.add_texts(texts) + + expected_query = { + "query": { + "script_score": { + "query": {"match_all": {}}, + "script": { + "source": "cosineSimilarity(params.query_vector, 'vector_field') + 1.0", # noqa: E501 + "params": { + "query_vector": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 0.0, + ] + }, + }, + } + } + } + + def assert_query(query_body: dict, query: Optional[str]) -> dict: + assert query_body == expected_query + return query_body + + output = store.search(query="foo", k=1, custom_query=assert_query) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + + def test_search_script_score_with_filter( + self, sync_client: Elasticsearch, index: str + ) -> None: + """Test end to end construction and search with metadata.""" + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorScriptScoreStrategy(), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + + texts = ["foo", "bar", "baz"] + metadatas = [{"page": i} for i in range(len(texts))] + store.add_texts(texts=texts, metadatas=metadatas) + + def assert_query(query_body: dict, query: Optional[str]) -> dict: + expected_query = { + "query": { + "script_score": { + "query": {"bool": {"filter": [{"term": {"metadata.page": 0}}]}}, + "script": { + "source": "cosineSimilarity(params.query_vector, 'vector_field') + 1.0", # noqa: E501 + "params": { + "query_vector": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 0.0, + ] + }, + }, + } + } + } + assert query_body == expected_query + return query_body + + output = store.search( + query="foo", + k=1, + custom_query=assert_query, + filter=[{"term": {"metadata.page": 0}}], + ) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + assert [doc["_source"]["metadata"]["page"] for doc in output] == [0] + + def test_search_script_score_distance_dot_product( + self, sync_client: Elasticsearch, index: str + ) -> None: + """Test end to end construction and search with metadata.""" + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorScriptScoreStrategy( + distance=DistanceMetric.DOT_PRODUCT, + ), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + + texts = ["foo", "bar", "baz"] + store.add_texts(texts) + + def assert_query(query_body: dict, query: Optional[str]) -> dict: + assert query_body == { + "query": { + "script_score": { + "query": {"match_all": {}}, + "script": { + "source": """ + double value = dotProduct(params.query_vector, 'vector_field'); + return sigmoid(1, Math.E, -value); + """, + "params": { + "query_vector": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 0.0, + ] + }, + }, + } + } + } + return query_body + + output = store.search(query="foo", k=1, custom_query=assert_query) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + + def test_search_knn_with_hybrid_search( + self, sync_client: Elasticsearch, index: str + ) -> None: + """Test end to end construction and search with metadata.""" + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(hybrid=True), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + + texts = ["foo", "bar", "baz"] + store.add_texts(texts) + + def assert_query(query_body: dict, query: Optional[str]) -> dict: + assert query_body == { + "knn": { + "field": "vector_field", + "filter": [], + "k": 1, + "num_candidates": 50, + "query_vector": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0], + }, + "query": { + "bool": { + "filter": [], + "must": [{"match": {"text_field": {"query": "foo"}}}], + } + }, + "rank": {"rrf": {}}, + } + return query_body + + output = store.search(query="foo", k=1, custom_query=assert_query) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + + def test_search_knn_with_hybrid_search_rrf( + self, sync_client: Elasticsearch, index: str + ) -> None: + """Test end to end construction and rrf hybrid search with metadata.""" + texts = ["foo", "bar", "baz"] + + def assert_query( + query_body: dict, + query: Optional[str], + expected_rrf: Union[dict, bool], + ) -> dict: + cmp_query_body = { + "knn": { + "field": "vector_field", + "filter": [], + "k": 3, + "num_candidates": 50, + "query_vector": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 0.0, + ], + }, + "query": { + "bool": { + "filter": [], + "must": [{"match": {"text_field": {"query": "foo"}}}], + } + }, + } + + if isinstance(expected_rrf, dict): + cmp_query_body["rank"] = {"rrf": expected_rrf} + elif isinstance(expected_rrf, bool) and expected_rrf is True: + cmp_query_body["rank"] = {"rrf": {}} + + assert query_body == cmp_query_body + + return query_body + + # 1. check query_body is okay + rrf_test_cases: List[Union[dict, bool]] = [ + True, + False, + {"rank_constant": 1, "window_size": 5}, + ] + for rrf_test_case in rrf_test_cases: + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(hybrid=True, rrf=rrf_test_case), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + store.add_texts(texts) + + ## without fetch_k parameter + output = store.search( + query="foo", + k=3, + custom_query=partial(assert_query, expected_rrf=rrf_test_case), + ) + + # 2. check query result is okay + es_output = store.client.search( + index=index, + query={ + "bool": { + "filter": [], + "must": [{"match": {"text_field": {"query": "foo"}}}], + } + }, + knn={ + "field": "vector_field", + "filter": [], + "k": 3, + "num_candidates": 50, + "query_vector": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0], + }, + size=3, + rank={"rrf": {"rank_constant": 1, "window_size": 5}}, + ) + + assert [o["_source"]["text_field"] for o in output] == [ + e["_source"]["text_field"] for e in es_output["hits"]["hits"] + ] + + # 3. check rrf default option is okay + store = VectorStore( + index=f"{index}_default", + retrieval_strategy=DenseVectorStrategy(hybrid=True), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + store.add_texts(texts) + + ## with fetch_k parameter + output = store.search( + query="foo", + k=3, + num_candidates=50, + custom_query=partial(assert_query, expected_rrf={}), + ) + + def test_search_knn_with_custom_query_fn( + self, sync_client: Elasticsearch, index: str + ) -> None: + """test that custom query function is called + with the query string and query body""" + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + + def my_custom_query(query_body: dict, query: Optional[str]) -> dict: + assert query == "foo" + assert query_body == { + "knn": { + "field": "vector_field", + "filter": [], + "k": 1, + "num_candidates": 50, + "query_vector": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0], + } + } + return {"query": {"match": {"text_field": {"query": "bar"}}}} + + """Test end to end construction and search with metadata.""" + texts = ["foo", "bar", "baz"] + store.add_texts(texts) + + output = store.search(query="foo", k=1, custom_query=my_custom_query) + assert [doc["_source"]["text_field"] for doc in output] == ["bar"] + + def test_search_with_knn_infer_instack( + self, sync_client: Elasticsearch, index: str + ) -> None: + """test end to end with knn retrieval strategy and inference in-stack""" + + if not model_is_deployed(sync_client, TRANSFORMER_MODEL_ID): + pytest.skip( + f"{TRANSFORMER_MODEL_ID} model not deployed in ML Node skipping test" + ) + + text_field = "text_field" + + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(model_id=TRANSFORMER_MODEL_ID), + client=sync_client, + ) + + # setting up the pipeline for inference + store.client.ingest.put_pipeline( + id="test_pipeline", + processors=[ + { + "inference": { + "model_id": TRANSFORMER_MODEL_ID, + "field_map": {"query_field": text_field}, + "target_field": "vector_query_field", + } + } + ], + ) + + # creating a new index with the pipeline, + # not relying on langchain to create the index + store.client.indices.create( + index=index, + mappings={ + "properties": { + text_field: {"type": "text_field"}, + "vector_query_field": { + "properties": { + "predicted_value": { + "type": "dense_vector", + "dims": 384, + "index": True, + "similarity": "l2_norm", + } + } + }, + } + }, + settings={"index": {"default_pipeline": "test_pipeline"}}, + ) + + # adding documents to the index + texts = ["foo", "bar", "baz"] + + for i, text in enumerate(texts): + store.client.create( + index=index, + id=str(i), + document={text_field: text, "metadata": {}}, + ) + + store.client.indices.refresh(index=index) + + def assert_query(query_body: dict, query: Optional[str]) -> dict: + assert query_body == { + "knn": { + "filter": [], + "field": "vector_query_field.predicted_value", + "k": 1, + "num_candidates": 50, + "query_vector_builder": { + "text_embedding": { + "model_id": TRANSFORMER_MODEL_ID, + "model_text": "foo", + } + }, + } + } + return query_body + + output = store.search(query="foo", k=1, custom_query=assert_query) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + + output = store.search(query="bar", k=1) + assert [doc["_source"]["text_field"] for doc in output] == ["bar"] + + def test_search_with_sparse_infer_instack( + self, sync_client: Elasticsearch, index: str + ) -> None: + """test end to end with sparse retrieval strategy and inference in-stack""" + + if not model_is_deployed(sync_client, ELSER_MODEL_ID): + reason = f"{ELSER_MODEL_ID} model not deployed in ML Node, skipping test" + pytest.skip(reason) + + store = VectorStore( + index=index, + retrieval_strategy=SparseVectorStrategy(model_id=ELSER_MODEL_ID), + client=sync_client, + ) + + texts = ["foo", "bar", "baz"] + store.add_texts(texts) + + output = store.search(query="foo", k=1) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + + def test_deployed_model_check_fails_semantic( + self, sync_client: Elasticsearch, index: str + ) -> None: + """test that exceptions are raised if a specified model is not deployed""" + with pytest.raises(NotFoundError): + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy( + model_id="non-existing model ID" + ), + client=sync_client, + ) + store.add_texts(["foo", "bar", "baz"]) + + def test_search_bm25(self, sync_client: Elasticsearch, index: str) -> None: + """Test end to end using the BM25Strategy retrieval strategy.""" + store = VectorStore( + index=index, + retrieval_strategy=BM25Strategy(), + client=sync_client, + ) + + texts = ["foo", "bar", "baz"] + store.add_texts(texts) + + def assert_query(query_body: dict, query: Optional[str]) -> dict: + assert query_body == { + "query": { + "bool": { + "must": [{"match": {"text_field": {"query": "foo"}}}], + "filter": [], + } + } + } + return query_body + + output = store.search(query="foo", k=1, custom_query=assert_query) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + + def test_search_bm25_with_filter( + self, sync_client: Elasticsearch, index: str + ) -> None: + """Test end to using the BM25Strategy retrieval strategy with metadata.""" + store = VectorStore( + index=index, + retrieval_strategy=BM25Strategy(), + client=sync_client, + ) + + texts = ["foo", "foo", "foo"] + metadatas = [{"page": i} for i in range(len(texts))] + store.add_texts(texts=texts, metadatas=metadatas) + + def assert_query(query_body: dict, query: Optional[str]) -> dict: + assert query_body == { + "query": { + "bool": { + "must": [{"match": {"text_field": {"query": "foo"}}}], + "filter": [{"term": {"metadata.page": 1}}], + } + } + } + return query_body + + output = store.search( + query="foo", + k=3, + custom_query=assert_query, + filter=[{"term": {"metadata.page": 1}}], + ) + assert [doc["_source"]["text_field"] for doc in output] == ["foo"] + assert [doc["_source"]["metadata"]["page"] for doc in output] == [1] + + def test_delete(self, sync_client: Elasticsearch, index: str) -> None: + """Test delete methods from vector store.""" + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(), + embedding_service=FakeEmbeddings(), + client=sync_client, + ) + + texts = ["foo", "bar", "baz", "gni"] + metadatas = [{"page": i} for i in range(len(texts))] + ids = store.add_texts(texts=texts, metadatas=metadatas) + + output = store.search(query="foo", k=10) + assert len(output) == 4 + + store.delete(ids=ids[1:3]) + output = store.search(query="foo", k=10) + assert len(output) == 2 + + store.delete(ids=["not-existing"]) + output = store.search(query="foo", k=10) + assert len(output) == 2 + + store.delete(ids=[ids[0]]) + output = store.search(query="foo", k=10) + assert len(output) == 1 + + store.delete(ids=[ids[3]]) + output = store.search(query="gni", k=10) + assert len(output) == 0 + + def test_indexing_exception_error( + self, + sync_client: Elasticsearch, + index: str, + caplog: pytest.LogCaptureFixture, + ) -> None: + """Test bulk exception logging is giving better hints.""" + store = VectorStore( + index=index, + retrieval_strategy=BM25Strategy(), + client=sync_client, + ) + + store.client.indices.create( + index=index, + mappings={"properties": {}}, + settings={"index": {"default_pipeline": "not-existing-pipeline"}}, + ) + + texts = ["foo"] + + with pytest.raises(BulkIndexError): + store.add_texts(texts) + + error_reason = "pipeline with id [not-existing-pipeline] does not exist" + log_message = f"First error reason: {error_reason}" + + assert log_message in caplog.text + + def test_user_agent_default( + self, sync_client_request_saving: Elasticsearch, index: str + ) -> None: + """Test to make sure the user-agent is set correctly.""" + store = VectorStore( + index=index, + retrieval_strategy=BM25Strategy(), + client=sync_client_request_saving, + ) + expected_pattern = r"^elasticsearch-py-vs/\d+\.\d+\.\d+$" + + got_agent = store.client._headers["User-Agent"] + assert ( + re.match(expected_pattern, got_agent) is not None + ), f"The user agent '{got_agent}' does not match the expected pattern." + + texts = ["foo", "bob", "baz"] + store.add_texts(texts) + + for request in store.client.transport.requests: # type: ignore + agent = request["headers"]["User-Agent"] + assert ( + re.match(expected_pattern, agent) is not None + ), f"The user agent '{agent}' does not match the expected pattern." + + def test_user_agent_custom( + self, sync_client_request_saving: Elasticsearch, index: str + ) -> None: + """Test to make sure the user-agent is set correctly.""" + user_agent = "this is THE user_agent!" + + store = VectorStore( + user_agent=user_agent, + index=index, + retrieval_strategy=BM25Strategy(), + client=sync_client_request_saving, + ) + + assert store.client._headers["User-Agent"] == user_agent + + texts = ["foo", "bob", "baz"] + store.add_texts(texts) + + for request in store.client.transport.requests: # type: ignore + assert request["headers"]["User-Agent"] == user_agent + + def test_bulk_args(self, sync_client_request_saving: Any, index: str) -> None: + """Test to make sure the bulk arguments work as expected.""" + store = VectorStore( + index=index, + retrieval_strategy=BM25Strategy(), + client=sync_client_request_saving, + ) + + texts = ["foo", "bob", "baz"] + store.add_texts(texts, bulk_kwargs={"chunk_size": 1}) + + # 1 for index exist, 1 for index create, 3 to index docs + assert len(store.client.transport.requests) == 5 # type: ignore + + def test_max_marginal_relevance_search( + self, sync_client: Elasticsearch, index: str + ) -> None: + """Test max marginal relevance search.""" + texts = ["foo", "bar", "baz"] + vector_field = "vector_field" + text_field = "text_field" + embedding_service = ConsistentFakeEmbeddings() + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorScriptScoreStrategy(), + embedding_service=embedding_service, + vector_field=vector_field, + text_field=text_field, + client=sync_client, + ) + store.add_texts(texts) + + mmr_output = store.max_marginal_relevance_search( + embedding_service=embedding_service, + query=texts[0], + vector_field=vector_field, + k=3, + num_candidates=3, + ) + sim_output = store.search(query=texts[0], k=3) + assert mmr_output == sim_output + + mmr_output = store.max_marginal_relevance_search( + embedding_service=embedding_service, + query=texts[0], + vector_field=vector_field, + k=2, + num_candidates=3, + ) + assert len(mmr_output) == 2 + assert mmr_output[0]["_source"][text_field] == texts[0] + assert mmr_output[1]["_source"][text_field] == texts[1] + + mmr_output = store.max_marginal_relevance_search( + embedding_service=embedding_service, + query=texts[0], + vector_field=vector_field, + k=2, + num_candidates=3, + lambda_mult=0.1, # more diversity + ) + assert len(mmr_output) == 2 + assert mmr_output[0]["_source"][text_field] == texts[0] + assert mmr_output[1]["_source"][text_field] == texts[2] + + # if fetch_k < k, then the output will be less than k + mmr_output = store.max_marginal_relevance_search( + embedding_service=embedding_service, + query=texts[0], + vector_field=vector_field, + k=3, + num_candidates=2, + ) + assert len(mmr_output) == 2 + + def test_metadata_mapping(self, sync_client: Elasticsearch, index: str) -> None: + """Test that the metadata mapping is applied.""" + test_mappings = { + "my_field": {"type": "keyword"}, + "another_field": {"type": "text"}, + } + store = VectorStore( + index=index, + retrieval_strategy=DenseVectorStrategy(distance=DistanceMetric.COSINE), + embedding_service=FakeEmbeddings(), + num_dimensions=10, + client=sync_client, + metadata_mappings=test_mappings, + ) + + texts = ["foo", "foo", "foo"] + metadatas = [{"my_field": str(i)} for i in range(len(texts))] + store.add_texts(texts=texts, metadatas=metadatas) + + mapping_response = sync_client.indices.get_mapping(index=index) + mapping_properties = mapping_response[index]["mappings"]["properties"] + assert mapping_properties["vector_field"] == { + "type": "dense_vector", + "dims": 10, + "index": True, + "similarity": "cosine", + } + + assert "metadata" in mapping_properties + for key, val in test_mappings.items(): + assert mapping_properties["metadata"]["properties"][key] == val diff --git a/test_elasticsearch/test_strategies.py b/test_elasticsearch/test_strategies.py new file mode 100644 index 000000000..36ce63e9f --- /dev/null +++ b/test_elasticsearch/test_strategies.py @@ -0,0 +1,90 @@ +# 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. + +import pytest + +from elasticsearch.helpers.vectorstore import ( + DenseVectorScriptScoreStrategy, + DenseVectorStrategy, + SparseVectorStrategy, +) + + +def test_sparse_vector_strategy_raises_errors(): + strategy = SparseVectorStrategy("my_model_id") + + with pytest.raises(ValueError): + # missing query + strategy.es_query( + query=None, + query_vector=None, + text_field="text_field", + vector_field="vector_field", + k=10, + num_candidates=20, + filter=[], + ) + + with pytest.raises(ValueError): + # query vector not allowed + strategy.es_query( + query="hi", + query_vector=[1, 2, 3], + text_field="text_field", + vector_field="vector_field", + k=10, + num_candidates=20, + filter=[], + ) + + +def test_dense_vector_strategy_raises_error(): + with pytest.raises(ValueError): + # unknown distance + DenseVectorStrategy(hybrid=True, text_field=None) + + with pytest.raises(ValueError): + # unknown distance + DenseVectorStrategy(distance="unknown distance").es_mappings_settings( + text_field="text_field", vector_field="vector_field", num_dimensions=10 + ) + + +def test_dense_vector_script_score_strategy_raises_error(): + with pytest.raises(ValueError): + # missing query vector + DenseVectorScriptScoreStrategy().es_query( + query=None, + query_vector=None, + text_field="text_field", + vector_field="vector_field", + k=10, + num_candidates=20, + filter=[], + ) + + with pytest.raises(ValueError): + # unknown distance + DenseVectorScriptScoreStrategy(distance="unknown distance").es_query( + query=None, + query_vector=[1, 2, 3], + text_field="text_field", + vector_field="vector_field", + k=10, + num_candidates=20, + filter=[], + ) diff --git a/utils/run-unasync.py b/utils/run-unasync.py index 122ba621f..4a943c10f 100644 --- a/utils/run-unasync.py +++ b/utils/run-unasync.py @@ -16,42 +16,84 @@ # under the License. import os +import subprocess +from glob import glob from pathlib import Path import unasync -def main(): - # Unasync all the generated async code - additional_replacements = { - # We want to rewrite to 'Transport' instead of 'SyncTransport', etc - "AsyncTransport": "Transport", - "AsyncElasticsearch": "Elasticsearch", - # We don't want to rewrite this class - "AsyncSearchClient": "AsyncSearchClient", - # Handling typing.Awaitable[...] isn't done yet by unasync. - "_TYPE_ASYNC_SNIFF_CALLBACK": "_TYPE_SYNC_SNIFF_CALLBACK", - } - rules = [ - unasync.Rule( - fromdir="/elasticsearch/_async/client/", - todir="/elasticsearch/_sync/client/", - additional_replacements=additional_replacements, - ), - ] +def cleanup(source_dir: Path, output_dir: Path, patterns: list[str]): + for file in glob("*.py", root_dir=source_dir): + path = output_dir / file + for pattern in patterns: + subprocess.check_call(["sed", "-i.bak", pattern, str(path)]) + subprocess.check_call(["rm", f"{path}.bak"]) + + +def run( + rule: unasync.Rule, + cleanup_patterns: list[str] = [], +): + root_dir = Path(__file__).absolute().parent.parent + source_dir = root_dir / rule.fromdir.lstrip("/") + output_dir = root_dir / rule.todir.lstrip("/") filepaths = [] - for root, _, filenames in os.walk( - Path(__file__).absolute().parent.parent / "elasticsearch/_async" - ): + for root, _, filenames in os.walk(source_dir): for filename in filenames: - if filename.rpartition(".")[-1] in ( + if filename.rpartition(".")[-1] in { "py", "pyi", - ) and not filename.startswith("utils.py"): + } and not filename.startswith("utils.py"): filepaths.append(os.path.join(root, filename)) - unasync.unasync_files(filepaths, rules) + unasync.unasync_files(filepaths, [rule]) + + if cleanup_patterns: + cleanup(source_dir, output_dir, cleanup_patterns) + + +def main(): + run( + rule=unasync.Rule( + fromdir="/elasticsearch/_async/client/", + todir="/elasticsearch/_sync/client/", + additional_replacements={ + # We want to rewrite to 'Transport' instead of 'SyncTransport', etc + "AsyncTransport": "Transport", + "AsyncElasticsearch": "Elasticsearch", + # We don't want to rewrite this class + "AsyncSearchClient": "AsyncSearchClient", + # Handling typing.Awaitable[...] isn't done yet by unasync. + "_TYPE_ASYNC_SNIFF_CALLBACK": "_TYPE_SYNC_SNIFF_CALLBACK", + }, + ), + ) + + run( + rule=unasync.Rule( + fromdir="elasticsearch/helpers/vectorstore/_async/", + todir="elasticsearch/helpers/vectorstore/_sync/", + additional_replacements={ + "AsyncBM25Strategy": "BM25Strategy", + "AsyncDenseVectorStrategy": "DenseVectorStrategy", + "AsyncDenseVectorScriptScoreStrategy": "DenseVectorScriptScoreStrategy", + "AsyncElasticsearch": "Elasticsearch", + "AsyncElasticsearchEmbeddings": "ElasticsearchEmbeddings", + "AsyncEmbeddingService": "EmbeddingService", + "AsyncRetrievalStrategy": "RetrievalStrategy", + "AsyncSparseVectorStrategy": "SparseVectorStrategy", + "AsyncTransport": "Transport", + "AsyncVectorStore": "VectorStore", + "async_bulk": "bulk", + "_async": "_sync", + }, + ), + cleanup_patterns=[ + "/^import asyncio$/d", + ], + ) if __name__ == "__main__":