diff --git a/py/packages/genkit/src/genkit/_ai/_aio.py b/py/packages/genkit/src/genkit/_ai/_aio.py index 84ddab6f5b..fc74f8453b 100644 --- a/py/packages/genkit/src/genkit/_ai/_aio.py +++ b/py/packages/genkit/src/genkit/_ai/_aio.py @@ -92,9 +92,9 @@ CancelModelOpFn, CheckModelOpFn, StartModelOpFn, + cancel_operation, check_operation, define_background_model, - lookup_background_action, ) from genkit._core._channel import Channel, run_loop from genkit._core._dap import ( @@ -150,18 +150,10 @@ def _model_supports_long_running(model_action: Action) -> bool: """Check if a model action supports long-running operations.""" model_info = model_action.metadata.get('model') if model_action.metadata else None - if not model_info: + if not isinstance(model_info, dict): return False - # Handle ModelInfo object - if hasattr(model_info, 'supports'): - supports = getattr(model_info, 'supports', None) - return bool(getattr(supports, 'long_running', False)) if supports else False - # Handle dict (cast needed because isinstance narrows too much for type checkers) - if isinstance(model_info, dict): - model_dict = cast(dict[str, Any], model_info) - supports = model_dict.get('supports') - return bool(supports.get('longRunning', False)) if isinstance(supports, dict) else False - return False + supports = cast(dict[str, Any], model_info).get('supports') + return bool(supports.get('longRunning', False)) if isinstance(supports, dict) else False class Genkit: @@ -1633,14 +1625,7 @@ async def check_operation(self, operation: Operation) -> Operation: async def cancel_operation(self, operation: Operation) -> Operation: """Cancel a long-running background operation.""" - if not operation.action: - raise ValueError('Provided operation is missing original request information') - - background_action = await lookup_background_action(self.registry, operation.action) - if background_action is None: - raise ValueError(f'Failed to resolve background action from original request: {operation.action}') - - return await background_action.cancel(operation) + return await cancel_operation(self.registry, operation) @overload async def generate_operation( @@ -1752,8 +1737,7 @@ async def generate_operation( docs=docs, ) - # Extract operation from response - if not hasattr(response, 'operation') or not response.operation: + if not response.operation: raise GenkitError( status='FAILED_PRECONDITION', message=f"Model '{model_action.name}' did not return an operation.", diff --git a/py/packages/genkit/src/genkit/_core/_background.py b/py/packages/genkit/src/genkit/_core/_background.py index e9c23329be..832afc8a8d 100644 --- a/py/packages/genkit/src/genkit/_core/_background.py +++ b/py/packages/genkit/src/genkit/_core/_background.py @@ -19,12 +19,11 @@ from __future__ import annotations import time -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping from typing import Any, Generic, TypeVar -from pydantic import BaseModel - from genkit._core._action import Action, ActionKind, ActionRunContext +from genkit._core._error import GenkitError from genkit._core._model import ModelRequest, ModelResponse from genkit._core._registry import Registry from genkit._core._schema import to_json_schema @@ -119,7 +118,7 @@ async def start( An Operation with an ID to track the job. """ result = await self.start_action.run(input) - return _ensure_operation(result.response) + return result.response async def check(self, operation: Operation) -> Operation: """Check the status of a background operation. @@ -130,52 +129,30 @@ async def check(self, operation: Operation) -> Operation: Returns: Updated Operation with current status. """ + operation = require_operation(value=operation) result = await self.check_action.run(operation) - return _ensure_operation(result.response) + return result.response async def cancel(self, operation: Operation) -> Operation: """Cancel a background operation. - If cancellation is not supported, returns the operation unchanged. - Args: operation: The operation to cancel. Returns: Updated Operation reflecting cancellation attempt. + + Raises: + GenkitError: If this action does not implement cancel. """ + operation = require_operation(value=operation) if self.cancel_action is None: - # Return operation unchanged if cancel not supported - return operation + raise GenkitError( + status='UNIMPLEMENTED', + message=f'Background action {operation.action} does not support cancellation.', + ) result = await self.cancel_action.run(operation) - return _ensure_operation(result.response) - - -def _ensure_operation(response: Any) -> Operation: # noqa: ANN401 - """Convert response to Operation type.""" - if isinstance(response, Operation): - return response - if isinstance(response, dict): - return Operation.model_validate(response) - raise TypeError(f'Expected Operation, got {type(response)}') - - -class DefineBackgroundModelOptions(BaseModel): - """Options for defining a background model. - - Attributes: - name: Unique name for this background model. - label: Human-readable label (defaults to name). - versions: Known version names for this model. - supports: Model capability information. - config_schema: Custom options schema for this model. - """ - - name: str - label: str | None = None - versions: list[str] | None = None - supports: dict[str, Any] | None = None - config_schema: type | dict[str, Any] | None = None + return result.response def define_background_model( @@ -365,29 +342,89 @@ async def lookup_background_action( ) +def require_operation(*, value: object) -> Operation: + """A poll handle is an Operation. A dump or generate() box is not.""" + if isinstance(value, Operation): + return value + if isinstance(value, ModelResponse): + raise GenkitError( + status='INVALID_ARGUMENT', + message='got ModelResponse; pass response.operation', + ) + if isinstance(value, Mapping): + raise GenkitError( + status='INVALID_ARGUMENT', + message='got a dump; pass Operation.model_validate(...)', + ) + raise GenkitError( + status='INVALID_ARGUMENT', + message=f'got {type(value).__name__}, expected Operation', + ) + + +async def resolve_operation_action( + registry: Registry, + operation: Operation, +) -> tuple[Operation, BackgroundAction]: + """Turn a poll handle into the background action that owns it.""" + operation = require_operation(value=operation) + if not operation.action: + raise GenkitError( + status='INVALID_ARGUMENT', + message='Provided operation is missing original request information', + ) + + background_action = await lookup_background_action(registry, operation.action) + if background_action is None: + raise GenkitError( + status='INVALID_ARGUMENT', + message=f'Failed to resolve background action from original request: {operation.action}', + ) + return operation, background_action + + async def check_operation( registry: Registry, operation: Operation, ) -> Operation: """Check the status of a background operation. - Matches JS checkOperation from js/ai/src/check-operation.ts. - Args: registry: The registry to look up actions from. - operation: The operation to check. + operation: The poll handle. Returns: Updated Operation with current status. Raises: - ValueError: If operation is missing action or action not found. + GenkitError: If the handle is missing action, or the action is + not found. """ - if not operation.action: - raise ValueError('Provided operation is missing original request information') + resolved, background_action = await resolve_operation_action(registry, operation) + return await background_action.check(resolved) - background_action = await lookup_background_action(registry, operation.action) - if background_action is None: - raise ValueError(f'Failed to resolve background action from original request: {operation.action}') - return await background_action.check(operation) +async def cancel_operation( + registry: Registry, + operation: Operation, +) -> Operation: + """Cancel a background operation. + + Args: + registry: The registry to look up actions from. + operation: The poll handle. + + Returns: + Updated Operation reflecting the cancel attempt. + + Raises: + GenkitError: If the handle is missing action, the action is not + found, or cancel is not implemented. + """ + resolved, background_action = await resolve_operation_action(registry, operation) + if not background_action.supports_cancel: + raise GenkitError( + status='UNIMPLEMENTED', + message=f'Background action {resolved.action} does not support cancellation.', + ) + return await background_action.cancel(resolved) diff --git a/py/packages/genkit/src/genkit/_core/_typing.py b/py/packages/genkit/src/genkit/_core/_typing.py index b283404ea5..b02d49ebd6 100644 --- a/py/packages/genkit/src/genkit/_core/_typing.py +++ b/py/packages/genkit/src/genkit/_core/_typing.py @@ -567,7 +567,7 @@ class MultipartToolResponse(GenkitModel): class Operation(GenkitModel): """Model for operation data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='ignore', populate_by_name=True) action: str | None = None id: str = Field(...) done: bool | None = None diff --git a/py/packages/genkit/tests/genkit/ai/background_generate_test.py b/py/packages/genkit/tests/genkit/ai/background_generate_test.py index deb12c203b..f522dc7359 100644 --- a/py/packages/genkit/tests/genkit/ai/background_generate_test.py +++ b/py/packages/genkit/tests/genkit/ai/background_generate_test.py @@ -239,3 +239,21 @@ def test_model_response_eq_includes_operation() -> None: b = ModelResponse(operation=Operation(id='unique-b')) assert a != b assert a == ModelResponse(operation=Operation(id='unique-a')) + + +@pytest.mark.asyncio +async def test_check_action_accepts_dumped_operation_with_extra_keys(ai: Genkit) -> None: + """A persisted dump still checks, even with leftover keys like latencyMs.""" + action = await register_bg_model(ai) + dumped = { + 'id': 'bg-op-123', + 'done': False, + 'action': '/background-model/bg-model', + 'latencyMs': 42, + } + + result = await action.check_action.run(dumped) + + assert result.response.id == 'bg-op-123' + assert result.response.action == '/background-model/bg-model' + assert 'latencyMs' not in result.response.model_dump() diff --git a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py index ee3ae2d601..897564171b 100644 --- a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py +++ b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py @@ -11,7 +11,9 @@ import pytest from genkit import Genkit -from genkit._core._action import _action_context +from genkit._core._action import ActionRunContext, _action_context +from genkit._core._error import GenkitError +from genkit._core._model import ModelRequest, ModelResponse from genkit._core._typing import Operation @@ -67,8 +69,9 @@ async def test_genkit_check_operation_no_action() -> None: ai = Genkit() op = Operation(id='123', done=False) # action is None - with pytest.raises(ValueError, match='Provided operation is missing original request information'): + with pytest.raises(GenkitError, match='Provided operation is missing original request information') as exc_info: await ai.check_operation(op) + assert exc_info.value.status == 'INVALID_ARGUMENT' @pytest.mark.asyncio @@ -78,8 +81,189 @@ async def test_genkit_check_operation_not_found() -> None: op = Operation(id='123', done=False, action='missing') ai.registry.resolve_action_by_key = AsyncMock(return_value=None) # type: ignore[assignment] - with pytest.raises(ValueError, match='Failed to resolve background action from original request: missing'): + with pytest.raises( + GenkitError, match='Failed to resolve background action from original request: missing' + ) as exc_info: await ai.check_operation(op) + assert exc_info.value.status == 'INVALID_ARGUMENT' + + +@pytest.mark.asyncio +async def test_check_operation_accepts_dumped_operation() -> None: + """A persisted dump still polls, even with leftover keys like latencyMs.""" + ai = Genkit() + dumped = { + 'id': '123', + 'done': False, + 'action': '/background-model/test_action', + 'latencyMs': 42, + } + mock_background_action = MagicMock() + mock_background_action.check = AsyncMock(return_value=Operation(id='123', done=True)) + + with mock.patch( + 'genkit._core._background.lookup_background_action', + new=AsyncMock(return_value=mock_background_action), + ): + updated = await ai.check_operation(Operation.model_validate(dumped)) + + assert updated.done is True + + +@pytest.mark.asyncio +async def test_check_operation_dump_is_invalid_argument() -> None: + """A saved dict is not an Operation until model_validate.""" + ai = Genkit() + dumped = { + 'id': '123', + 'done': False, + 'action': '/background-model/test_action', + 'latencyMs': 42, + } + + with pytest.raises(GenkitError, match='got a dump; pass Operation.model_validate') as exc_info: + await ai.check_operation(dumped) # type: ignore[arg-type] + assert exc_info.value.status == 'INVALID_ARGUMENT' + + +@pytest.mark.asyncio +async def test_check_operation_boxed_response_is_invalid_argument() -> None: + """generate() returns a ModelResponse; the handle is response.operation.""" + ai = Genkit() + boxed = ModelResponse(operation=Operation(id='123', action='/background-model/test_action')) + + with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as exc_info: + await ai.check_operation(boxed) # type: ignore[arg-type] + assert exc_info.value.status == 'INVALID_ARGUMENT' + + +@pytest.mark.asyncio +async def test_check_operation_str_is_invalid_argument() -> None: + ai = Genkit() + + with pytest.raises(GenkitError, match='got str, expected Operation') as exc_info: + await ai.check_operation('not-an-op') # type: ignore[arg-type] + assert exc_info.value.status == 'INVALID_ARGUMENT' + + +@pytest.mark.asyncio +async def test_cancel_operation_accepts_dumped_operation() -> None: + """Cancel accepts the same persisted Operation dump as check.""" + ai = Genkit() + dumped = { + 'id': '123', + 'done': False, + 'action': '/background-model/test_action', + 'latencyMs': 42, + } + mock_background_action = MagicMock() + mock_background_action.cancel = AsyncMock(return_value=Operation(id='123', done=True)) + + with mock.patch( + 'genkit._core._background.lookup_background_action', + new=AsyncMock(return_value=mock_background_action), + ): + updated = await ai.cancel_operation(Operation.model_validate(dumped)) + + assert updated.done is True + + +@pytest.mark.asyncio +async def test_cancel_operation_dump_is_invalid_argument() -> None: + ai = Genkit() + dumped = { + 'id': '123', + 'done': False, + 'action': '/background-model/test_action', + 'latencyMs': 42, + } + + with pytest.raises(GenkitError, match='got a dump; pass Operation.model_validate') as exc_info: + await ai.cancel_operation(dumped) # type: ignore[arg-type] + assert exc_info.value.status == 'INVALID_ARGUMENT' + + +@pytest.mark.asyncio +async def test_cancel_operation_without_cancel_is_unimplemented() -> None: + ai = Genkit() + op = Operation(id='123', done=False, action='/background-model/test_action') + mock_background_action = MagicMock() + mock_background_action.supports_cancel = False + mock_background_action.cancel = AsyncMock() + + with mock.patch( + 'genkit._core._background.lookup_background_action', + new=AsyncMock(return_value=mock_background_action), + ): + with pytest.raises(GenkitError, match='does not support cancellation') as exc_info: + await ai.cancel_operation(op) + assert exc_info.value.status == 'UNIMPLEMENTED' + mock_background_action.cancel.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_background_action_cancel_without_fn_is_unimplemented() -> None: + """A real no-cancel BackgroundAction raises UNIMPLEMENTED from .cancel.""" + + async def start(_request: ModelRequest, _ctx: ActionRunContext) -> Operation: + return Operation(id='1', done=False) + + async def check(op: Operation) -> Operation: + return op + + ai = Genkit() + action = ai.define_background_model(name='no-cancel', start=start, check=check) + op = Operation(id='1', action='/background-model/no-cancel') + + with pytest.raises(GenkitError, match='does not support cancellation') as exc_info: + await action.cancel(op) + assert exc_info.value.status == 'UNIMPLEMENTED' + + +@pytest.mark.asyncio +async def test_background_action_check_rejects_non_operation() -> None: + """BackgroundAction.check uses the same require_operation gate as the veneer.""" + + async def start(_request: ModelRequest, _ctx: ActionRunContext) -> Operation: + return Operation(id='1', done=False) + + async def check(op: Operation) -> Operation: + return op + + ai = Genkit() + action = ai.define_background_model(name='bg-check', start=start, check=check) + dumped = {'id': '1', 'action': '/background-model/bg-check'} + boxed = ModelResponse(operation=Operation(id='1', action='/background-model/bg-check')) + + with pytest.raises(GenkitError, match='got a dump; pass Operation.model_validate') as dump_exc: + await action.check(dumped) # type: ignore[arg-type] + assert dump_exc.value.status == 'INVALID_ARGUMENT' + + with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as box_exc: + await action.check(boxed) # type: ignore[arg-type] + assert box_exc.value.status == 'INVALID_ARGUMENT' + + with pytest.raises(GenkitError, match='got str, expected Operation') as str_exc: + await action.check('not-an-op') # type: ignore[arg-type] + assert str_exc.value.status == 'INVALID_ARGUMENT' + + +@pytest.mark.asyncio +async def test_background_action_cancel_rejects_non_operation() -> None: + """A dump must not AttributeError on .action before UNIMPLEMENTED.""" + + async def start(_request: ModelRequest, _ctx: ActionRunContext) -> Operation: + return Operation(id='1', done=False) + + async def check(op: Operation) -> Operation: + return op + + ai = Genkit() + action = ai.define_background_model(name='no-cancel', start=start, check=check) + + with pytest.raises(GenkitError, match='got a dump; pass Operation.model_validate') as exc_info: + await action.cancel({'id': '1', 'action': '/background-model/no-cancel'}) # type: ignore[arg-type] + assert exc_info.value.status == 'INVALID_ARGUMENT' @pytest.mark.asyncio diff --git a/py/scripts/schema_to_typing.py b/py/scripts/schema_to_typing.py index 7bffe66dfb..ae2bea59e1 100644 --- a/py/scripts/schema_to_typing.py +++ b/py/scripts/schema_to_typing.py @@ -133,6 +133,16 @@ def _models_allowing_extra(schema: dict) -> set[str]: return result +def _extra_policy(name: str, allow: set[str]) -> str: + if name in allow: + return 'allow' + # A saved handle dump can carry leftover keys that aren't Operation + # fields. Drop them so Operation.model_validate(saved) does not 500. + if name == 'Operation': + return 'ignore' + return 'forbid' + + def _typed_map_aliases(defs: dict) -> dict[str, str]: """Inline object schemas with typed scalar ``additionalProperties`` -> Python dict alias. @@ -274,7 +284,8 @@ def _emit_model( req = req - omit - {_camel_to_snake(k) for k in omit} ext = ', protected_namespaces=()' if any(_camel_to_snake(k) in ('schema', 'schema_') for k in props) else '' frz = ', frozen=True' if name == 'PathMetadata' else '' - cfg = f"ConfigDict(alias_generator=to_camel, extra='{'allow' if name in allow else 'forbid'}', populate_by_name=True{ext}{frz})" + extra = _extra_policy(name, allow) + cfg = f"ConfigDict(alias_generator=to_camel, extra='{extra}', populate_by_name=True{ext}{frz})" lines = [ f'class {name}(GenkitModel):', f' """Model for {name.lower().replace("_", " ")} data."""',