diff --git a/tests/compose/docker-compose.test.yaml b/tests/compose/docker-compose.test.yaml index a5cbbee0c3..b1e016e416 100644 --- a/tests/compose/docker-compose.test.yaml +++ b/tests/compose/docker-compose.test.yaml @@ -26,3 +26,14 @@ services: image: unstract/runner:${UNSTRACT_TEST_VERSION:-latest} environment: - ENVIRONMENT=test + + # LLM completions run inside the worker-unified image; passing the mock env + # here lets execute-path e2e tests run hermetically (no real LLM/secret). + # Empty unless the rig/CI sets it, so production-like runs are unaffected. + worker-executor-v2: + environment: + - UNSTRACT_LLM_MOCK_RESPONSE=${UNSTRACT_LLM_MOCK_RESPONSE:-} + + worker-file-processing-v2: + environment: + - UNSTRACT_LLM_MOCK_RESPONSE=${UNSTRACT_LLM_MOCK_RESPONSE:-} diff --git a/tests/critical_paths.yaml b/tests/critical_paths.yaml index 2a2ff8c06b..bdd87a12d3 100644 --- a/tests/critical_paths.yaml +++ b/tests/critical_paths.yaml @@ -43,6 +43,7 @@ paths: - id: workflow-create-execute description: "Create a workflow, configure source+destination, execute, poll, fetch result." covered_by: [e2e-workflow] + proof: marker - id: api-deployment-provision description: "Deploying a workflow as an API mints a usable key and a resolvable endpoint." @@ -57,6 +58,7 @@ paths: - id: api-deployment-run description: "Deploy a workflow as an API, POST a document, receive structured JSON." covered_by: [e2e-api-deployment] + proof: marker - id: prompt-studio-author description: "Create a Prompt Studio project and add a prompt to it." @@ -83,7 +85,8 @@ paths: - id: usage-token-tracking description: "Per-execution token usage is recorded and retrievable." - covered_by: [] # gap + covered_by: [e2e-api-deployment] + proof: marker - id: workflow-execution-fan-out description: "Multi-file workflow execution fans out to file-processing workers and rejoins." diff --git a/tests/e2e/api_deployment/__init__.py b/tests/e2e/api_deployment/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/e2e/api_deployment/test_api_deployment_run.py b/tests/e2e/api_deployment/test_api_deployment_run.py new file mode 100644 index 0000000000..f33d9e002f --- /dev/null +++ b/tests/e2e/api_deployment/test_api_deployment_run.py @@ -0,0 +1,77 @@ +"""E2E: deploy a workflow as an API, POST a document, get structured JSON back. + +Covers two critical paths in one hermetic run (the LLM is mocked via +UNSTRACT_LLM_MOCK_RESPONSE, so no real provider/secret is touched): + + • api-deployment-run — the public deployment endpoint executes and returns + the mocked answer as structured JSON. + • usage-token-tracking — per-execution token usage is recorded and returned + (litellm stamps a fixed 10/20/30 on the mock). + +Synchronous execution (timeout > 0 blocks) so the result comes back on the POST +itself — no polling and no out-of-band result store to read. +""" + +from __future__ import annotations + +import io +import uuid + +import pytest + +from tests.e2e.conftest import ProvisionedWorkflow + +pytestmark = [pytest.mark.e2e, pytest.mark.critical] + + +@pytest.mark.critical_path("api-deployment-run") +@pytest.mark.critical_path("usage-token-tracking") +def test_api_deployment_returns_mocked_answer( + provisioned_workflow: ProvisionedWorkflow, llm_mock_response: str +) -> None: + pw = provisioned_workflow + session = pw.session + csrf = {"X-CSRFToken": session.cookies.get("csrftoken", "")} + + api_name = f"e2edep{uuid.uuid4().hex[:8]}" + resp = session.post( + f"{pw.prefix}/api/deployment/", + headers=csrf, + json={ + "workflow": pw.workflow_id, + "display_name": f"e2e {api_name}", + "description": "e2e api deployment", + "api_name": api_name, + "is_active": True, + }, + timeout=30, + ) + assert resp.status_code == 201, f"deploy: {resp.text}" + deployment = resp.json() + api_key = deployment["api_key"] + endpoint = deployment["api_endpoint"] + exec_url = endpoint if endpoint.startswith("http") else f"{pw.base}/{endpoint.lstrip('/')}" + + document = io.BytesIO(b"Hello invoice 123. This is a test document about widgets.") + resp = session.post( + exec_url, + headers={"Authorization": f"Bearer {api_key}"}, + data={"timeout": 300, "include_metadata": True}, + files={"files": ("probe.txt", document, "text/plain")}, + timeout=310, + ) + assert resp.status_code == 200, f"execute: HTTP {resp.status_code}: {resp.text}" + message = resp.json()["message"] + assert message["execution_status"] == "COMPLETED", message + + file_result = message["result"][0] + assert file_result["status"] == "Success", file_result + # api-deployment-run: the mocked completion surfaces as the prompt's answer. + assert file_result["result"]["output"]["answer"] == llm_mock_response + + # usage-token-tracking: litellm stamps a deterministic 10/20/30 on the mock, + # recorded per-execution and returned under the file's metadata. + usage = file_result["result"]["metadata"]["extraction_llm"][0] + assert usage["input_tokens"] == 10, usage + assert usage["output_tokens"] == 20, usage + assert usage["total_tokens"] == 30, usage diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index d4ba147507..bd0ca9e122 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -9,11 +9,24 @@ from __future__ import annotations import os +import uuid +from dataclasses import dataclass import pytest import requests -from tests.rig.runtime import PlatformEndpoints +from tests.rig.runtime import LLM_MOCK_RESPONSE_ENV, PlatformEndpoints + +# Static adapter registry ids from unstract.sdk1 (see adapters/*/). The NoOp +# x2text/vectordb return canned output so extraction/indexing need no real +# service; the OpenAI LLM/embedding rows never reach a provider — the LLM is +# mocked via UNSTRACT_LLM_MOCK_RESPONSE and, with chunk_size=0, embedding + +# vectordb are never invoked. Adapter create does not validate connectivity, +# so fake creds persist. +_LLM_ADAPTER = "openai|502ecf49-e47c-445c-9907-6d4b90c5cd17" +_EMBED_ADAPTER = "openai|717a0b0e-3bbc-41dc-9f0c-5689437a1151" +_VECTORDB_ADAPTER = "noOpVectorDb|ca4d6056-4971-4bc8-97e3-9e36290b5bc0" +_X2TEXT_ADAPTER = "noOpX2text|mp66d1op-7100-d340-9101-846fc7115676" @pytest.fixture(scope="session") @@ -69,3 +82,192 @@ def authed_session(platform: PlatformEndpoints) -> requests.Session: ) resp.raise_for_status() return session + + +@pytest.fixture(scope="session") +def llm_mock_response() -> str: + """The exact string the workers were told to return for every completion. + + The rig sets this when it boots the platform (ComposeRuntime); a manual/local + run must export it to match what the workers got. Absent it, execute-path + tests can't assert a deterministic answer, so skip rather than guess. + """ + value = os.environ.get(LLM_MOCK_RESPONSE_ENV) + if not value: + pytest.skip( + f"{LLM_MOCK_RESPONSE_ENV} not set — execute-path e2e needs the LLM " + "mock; the rig sets it when it boots the platform." + ) + return value + + +@dataclass(frozen=True) +class ProvisionedWorkflow: + """Handles for a hermetic, execute-ready workflow (one Prompt Studio tool).""" + + session: requests.Session + base: str # backend root, e.g. http://localhost:8000 + prefix: str # tenant-scoped API root: {base}/api/v1/unstract/{org_id} + org_id: str + workflow_id: str + tool_id: str + prompt_registry_id: str + + +def _org_id(session: requests.Session, base: str) -> str: + orgs = session.get(f"{base}/api/v1/organization", timeout=10) + orgs.raise_for_status() + return orgs.json()["organizations"][0]["id"] + + +def _post(session: requests.Session, url: str, **kw: object) -> requests.Response: + headers = dict(kw.pop("headers", {})) + headers["X-CSRFToken"] = session.cookies.get("csrftoken", "") + return session.post(url, headers=headers, timeout=60, **kw) + + +def _patch(session: requests.Session, url: str, **kw: object) -> requests.Response: + headers = dict(kw.pop("headers", {})) + headers["X-CSRFToken"] = session.cookies.get("csrftoken", "") + return session.patch(url, headers=headers, timeout=60, **kw) + + +@pytest.fixture(scope="session") +def provisioned_workflow( + platform: PlatformEndpoints, authed_session: requests.Session +) -> ProvisionedWorkflow: + """Stand up an API workflow backed by a single Prompt Studio tool. + + Chain (all HTTP, mirrors the app): create 4 adapters -> Prompt Studio project + (auto-creates a default profile) -> pin adapters + chunk_size=0 on that + profile -> add one text prompt -> export (mints a PromptStudioRegistry row) -> + workflow -> point both endpoints at API -> attach the exported tool. The + resulting workflow executes hermetically because the LLM is mocked and, with + chunk_size=0, embedding + vectordb are skipped. + """ + s = authed_session + base = platform.backend_url.rstrip("/") + org_id = _org_id(s, base) + prefix = f"{base}/api/v1/unstract/{org_id}" + sfx = uuid.uuid4().hex[:8] # adapter/tool names are unique per org + + def create_adapter(adapter_id: str, adapter_type: str, metadata: dict) -> str: + name = f"e2e-{adapter_type.lower()}-{sfx}" + resp = _post( + s, + f"{prefix}/adapter/", + json={ + "adapter_id": adapter_id, + "adapter_name": name, + "adapter_type": adapter_type, + "adapter_metadata": {"adapter_name": name, **metadata}, + }, + ) + assert resp.status_code == 201, f"adapter {adapter_type}: {resp.text}" + return resp.json()["id"] + + # api_base is required by the SDK's OpenAI params even though completion is + # mocked (validation runs before the litellm call). + llm_id = create_adapter( + _LLM_ADAPTER, + "LLM", + {"api_key": "sk-test", "model": "gpt-4o", "api_base": "https://api.openai.com/v1"}, + ) + embed_id = create_adapter( + _EMBED_ADAPTER, "EMBEDDING", {"api_key": "sk-test", "model": "text-embedding-3-small"} + ) + vdb_id = create_adapter(_VECTORDB_ADAPTER, "VECTOR_DB", {"wait_time": 0}) + x2t_id = create_adapter(_X2TEXT_ADAPTER, "X2TEXT", {"wait_time": 0}) + + resp = _post( + s, + f"{prefix}/prompt-studio/", + json={"tool_name": f"e2e-{sfx}", "description": "e2e execute", "author": "e2e"}, + ) + assert resp.status_code == 201, f"prompt-studio project: {resp.text}" + tool_id = resp.json()["tool_id"] + + # Project creation auto-makes the default profile; creating a second yields + # two is_default rows and breaks export. Patch the auto one instead. + resp = s.get(f"{prefix}/prompt-studio/prompt-studio-profile/{tool_id}/", timeout=30) + resp.raise_for_status() + profile_id = next(p["profile_id"] for p in resp.json() if p.get("is_default")) + resp = _patch( + s, + f"{prefix}/prompt-studio/profile-manager/{profile_id}/", + json={ + "llm": llm_id, + "embedding_model": embed_id, + "vector_store": vdb_id, + "x2text": x2t_id, + "chunk_size": 0, + "chunk_overlap": 0, + }, + ) + assert resp.status_code == 200, f"profile patch: {resp.text}" + + resp = _post( + s, + f"{prefix}/prompt-studio/prompt-studio-prompt/{tool_id}/", + json={ + "tool_id": tool_id, + "prompt_key": "answer", + "prompt": "What is this document about?", + "prompt_type": "PROMPT", + "enforce_type": "text", + "sequence_number": 1, + "active": True, + "profile_manager": profile_id, + }, + ) + assert resp.status_code == 201, f"add prompt: {resp.text}" + + resp = _post( + s, + f"{prefix}/prompt-studio/export/{tool_id}", + json={"force_export": True, "is_shared_with_org": True}, + ) + assert resp.status_code == 200, f"export: {resp.text}" + # get_queryset returns None unless a filter arg is present -> pass custom_tool. + resp = s.get( + f"{prefix}/prompt-studio/registry/", params={"custom_tool": tool_id}, timeout=30 + ) + resp.raise_for_status() + regs = resp.json() + reg_list = regs if isinstance(regs, list) else regs.get("results", []) + assert reg_list, f"registry empty for tool {tool_id}" + prompt_registry_id = reg_list[0]["prompt_registry_id"] + + resp = _post(s, f"{prefix}/workflow/", json={"workflow_name": f"e2e-wf-{sfx}"}) + assert resp.status_code == 201, f"create workflow: {resp.text}" + workflow_id = resp.json()["id"] + + resp = s.get(f"{prefix}/workflow/endpoint/", params={"workflow": workflow_id}, timeout=30) + resp.raise_for_status() + eps = resp.json() + eps = eps if isinstance(eps, list) else eps.get("results", []) + for endpoint in eps: + if endpoint.get("workflow") == workflow_id: + resp = _patch( + s, + f"{prefix}/workflow/endpoint/{endpoint['id']}/", + json={"connection_type": "API"}, + ) + assert resp.status_code == 200, f"patch endpoint: {resp.text}" + + resp = _post( + s, + f"{prefix}/tool_instance/", + json={"workflow_id": workflow_id, "tool_id": prompt_registry_id}, + ) + assert resp.status_code == 201, f"attach tool: {resp.text}" + + return ProvisionedWorkflow( + session=s, + base=base, + prefix=prefix, + org_id=org_id, + workflow_id=workflow_id, + tool_id=tool_id, + prompt_registry_id=prompt_registry_id, + ) diff --git a/tests/e2e/workflows/__init__.py b/tests/e2e/workflows/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/e2e/workflows/test_workflow_execute.py b/tests/e2e/workflows/test_workflow_execute.py new file mode 100644 index 0000000000..c9dc53efdd --- /dev/null +++ b/tests/e2e/workflows/test_workflow_execute.py @@ -0,0 +1,70 @@ +"""E2E: create a workflow, execute a document through it, poll to COMPLETED. + +Exercises the app-facing ``/workflow/execute/`` path (distinct from the public +API-deployment endpoint). The LLM is mocked via UNSTRACT_LLM_MOCK_RESPONSE, so a +COMPLETED status with a successful file is itself the hermetic proof: without a +real key the completion would fail and the file would not succeed. + +The per-file answer is not exposed over HTTP for a manual execute (it lands in a +worker-side result store), so this asserts on the execution status the status +endpoint does expose. The exact mocked answer is asserted by the API-deployment +test, whose endpoint returns it as JSON. +""" + +from __future__ import annotations + +import io +import time + +import pytest + +from tests.e2e.conftest import ProvisionedWorkflow + +pytestmark = [pytest.mark.e2e, pytest.mark.critical] + +_TERMINAL = {"COMPLETED", "ERROR", "STOPPED"} + + +@pytest.mark.critical_path("workflow-create-execute") +def test_workflow_execute_completes( + provisioned_workflow: ProvisionedWorkflow, llm_mock_response: str +) -> None: + pw = provisioned_workflow + session = pw.session + csrf = {"X-CSRFToken": session.cookies.get("csrftoken", "")} + + # Two-step contract: call 1 (no files) creates a PENDING execution without + # dispatching; call 2 (with that execution_id + files) uploads and dispatches. + resp = session.post( + f"{pw.prefix}/workflow/execute/", + headers=csrf, + data={"workflow_id": pw.workflow_id}, + timeout=60, + ) + assert resp.status_code == 200, f"create execution: {resp.text}" + execution_id = resp.json()["execution_id"] + + document = io.BytesIO(b"Hello invoice 123. This is a test document about widgets.") + resp = session.post( + f"{pw.prefix}/workflow/execute/", + headers=csrf, + data={"workflow_id": pw.workflow_id, "execution_id": execution_id}, + files={"files": ("probe.txt", document, "text/plain")}, + timeout=120, + ) + assert resp.status_code == 200, f"dispatch execution: {resp.text}" + + # Poll the top-level execution app (retrieve by execution_id). The + # workflow/execution// route filters by workflow_id and 404s here. + final = {} + deadline = time.monotonic() + 180 + while time.monotonic() < deadline: + resp = session.get(f"{pw.prefix}/execution/{execution_id}/", timeout=30) + resp.raise_for_status() + final = resp.json() + if final.get("status") in _TERMINAL: + break + time.sleep(2) + + assert final.get("status") == "COMPLETED", final + assert final.get("successful_files") == 1, final diff --git a/tests/rig/runtime.py b/tests/rig/runtime.py index 2643d9da99..c02fcf8b08 100644 --- a/tests/rig/runtime.py +++ b/tests/rig/runtime.py @@ -33,6 +33,12 @@ COMPOSE_OVERLAY = REPO_ROOT / "tests" / "compose" / "docker-compose.test.yaml" BASE_COMPOSE = REPO_ROOT / "docker" / "docker-compose.yaml" +# Deterministic completion the execute-path e2e tests provision against. The +# overlay forwards UNSTRACT_LLM_MOCK_RESPONSE into the workers; the same value +# reaches pytest via os.environ so tests can assert the exact string back. +LLM_MOCK_RESPONSE_ENV = "UNSTRACT_LLM_MOCK_RESPONSE" +DEFAULT_LLM_MOCK_RESPONSE = "MOCK_LLM_OK" + @dataclass(frozen=True) class InfraEndpoints: @@ -130,6 +136,9 @@ def __init__(self, *, project_name: str = "unstract-test") -> None: def up(self) -> PlatformEndpoints: if shutil.which("docker") is None: raise RuntimeError("ComposeRuntime requires the `docker` CLI on PATH") + # setdefault so a CI/dev override wins; copied into os.environ so both + # the compose subprocess (worker env) and the pytest groups see it. + os.environ.setdefault(LLM_MOCK_RESPONSE_ENV, DEFAULT_LLM_MOCK_RESPONSE) files = ["-f", str(BASE_COMPOSE)] if COMPOSE_OVERLAY.exists(): files += ["-f", str(COMPOSE_OVERLAY)] diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index b45b856239..fde9cf321c 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -31,6 +31,20 @@ logger = logging.getLogger(__name__) +# Test-only escape hatch: when UNSTRACT_LLM_MOCK_RESPONSE is set, litellm returns +# it as the completion instead of calling a provider, so execute-path tests run +# hermetically with no real LLM/secret. litellm stamps fixed usage on the mock +# (tune via DEFAULT_MOCK_RESPONSE_PROMPT/COMPLETION_TOKEN_COUNT). Sentinels like +# "litellm.RateLimitError" force error paths. Unset in production => no-op. +_MOCK_RESPONSE_ENV = "UNSTRACT_LLM_MOCK_RESPONSE" + + +def _inject_mock_response(completion_kwargs: dict[str, object]) -> None: + mock = os.getenv(_MOCK_RESPONSE_ENV) + if mock and "mock_response" not in completion_kwargs: + completion_kwargs["mock_response"] = mock + + # Drop unsupported params rather than raising errors. # Set once at module level instead of per-call to avoid repeated # global mutation in concurrent environments. @@ -327,6 +341,7 @@ def complete(self, prompt: str, **kwargs: object) -> dict[str, object]: ) completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) + _inject_mock_response(completion_kwargs) completion_kwargs.pop("cost_model", None) # if hasattr(self, "model") and self.model not in O1_MODELS: @@ -450,6 +465,7 @@ def complete_vision( ) completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) + _inject_mock_response(completion_kwargs) completion_kwargs.pop("cost_model", None) response: dict[str, object] = litellm.completion( @@ -517,6 +533,7 @@ def stream_complete( ) completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) + _inject_mock_response(completion_kwargs) completion_kwargs.pop("cost_model", None) max_retries = pop_litellm_retry_kwargs( @@ -588,6 +605,7 @@ async def acomplete(self, prompt: str, **kwargs: object) -> dict[str, object]: ) completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) + _inject_mock_response(completion_kwargs) completion_kwargs.pop("cost_model", None) max_retries = pop_litellm_retry_kwargs( diff --git a/unstract/sdk1/tests/test_mock_response.py b/unstract/sdk1/tests/test_mock_response.py new file mode 100644 index 0000000000..902407cabe --- /dev/null +++ b/unstract/sdk1/tests/test_mock_response.py @@ -0,0 +1,73 @@ +"""The UNSTRACT_LLM_MOCK_RESPONSE escape hatch and the litellm mock contract it +relies on. Execute-path critical-path tests run hermetically only if both hold: +our injector wires mock_response when (and only when) the env is set, and litellm +returns that string with deterministic usage so token-tracking assertions are exact. +""" + +from __future__ import annotations + +from functools import lru_cache +from importlib import import_module + +import litellm +import pytest + + +@lru_cache(maxsize=1) +def _load_llm_module() -> object: + import sys + from types import ModuleType + + # Stub python-magic so importing LLM does not depend on libmagic. + sys.modules.setdefault("magic", ModuleType("magic")) + return import_module("unstract.sdk1.llm") + + +def _inject(kwargs: dict[str, object]) -> dict[str, object]: + _load_llm_module()._inject_mock_response(kwargs) + return kwargs + + +def test_inject_is_noop_when_env_unset(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("UNSTRACT_LLM_MOCK_RESPONSE", raising=False) + assert _inject({"model": "gpt-4o"}) == {"model": "gpt-4o"} + + +def test_inject_sets_mock_response_when_env_set( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("UNSTRACT_LLM_MOCK_RESPONSE", "canned answer") + assert _inject({"model": "gpt-4o"})["mock_response"] == "canned answer" + + +def test_inject_does_not_clobber_explicit_mock_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("UNSTRACT_LLM_MOCK_RESPONSE", "from-env") + assert _inject({"mock_response": "explicit"})["mock_response"] == "explicit" + + +def test_litellm_mock_contract_returns_string_and_fixed_usage() -> None: + # The whole hermetic-execute design leans on this: mock_response yields the + # string verbatim plus deterministic usage (litellm defaults 10/20/30), + # letting usage-token-tracking assert exact counts. Guards a litellm bump. + resp = litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "anything"}], + mock_response="canned answer", + ) + assert resp["choices"][0]["message"]["content"] == "canned answer" + assert resp["usage"]["prompt_tokens"] == 10 + assert resp["usage"]["completion_tokens"] == 20 + assert resp["usage"]["total_tokens"] == 30 + + +def test_litellm_mock_error_sentinel_raises() -> None: + # Error-path critical paths depend on the sentinel forcing a real + # litellm error type rather than a normal completion. + with pytest.raises(litellm.RateLimitError): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "anything"}], + mock_response="litellm.RateLimitError", + )