ENG-1245: transition from update to save

.gitignore

@@ -2,9 +2,12 @@
 __pycache__/
 *.py[cod]
 *$py.class
+.aixplain_cache/
 
+setup_env_ahmet.sh
 # C extensions
 *.so
+*.ipynb
 
 # Distribution / packaging
 .Python

aixplain/modules/agent/__init__.py

@@ -305,9 +305,20 @@ def delete(self) -> None:
                 message = f"Agent Deletion Error (HTTP {r.status_code}): There was an error in deleting the agent."
             logging.error(message)
             raise Exception(f"{message}")
-
+        
     def update(self) -> None:
         """Update agent."""
+        import warnings
+        import inspect
+        # Get the current call stack
+        stack = inspect.stack()
+        if len(stack) > 2 and stack[1].function != 'save':
+            warnings.warn(
+                "update() is deprecated and will be removed in a future version. "
+                "Please use save() instead.",
+                DeprecationWarning,
+                stacklevel=2
+            )
         from aixplain.factories.agent_factory.utils import build_agent
 
         self.validate()
@@ -330,6 +341,11 @@ def update(self) -> None:
             error_msg = f"Agent Update Error (HTTP {r.status_code}): {resp}"
             raise Exception(error_msg)
 
+
+    def save(self) -> None:
+        """Save the Agent."""
+        self.update() 
+
     def deploy(self) -> None:
         assert self.status == AssetStatus.DRAFT, "Agent must be in draft status to be deployed."
         assert self.status != AssetStatus.ONBOARDED, "Agent is already deployed."

aixplain/modules/model/utility_model.py

@@ -138,6 +138,18 @@ def to_dict(self):
         }
 
     def update(self):
+        """Update the Utility Model."""
+        import warnings
+        import inspect
+        # Get the current call stack
+        stack = inspect.stack()
+        if len(stack) > 2 and stack[1].function != 'save':
+            warnings.warn(
+                "update() is deprecated and will be removed in a future version. "
+                "Please use save() instead.",
+                DeprecationWarning,
+                stacklevel=2
+            )
         self.validate()
         url = urljoin(self.backend_url, f"sdk/utilities/{self.id}")
         headers = {"x-api-key": f"{self.api_key}", "Content-Type": "application/json"}
@@ -156,7 +168,12 @@ def update(self):
             logging.error(message)
             raise Exception(f"{message}")
 
+    def save(self):
+        """Save the Utility Model."""
+        self.update()
+
     def delete(self):
+        """Delete the Utility Model."""
         url = urljoin(self.backend_url, f"sdk/utilities/{self.id}")
         headers = {"x-api-key": f"{self.api_key}", "Content-Type": "application/json"}
         try:

aixplain/modules/pipeline/asset.py

@@ -384,6 +384,17 @@ def update(
         Raises:
             Exception: Make sure the pipeline to be save is in a JSON file.
         """
+        import warnings
+        import inspect
+        # Get the current call stack
+        stack = inspect.stack()
+        if len(stack) > 2 and stack[1].function != 'save':
+            warnings.warn(
+                "update() is deprecated and will be removed in a future version. "
+                "Please use save() instead.",
+                DeprecationWarning,
+                stacklevel=2
+            )
         try:
             if isinstance(pipeline, str) is True:
                 _, ext = os.path.splitext(pipeline)
@@ -437,18 +448,29 @@ def delete(self) -> None:
             logging.error(message)
             raise Exception(f"{message}")
 
-    def save(self, save_as_asset: bool = False, api_key: Optional[Text] = None):
-        """Save Pipeline
+    def save(self, pipeline: Optional[Union[Text, Dict]] = None, save_as_asset: bool = False, api_key: Optional[Text] = None):
+        """Update and Save Pipeline
 
         Args:
+            pipeline (Optional[Union[Text, Dict]]): Pipeline as a Python dictionary or in a JSON file
             save_as_asset (bool, optional): Save as asset (True) or draft (False). Defaults to False.
             api_key (Optional[Text], optional): Team API Key to create the Pipeline. Defaults to None.
 
         Raises:
             Exception: Make sure the pipeline to be save is in a JSON file.
         """
         try:
-            pipeline = self.to_dict()
+            if pipeline is None:
+                pipeline = self.to_dict()
+            else:
+                if isinstance(pipeline, str) is True:
+                    _, ext = os.path.splitext(pipeline)
+                    assert (
+                        os.path.exists(pipeline) and ext == ".json"
+                    ), "Pipeline Update Error: Make sure the pipeline to be saved is in a JSON file."
+                    with open(pipeline) as f:
+                        pipeline = json.load(f)
+                self.update(pipeline=pipeline, save_as_asset=save_as_asset, api_key=api_key)
 
             for i, node in enumerate(pipeline["nodes"]):
                 if "functionType" in node:
@@ -463,19 +485,14 @@ def save(self, save_as_asset: bool = False, api_key: Optional[Text] = None):
                 "architecture": pipeline,
             }
 
