diff --git a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/embedding_model_cross_language_test.py b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/embedding_model_cross_language_test.py index a0afc56f2..1e5de5268 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/embedding_model_cross_language_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/embedding_model_cross_language_test.py @@ -41,7 +41,6 @@ current_dir = Path(__file__).parent OLLAMA_MODEL = os.environ.get("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text:latest") -os.environ["OLLAMA_EMBEDDING_MODEL"] = OLLAMA_MODEL client = pull_model(OLLAMA_MODEL) @@ -51,7 +50,10 @@ @pytest.mark.skipif( client is None, reason="Ollama client is not available or test model is missing." ) -def test_java_embedding_model_integration(tmp_path: Path) -> None: +def test_java_embedding_model_integration( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("OLLAMA_EMBEDDING_MODEL", OLLAMA_MODEL) env = StreamExecutionEnvironment.get_execution_environment() env.set_runtime_mode(RuntimeExecutionMode.STREAMING) env.set_parallelism(1) diff --git a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/vector_store_cross_language_test.py b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/vector_store_cross_language_test.py index ac49eecdf..da987e011 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/vector_store_cross_language_test.py +++ b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/vector_store_cross_language_test.py @@ -43,7 +43,6 @@ current_dir = Path(__file__).parent OLLAMA_MODEL = os.environ.get("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text:latest") -os.environ["OLLAMA_EMBEDDING_MODEL"] = OLLAMA_MODEL ES_HOST = os.environ.get("ES_HOST") MILVUS_URI = os.environ.get("MILVUS_URI") @@ -61,6 +60,7 @@ def _run_vector_store_integration( backend: str, ) -> None: print(f"[TEST][{backend}] Vector store e2e embedding={embedding_type}") + monkeypatch.setenv("OLLAMA_EMBEDDING_MODEL", OLLAMA_MODEL) monkeypatch.setenv("EMBEDDING_TYPE", embedding_type) monkeypatch.setenv("VECTOR_STORE_BACKEND", backend) suffix = uuid.uuid4().hex diff --git a/python/flink_agents/integrations/embedding_models/local/tests/start_ollama_server.sh b/python/flink_agents/integrations/embedding_models/local/tests/start_ollama_server.sh index 674ef6e70..92da52d79 100755 --- a/python/flink_agents/integrations/embedding_models/local/tests/start_ollama_server.sh +++ b/python/flink_agents/integrations/embedding_models/local/tests/start_ollama_server.sh @@ -36,5 +36,4 @@ if [[ $os == "Linux" ]]; then fi ollama pull all-minilm:22m - ollama run all-minilm:22m fi \ No newline at end of file diff --git a/python/flink_agents/integrations/embedding_models/local/tests/test_ollama_embedding_model.py b/python/flink_agents/integrations/embedding_models/local/tests/test_ollama_embedding_model.py index ab983cfa9..b770eea3f 100644 --- a/python/flink_agents/integrations/embedding_models/local/tests/test_ollama_embedding_model.py +++ b/python/flink_agents/integrations/embedding_models/local/tests/test_ollama_embedding_model.py @@ -78,6 +78,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: truncate=True, resource_context=mock_ctx, ) + setup.open() # Test embedding through setup embedding = setup.embed("This is a test sentence for embedding.") diff --git a/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py b/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py index 753820358..e76cc3faa 100644 --- a/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py +++ b/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py @@ -48,6 +48,7 @@ def get_resource(name: str, type: ResourceType) -> Resource: embedding_model = OpenAIEmbeddingModelSetup( name="openai", model=test_model, connection="openai", resource_context=mock_ctx ) + embedding_model.open() response = embedding_model.embed("Hello, Flink Agent!") assert response is not None