diff --git a/nua/e2e/conftest.py b/nua/e2e/conftest.py index 519f232..fe1a8e6 100644 --- a/nua/e2e/conftest.py +++ b/nua/e2e/conftest.py @@ -4,9 +4,10 @@ import nuclia import pytest -from nuclia.config import reset_config_file, set_config_file from nuclia.sdk.auth import NucliaAuth +from nuclia.lib.nua import AsyncNuaClient from dataclasses import dataclass +from typing import AsyncGenerator logger = logging.getLogger("e2e") @@ -22,17 +23,17 @@ class Tokens: PROD_ACCOUNT_ID = "5cec111b-ea23-4b0c-a82a-d1a666dd1fd2" TOKENS: dict[str, Tokens] = { - "europe-1.stashify.cloud": Tokens( + "https://europe-1.stashify.cloud": Tokens( nua_key=os.environ.get("TEST_EUROPE1_STASHIFY_NUA"), pat_key=os.environ.get("STAGE_PERMAMENT_ACCOUNT_OWNER_PAT_TOKEN"), account_id=os.environ.get("TEST_EUROPE1_STASHIFY_ACCOUNT", STAGE_ACCOUNT_ID), ), - "europe-1.nuclia.cloud": Tokens( + "https://europe-1.nuclia.cloud": Tokens( nua_key=os.environ.get("TEST_EUROPE1_NUCLIA_NUA"), pat_key=os.environ.get("PROD_PERMAMENT_ACCOUNT_OWNER_PAT_TOKEN"), account_id=PROD_ACCOUNT_ID, ), - "aws-us-east-2-1.nuclia.cloud": Tokens( + "https://aws-us-east-2-1.nuclia.cloud": Tokens( nua_key=os.environ.get("TEST_AWS_US_EAST_2_1_NUCLIA_NUA"), pat_key=os.environ.get("PROD_PERMAMENT_ACCOUNT_OWNER_PAT_TOKEN"), account_id=PROD_ACCOUNT_ID, @@ -43,12 +44,12 @@ class Tokens: @pytest.fixture( scope="function", params=[ - "europe-1.stashify.cloud", - "europe-1.nuclia.cloud", - "aws-us-east-2-1.nuclia.cloud", + "https://europe-1.stashify.cloud", + "https://europe-1.nuclia.cloud", + "https://aws-us-east-2-1.nuclia.cloud", ], ) -async def nua_config(request): +async def nua_config(request) -> AsyncGenerator[AsyncNuaClient, None]: if ( os.environ.get("TEST_ENV") == "stage" and "stashify.cloud" not in request.param # noqa ): # noqa @@ -63,13 +64,6 @@ async def nua_config(request): assert token.pat_key assert token.account_id - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(b"{}") - temp_file.flush() - set_config_file(temp_file.name) - nuclia_auth = NucliaAuth() - client_id = nuclia_auth.nua(token.nua_key) - assert client_id - nuclia_auth._config.set_default_nua(client_id) - yield request.param - reset_config_file() + yield AsyncNuaClient( + region=request.param, account=token.account_id, token=token.nua_key + ) diff --git a/nua/e2e/regional/test_da_tasks.py b/nua/e2e/regional/test_da_tasks.py index b40b866..7a70593 100644 --- a/nua/e2e/regional/test_da_tasks.py +++ b/nua/e2e/regional/test_da_tasks.py @@ -28,6 +28,7 @@ import base64 from typing import Optional import aiohttp +from nuclia.lib.nua import AsyncNuaClient @dataclass @@ -510,17 +511,20 @@ def validate_labeler_output_text_block(msg: BrokerMessage): @pytest.fixture async def tmp_nua_key( - nua_config: str, aiohttp_client: AsyncGenerator[aiohttp.ClientSession, None] + nua_config: AsyncNuaClient, + aiohttp_client: AsyncGenerator[aiohttp.ClientSession, None], ) -> AsyncGenerator[str, None]: - account_id = TOKENS[nua_config].account_id + account_id = TOKENS[nua_config.region].account_id pat_client_generator = aiohttp_client( - base_url=f"https://{nua_config}", pat_key=TOKENS[nua_config].pat_key, timeout=30 + base_url=nua_config.url, + pat_key=TOKENS[nua_config.url].pat_key, + timeout=30, ) pat_client = await anext(pat_client_generator) nua_client_id, nua_key = await create_nua_key( client=pat_client, account_id=account_id, - title=f"E2E DA AGENTS - {nua_config}", + title=f"E2E DA AGENTS - {nua_config.region}", ) try: yield nua_key @@ -535,7 +539,7 @@ async def tmp_nua_key( "test_input", DA_TEST_INPUTS, ids=lambda test_input: test_input.parameters.name ) async def test_da_agent_tasks( - nua_config: str, + nua_config: AsyncNuaClient, aiohttp_client: AsyncGenerator[aiohttp.ClientSession, None], tmp_nua_key: str, test_input: TestInput, @@ -545,7 +549,7 @@ async def test_da_agent_tasks( start_time = asyncio.get_event_loop().time() try: nua_client_generator = aiohttp_client( - base_url=f"https://{nua_config}", nua_key=tmp_nua_key, timeout=30 + base_url=nua_config.url, nua_key=tmp_nua_key, timeout=30 ) nua_client = await anext(nua_client_generator) diff --git a/nua/e2e/regional/test_llm_chat.py b/nua/e2e/regional/test_llm_chat.py index ed386cb..c5d4f1b 100644 --- a/nua/e2e/regional/test_llm_chat.py +++ b/nua/e2e/regional/test_llm_chat.py @@ -1,12 +1,12 @@ from nuclia.lib.nua_responses import ChatModel, UserPrompt from nuclia.sdk.predict import AsyncNucliaPredict - +from nuclia.lib.nua import AsyncNuaClient from regional.models import ALL_LLMS import pytest @pytest.mark.asyncio_cooperative -async def test_llm_chat(nua_config): +async def test_llm_chat(nua_config: AsyncNuaClient): # Validate that other features such as # * citations # * custom prompts @@ -29,6 +29,7 @@ async def test_llm_chat(nua_config): generated = await np.generate( text=chat_model, model=ALL_LLMS[0], + nc=nua_config, ) # Check that system + user prompt worked assert generated.answer.startswith("ITALIAN") diff --git a/nua/e2e/regional/test_llm_config.py b/nua/e2e/regional/test_llm_config.py index 77412e0..531ad09 100644 --- a/nua/e2e/regional/test_llm_config.py +++ b/nua/e2e/regional/test_llm_config.py @@ -1,25 +1,26 @@ import pytest from nuclia.exceptions import NuaAPIException from nuclia.lib.nua_responses import LearningConfigurationCreation +from nuclia.lib.nua import AsyncNuaClient from nuclia.sdk.predict import AsyncNucliaPredict @pytest.mark.asyncio_cooperative -async def test_llm_config_nua(nua_config): +async def test_llm_config_nua(nua_config: AsyncNuaClient): np = AsyncNucliaPredict() try: - await np.del_config("kbid") + await np.del_config("kbid", nc=nua_config) except NuaAPIException: pass with pytest.raises(NuaAPIException): - config = await np.config("kbid") + config = await np.config("kbid", nc=nua_config) lcc = LearningConfigurationCreation() - await np.set_config("kbid", lcc) + await np.set_config("kbid", lcc, nc=nua_config) - config = await np.config("kbid") + config = await np.config("kbid", nc=nua_config) assert config.resource_labelers_models is None assert config.ner_model == "multilingual" diff --git a/nua/e2e/regional/test_llm_generate.py b/nua/e2e/regional/test_llm_generate.py index 8a43b1d..143b6e4 100644 --- a/nua/e2e/regional/test_llm_generate.py +++ b/nua/e2e/regional/test_llm_generate.py @@ -1,12 +1,14 @@ import pytest from nuclia.sdk.predict import AsyncNucliaPredict - +from nuclia.lib.nua import AsyncNuaClient from regional.models import ALL_LLMS @pytest.mark.asyncio_cooperative @pytest.mark.parametrize("model", ALL_LLMS) -async def test_llm_generate(nua_config, model): +async def test_llm_generate(nua_config: AsyncNuaClient, model): np = AsyncNucliaPredict() - generated = await np.generate("Which is the capital of Catalonia?", model=model) + generated = await np.generate( + "Which is the capital of Catalonia?", model=model, nc=nua_config + ) assert "Barcelona" in generated.answer diff --git a/nua/e2e/regional/test_llm_json.py b/nua/e2e/regional/test_llm_json.py index 2ecf577..2c5af5f 100644 --- a/nua/e2e/regional/test_llm_json.py +++ b/nua/e2e/regional/test_llm_json.py @@ -1,7 +1,7 @@ import pytest from nuclia.lib.nua_responses import ChatModel, UserPrompt from nuclia.sdk.predict import AsyncNucliaPredict - +from nuclia.lib.nua import AsyncNuaClient from regional.models import LLM_WITH_JSON_OUTPUT_SUPPORT SCHEMA = { @@ -36,7 +36,7 @@ @pytest.mark.asyncio_cooperative @pytest.mark.parametrize("model_name", LLM_WITH_JSON_OUTPUT_SUPPORT) -async def test_llm_json(nua_config, model_name): +async def test_llm_json(nua_config: AsyncNuaClient, model_name): np = AsyncNucliaPredict() results = await np.generate( text=ChatModel( @@ -47,5 +47,6 @@ async def test_llm_json(nua_config, model_name): json_schema=SCHEMA, ), model=model_name, + nc=nua_config, ) assert "SPORTS" in results.object["document_type"] diff --git a/nua/e2e/regional/test_llm_rag.py b/nua/e2e/regional/test_llm_rag.py index 0555a18..cc5222d 100644 --- a/nua/e2e/regional/test_llm_rag.py +++ b/nua/e2e/regional/test_llm_rag.py @@ -1,12 +1,12 @@ import pytest from nuclia.sdk.predict import AsyncNucliaPredict - +from nuclia.lib.nua import AsyncNuaClient from regional.models import ALL_LLMS @pytest.mark.asyncio_cooperative @pytest.mark.parametrize("model", ALL_LLMS) -async def test_llm_rag(nua_config, model): +async def test_llm_rag(nua_config: AsyncNuaClient, model): np = AsyncNucliaPredict() generated = await np.rag( question="Which is the CEO of Nuclia?", @@ -15,6 +15,7 @@ async def test_llm_rag(nua_config, model): "Eudald Camprubí is CEO at the same company as Ramon Navarro", ], model=model, + nc=nua_config, ) assert "Eudald" in generated.answer diff --git a/nua/e2e/regional/test_llm_schema.py b/nua/e2e/regional/test_llm_schema.py index 1ab98ee..616a174 100644 --- a/nua/e2e/regional/test_llm_schema.py +++ b/nua/e2e/regional/test_llm_schema.py @@ -1,19 +1,20 @@ from nuclia.sdk.predict import AsyncNucliaPredict import pytest +from nuclia.lib.nua import AsyncNuaClient @pytest.mark.asyncio_cooperative -async def test_llm_schema_nua(nua_config): +async def test_llm_schema_nua(nua_config: AsyncNuaClient): np = AsyncNucliaPredict() - config = await np.schema() + config = await np.schema(nc=nua_config) assert len(config.ner_model.options) == 1 assert len(config.generative_model.options) >= 5 @pytest.mark.asyncio_cooperative -async def test_llm_schema_kbid(nua_config): +async def test_llm_schema_kbid(nua_config: AsyncNuaClient): np = AsyncNucliaPredict() - config = await np.schema("fake_kbid") + config = await np.schema("fake_kbid", nc=nua_config) assert len(config.ner_model.options) == 1 assert len(config.generative_model.options) >= 5 diff --git a/nua/e2e/regional/test_llm_summarize.py b/nua/e2e/regional/test_llm_summarize.py index e8fd976..75cfd95 100644 --- a/nua/e2e/regional/test_llm_summarize.py +++ b/nua/e2e/regional/test_llm_summarize.py @@ -1,5 +1,6 @@ from nuclia.sdk.predict import AsyncNucliaPredict import pytest +from nuclia.lib.nua import AsyncNuaClient DATA = { "barcelona": "Barcelona (pronunciat en català central, [bərsəˈlonə]) és una ciutat i metròpoli a la costa mediterrània de la península Ibèrica. És la capital de Catalunya,[1] així com de la comarca del Barcelonès i de la província de Barcelona, i la segona ciutat en població i pes econòmic de la península Ibèrica,[2][3] després de Madrid. El municipi creix sobre una plana encaixada entre la serralada Litoral, el mar Mediterrani, el riu Besòs i la muntanya de Montjuïc. La ciutat acull les seus de les institucions d'autogovern més importants de la Generalitat de Catalunya: el Parlament de Catalunya, el President i el Govern de la Generalitat. Pel fet d'haver estat capital del Comtat de Barcelona, rep sovint el sobrenom de Ciutat Comtal. També, com que ha estat la ciutat més important del Principat de Catalunya des d'època medieval, rep sovint el sobrenom o títol de cap i casal.[4]", # noqa @@ -27,25 +28,25 @@ @pytest.mark.asyncio_cooperative -async def test_summarize_chatgpt(nua_config): +async def test_summarize_chatgpt(nua_config: AsyncNuaClient): np = AsyncNucliaPredict() - embed = await np.summarize(DATA, model="chatgpt4o") + embed = await np.summarize(DATA, model="chatgpt4o", nc=nua_config) assert "Manresa" in embed.summary assert "Barcelona" in embed.summary @pytest.mark.asyncio_cooperative -async def test_summarize_azure_chatgpt(nua_config): +async def test_summarize_azure_chatgpt(nua_config: AsyncNuaClient): np = AsyncNucliaPredict() - embed = await np.summarize(DATA, model="chatgpt-azure-4o") + embed = await np.summarize(DATA, model="chatgpt-azure-4o", nc=nua_config) assert "Manresa" in embed.summary assert "Barcelona" in embed.summary @pytest.mark.asyncio_cooperative -async def test_summarize_claude(nua_config): +async def test_summarize_claude(nua_config: AsyncNuaClient): np = AsyncNucliaPredict() - embed = await np.summarize(DATA_COFFEE, model="claude-3-fast") + embed = await np.summarize(DATA_COFFEE, model="claude-3-fast", nc=nua_config) # changed to partial summaries since anthropic is not consistent in the global summary at all assert "flat white" in embed.resources["Flat white"].summary.lower() assert "macchiato" in embed.resources["Macchiato"].summary.lower() diff --git a/nua/e2e/regional/test_predict.py b/nua/e2e/regional/test_predict.py index 60e0a6c..07da311 100644 --- a/nua/e2e/regional/test_predict.py +++ b/nua/e2e/regional/test_predict.py @@ -3,13 +3,14 @@ from regional.models import ALL_ENCODERS, ALL_LLMS from nuclia_models.predict.remi import RemiRequest +from nuclia.lib.nua import AsyncNuaClient @pytest.mark.asyncio_cooperative @pytest.mark.parametrize("model", ALL_ENCODERS.keys()) -async def test_predict_sentence(nua_config, model): +async def test_predict_sentence(nua_config: AsyncNuaClient, model): np = AsyncNucliaPredict() - embed = await np.sentence(text="This is my text", model=model) + embed = await np.sentence(text="This is my text", model=model, nc=nua_config) assert embed.time > 0 # Deprecated field (data) assert len(embed.data) == ALL_ENCODERS[model] @@ -20,9 +21,9 @@ async def test_predict_sentence(nua_config, model): @pytest.mark.asyncio_cooperative -async def test_predict_query(nua_config): +async def test_predict_query(nua_config: AsyncNuaClient): np = AsyncNucliaPredict() - embed = await np.query(text="I love Barcelona") + embed = await np.query(text="I love Barcelona", nc=nua_config) # Semantic assert embed.semantic_threshold > 0 assert len(embed.sentence.data) > 128 @@ -39,9 +40,9 @@ async def test_predict_query(nua_config): @pytest.mark.asyncio_cooperative -async def test_predict_tokens(nua_config): +async def test_predict_tokens(nua_config: AsyncNuaClient): np = AsyncNucliaPredict() - embed = await np.tokens(text="I love Barcelona") + embed = await np.tokens(text="I love Barcelona", nc=nua_config) assert embed.tokens[0].text == "Barcelona" assert embed.tokens[0].start == 7 assert embed.tokens[0].end == 16 @@ -50,16 +51,18 @@ async def test_predict_tokens(nua_config): @pytest.mark.asyncio_cooperative @pytest.mark.parametrize("model", ALL_LLMS) -async def test_predict_rephrase(nua_config, model): +async def test_predict_rephrase(nua_config: AsyncNuaClient, model): # Check that rephrase is working for all models np = AsyncNucliaPredict() # TODO: Test that custom rephrase prompt works once SDK supports it - rephrased = await np.rephrase(question="Barcelona best coffe", model=model) + rephrased = await np.rephrase( + question="Barcelona best coffe", model=model, nc=nua_config + ) assert rephrased != "Barcelona best coffe" and rephrased != "" @pytest.mark.asyncio_cooperative -async def test_predict_remi(nua_config): +async def test_predict_remi(nua_config: AsyncNuaClient): np = AsyncNucliaPredict() results = await np.remi( RemiRequest( @@ -70,7 +73,8 @@ async def test_predict_remi(nua_config): "Paris is the capital of France.", "Berlin is the capital of Germany.", ], - ) + ), + nc=nua_config, ) assert results.answer_relevance.score >= 4 diff --git a/nua/e2e/regional/test_processor.py b/nua/e2e/regional/test_processor.py index 0af2815..4a3e25c 100644 --- a/nua/e2e/regional/test_processor.py +++ b/nua/e2e/regional/test_processor.py @@ -1,14 +1,15 @@ import pytest from nuclia.sdk.process import AsyncNucliaProcessing from regional.utils import define_path +from nuclia.lib.nua import AsyncNuaClient @pytest.mark.asyncio_cooperative @pytest.mark.timeout(660) -async def test_pdf(nua_config): +async def test_pdf(nua_config: AsyncNuaClient): path = define_path("2310.14587.pdf") nc = AsyncNucliaProcessing() - payload = await nc.process_file(path, kbid="kbid", timeout=300) + payload = await nc.process_file(path, kbid="kbid", timeout=300, nc=nua_config) assert payload assert ( "As the training data of LLMs often contains undesirable" @@ -18,10 +19,10 @@ async def test_pdf(nua_config): @pytest.mark.asyncio_cooperative @pytest.mark.timeout(360) -async def test_video(nua_config): +async def test_video(nua_config: AsyncNuaClient): path = define_path("simple_video.mp4") nc = AsyncNucliaProcessing() - payload = await nc.process_file(path, kbid="kbid", timeout=300) + payload = await nc.process_file(path, kbid="kbid", timeout=300, nc=nua_config) assert payload assert ( "This is one of the most reflective mirrors" @@ -31,18 +32,18 @@ async def test_video(nua_config): @pytest.mark.asyncio_cooperative @pytest.mark.timeout(660) -async def test_vude_1(nua_config): +async def test_vude_1(nua_config: AsyncNuaClient): path = define_path( "y2mate_is_Stone_1_Minute_Short_Film_Hot_Shot_5hPtU8Jbpg0_720p_1701938639.mp4" # noqa ) nc = AsyncNucliaProcessing() - payload = await nc.process_file(path, kbid="kbid", timeout=300) + payload = await nc.process_file(path, kbid="kbid", timeout=300, nc=nua_config) assert payload print(payload.extracted_text[0].body) assert "harmful" in payload.extracted_text[0].body.text @pytest.mark.asyncio_cooperative -async def test_activity(nua_config): +async def test_activity(nua_config: AsyncNuaClient): nc = AsyncNucliaProcessing() - await nc.status() + await nc.status(nc=nua_config) diff --git a/nua/e2e/regional/utils.py b/nua/e2e/regional/utils.py index 4e9331a..d0a6884 100644 --- a/nua/e2e/regional/utils.py +++ b/nua/e2e/regional/utils.py @@ -1,7 +1,14 @@ from pathlib import Path +from dataclasses import dataclass FILE_PATH = f"{Path(__file__).parent.parent}/assets/" def define_path(file: str): return FILE_PATH + file + + +@dataclass +class TestConfig: + domain: str + config_path: str diff --git a/nua/requirements.txt b/nua/requirements.txt index ca41d0b..5ccf2ba 100644 --- a/nua/requirements.txt +++ b/nua/requirements.txt @@ -1,4 +1,4 @@ -nuclia>=4.3.11 +nuclia>=4.4.2 pytest pytest-timeout aiofiles