add better input object validation

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,

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]
 

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)
 

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()