diff --git a/narrative_llm_agent/crews/job_crew.py b/narrative_llm_agent/crews/job_crew.py index 269f1f8..e091a84 100644 --- a/narrative_llm_agent/crews/job_crew.py +++ b/narrative_llm_agent/crews/job_crew.py @@ -14,11 +14,11 @@ from narrative_llm_agent.kbase.objects.workspace import ObjectInfo from narrative_llm_agent.tools.app_tools import app_params_pydantic, get_app_params from narrative_llm_agent.tools.job_tools import CompletedJob, CreatedObject -import logging +import logging workflow_logger = logging.getLogger("WorkflowExecution") class JobCrew: - """ + """ƒ Initializes and runs a CrewAI Crew that will run a single KBase job from start to finish, analyze, and interpret the results, saving a summary in a Narrative. @@ -59,7 +59,7 @@ def start_job(self, app_name: str, input_object_upa: str, narrative_id: int, app """ ws = Workspace(token=self._token) object_info = ws.get_object_info(input_object_upa) - self._tasks = self.build_tasks(app_id, narrative_id, object_info) + self._tasks = self.build_tasks(app_id, narrative_id, object_info, ws) crew = Crew( agents=self._agents, tasks=self._tasks, @@ -130,33 +130,40 @@ def start_job_debug_skip(self, app_name: str, input_object_upa: str, narrative_i ) def build_tasks( - self, app_id: str, narrative_id: int, object_info: ObjectInfo + self, app_id: str, narrative_id: int, object_info: ObjectInfo, ws: Workspace ) -> list[Task]: # TODO: make sure that input objects are ALWAYS UPAs param_template = get_app_params(app_id, self._nms) spec = AppSpec(**self._nms.get_app_spec(app_id)) - param_model = app_params_pydantic(spec) + param_model = app_params_pydantic(spec, ws) build_params_task = Task( name=f"1. Build the parameters for {app_id}", description=f""" - From the given KBase app id, {app_id}, fetch the list of parameters needed to run it. Use the App and Job manager agent - for assistance. Using the data object with UPA "{object_info.upa}" and name "{object_info.name}", populate a dictionary - with the parameters where the keys are parameter ids, and values are the proper parameter values, or their - default values if no value can be found or calculated. + From the given KBase app id, {app_id}, fetch the list of parameters needed to run it. + Use the App and Job manager agent for assistance. Using the data object with UPA "{object_info.upa}" + and name "{object_info.name}", populate a dictionary with the parameters where the keys are parameter + ids, and values are the proper parameter values or their default values if no value can be found or + calculated. + Here is the parameter information you must use: {param_template} + Any input object parameter must be the input object UPA. - Be sure to make sure there is a non-null value for any parameter that is not optional. + Ensure there is a non-null value for any parameter that is not optional. Any parameter that has a true value for "is_output_object" must have a valid name for the new object. - The new object name should be based on the input object name, not its UPA. But it must NEVER be identical to the input object name, - always create a new name. - If the input object name is not available, the Workspace Manager can assist. - If the parameter type is 'dropdown', use the allowed 'name' option to determine what should be used, but only set the - 'value' associated with that name, or the default value if any. + The new object name should be based on the input object name, not its UPA. It must NEVER be identical to the input object name; + always create a new name. If the input object name is not available, the Workspace Manager can assist. + + If the parameter type is 'dropdown', use the allowed 'name' option to determine what should be used, + but only set the 'value' associated with that name, or the default value if any. Only alphanumeric characters and underscores are allowed in new object names. - Return the dictionary of inputs, the app id, and the - narrative id {narrative_id} for use in the next task. Do not add comments or other text. If the parameters are rejected, examine the reason why and - reform them. The dictionary of inputs and the app id must not be combined into a single dictionary. + + After populating the dictionary, validate the parameters using the `validate_parameters()` tool. + If the parameters are incorrect, reform them accordingly and ensure that they comply with the specified requirements. + + Return the dictionary of inputs, the app id, and the narrative id {narrative_id} for use in the next task. + Do not add comments or other text. If the parameters are rejected, examine the reason why and reform them. + The dictionary of inputs and the app id must not be combined into a single dictionary. """, expected_output="A dictionary of parameters used to run the app with the given id along with the narrative id.", output_pydantic=param_model, diff --git a/narrative_llm_agent/tools/app_tools.py b/narrative_llm_agent/tools/app_tools.py index 892c17c..d6a503c 100644 --- a/narrative_llm_agent/tools/app_tools.py +++ b/narrative_llm_agent/tools/app_tools.py @@ -1,9 +1,11 @@ -from typing import Annotated, Literal -from pydantic import BaseModel, ConfigDict, Field, create_model +from typing import Annotated, Callable, Literal +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, create_model from narrative_llm_agent.kbase.clients.narrative_method_store import ( NarrativeMethodStore, ) +from narrative_llm_agent.kbase.clients.workspace import Workspace from narrative_llm_agent.kbase.objects.app_spec import AppParameter, AppSpec +from narrative_llm_agent.kbase.service_client import ServerError from narrative_llm_agent.util.app import get_processed_app_spec_params @@ -12,7 +14,20 @@ def get_app_params(app_id: str, nms: NarrativeMethodStore) -> dict: return get_processed_app_spec_params(AppSpec(**spec)) -def app_params_pydantic(app_spec: AppSpec) -> BaseModel: +def make_input_upa_validator(ws_client: Workspace, valid_types: list[str]) -> Callable: + def validator(upa: str) -> str: + print(f"running input upa validator on {upa}") + try: + info = ws_client.get_object_info(upa) + except ServerError as e: + raise ValueError(e) + for obj_type in valid_types: + if info.type.lower().startswith(obj_type.lower()): + return upa + raise ValueError(f"Object {upa} is of type {info.type}, it must be one of types {valid_types}") + return validator + +def app_params_pydantic(app_spec: AppSpec, ws_client: Workspace) -> BaseModel: model_atts = {} proc = get_processed_app_spec_params(app_spec) params_dict = {} @@ -28,7 +43,8 @@ def app_params_pydantic(app_spec: AppSpec) -> BaseModel: for param_id in param_group.parameter_ids: group_model_atts[param_id] = _param_to_model_attribute( params_dict[param_id], - proc[param_group.id]["params"][param_id]["type"] + proc[param_group.id]["params"][param_id]["type"], + ws_client ) pg_model = create_model( f"{param_group.id}Model", @@ -44,7 +60,7 @@ def app_params_pydantic(app_spec: AppSpec) -> BaseModel: for param in app_spec.parameters: if param.id not in param_group_params: - model_atts[param.id] = _param_to_model_attribute(param, proc[param.id].get("type", "string")) + model_atts[param.id] = _param_to_model_attribute(param, proc[param.id].get("type", "string"), ws_client) model_atts = model_atts | param_group_model_atts return create_model( @@ -53,7 +69,7 @@ def app_params_pydantic(app_spec: AppSpec) -> BaseModel: __config__=ConfigDict(regex_engine="python-re") ) -def _param_to_model_attribute(param: AppParameter, param_type: str): +def _param_to_model_attribute(param: AppParameter, param_type: str, ws_client: Workspace): """ `param_type` is figured out from the processed version - turns "text" into "data_object", for example. @@ -110,6 +126,13 @@ def _param_to_model_attribute(param: AppParameter, param_type: str): pattern=r"^(?!\d+$)[A-Za-z0-9|_\.-]+$" ) ] + elif param_type == "data_object" and param.text_options.is_output_name != 1: + # This is a data object input, and the parameter string must be an UPA. + param_attr = Annotated[ + str, + Field(strict=True, pattern=r"^\d+/\d+/\d+$"), + AfterValidator(make_input_upa_validator(ws_client, param.text_options.valid_ws_types)) + ] if param.allow_multiple == 1: param_attr = list[param_attr] diff --git a/tests/crews/test_job_crew.py b/tests/crews/test_job_crew.py index fae9d87..21bd57f 100644 --- a/tests/crews/test_job_crew.py +++ b/tests/crews/test_job_crew.py @@ -1,6 +1,7 @@ from pathlib import Path import pytest from narrative_llm_agent.crews.job_crew import JobCrew +from narrative_llm_agent.kbase.clients.workspace import Workspace from narrative_llm_agent.tools.job_tools import CompletedJob from narrative_llm_agent.kbase.objects.workspace import ObjectInfo from narrative_llm_agent.kbase.clients.narrative_method_store import NarrativeMethodStore @@ -33,8 +34,9 @@ def test_build_tasks_returns_list(job_crew, mocker): ) job_crew._nms = mocker.Mock(spec=NarrativeMethodStore) job_crew._nms.get_app_spec.return_value = load_test_data_json(Path("app_spec_data") / "test_app_spec.json") + ws_client = mocker.Mock(spec=Workspace) - tasks = job_crew.build_tasks("prokka/annotate_contigs", 123, obj_info) + tasks = job_crew.build_tasks("prokka/annotate_contigs", 123, obj_info, ws_client) assert isinstance(tasks, list) assert all(isinstance(t, Task) for t in tasks) diff --git a/tests/tools/test_app_tools.py b/tests/tools/test_app_tools.py index 99ce0b2..adbedae 100644 --- a/tests/tools/test_app_tools.py +++ b/tests/tools/test_app_tools.py @@ -20,8 +20,12 @@ ) from narrative_llm_agent.tools.app_tools import ( get_app_params, - app_params_pydantic + app_params_pydantic, + make_input_upa_validator ) +from narrative_llm_agent.kbase.clients.workspace import Workspace +from narrative_llm_agent.kbase.service_client import ServerError +from narrative_llm_agent.kbase.objects.workspace import ObjectInfo from tests.test_data.test_data import load_test_data_json @@ -59,7 +63,8 @@ def widget_spec(): def behavior(): return AppBehavior() -def test_single_required_string_param(base_info, widget_spec, behavior): +def test_single_required_string_param(base_info, widget_spec, behavior, mocker: MockerFixture): + mock_ws = mocker.Mock(spec=Workspace) param = AppParameter( id="param1", ui_name="Param 1", @@ -80,11 +85,12 @@ def test_single_required_string_param(base_info, widget_spec, behavior): behavior=behavior, parameters=[param] ) - model_cls = app_params_pydantic(app_spec) + model_cls = app_params_pydantic(app_spec, mock_ws) model = model_cls(param1="value") assert model.param1 == "value" -def test_optional_int_with_default(base_info, widget_spec, behavior): +def test_optional_int_with_default(base_info, widget_spec, behavior, mocker: MockerFixture): + mock_ws = mocker.Mock(spec=Workspace) param = AppParameter( id="int_param", ui_name="Int Param", @@ -105,11 +111,12 @@ def test_optional_int_with_default(base_info, widget_spec, behavior): behavior=behavior, parameters=[param] ) - model_cls = app_params_pydantic(app_spec) + model_cls = app_params_pydantic(app_spec, mock_ws) model = model_cls() assert model.int_param == 5 -def test_checkbox_param(base_info, widget_spec, behavior): +def test_checkbox_param(base_info, widget_spec, behavior, mocker: MockerFixture): + mock_ws = mocker.Mock(spec=Workspace) param = AppParameter( id="check", ui_name="Checkbox Param", @@ -130,11 +137,12 @@ def test_checkbox_param(base_info, widget_spec, behavior): behavior=behavior, parameters=[param] ) - model_cls = app_params_pydantic(app_spec) + model_cls = app_params_pydantic(app_spec, mock_ws) model = model_cls(check=1) assert model.check == 1 -def test_dropdown_param(base_info, widget_spec, behavior): +def test_dropdown_param(base_info, widget_spec, behavior, mocker: MockerFixture): + mock_ws = mocker.Mock(spec=Workspace) param = AppParameter( id="select", ui_name="Dropdown Param", @@ -157,11 +165,12 @@ def test_dropdown_param(base_info, widget_spec, behavior): behavior=behavior, parameters=[param] ) - model_cls = app_params_pydantic(app_spec) + model_cls = app_params_pydantic(app_spec, mock_ws) model = model_cls(select="opt1") assert model.select == "opt1" -def test_parameter_group_optional_multiple(base_info, widget_spec, behavior): +def test_parameter_group_optional_multiple(base_info, widget_spec, behavior, mocker: MockerFixture): + mock_ws = mocker.Mock(spec=Workspace) param1 = AppParameter( id="g1p1", ui_name="Group Param 1", @@ -208,8 +217,124 @@ def test_parameter_group_optional_multiple(base_info, widget_spec, behavior): parameters=[param1, param2], parameter_groups=[group] ) - model_cls = app_params_pydantic(app_spec) + model_cls = app_params_pydantic(app_spec, mock_ws) model = model_cls(group1=[{"g1p1": "A", "g1p2": "B"}]) assert isinstance(model.group1, list) assert model.group1[0].g1p1 == "A" assert model.group1[0].g1p2 == "B" + + +def test_input_upa_validator_successful(mocker: MockerFixture): + """Test successful UPA validation when object type matches""" + mock_ws = mocker.Mock(spec=Workspace) + mock_info = mocker.Mock(spec=ObjectInfo) + mock_info.type = "KBaseFetch.SingleEndLibrary" + mock_ws.get_object_info.return_value = mock_info + + validator = make_input_upa_validator(mock_ws, ["KBaseFetch.SingleEndLibrary"]) + result = validator("124/5/4") + + assert result == "124/5/4" + mock_ws.get_object_info.assert_called_once_with("124/5/4") + + +def test_input_upa_validator_bad_auth_token(mocker: MockerFixture): + """Test UPA validation fails with authentication error""" + mock_ws = mocker.Mock(spec=Workspace) + mock_ws.get_object_info.side_effect = ServerError("Unauthorized", 401, "Invalid authentication token") + + validator = make_input_upa_validator(mock_ws, ["KBaseFetch.SingleEndLibrary"]) + + with pytest.raises(ValueError): + validator("124/5/4") + + +def test_input_upa_validator_invalid_access(mocker: MockerFixture): + """Test UPA validation fails with access denied error""" + mock_ws = mocker.Mock(spec=Workspace) + mock_ws.get_object_info.side_effect = ServerError("AccessDenied", 403, "User does not have read access") + + validator = make_input_upa_validator(mock_ws, ["KBaseFetch.SingleEndLibrary"]) + + with pytest.raises(ValueError): + validator("124/5/4") + + +def test_input_upa_validator_wrong_type(mocker: MockerFixture): + """Test UPA validation fails when object type isn't in valid_types""" + mock_ws = mocker.Mock(spec=Workspace) + mock_info = mocker.Mock(spec=ObjectInfo) + mock_info.type = "KBaseGenomes.Genome" + mock_ws.get_object_info.return_value = mock_info + + validator = make_input_upa_validator(mock_ws, ["KBaseFetch.SingleEndLibrary", "KBaseFetch.PairedEndLibrary"]) + + with pytest.raises(ValueError, match="Object .* is of type KBaseGenomes.Genome"): + validator("124/5/4") + + +def test_input_upa_in_pydantic_model(mocker: MockerFixture, base_info, widget_spec, behavior): + """Test UPA validation integrated into Pydantic model""" + mock_ws = mocker.Mock(spec=Workspace) + mock_info = mocker.Mock(spec=ObjectInfo) + mock_info.type = "KBaseFetch.SingleEndLibrary" + mock_ws.get_object_info.return_value = mock_info + + param = AppParameter( + id="input_reads", + ui_name="Input Reads", + short_hint="hint", + description="desc", + field_type="text", + allow_multiple=0, + optional=0, + advanced=0, + disabled=0, + default_values=[], + ui_class="input", + text_options=TextOptions(is_output_name=0, valid_ws_types=["KBaseFetch.SingleEndLibrary"]) + ) + app_spec = AppSpec( + info=base_info, + widgets=widget_spec, + behavior=behavior, + parameters=[param] + ) + model_cls = app_params_pydantic(app_spec, mock_ws) + model = model_cls(input_reads="124/5/4") + + assert model.input_reads == "124/5/4" + mock_ws.get_object_info.assert_called_once_with("124/5/4") + + +def test_input_upa_invalid_format_rejected_before_validator(mocker: MockerFixture, base_info, widget_spec, behavior): + """Test that invalid UPA format is rejected by regex before hitting validator""" + mock_ws = mocker.Mock(spec=Workspace) + + param = AppParameter( + id="input_reads", + ui_name="Input Reads", + short_hint="hint", + description="desc", + field_type="text", + allow_multiple=0, + optional=0, + advanced=0, + disabled=0, + default_values=[], + ui_class="input", + text_options=TextOptions(is_output_name=0, valid_ws_types=["KBaseFetch.SingleEndLibrary"]) + ) + app_spec = AppSpec( + info=base_info, + widgets=widget_spec, + behavior=behavior, + parameters=[param] + ) + model_cls = app_params_pydantic(app_spec, mock_ws) + + with pytest.raises(ValueError): + model_cls(input_reads="invalid-upa-format") + + # validator should not be called since regex validation fails first + mock_ws.get_object_info.assert_not_called()