diff --git a/.github/workflows/_compile_integration_test.yml b/.github/workflows/_integration_test.yml similarity index 55% rename from .github/workflows/_compile_integration_test.yml rename to .github/workflows/_integration_test.yml index d8b5b9b..0772b99 100644 --- a/.github/workflows/_compile_integration_test.yml +++ b/.github/workflows/_integration_test.yml @@ -1,4 +1,4 @@ -name: compile-integration-test +name: integration-test on: workflow_call: @@ -13,6 +13,7 @@ env: jobs: build: + name: "make integration_tests" defaults: run: working-directory: ${{ inputs.working-directory }} @@ -20,11 +21,24 @@ jobs: strategy: matrix: python-version: - - "3.8" - "3.9" - "3.10" - "3.11" - name: "poetry run pytest -m compile tests/integration_tests #${{ matrix.python-version }}" + services: + elasticsearch: + image: elasticsearch:8.13.0 + env: + 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 + options: >- + --health-cmd "curl --fail http://localhost:9200/_cluster/health" + --health-start-period 10s + --health-timeout 3s + --health-interval 3s + --health-retries 10 steps: - uses: actions/checkout@v4 @@ -34,24 +48,12 @@ jobs: python-version: ${{ matrix.python-version }} poetry-version: ${{ env.POETRY_VERSION }} working-directory: ${{ inputs.working-directory }} - cache-key: compile-integration + cache-key: integration-tests - - name: Install integration dependencies + - name: Install dependencies shell: bash run: poetry install --with=test_integration,test - - name: Check integration tests compile + - name: Run integration tests shell: bash - run: poetry run pytest -m compile tests/integration_tests - - - name: Ensure the tests did not create any additional files - shell: bash - run: | - set -eu - - STATUS="$(git status)" - echo "$STATUS" - - # grep will exit non-zero if the target message isn't found, - # and `set -e` above will cause the step to fail. - echo "$STATUS" | grep 'nothing to commit, working tree clean' + run: make integration_test diff --git a/.github/workflows/check_diffs.yml b/.github/workflows/check_diffs.yml index b6f4891..2137436 100644 --- a/.github/workflows/check_diffs.yml +++ b/.github/workflows/check_diffs.yml @@ -59,20 +59,21 @@ jobs: working-directory: ${{ matrix.working-directory }} secrets: inherit - compile-integration-tests: + integration-test: name: cd ${{ matrix.working-directory }} needs: [ build ] if: ${{ needs.build.outputs.dirs-to-test != '[]' }} strategy: matrix: working-directory: ${{ fromJson(needs.build.outputs.dirs-to-test) }} - uses: ./.github/workflows/_compile_integration_test.yml + uses: ./.github/workflows/_integration_test.yml with: working-directory: ${{ matrix.working-directory }} secrets: inherit + ci_success: name: "CI Success" - needs: [build, lint, test, compile-integration-tests] + needs: [build, lint, test, integration-test] if: | always() runs-on: ubuntu-latest diff --git a/libs/elasticsearch/langchain_elasticsearch/_utilities.py b/libs/elasticsearch/langchain_elasticsearch/_utilities.py index 33b3022..79d1d12 100644 --- a/libs/elasticsearch/langchain_elasticsearch/_utilities.py +++ b/libs/elasticsearch/langchain_elasticsearch/_utilities.py @@ -90,7 +90,7 @@ def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray: return similarity -def check_if_model_deployed(client: Elasticsearch, model_id: str) -> None: +def model_must_be_deployed(client: Elasticsearch, model_id: str) -> None: try: dummy = {"x": "y"} client.ml.infer_trained_model(model_id=model_id, docs=[dummy]) @@ -106,3 +106,11 @@ def check_if_model_deployed(client: Elasticsearch, model_id: str) -> None: # This error is expected because we do not know the expected document # shape and just use a dummy doc above. pass + + +def model_is_deployed(es_client: Elasticsearch, model_id: str) -> bool: + try: + model_must_be_deployed(es_client, model_id) + return True + except NotFoundError: + return False diff --git a/libs/elasticsearch/langchain_elasticsearch/vectorstores.py b/libs/elasticsearch/langchain_elasticsearch/vectorstores.py index b3c0b5b..cc01c92 100644 --- a/libs/elasticsearch/langchain_elasticsearch/vectorstores.py +++ b/libs/elasticsearch/langchain_elasticsearch/vectorstores.py @@ -22,8 +22,8 @@ from langchain_elasticsearch._utilities import ( DistanceStrategy, - check_if_model_deployed, maximal_marginal_relevance, + model_must_be_deployed, with_user_agent_header, ) @@ -204,7 +204,7 @@ def before_index_setup( self, client: "Elasticsearch", text_field: str, vector_query_field: str ) -> None: if self.query_model_id: - check_if_model_deployed(client, self.query_model_id) + model_must_be_deployed(client, self.query_model_id) def index( self, @@ -348,7 +348,7 @@ def before_index_setup( self, client: "Elasticsearch", text_field: str, vector_query_field: str ) -> None: if self.model_id: - check_if_model_deployed(client, self.model_id) + model_must_be_deployed(client, self.model_id) # Create a pipeline for the model client.ingest.put_pipeline( diff --git a/libs/elasticsearch/tests/integration_tests/_test_utilities.py b/libs/elasticsearch/tests/integration_tests/_test_utilities.py index 742a0a5..f191c39 100644 --- a/libs/elasticsearch/tests/integration_tests/_test_utilities.py +++ b/libs/elasticsearch/tests/integration_tests/_test_utilities.py @@ -1,10 +1,33 @@ import os -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from elastic_transport import Transport from elasticsearch import Elasticsearch +def read_es_params() -> 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} + + +def create_es_client( + es_params: Optional[Dict[str, str]] = None, **kwargs: Transport +) -> Elasticsearch: + if es_params is None: + es_params = read_es_params() + + if "es_cloud_id" in es_params: + return Elasticsearch( + cloud_id=es_params["es_cloud_id"], api_key=es_params["es_api_key"] + ) + return Elasticsearch(hosts=[es_params["es_url"]], **kwargs) + + def clear_test_indices(es: Elasticsearch) -> None: index_names = es.indices.get(index="_all").keys() for index_name in index_names: @@ -23,20 +46,4 @@ def perform_request(self, *args, **kwargs): # type: ignore self.requests.append(kwargs) return super().perform_request(*args, **kwargs) - es_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: - # Running this integration test with Elastic Cloud - # Required for in-stack inference testing (ELSER + model_id) - es = Elasticsearch( - cloud_id=cloud_id, - api_key=api_key, - transport_class=CustomTransport, - ) - else: - # Running this integration test with local docker instance - es = Elasticsearch(hosts=[es_url], transport_class=CustomTransport) - - return es + return create_es_client(transport_class=CustomTransport) diff --git a/libs/elasticsearch/tests/integration_tests/docker-compose.yml b/libs/elasticsearch/tests/integration_tests/docker-compose.yml index b39daa6..b0e832e 100644 --- a/libs/elasticsearch/tests/integration_tests/docker-compose.yml +++ b/libs/elasticsearch/tests/integration_tests/docker-compose.yml @@ -2,12 +2,11 @@ version: "3" services: elasticsearch: - image: docker.elastic.co/elasticsearch/elasticsearch:8.12.1 # https://www.docker.elastic.co/r/elasticsearch/elasticsearch + image: elasticsearch:8.13.0 environment: - discovery.type=single-node - - xpack.security.enabled=false # security has been disabled, so no login or password is required. - - xpack.security.http.ssl.enabled=false - xpack.license.self_generated.type=trial + - xpack.security.enabled=false # disable password and TLS; never do this in production! ports: - "9200:9200" healthcheck: @@ -20,7 +19,7 @@ services: retries: 60 kibana: - image: docker.elastic.co/kibana/kibana:8.12.1 + image: kibana:8.13.0 environment: - ELASTICSEARCH_URL=http://elasticsearch:9200 ports: diff --git a/libs/elasticsearch/tests/integration_tests/test_embeddings.py b/libs/elasticsearch/tests/integration_tests/test_embeddings.py index c512bb7..9519bbc 100644 --- a/libs/elasticsearch/tests/integration_tests/test_embeddings.py +++ b/libs/elasticsearch/tests/integration_tests/test_embeddings.py @@ -1,48 +1,44 @@ """Test elasticsearch_embeddings embeddings.""" +import os + import pytest -from langchain_core.utils import get_from_env +from elasticsearch import Elasticsearch +from langchain_elasticsearch._utilities import model_is_deployed from langchain_elasticsearch.embeddings import ElasticsearchEmbeddings # deployed with # https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-text-emb-vector-search-example.html -DEFAULT_MODEL = "sentence-transformers__msmarco-minilm-l-12-v3" -DEFAULT_NUM_DIMENSIONS = "384" - - -@pytest.fixture -def model_id() -> str: - return get_from_env("model_id", "MODEL_ID", DEFAULT_MODEL) - +MODEL_ID = os.getenv("MODEL_ID", "sentence-transformers__msmarco-minilm-l-12-v3") +NUM_DIMENSIONS = int(os.getenv("NUM_DIMENTIONS", "384")) -@pytest.fixture -def expected_num_dimensions() -> int: - return int( - get_from_env( - "expected_num_dimensions", "EXPECTED_NUM_DIMENSIONS", DEFAULT_NUM_DIMENSIONS - ) - ) +ES_URL = os.environ.get("ES_URL", "http://localhost:9200") +ES_CLIENT = Elasticsearch(hosts=[ES_URL]) -def test_elasticsearch_embedding_documents( - model_id: str, expected_num_dimensions: int -) -> None: +@pytest.mark.skipif( + not model_is_deployed(ES_CLIENT, MODEL_ID), + reason=f"{MODEL_ID} model is not deployed in ML Node, skipping test", +) +def test_elasticsearch_embedding_documents() -> None: """Test Elasticsearch embedding documents.""" documents = ["foo bar", "bar foo", "foo"] - embedding = ElasticsearchEmbeddings.from_credentials(model_id) + embedding = ElasticsearchEmbeddings.from_credentials(MODEL_ID) output = embedding.embed_documents(documents) assert len(output) == 3 - assert len(output[0]) == expected_num_dimensions - assert len(output[1]) == expected_num_dimensions - assert len(output[2]) == expected_num_dimensions + assert len(output[0]) == NUM_DIMENSIONS + assert len(output[1]) == NUM_DIMENSIONS + assert len(output[2]) == NUM_DIMENSIONS -def test_elasticsearch_embedding_query( - model_id: str, expected_num_dimensions: int -) -> None: +@pytest.mark.skipif( + not model_is_deployed(ES_CLIENT, MODEL_ID), + reason=f"{MODEL_ID} model is not deployed in ML Node, skipping test", +) +def test_elasticsearch_embedding_query() -> None: """Test Elasticsearch embedding query.""" document = "foo bar" - embedding = ElasticsearchEmbeddings.from_credentials(model_id) + embedding = ElasticsearchEmbeddings.from_credentials(MODEL_ID) output = embedding.embed_query(document) - assert len(output) == expected_num_dimensions + assert len(output) == NUM_DIMENSIONS diff --git a/libs/elasticsearch/tests/integration_tests/test_vectorstores.py b/libs/elasticsearch/tests/integration_tests/test_vectorstores.py index 3ac6d45..d2941d1 100644 --- a/libs/elasticsearch/tests/integration_tests/test_vectorstores.py +++ b/libs/elasticsearch/tests/integration_tests/test_vectorstores.py @@ -1,23 +1,28 @@ """Test ElasticsearchStore functionality.""" import logging -import os import re import uuid from typing import Any, Dict, Generator, List, Union import pytest -from elasticsearch import Elasticsearch, NotFoundError +from elasticsearch import NotFoundError from elasticsearch.helpers import BulkIndexError from langchain_core.documents import Document +from langchain_elasticsearch._utilities import model_is_deployed from langchain_elasticsearch.vectorstores import ElasticsearchStore from ..fake_embeddings import ( ConsistentFakeEmbeddings, FakeEmbeddings, ) -from ._test_utilities import clear_test_indices, requests_saving_es_client +from ._test_utilities import ( + clear_test_indices, + create_es_client, + read_es_params, + requests_saving_es_client, +) logging.basicConfig(level=logging.DEBUG) @@ -32,47 +37,24 @@ 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 API, loaded via eland) +- 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. """ -modelsDeployed: List[str] = [ - # ".elser_model_1", - # "sentence-transformers__all-minilm-l6-v2", -] +ELSER_MODEL_ID = ".elser_model_2" +TRANSFORMER_MODEL_ID = "sentence-transformers__all-minilm-l6-v2" class TestElasticsearch: - @classmethod - def setup_class(cls) -> None: - if not os.getenv("OPENAI_API_KEY"): - raise ValueError("OPENAI_API_KEY environment variable is not set") - @pytest.fixture(scope="class", autouse=True) def elasticsearch_connection(self) -> Union[dict, Generator[dict, None, None]]: - es_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: - # Running this integration test with Elastic Cloud - # Required for in-stack inference testing (ELSER + model_id) - es = Elasticsearch( - cloud_id=cloud_id, - api_key=api_key, - ) - yield { - "es_cloud_id": cloud_id, - "es_api_key": api_key, - } + params = read_es_params() + es = create_es_client(params) - else: - # Running this integration test with local docker instance - es = Elasticsearch(hosts=es_url) - yield {"es_url": es_url} + yield params # clear indices clear_test_indices(es) @@ -620,8 +602,9 @@ def my_custom_query(query_body: dict, query: str) -> dict: assert output == [Document(page_content="bar")] @pytest.mark.skipif( - "sentence-transformers__all-minilm-l6-v2" not in modelsDeployed, - reason="Sentence Transformers model not deployed in ML Node, skipping test", + not model_is_deployed(create_es_client(), TRANSFORMER_MODEL_ID), + reason=f"{TRANSFORMER_MODEL_ID} model not deployed in ML Node, " + "skipping test", ) def test_similarity_search_with_approx_infer_instack( self, elasticsearch_connection: dict, index_name: str @@ -643,7 +626,7 @@ def test_similarity_search_with_approx_infer_instack( processors=[ { "inference": { - "model_id": "sentence-transformers__all-minilm-l6-v2", + "model_id": TRANSFORMER_MODEL_ID, "field_map": {"query_field": "text_field"}, "target_field": "vector_query_field", } @@ -694,7 +677,7 @@ def assert_query(query_body: dict, query: str) -> dict: "num_candidates": 50, "query_vector_builder": { "text_embedding": { - "model_id": "sentence-transformers__all-minilm-l6-v2", + "model_id": TRANSFORMER_MODEL_ID, "model_text": "foo", } }, @@ -709,8 +692,8 @@ def assert_query(query_body: dict, query: str) -> dict: assert output == [Document(page_content="bar")] @pytest.mark.skipif( - ".elser_model_1" not in modelsDeployed, - reason="ELSER not deployed in ML Node, skipping test", + not model_is_deployed(create_es_client(), ELSER_MODEL_ID), + reason=f"{ELSER_MODEL_ID} model not deployed in ML Node, skipping test", ) def test_similarity_search_with_sparse_infer_instack( self, elasticsearch_connection: dict, index_name: str @@ -721,7 +704,7 @@ def test_similarity_search_with_sparse_infer_instack( texts, **elasticsearch_connection, index_name=index_name, - strategy=ElasticsearchStore.SparseVectorRetrievalStrategy(), + strategy=ElasticsearchStore.SparseVectorRetrievalStrategy(ELSER_MODEL_ID), ) output = docsearch.similarity_search("foo", k=1) assert output == [Document(page_content="foo")]