use elasticsearch_url fixture; create conftest.py

Author: maxjakob

elasticsearch/vectorstore/_async/strategies.py

@@ -17,7 +17,7 @@
 
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 from elasticsearch import AsyncElasticsearch
 from elasticsearch.vectorstore._async._utils import model_must_be_deployed
@@ -250,7 +250,6 @@ class DenseVector(RetrievalStrategy):
 
     def __init__(
         self,
-        knn_type: Literal["hnsw", "int8_hnsw", "flat", "int8_flat"] = "hnsw",
         distance: DistanceMetric = DistanceMetric.COSINE,
         model_id: Optional[str] = None,
         hybrid: bool = False,
@@ -262,7 +261,6 @@ def __init__(
                 "to enable hybrid you have to specify a text_field (for BM25 matching)"
             )
 
-        self.knn_type = knn_type
         self.distance = distance
         self.model_id = model_id
         self.hybrid = hybrid

elasticsearch/vectorstore/_sync/strategies.py

@@ -17,7 +17,7 @@
 
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 from elasticsearch import Elasticsearch
 from elasticsearch.vectorstore._sync._utils import model_must_be_deployed
@@ -250,7 +250,6 @@ class DenseVector(RetrievalStrategy):
 
     def __init__(
         self,
-        knn_type: Literal["hnsw", "int8_hnsw", "flat", "int8_flat"] = "hnsw",
         distance: DistanceMetric = DistanceMetric.COSINE,
         model_id: Optional[str] = None,
         hybrid: bool = False,
@@ -262,7 +261,6 @@ def __init__(
                 "to enable hybrid you have to specify a text_field (for BM25 matching)"
             )
 
-        self.knn_type = knn_type
         self.distance = distance
         self.model_id = model_id
         self.hybrid = hybrid

elasticsearch/vectorstore/_utils.py

@@ -23,8 +23,8 @@
 
 
 def maximal_marginal_relevance(
-    query_embedding: list[float],
-    embedding_list: list[list[float]],
+    query_embedding: List[float],
+    embedding_list: List[List[float]],
     lambda_mult: float = 0.5,
     k: int = 4,
 ) -> List[int]:

setup.py

@@ -94,77 +94,3 @@
         "orjson": ["orjson>=3"],
     },
 )
-
-vectorstore_package_name = "elasticsearch[vectorstore]"
-base_dir = abspath(dirname(__file__))
-
-with open(join(base_dir, package_name, "_version.py")) as f:
-    package_version = re.search(
-        r"__versionstr__\s+=\s+[\"\']([^\"\']+)[\"\']", f.read()
-    ).group(1)
-
-with open(join(base_dir, "README.rst")) as f:
-    # Remove reST raw directive from README as they're not allowed on PyPI
-    # Those blocks start with a newline and continue until the next newline
-    mode = None
-    lines = []
-    for line in f:
-        if line.startswith(".. raw::"):
-            mode = "ignore_nl"
-        elif line == "\n":
-            mode = "wait_nl" if mode == "ignore_nl" else None
-        if mode is None:
-            lines.append(line)
-
-    long_description = "".join(lines)
-
-
-packages = [
-    package
-    for package in find_packages(where=".", exclude=("test_elasticsearch*",))
-    if package == package_name or package.startswith(package_name + ".")
-]
-
-setup(
-    name=package_name,
-    description="Python client for Elasticsearch",
-    license="Apache-2.0",
-    url="https://github.com/elastic/elasticsearch-py",
-    long_description=long_description,
-    long_description_content_type="text/x-rst",
-    version=package_version,
-    author="Elastic Client Library Maintainers",
-    author_email="[email protected]",
-    project_urls={
-        "Documentation": "https://elasticsearch-py.readthedocs.io",
-        "Source Code": "https://github.com/elastic/elasticsearch-py",
-        "Issue Tracker": "https://github.com/elastic/elasticsearch-py/issues",
-    },
-    packages=packages,
-    package_data={"elasticsearch": ["py.typed", "*.pyi"]},
-    include_package_data=True,
-    zip_safe=False,
-    classifiers=[
-        "Development Status :: 5 - Production/Stable",
-        "License :: OSI Approved :: Apache Software License",
-        "Intended Audience :: Developers",
-        "Operating System :: OS Independent",
-        "Programming Language :: Python",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.7",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.10",
-        "Programming Language :: Python :: 3.11",
-        "Programming Language :: Python :: 3.12",
-        "Programming Language :: Python :: Implementation :: CPython",
-        "Programming Language :: Python :: Implementation :: PyPy",
-    ],
-    python_requires=">=3.7",
-    install_requires=["elastic-transport>=8.13,<9"],
-    extras_require={
-        "requests": ["requests>=2.4.0, <3.0.0"],
-        "async": ["aiohttp>=3,<4"],
-        "orjson": ["orjson>=3"],
-    },
-)