-            if self.id != "":
-                method = "put"
-                url = urljoin(config.BACKEND_URL, f"sdk/pipelines/{self.id}")
-            else:
-                method = "post"
-                url = urljoin(config.BACKEND_URL, "sdk/pipelines")
+            url = urljoin(config.BACKEND_URL, "sdk/pipelines")
             api_key = api_key if api_key is not None else config.TEAM_API_KEY
             headers = {
                 "Authorization": f"Token {api_key}",
                 "Content-Type": "application/json",
             }
             logging.info(f"Start service for Save Pipeline - {url} - {headers} - {json.dumps(payload)}")
-            r = _request_with_retry(method, url, headers=headers, json=payload)
+            r = _request_with_retry("post", url, headers=headers, json=payload)
             response = r.json()
             self.id = response["id"]
             logging.info(f"Pipeline {response['id']} Saved.")

aixplain/modules/team_agent/__init__.py

@@ -310,6 +310,17 @@ def validate(self) -> None:
 
     def update(self) -> None:
         """Update the Team Agent."""
+        import warnings
+        import inspect
+        # Get the current call stack
+        stack = inspect.stack()
+        if len(stack) > 2 and stack[1].function != 'save':
+            warnings.warn(
+                "update() is deprecated and will be removed in a future version. "
+                "Please use save() instead.",
+                DeprecationWarning,
+                stacklevel=2
+            )
         from aixplain.factories.team_agent_factory.utils import build_team_agent
 
         self.validate()
@@ -332,6 +343,10 @@ def update(self) -> None:
             error_msg = f"Team Agent Update Error (HTTP {r.status_code}): {resp}"
             raise Exception(error_msg)
 
+    def save(self) -> None:
+        """Save the Team Agent."""
+        self.update()
+
     def deploy(self) -> None:
         """Deploy the Team Agent."""
         assert self.status == AssetStatus.DRAFT, "Team Agent Deployment Error: Team Agent must be in draft status."

tests/unit/agent_test.py

@@ -8,9 +8,11 @@
 from aixplain.modules.agent import PipelineTool, ModelTool
 from aixplain.modules.agent.utils import process_variables
 from urllib.parse import urljoin
+import warnings
 from aixplain.enums.function import Function
 
 
+
 def test_fail_no_data_query():
     agent = Agent("123", "Test Agent", "Sample Description")
     with pytest.raises(Exception) as exc_info:
@@ -222,14 +224,78 @@ def test_update_success():
         }
         mock.get(url, headers=headers, json=model_ref_response)
 
-        agent.update()
+        # Capture warnings
+        with pytest.warns(DeprecationWarning, match="update\(\) is deprecated and will be removed in a future version. Please use save\(\) instead."):
+            agent.update()
 
     assert agent.id == ref_response["id"]
     assert agent.name == ref_response["name"]
     assert agent.description == ref_response["description"]
     assert agent.llm_id == ref_response["llmId"]
     assert agent.tools[0].function.value == ref_response["assets"][0]["function"]
 
+def test_save_success():
+    agent = Agent(
+        id="123",
+        name="Test Agent",
+        description="Test Agent Description",
+        llm_id="6646261c6eb563165658bbb1",
+        tools=[AgentFactory.create_model_tool(function="text-generation")],
+    )
+
+    with requests_mock.Mocker() as mock:
+        url = urljoin(config.BACKEND_URL, f"sdk/agents/{agent.id}")
+        headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"}
+        ref_response = {
+            "id": "123",
+            "name": "Test Agent",
+            "description": "Test Agent Description",
+            "teamId": "123",
+            "version": "1.0",
+            "status": "onboarded",
+            "llmId": "6646261c6eb563165658bbb1",
+            "#": {"currency": "USD", "value": 0.0},
+            "assets": [
+                {
+                    "type": "model",
+                    "supplier": "openai",
+                    "version": "1.0",
+                    "assetId": "6646261c6eb563165658bbb1",
+                    "function": "text-generation",
+                }
+            ],
+        }
+        mock.put(url, headers=headers, json=ref_response)
+
+        url = urljoin(config.BACKEND_URL, "sdk/models/6646261c6eb563165658bbb1")
+        model_ref_response = {
+            "id": "6646261c6eb563165658bbb1",
+            "name": "Test LLM",
+            "description": "Test LLM Description",
+            "function": {"id": "text-generation"},
+            "supplier": "openai",
+            "version": {"id": "1.0"},
+            "status": "onboarded",
+            "#": {"currency": "USD", "value": 0.0},
+        }
+        mock.get(url, headers=headers, json=model_ref_response)
+
+        import warnings
+        # Capture warnings
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")  # Trigger all warnings
+
+            # Call the save method
+            agent.save()
+
+            # Assert no warnings were triggered
+            assert len(w) == 0, f"Warnings were raised: {[str(warning.message) for warning in w]}"
+
+    assert agent.id == ref_response["id"]
+    assert agent.name == ref_response["name"]
+    assert agent.description == ref_response["description"]
+    assert agent.llm_id == ref_response["llmId"]
+    assert agent.tools[0].function.value == ref_response["assets"][0]["function"]
 
 def test_run_success():
     agent = Agent("123", "Test Agent", "Sample Description")

tests/unit/utility_test.py

@@ -106,12 +106,49 @@ def test_update_utility_model():
                         function=Function.UTILITIES,
                         api_key=config.TEAM_API_KEY,
                     )
-                    utility_model.description = "updated_description"
-                    utility_model.update()
+
+                    with pytest.warns(DeprecationWarning, match="update\(\) is deprecated and will be removed in a future version. Please use save\(\) instead."):
+                        utility_model.description = "updated_description"
+                        utility_model.update()
 
                     assert utility_model.id == "123"
                     assert utility_model.description == "updated_description"
 
+def test_save_utility_model():
+    with requests_mock.Mocker() as mock:
+        with patch("aixplain.factories.file_factory.FileFactory.to_link", return_value="def main(originCode: str)"):
+            with patch("aixplain.factories.file_factory.FileFactory.upload", return_value="def main(originCode: str)"):
+                with patch(
+                    "aixplain.modules.model.utils.parse_code",
+                    return_value=(
+                        "def main(originCode: str)",
+                        [UtilityModelInput(name="originCode", description="originCode", type=DataType.TEXT)],
+                        "utility_model_test",
+                    ),
+                ):
+                    mock.put(urljoin(config.BACKEND_URL, "sdk/utilities/123"), json={"id": "123"})
+                    utility_model = UtilityModel(
+                        id="123",
+                        name="utility_model_test",
+                        description="utility_model_test",
+                        code="def main(originCode: str)",
+                        output_examples="output_description",
+                        inputs=[UtilityModelInput(name="originCode", description="originCode", type=DataType.TEXT)],
+                        function=Function.UTILITIES,
+                        api_key=config.TEAM_API_KEY,
+                    )
+                    import warnings
+                    # it should not trigger any warning
+                    with warnings.catch_warnings(record=True) as w:
+                        warnings.simplefilter("always")  # Trigger all warnings
+                        utility_model.description = "updated_description"
+                        utility_model.save()
+
+                        assert len(w) == 0
+
+                    assert utility_model.id == "123"
+                    assert utility_model.description == "updated_description"
+
 
 def test_delete_utility_model():
     with requests_mock.Mocker() as mock: