diff --git a/autoblocks/_impl/prompts/v2/manager.py b/autoblocks/_impl/prompts/v2/manager.py index 3334a611..c79c0df2 100644 --- a/autoblocks/_impl/prompts/v2/manager.py +++ b/autoblocks/_impl/prompts/v2/manager.py @@ -1,8 +1,10 @@ import abc import asyncio import contextlib +import json import logging from datetime import timedelta +from http import HTTPStatus from typing import Any from typing import Dict from typing import Generic @@ -11,13 +13,20 @@ from typing import TypeVar from autoblocks._impl import global_state +from autoblocks._impl.config.constants import API_ENDPOINT_V2 from autoblocks._impl.config.constants import REVISION_LATEST from autoblocks._impl.config.constants import REVISION_UNDEPLOYED +from autoblocks._impl.context_vars import RevisionType +from autoblocks._impl.context_vars import register_revision_usage +from autoblocks._impl.prompts.error import IncompatiblePromptRevisionError from autoblocks._impl.prompts.v2.client import PromptsAPIClient from autoblocks._impl.prompts.v2.context import PromptExecutionContext from autoblocks._impl.prompts.v2.models import PromptMinorVersion from autoblocks._impl.util import AnyTask +from autoblocks._impl.util import AutoblocksEnvVar +from autoblocks._impl.util import encode_uri_component from autoblocks._impl.util import get_running_loop +from autoblocks._impl.util import parse_autoblocks_overrides log = logging.getLogger(__name__) @@ -28,6 +37,34 @@ ExecutionContextType = TypeVar("ExecutionContextType", bound=PromptExecutionContext[Any, Any, Any]) +def is_testing_context() -> bool: + """ + Note that we check for the presence of V2_CI_TEST_RUN_BUILD_ID + """ + return bool(AutoblocksEnvVar.V2_CI_TEST_RUN_BUILD_ID.get()) + + +def prompt_revisions_map() -> dict[str, str]: + """ + The AUTOBLOCKS_OVERRIDES environment variable is a JSON-stringified + object that can contain promptRevisions. This is set in CI test runs triggered + from the UI. Falls back to legacy AUTOBLOCKS_OVERRIDES_PROMPT_REVISIONS. + """ + if not is_testing_context(): + return {} + + # Try new unified format first + overrides = parse_autoblocks_overrides() + if overrides.prompt_revisions: + return overrides.prompt_revisions + + # Fallback to legacy format + prompt_revisions_raw = AutoblocksEnvVar.OVERRIDES_PROMPT_REVISIONS.get() + if not prompt_revisions_raw: + return {} + return json.loads(prompt_revisions_raw) # type: ignore + + class AutoblocksPromptManager( abc.ABC, Generic[ExecutionContextType], @@ -82,6 +119,8 @@ def __init__( # Initialize API client self._client = PromptsAPIClient(api_key=api_key) + self._prompt_revision_override: Optional[Dict[str, Any]] = None + # Validate refresh interval refresh_seconds = refresh_interval.total_seconds() if refresh_seconds < 1: @@ -92,6 +131,10 @@ def __init__( # Set up periodic refresh for latest versions if self._minor_version.version == REVISION_LATEST or self.__prompt_major_version__ == REVISION_UNDEPLOYED: + if is_testing_context(): + log.info("Prompt refreshing is disabled when in a testing context.") + return + log.info(f"Refreshing latest prompt every {refresh_seconds} seconds") if running_loop := get_running_loop(): running_loop.create_task(self._refresh_loop()) @@ -126,7 +169,12 @@ async def _get_prompt( async def _init_async(self) -> None: """Initialize the prompt manager asynchronously.""" - # No longer need to extract from all_minor_versions since we only have a single version + # Check for revision override first + if is_testing_context() and (revision_id := prompt_revisions_map().get(self.__prompt_id__)): + await self._set_prompt_revision(revision_id) + return + + # Normal initialization logic try: prompt = await self._get_prompt(self._minor_version.version, self._init_timeout) except Exception as err: @@ -189,10 +237,83 @@ def exec(self) -> Iterator[ExecutionContextType]: Returns: A context manager that yields an execution context """ - version = self._minor_version.version - prompt = self._minor_version_to_prompt[version] + # Choose prompt with override priority + if is_testing_context() and self._prompt_revision_override: + prompt = self._prompt_revision_override + else: + version = self._minor_version.version + prompt = self._minor_version_to_prompt[version] + + # Register revision usage for tracking + if is_testing_context(): + entity_id = prompt.get("id") + revision_id = prompt.get("revisionId") + if entity_id and revision_id: + register_revision_usage( + entity_id=entity_id, + entity_type=RevisionType.PROMPT, + revision_id=revision_id, + ) + context = self.__execution_context_class__(prompt) try: yield context finally: pass + + def _make_revision_validate_override_request_url(self, revision_id: str) -> str: + """ + Construct the V2 validation endpoint URL. + V2 pattern: /apps/:appId/prompts/:externalId/revisions/:revisionId/validate + """ + app_id = encode_uri_component(self.__app_id__) + prompt_id = encode_uri_component(self.__prompt_id__) + revision_id = encode_uri_component(revision_id) + return f"{API_ENDPOINT_V2}/apps/{app_id}/prompts/{prompt_id}/revisions/{revision_id}/validate" + + async def _set_prompt_revision(self, revision_id: str) -> None: + """ + If this prompt has a revision override set, use the /validate endpoint to + check if the major version this prompt manager is configured to use is compatible + to be overridden with the revision. + """ + # Double check we're in a testing context + if not is_testing_context(): + log.error("Can't set prompt revision unless in a testing context.") + return + + # Double check the given revision_id belongs to this prompt manager + expected_revision_id = prompt_revisions_map()[self.__prompt_id__] + if revision_id != expected_revision_id: + raise RuntimeError( + f"Revision ID '{revision_id}' does not match the revision ID " + f"for this prompt manager '{expected_revision_id}'." + ) + + resp = await global_state.http_client().post( + self._make_revision_validate_override_request_url(revision_id), + timeout=self._init_timeout.total_seconds(), + headers={"Authorization": f"Bearer {self._client._api_key}"}, + json=dict( + majorVersion=int(self.__prompt_major_version__), + ), + ) + + if resp.status_code == HTTPStatus.CONFLICT: + # The /validate endpoint returns this status code when the revision is + # not compatible with the major version this prompt manager + # is configured to use. + raise IncompatiblePromptRevisionError( + f"Can't override prompt '{self._class_name}' with revision '{revision_id}' because it is not " + f"compatible with major version '{self.__prompt_major_version__}'." + ) + + # Raise for any unexpected errors + resp.raise_for_status() + + # Set the prompt revision override + log.warning(f"Overriding prompt '{self._class_name}' with revision '{revision_id}'!") + response_data = resp.json() + # Ensure appId is included in the response data + response_data["appId"] = self.__app_id__ + self._prompt_revision_override = response_data diff --git a/tests/autoblocks/test_prompts_v2_manager.py b/tests/autoblocks/test_prompts_v2_manager.py new file mode 100644 index 00000000..f6b6035d --- /dev/null +++ b/tests/autoblocks/test_prompts_v2_manager.py @@ -0,0 +1,302 @@ +import json +import os +from http import HTTPStatus +from typing import Any +from unittest import mock + +import pytest + +from autoblocks._impl.config.constants import API_ENDPOINT_V2 +from autoblocks._impl.prompts.error import IncompatiblePromptRevisionError +from autoblocks.prompts.v2.context import PromptExecutionContext +from autoblocks.prompts.v2.manager import AutoblocksPromptManager +from autoblocks.prompts.v2.models import FrozenModel +from autoblocks.prompts.v2.renderer import TemplateRenderer +from autoblocks.prompts.v2.renderer import ToolRenderer +from tests.util import make_expected_body + + +class MyV2Params(FrozenModel): + pass + + +class MyV2TemplateRenderer(TemplateRenderer): + __name_mapper__ = { + "name": "name", + "weather": "weather", + } + + def my_template( + self, + *, + name: str, + weather: str, + ) -> str: + return self._render( + "my-template", + name=name, + weather=weather, + ) + + +class MyV2ToolRenderer(ToolRenderer): + __name_mapper__ = { + "description": "description", + } + + def my_tool( + self, + *, + description: str, + ) -> dict[str, Any]: + return self._render( + "my-tool", + description=description, + ) + + +class MyV2ExecutionContext( + PromptExecutionContext[MyV2Params, MyV2TemplateRenderer, MyV2ToolRenderer], +): + __params_class__ = MyV2Params + __template_renderer_class__ = MyV2TemplateRenderer + __tool_renderer_class__ = MyV2ToolRenderer + + +class MyV2PromptManager( + AutoblocksPromptManager[MyV2ExecutionContext], +): + __app_id__ = "test-app-id" + __prompt_id__ = "my-prompt-id" + __prompt_major_version__ = "1" + __execution_context_class__ = MyV2ExecutionContext + + +@mock.patch.dict( + os.environ, + { + "AUTOBLOCKS_V2_API_KEY": "mock-api-key", + "AUTOBLOCKS_V2_CI_TEST_RUN_BUILD_ID": "mock-build-id", + "AUTOBLOCKS_OVERRIDES": json.dumps({"promptRevisions": {MyV2PromptManager.__prompt_id__: "mock-revision-id"}}), + }, +) +def test_uses_prompt_revision(httpx_mock): + httpx_mock.add_response( + url=f"{API_ENDPOINT_V2}/apps/test-app-id/prompts/my-prompt-id/revisions/mock-revision-id/validate", + method="POST", + match_headers={"Authorization": "Bearer mock-api-key"}, + match_content=make_expected_body( + dict( + majorVersion=1, + ), + ), + json=dict( + id="my-prompt-id", + version="revision:mock-revision-id", + revisionId="mock-revision-id", + templates=[ + dict( + id="my-template", + template="Hello, {{ name }}! The weather is {{ weather }} today!!!", + ), + ], + params=dict(params={}), + tools=[dict(name="my-tool", description="{{ description }}")], + toolsParams=[dict(name="my-tool", params=["description"])], + ), + ) + + mgr = MyV2PromptManager(minor_version="0") + with mgr.exec() as p: + rendered = p.render_template.my_template(name="Nicole", weather="sunny") + assert rendered == "Hello, Nicole! The weather is sunny today!!!" + # Test that the override prompt data is being used + assert p.track()["revisionId"] == "mock-revision-id" + assert p.track()["version"] == "revision:mock-revision-id" + + assert len(httpx_mock.get_requests()) == 1 + + +@mock.patch.dict( + os.environ, + { + "AUTOBLOCKS_V2_API_KEY": "mock-api-key", + "AUTOBLOCKS_V2_CI_TEST_RUN_BUILD_ID": "mock-build-id", + "AUTOBLOCKS_OVERRIDES": json.dumps({"promptRevisions": {MyV2PromptManager.__prompt_id__: "mock-revision-id"}}), + }, +) +def test_uses_prompt_revision_when_version_is_latest(httpx_mock): + httpx_mock.add_response( + url=f"{API_ENDPOINT_V2}/apps/test-app-id/prompts/my-prompt-id/revisions/mock-revision-id/validate", + method="POST", + match_headers={"Authorization": "Bearer mock-api-key"}, + match_content=make_expected_body( + dict( + majorVersion=1, + ), + ), + json=dict( + id="my-prompt-id", + version="revision:mock-revision-id", + revisionId="mock-revision-id", + templates=[ + dict( + id="my-template", + template="Hello, {{ name }}! The weather is {{ weather }} today!!!", + ), + ], + params=dict(params={}), + tools=[dict(name="my-tool", description="{{ description }}")], + toolsParams=[dict(name="my-tool", params=["description"])], + ), + ) + + mgr = MyV2PromptManager(minor_version="latest") + with mgr.exec() as p: + rendered = p.render_template.my_template(name="Nicole", weather="sunny") + assert rendered == "Hello, Nicole! The weather is sunny today!!!" + + assert len(httpx_mock.get_requests()) == 1 + + +@mock.patch.dict( + os.environ, + { + "AUTOBLOCKS_V2_API_KEY": "mock-api-key", + "AUTOBLOCKS_CI_TEST_RUN_BUILD_ID": "mock-build-id", + "AUTOBLOCKS_OVERRIDES": json.dumps({"promptRevisions": {"some-other-prompt-id": "mock-revision-id"}}), + }, +) +def test_uses_configured_version_if_revision_is_for_different_prompt(httpx_mock): + httpx_mock.add_response( + url=f"{API_ENDPOINT_V2}/apps/test-app-id/prompts/my-prompt-id/major/1/minor/0", + method="GET", + match_headers={"Authorization": "Bearer mock-api-key"}, + json=dict( + id="my-prompt-id", + version="1.0", + revisionId="mock-revision-id", + appId="test-app-id", + templates=[ + dict( + id="my-template", + template="Hello, {{ name }}! The weather is {{ weather }} today.", + ), + ], + ), + ) + + mgr = MyV2PromptManager(minor_version="0") + with mgr.exec() as p: + rendered = p.render_template.my_template(name="Nicole", weather="sunny") + assert rendered == "Hello, Nicole! The weather is sunny today." + + assert len(httpx_mock.get_requests()) == 1 + + +@mock.patch.dict( + os.environ, + { + "AUTOBLOCKS_V2_API_KEY": "mock-api-key", + "AUTOBLOCKS_V2_CI_TEST_RUN_BUILD_ID": "mock-build-id", + "AUTOBLOCKS_OVERRIDES": json.dumps({"promptRevisions": {MyV2PromptManager.__prompt_id__: "mock-revision-id"}}), + }, +) +def test_raises_if_prompt_is_incompatible(httpx_mock): + httpx_mock.add_response( + url=f"{API_ENDPOINT_V2}/apps/test-app-id/prompts/my-prompt-id/revisions/mock-revision-id/validate", + method="POST", + match_headers={"Authorization": "Bearer mock-api-key"}, + match_content=make_expected_body( + dict( + majorVersion=1, + ), + ), + # If the revision's prompt ID matches the manager's prompt ID but the revision is + # incompatible, the endpoint will return a 409. + status_code=HTTPStatus.CONFLICT, + ) + + with pytest.raises(IncompatiblePromptRevisionError) as exc: + MyV2PromptManager(minor_version="0") + + assert str(exc.value) == ( + "Can't override prompt 'MyV2PromptManager' with revision 'mock-revision-id' " + "because it is not compatible with major version '1'." + ) + + assert len(httpx_mock.get_requests()) == 1 + + +@mock.patch.dict( + os.environ, + { + "AUTOBLOCKS_V2_API_KEY": "mock-api-key", + # Note that AUTOBLOCKS_CLI_SERVER_ADDRESS is not set + "AUTOBLOCKS_OVERRIDES": json.dumps({"promptRevisions": {MyV2PromptManager.__prompt_id__: "mock-revision-id"}}), + }, +) +def test_ignores_revision_id_if_not_in_test_run_context(httpx_mock): + httpx_mock.add_response( + url=f"{API_ENDPOINT_V2}/apps/test-app-id/prompts/my-prompt-id/major/1/minor/0", + method="GET", + match_headers={"Authorization": "Bearer mock-api-key"}, + json=dict( + id="my-prompt-id", + version="1.0", + revisionId="mock-revision-id", + appId="test-app-id", + templates=[ + dict( + id="my-template", + template="Hello, {{ name }}! The weather is {{ weather }} today.", + ), + ], + ), + ) + + mgr = MyV2PromptManager(minor_version="0") + with mgr.exec() as p: + rendered = p.render_template.my_template(name="Nicole", weather="sunny") + assert rendered == "Hello, Nicole! The weather is sunny today." + + assert len(httpx_mock.get_requests()) == 1 + + +@mock.patch.dict( + os.environ, + { + "AUTOBLOCKS_V2_API_KEY": "mock-api-key", + "AUTOBLOCKS_V2_CI_TEST_RUN_BUILD_ID": "mock-build-id", + "AUTOBLOCKS_OVERRIDES_PROMPT_REVISIONS": json.dumps({MyV2PromptManager.__prompt_id__: "legacy-revision-id"}), + }, +) +def test_falls_back_to_legacy_format(httpx_mock): + httpx_mock.add_response( + url=f"{API_ENDPOINT_V2}/apps/test-app-id/prompts/my-prompt-id/revisions/legacy-revision-id/validate", + method="POST", + match_headers={"Authorization": "Bearer mock-api-key"}, + match_content=make_expected_body( + dict( + majorVersion=1, + ), + ), + json=dict( + id="my-prompt-id", + version="revision:legacy-revision-id", + revisionId="legacy-revision-id", + templates=[ + dict( + id="my-template", + template="Hello, {{ name }}! Legacy format works!", + ), + ], + ), + ) + + mgr = MyV2PromptManager(minor_version="0") + with mgr.exec() as p: + rendered = p.render_template.my_template(name="Nicole", weather="sunny") + assert rendered == "Hello, Nicole! Legacy format works!" + + assert len(httpx_mock.get_requests()) == 1 diff --git a/tests/e2e/test_prompts_v2.py b/tests/e2e/test_prompts_v2.py index bd5bd7fb..9de036cc 100644 --- a/tests/e2e/test_prompts_v2.py +++ b/tests/e2e/test_prompts_v2.py @@ -60,3 +60,21 @@ def test_app_sdk_test_prompt_basic_v2(): assert track_info["appId"] == "b5sz3k5d61w9f8325fhxuxkr" assert "templates" in track_info assert len(track_info["templates"]) == 1 + + +# Note: E2E tests for revision overrides would require real revision IDs +# and would be better suited for integration tests with actual API endpoints +def test_prompt_revision_override_structure(): + """Test that V2 prompts have the same structure as V1 for revision overrides.""" + # This test verifies that the manager has the necessary attributes for overrides + mgr = ci_app.prompt_basic_prompt_manager( + major_version="2", + minor_version="0", + ) + + # Verify the manager has the required attributes for revision overrides + assert hasattr(mgr, "_prompt_revision_override") + assert hasattr(mgr, "__prompt_id__") + assert hasattr(mgr, "__app_id__") + assert mgr.__prompt_id__ == "prompt-basic" + assert mgr.__app_id__ == "b5sz3k5d61w9f8325fhxuxkr"