test_elasticsearch/test_server/test_vectorstore/_async/_test_utils.py

@@ -15,12 +15,10 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-import os
-from typing import Any, AsyncIterator, Dict, List, Optional
+from typing import Any, Dict, List
 
 from elastic_transport import AsyncTransport
 
-from elasticsearch import AsyncElasticsearch
 from elasticsearch.vectorstore._async.embedding_service import AsyncEmbeddingService
 
 
@@ -87,69 +85,3 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
     async def perform_request(self, *args, **kwargs):  # type: ignore
         self.requests.append(kwargs)
         return await super().perform_request(*args, **kwargs)
-
-
-def create_es_client(
-    es_params: Optional[Dict[str, str]] = None, es_kwargs: Dict = {}
-) -> AsyncElasticsearch:
-    if es_params is None:
-        es_params = read_env()
-    if not es_kwargs:
-        es_kwargs = {}
-
-    if "es_cloud_id" in es_params:
-        return AsyncElasticsearch(
-            cloud_id=es_params["es_cloud_id"],
-            api_key=es_params["es_api_key"],
-            **es_kwargs,
-        )
-    return AsyncElasticsearch(hosts=[es_params["es_url"]], **es_kwargs)
-
-
-def create_requests_saving_client() -> AsyncElasticsearch:
-    return create_es_client(es_kwargs={"transport_class": AsyncRequestSavingTransport})
-
-
-async def es_client_fixture() -> AsyncIterator[AsyncElasticsearch]:
-    params = read_env()
-    client = create_es_client(params)
-
-    yield client
-
-    # clear indices
-    await clear_test_indices(client)
-
-    # clear all test pipelines
-    try:
-        response = await client.ingest.get_pipeline(id="test_*,*_sparse_embedding")
-
-        for pipeline_id, _ in response.items():
-            try:
-                await client.ingest.delete_pipeline(id=pipeline_id)
-                print(f"Deleted pipeline: {pipeline_id}")  # noqa: T201
-            except Exception as e:
-                print(f"Pipeline error: {e}")  # noqa: T201
-
-    except Exception:
-        pass
-    finally:
-        await client.close()
-
-
-async def clear_test_indices(client: AsyncElasticsearch) -> None:
-    response = await client.indices.get(index="_all")
-    index_names = response.keys()
-    for index_name in index_names:
-        if index_name.startswith("test_"):
-            await client.indices.delete(index=index_name)
-    await client.indices.refresh(index="_all")
-
-
-def read_env() -> Dict:
-    url = os.environ.get("ES_URL", "http://localhost:9200")
-    cloud_id = os.environ.get("ES_CLOUD_ID")
-    api_key = os.environ.get("ES_API_KEY")
-
-    if cloud_id:
-        return {"es_cloud_id": cloud_id, "es_api_key": api_key}
-    return {"es_url": url}

test_elasticsearch/test_server/test_vectorstore/_async/conftest.py

@@ -0,0 +1,103 @@
+#  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 uuid
+from typing import AsyncIterator, Dict
+
+import pytest
+import pytest_asyncio
+
+from elasticsearch import AsyncElasticsearch
+
+from ._test_utils import AsyncRequestSavingTransport
+
+
+@pytest_asyncio.fixture
+async def es_client(elasticsearch_url: str) -> AsyncIterator[AsyncElasticsearch]:
+    client = _create_es_client(elasticsearch_url)
+
+    yield client
+
+    # clear indices
+    await _clear_test_indices(client)
+
+    # clear all test pipelines
+    try:
+        response = await client.ingest.get_pipeline(id="test_*,*_sparse_embedding")
+
+        for pipeline_id, _ in response.items():
+            try:
+                await client.ingest.delete_pipeline(id=pipeline_id)
+                print(f"Deleted pipeline: {pipeline_id}")  # noqa: T201
+            except Exception as e:
+                print(f"Pipeline error: {e}")  # noqa: T201
+
+    except Exception:
+        pass
+    finally:
+        await client.close()
+
+
+@pytest_asyncio.fixture
+async def requests_saving_client(
+    elasticsearch_url: str,
+) -> AsyncIterator[AsyncElasticsearch]:
+    client = _create_es_client(
+        elasticsearch_url, es_kwargs={"transport_class": AsyncRequestSavingTransport}
+    )
+
+    try:
+        yield client
+    finally:
+        await client.close()
+
+
+@pytest.fixture(scope="function")
+def index_name() -> str:
+    return f"test_{uuid.uuid4().hex}"
+
+
+async def _clear_test_indices(client: AsyncElasticsearch) -> None:
+    response = await client.indices.get(index="_all")
+    index_names = response.keys()
+    for index_name in index_names:
+        if index_name.startswith("test_"):
+            await client.indices.delete(index=index_name)
+    await client.indices.refresh(index="_all")
+
+
+def _create_es_client(
+    elasticsearch_url: str, es_kwargs: Dict = {}
+) -> AsyncElasticsearch:
+    if not elasticsearch_url:
+        elasticsearch_url = os.environ.get("ES_URL", "http://localhost:9200")
+    cloud_id = os.environ.get("ES_CLOUD_ID")
+    api_key = os.environ.get("ES_API_KEY")
+
+    if cloud_id:
+        es_params = {"es_cloud_id": cloud_id, "es_api_key": api_key}
+    else:
+        es_params = {"es_url": elasticsearch_url}
+
+    if "es_cloud_id" in es_params:
+        return AsyncElasticsearch(
+            cloud_id=es_params["es_cloud_id"],
+            api_key=es_params["es_api_key"],
+            **es_kwargs,
+        )
+    return AsyncElasticsearch(hosts=[es_params["es_url"]], **es_kwargs)

test_elasticsearch/test_server/test_vectorstore/_async/test_embedding_service.py

@@ -16,31 +16,21 @@
 #  under the License.
 
 import os
-from typing import AsyncIterator
 
 import pytest
-import pytest_asyncio
 
 from elasticsearch import AsyncElasticsearch
 from elasticsearch.vectorstore._async._utils import model_is_deployed
 from elasticsearch.vectorstore._async.embedding_service import (
     AsyncElasticsearchEmbeddings,
 )
 
-from ._test_utils import es_client_fixture
-
 # 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__msmarco-minilm-l-12-v3")
 NUM_DIMENSIONS = int(os.getenv("NUM_DIMENTIONS", "384"))
 
 
-@pytest_asyncio.fixture
-async def es_client() -> AsyncIterator[AsyncElasticsearch]:
-    async for x in es_client_fixture():
-        yield x
-
-
 @pytest.mark.asyncio
 async def test_elasticsearch_embedding_documents(es_client: AsyncElasticsearch) -> None:
     """Test Elasticsearch embedding documents."""