add cleanup step for _sync generation

Author: maxjakob

elasticsearch/vectorstore/_sync/_utils.py

@@ -1,9 +1,4 @@
-from elasticsearch import (
-    Elasticsearch,
-    BadRequestError,
-    ConflictError,
-    NotFoundError,
-)
+from elasticsearch import BadRequestError, ConflictError, Elasticsearch, NotFoundError
 
 
 def model_must_be_deployed(client: Elasticsearch, model_id: str) -> None:

elasticsearch/vectorstore/_sync/strategies.py

@@ -3,7 +3,6 @@
 from typing import Any, Dict, List, Literal, Optional, Union, cast
 
 from elasticsearch import Elasticsearch
-
 from elasticsearch.vectorstore._sync._utils import model_must_be_deployed
 from elasticsearch.vectorstore._sync.embedding_service import EmbeddingService
 
@@ -226,9 +225,7 @@ def create_index(
             mappings["properties"]["metadata"] = {"properties": metadata_mapping}
         settings = {"default_pipeline": pipeline_name}
 
-        client.indices.create(
-            index=index_name, mappings=mappings, settings=settings
-        )
+        client.indices.create(index=index_name, mappings=mappings, settings=settings)
 
         return None
 
@@ -287,9 +284,7 @@ def es_query(
         if query_vector:
             knn["query_vector"] = query_vector
         elif self.embedding_service:
-            knn["query_vector"] = self.embedding_service.embed_query(
-                cast(str, query)
-            )
+            knn["query_vector"] = self.embedding_service.embed_query(cast(str, query))
         else:
             # Inference in Elasticsearch. When initializing we make sure to always have
             # a model_id if don't have an embedding_service.
@@ -555,6 +550,4 @@ def create_index(
             }
         }
 
-        client.indices.create(
-            index=index_name, mappings=mappings, settings=settings
-        )
+        client.indices.create(index=index_name, mappings=mappings, settings=settings)

elasticsearch/vectorstore/_sync/vectorestore.py

@@ -4,12 +4,9 @@
 
 from elasticsearch import Elasticsearch
 from elasticsearch.helpers import BulkIndexError, bulk
-
-from elasticsearch.vectorstore._utils import (
-    maximal_marginal_relevance,
-)
 from elasticsearch.vectorstore._sync.embedding_service import EmbeddingService
 from elasticsearch.vectorstore._sync.strategies import RetrievalStrategy
+from elasticsearch.vectorstore._utils import maximal_marginal_relevance
 
 logger = logging.getLogger(__name__)
 

test_elasticsearch/test_server/test_vectorstore/_async/test_embedding_service.py

@@ -23,7 +23,7 @@
 NUM_DIMENSIONS = int(os.getenv("NUM_DIMENTIONS", "384"))
 
 
-@pytest_asyncio.fixture(autouse=True)
+@pytest_asyncio.fixture
 async def es_client() -> AsyncIterator[AsyncElasticsearch]:
     async for x in es_client_fixture():
         yield x

test_elasticsearch/test_server/test_vectorstore/_async/test_vectorestore.py

@@ -53,12 +53,12 @@
 
 
 class TestElasticsearch:
-    @pytest_asyncio.fixture(autouse=True)
+    @pytest_asyncio.fixture
     async def es_client(self) -> AsyncIterator[AsyncElasticsearch]:
         async for x in es_client_fixture():
             yield x
 
-    @pytest_asyncio.fixture(autouse=True)
+    @pytest_asyncio.fixture
     async def requests_saving_client(self) -> AsyncIterator[AsyncElasticsearch]:
         client = create_requests_saving_client()
         try:

test_elasticsearch/test_server/test_vectorstore/_sync/_test_utils.py

@@ -1,9 +1,9 @@
 import os
-from typing import Any, Dict, List, Optional, Iterator
+from typing import Any, Dict, Iterator, List, Optional
 
 from elastic_transport import Transport
-from elasticsearch import Elasticsearch
 
+from elasticsearch import Elasticsearch
 from elasticsearch.vectorstore._sync.embedding_service import EmbeddingService
 
 

test_elasticsearch/test_server/test_vectorstore/_sync/test_embedding_service.py

@@ -1,35 +1,26 @@
 import os
+from typing import Iterator
 
 import pytest
 
-import pytest_asyncio
 from elasticsearch import Elasticsearch
-
-from typing import Iterator
-
 from elasticsearch.vectorstore._sync._utils import model_is_deployed
+from elasticsearch.vectorstore._sync.embedding_service import ElasticsearchEmbeddings
 
-from ._test_utils import (
-    es_client_fixture,
-)
-
-from elasticsearch.vectorstore._sync.embedding_service import (
-    ElasticsearchEmbeddings,
-)
+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(autouse=True)
+@pytest.fixture
 def es_client() -> Iterator[Elasticsearch]:
     for x in es_client_fixture():
         yield x
 
 
-@pytest.mark.asyncio
 def test_elasticsearch_embedding_documents(es_client: Elasticsearch) -> None:
     """Test Elasticsearch embedding documents."""
 
@@ -47,7 +38,6 @@ def test_elasticsearch_embedding_documents(es_client: Elasticsearch) -> None:
     assert len(output[2]) == NUM_DIMENSIONS
 
 
-@pytest.mark.asyncio
 def test_elasticsearch_embedding_query(es_client: Elasticsearch) -> None:
     """Test Elasticsearch embedding query."""
 

test_elasticsearch/test_server/test_vectorstore/_sync/test_vectorestore.py

@@ -1,16 +1,12 @@
 import logging
 import uuid
-from typing import Iterator
-from typing import Any, List, Optional, Union, cast
 from functools import partial
+from typing import Any, Iterator, List, Optional, Union, cast
 
 import pytest
-import pytest_asyncio
-from elasticsearch import Elasticsearch
 
-from elasticsearch import NotFoundError
+from elasticsearch import Elasticsearch, NotFoundError
 from elasticsearch.helpers import BulkIndexError
-
 from elasticsearch.vectorstore._sync import VectorStore
 from elasticsearch.vectorstore._sync._utils import model_is_deployed
 from elasticsearch.vectorstore._sync.strategies import (
@@ -22,11 +18,11 @@
 )
 
 from ._test_utils import (
-    create_requests_saving_client,
-    es_client_fixture,
     ConsistentFakeEmbeddings,
     FakeEmbeddings,
     RequestSavingTransport,
+    create_requests_saving_client,
+    es_client_fixture,
 )
 
 logging.basicConfig(level=logging.DEBUG)
@@ -53,12 +49,12 @@
 
 
 class TestElasticsearch:
-    @pytest_asyncio.fixture(autouse=True)
+    @pytest.fixture
     def es_client(self) -> Iterator[Elasticsearch]:
         for x in es_client_fixture():
             yield x
 
-    @pytest_asyncio.fixture(autouse=True)
+    @pytest.fixture
     def requests_saving_client(self) -> Iterator[Elasticsearch]:
         client = create_requests_saving_client()
         try:
@@ -71,7 +67,6 @@ def index_name(self) -> str:
         """Return the index name."""
         return f"test_{uuid.uuid4().hex}"
 
-    @pytest.mark.asyncio
     def test_search_without_metadata(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -102,7 +97,6 @@ def assert_query(query_body: dict, query: Optional[str]) -> dict:
         output = store.search("foo", k=1, custom_query=assert_query)
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
 
-    @pytest.mark.asyncio
     def test_search_without_metadata_async(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -120,10 +114,7 @@ def test_search_without_metadata_async(
         output = store.search("foo", k=1)
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
 
-    @pytest.mark.asyncio
-    def test_add_vectors(
-        self, es_client: Elasticsearch, index_name: str
-    ) -> None:
+    def test_add_vectors(self, es_client: Elasticsearch, index_name: 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
@@ -145,14 +136,11 @@ def test_add_vectors(
             es_client=es_client,
         )
 
-        store.add_texts(
-            texts=texts, vectors=embedding_vectors, metadatas=metadatas
-        )
+        store.add_texts(texts=texts, vectors=embedding_vectors, metadatas=metadatas)
         output = store.search("foo1", k=1)
         assert [doc["_source"]["text_field"] for doc in output] == ["foo1"]
         assert [doc["_source"]["metadata"]["page"] for doc in output] == [0]
 
-    @pytest.mark.asyncio
     def test_search_with_metadata(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -178,7 +166,6 @@ def test_search_with_metadata(
         assert [doc["_source"]["text_field"] for doc in output] == ["bar"]
         assert [doc["_source"]["metadata"]["page"] for doc in output] == [1]
 
-    @pytest.mark.asyncio
     def test_search_with_filter(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -215,7 +202,6 @@ def assert_query(query_body: dict, query: Optional[str]) -> dict:
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
         assert [doc["_source"]["metadata"]["page"] for doc in output] == [1]
 
-    @pytest.mark.asyncio
     def test_search_script_score(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -264,7 +250,6 @@ def assert_query(query_body: dict, query: Optional[str]) -> dict:
         output = store.search("foo", k=1, custom_query=assert_query)
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
 
-    @pytest.mark.asyncio
     def test_search_script_score_with_filter(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -319,7 +304,6 @@ def assert_query(query_body: dict, query: Optional[str]) -> dict:
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
         assert [doc["_source"]["metadata"]["page"] for doc in output] == [0]
 
-    @pytest.mark.asyncio
     def test_search_script_score_distance_dot_product(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -370,7 +354,6 @@ def assert_query(query_body: dict, query: Optional[str]) -> dict:
         output = store.search("foo", k=1, custom_query=assert_query)
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
 
-    @pytest.mark.asyncio
     def test_search_knn_with_hybrid_search(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -410,7 +393,6 @@ def assert_query(query_body: dict, query: Optional[str]) -> dict:
         output = store.search("foo", k=1, custom_query=assert_query)
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
 
-    @pytest.mark.asyncio
     def test_search_knn_with_hybrid_search_rrf(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -528,7 +510,6 @@ def assert_query(
             custom_query=partial(assert_query, expected_rrf={}),
         )
 
-    @pytest.mark.asyncio
     def test_search_knn_with_custom_query_fn(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -561,7 +542,6 @@ def my_custom_query(query_body: dict, query: Optional[str]) -> dict:
         output = store.search("foo", k=1, custom_query=my_custom_query)
         assert [doc["_source"]["text_field"] for doc in output] == ["bar"]
 
-    @pytest.mark.asyncio
     def test_search_with_knn_infer_instack(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -655,7 +635,6 @@ def assert_query(query_body: dict, query: Optional[str]) -> dict:
         output = store.search("bar", k=1)
         assert [doc["_source"]["text_field"] for doc in output] == ["bar"]
 
-    @pytest.mark.asyncio
     def test_search_with_sparse_infer_instack(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -679,7 +658,6 @@ def test_search_with_sparse_infer_instack(
         output = store.search("foo", k=1)
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
 
-    @pytest.mark.asyncio
     def test_deployed_model_check_fails_semantic(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -693,10 +671,7 @@ def test_deployed_model_check_fails_semantic(
             )
             store.add_texts(["foo", "bar", "baz"])
 
-    @pytest.mark.asyncio
-    def test_search_bm25(
-        self, es_client: Elasticsearch, index_name: str
-    ) -> None:
+    def test_search_bm25(self, es_client: Elasticsearch, index_name: str) -> None:
         """Test end to end using the BM25 retrieval strategy."""
         store = VectorStore(
             user_agent="test",
@@ -722,7 +697,6 @@ def assert_query(query_body: dict, query: Optional[str]) -> dict:
         output = store.search("foo", k=1, custom_query=assert_query)
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
 
-    @pytest.mark.asyncio
     def test_search_bm25_with_filter(
         self, es_client: Elasticsearch, index_name: str
     ) -> None:
@@ -758,7 +732,6 @@ def assert_query(query_body: dict, query: Optional[str]) -> dict:
         assert [doc["_source"]["text_field"] for doc in output] == ["foo"]
         assert [doc["_source"]["metadata"]["page"] for doc in output] == [1]
 
-    @pytest.mark.asyncio
     def test_delete(self, es_client: Elasticsearch, index_name: str) -> None:
         """Test delete methods from vector store."""
         store = VectorStore(
@@ -791,7 +764,6 @@ def test_delete(self, es_client: Elasticsearch, index_name: str) -> None:
         output = store.search("gni", k=10)
         assert len(output) == 0
 
-    @pytest.mark.asyncio
     def test_indexing_exception_error(
         self,
         es_client: Elasticsearch,
@@ -822,7 +794,6 @@ def test_indexing_exception_error(
 
         assert log_message in caplog.text
 
-    @pytest.mark.asyncio
     def test_user_agent(
         self, requests_saving_client: Elasticsearch, index_name: str
     ) -> None:
@@ -845,10 +816,7 @@ def test_user_agent(
         for request in transport.requests:
             assert request["headers"]["User-Agent"] == user_agent
 
-    @pytest.mark.asyncio
-    def test_bulk_args(
-        self, requests_saving_client: Any, index_name: str
-    ) -> None:
+    def test_bulk_args(self, requests_saving_client: Any, index_name: str) -> None:
         """Test to make sure the bulk arguments work as expected."""
         store = VectorStore(
             user_agent="test",
@@ -863,7 +831,6 @@ def test_bulk_args(
         # 1 for index exist, 1 for index create, 3 to index docs
         assert len(store.es_client.transport.requests) == 5  # type: ignore
 
-    @pytest.mark.asyncio
     def test_max_marginal_relevance_search(
         self, es_client: Elasticsearch, index_name: str
     ) -> None: