diff --git a/aixplain/enums/__init__.py b/aixplain/enums/__init__.py index 725fdb90..6162ed86 100644 --- a/aixplain/enums/__init__.py +++ b/aixplain/enums/__init__.py @@ -20,3 +20,4 @@ from .asset_status import AssetStatus from .index_stores import IndexStores from .function_type import FunctionType +from .code_interpeter import CodeInterpreterModel diff --git a/aixplain/enums/code_interpeter.py b/aixplain/enums/code_interpeter.py new file mode 100644 index 00000000..9f0a14c6 --- /dev/null +++ b/aixplain/enums/code_interpeter.py @@ -0,0 +1,10 @@ +from enum import Enum + + +class CodeInterpreterModel(str, Enum): + """Code Interpreter Model IDs""" + + PYTHON_AZURE = "67476fa16eb563d00060ad62" + + def __str__(self): + return self._value_ diff --git a/aixplain/enums/function_type.py b/aixplain/enums/function_type.py index 514ff992..ae6f8e79 100644 --- a/aixplain/enums/function_type.py +++ b/aixplain/enums/function_type.py @@ -33,3 +33,4 @@ class FunctionType(Enum): SEARCH = "search" INTEGRATION = "connector" CONNECTION = "connection" + MCPSERVER = 'mcpserver' diff --git a/aixplain/factories/index_factory/__init__.py b/aixplain/factories/index_factory/__init__.py index 935907ab..2eacca55 100644 --- a/aixplain/factories/index_factory/__init__.py +++ b/aixplain/factories/index_factory/__init__.py @@ -43,6 +43,11 @@ def validate_embedding_model(model_id) -> bool: return model.function == Function.TEXT_EMBEDDING +def validate_embedding_model(model_id) -> bool: + model = ModelFactory.get(model_id) + return model.function == Function.TEXT_EMBEDDING + + class IndexFactory(ModelFactory, Generic[T]): @classmethod def create( diff --git a/aixplain/factories/model_factory/__init__.py b/aixplain/factories/model_factory/__init__.py index c7a2e164..2e1dc1da 100644 --- a/aixplain/factories/model_factory/__init__.py +++ b/aixplain/factories/model_factory/__init__.py @@ -30,8 +30,6 @@ from aixplain.factories.model_factory.mixins import ModelGetterMixin, ModelListMixin from typing import Callable, Dict, List, Optional, Text, Union - - class ModelFactory(ModelGetterMixin, ModelListMixin): """A static class for creating and exploring Model Objects. diff --git a/aixplain/factories/model_factory/utils.py b/aixplain/factories/model_factory/utils.py index dd71a0cf..ca3e1eec 100644 --- a/aixplain/factories/model_factory/utils.py +++ b/aixplain/factories/model_factory/utils.py @@ -118,7 +118,6 @@ def create_model_from_response(response: Dict) -> Model: supports_streaming=response.get("supportsStreaming", False), status=status, function_type=function_type, - **additional_kwargs, ) diff --git a/aixplain/modules/agent/tool/custom_python_code_tool.py b/aixplain/modules/agent/tool/custom_python_code_tool.py index 6433a408..05715b2a 100644 --- a/aixplain/modules/agent/tool/custom_python_code_tool.py +++ b/aixplain/modules/agent/tool/custom_python_code_tool.py @@ -25,6 +25,7 @@ from aixplain.modules.agent.tool import Tool import logging from aixplain.enums import AssetStatus +from aixplain.enums.code_interpeter import CodeInterpreterModel class CustomPythonCodeTool(Tool): @@ -37,11 +38,13 @@ def __init__( super().__init__(name=name or "", description=description, **additional_info) self.code = code self.status = AssetStatus.ONBOARDED # TODO: change to DRAFT when we have a way to onboard the tool + self.id = CodeInterpreterModel.PYTHON_AZURE self.validate() def to_dict(self): return { + "id": self.id, "name": self.name, "description": self.description, "type": "utility", diff --git a/aixplain/modules/model/index_model.py b/aixplain/modules/model/index_model.py index 72055d69..a240f3b2 100644 --- a/aixplain/modules/model/index_model.py +++ b/aixplain/modules/model/index_model.py @@ -7,7 +7,6 @@ from enum import Enum from typing import List from aixplain.enums.splitting_options import SplittingOptions - import os from urllib.parse import urljoin @@ -56,8 +55,6 @@ def __init__( self.split_length = split_length self.split_overlap = split_overlap - - class IndexModel(Model): def __init__( self, @@ -125,7 +122,6 @@ def to_dict(self) -> Dict: data["collection_type"] = self.version.split("-", 1)[0] return data - def search(self, query: str, top_k: int = 10, filters: List[IndexFilter] = []) -> ModelResponse: """Search for documents in the index diff --git a/aixplain/modules/pipeline/default.py b/aixplain/modules/pipeline/default.py index 2fa4e859..60d879f4 100644 --- a/aixplain/modules/pipeline/default.py +++ b/aixplain/modules/pipeline/default.py @@ -15,5 +15,4 @@ def save(self, *args, **kwargs): def to_dict(self) -> dict: return self.serialize() - - + \ No newline at end of file diff --git a/aixplain/modules/pipeline/designer/enums.py b/aixplain/modules/pipeline/designer/enums.py index fe4cbfed..b733265d 100644 --- a/aixplain/modules/pipeline/designer/enums.py +++ b/aixplain/modules/pipeline/designer/enums.py @@ -1,5 +1,5 @@ from enum import Enum - +from aixplain.enums import FunctionType class RouteType(str, Enum): CHECK_TYPE = "checkType" @@ -29,15 +29,6 @@ class NodeType(str, Enum): class AssetType(str, Enum): MODEL = "MODEL" - -class FunctionType(str, Enum): - AI = "ai" - SEGMENTOR = "segmentor" - RECONSTRUCTOR = "reconstructor" - UTILITY = "utility" - METRIC = "metric" - - class ParamType: INPUT = "INPUT" OUTPUT = "OUTPUT" diff --git a/aixplain/modules/pipeline/designer/nodes.py b/aixplain/modules/pipeline/designer/nodes.py index e81be81a..e92d5fe1 100644 --- a/aixplain/modules/pipeline/designer/nodes.py +++ b/aixplain/modules/pipeline/designer/nodes.py @@ -1,4 +1,5 @@ from typing import List, Union, Type, TYPE_CHECKING, Optional +from enum import Enum from aixplain.modules import Model from aixplain.enums import DataType, Function @@ -142,7 +143,11 @@ def serialize(self) -> dict: obj["supplier"] = self.supplier obj["version"] = self.version obj["assetType"] = self.assetType - obj["functionType"] = self.functionType + # Handle functionType as enum or string + if isinstance(self.functionType, Enum): + obj["functionType"] = self.functionType.value + else: + obj["functionType"] = self.functionType obj["type"] = self.type return obj diff --git a/aixplain/utils/asset_cache.py b/aixplain/utils/asset_cache.py index 8a693c26..357b70ef 100644 --- a/aixplain/utils/asset_cache.py +++ b/aixplain/utils/asset_cache.py @@ -20,7 +20,6 @@ CACHE_DURATION = 86400 - @dataclass class Store(Generic[T]): data: Dict[str, T] @@ -62,7 +61,6 @@ def compute_expiry(self): del os.environ["CACHE_EXPIRY_TIME"] expiry = CACHE_DURATION - return time.time() + int(expiry) def invalidate(self): diff --git a/aixplain/utils/cache_utils.py b/aixplain/utils/cache_utils.py index 464e916a..fcfe1cb6 100644 --- a/aixplain/utils/cache_utils.py +++ b/aixplain/utils/cache_utils.py @@ -14,7 +14,6 @@ def get_cache_expiry(): return int(os.getenv("CACHE_EXPIRY_TIME", CACHE_DURATION)) - def save_to_cache(cache_file, data, lock_file): try: os.makedirs(os.path.dirname(cache_file), exist_ok=True) diff --git a/tests/functional/agent/agent_functional_test.py b/tests/functional/agent/agent_functional_test.py index 08e6e095..2783ffc3 100644 --- a/tests/functional/agent/agent_functional_test.py +++ b/tests/functional/agent/agent_functional_test.py @@ -18,6 +18,7 @@ import copy import json import os +import re from dotenv import load_dotenv load_dotenv() @@ -234,7 +235,10 @@ def test_delete_agent_in_use(delete_agents_and_team_agents, AgentFactory): with pytest.raises(Exception) as exc_info: agent.delete() - assert str(exc_info.value) == "Agent Deletion Error (HTTP 403): err.agent_is_in_use." + assert re.match( + r"Error: Agent cannot be deleted\.\nReason: This agent is currently used by one or more team agents\.\n\nteam_agent_id: [a-f0-9]{24}\. To proceed, remove the agent from all team agents before deletion\.", + str(exc_info.value), + ) @pytest.mark.parametrize("AgentFactory", [AgentFactory, v2.Agent]) @@ -302,7 +306,7 @@ def test_update_tools_of_agent(run_input_map, delete_agents_and_team_agents, Age "type": "translation", "supplier": "Microsoft", "function": "translation", - "query": "Translate: Olá, como vai você?", + "query": "Translate: 'Olá, como vai você?'", "description": "Translation tool with target language", "expected_tool_input": "targetlanguage", }, @@ -334,7 +338,7 @@ def test_specific_model_parameters_e2e(tool_config, delete_agents_and_team_agent # Create and run agent agent = AgentFactory.create( name="Test Parameter Agent", - description="Test agent with parameterized tools. You MUST use a tool for the tasks.", + description="Test agent with parameterized tools. You MUST use a tool for the tasks. Do not directly answer the question.", tools=[tool], llm_id="6646261c6eb563165658bbb1", # Using LLM ID from test data ) @@ -351,8 +355,9 @@ def test_specific_model_parameters_e2e(tool_config, delete_agents_and_team_agent # Verify tool was used in execution assert len(response["data"]["intermediate_steps"]) > 0 tool_used = False + for step in response["data"]["intermediate_steps"]: - if tool_config["expected_tool_input"] in step["tool_steps"][0]["input"]: + if len(step["tool_steps"]) > 0 and tool_config["expected_tool_input"] in step["tool_steps"][0]["input"]: tool_used = True break assert tool_used, "Tool was not used in execution" @@ -643,6 +648,7 @@ def test_agent_llm_parameter_preservation(delete_agents_and_team_agents, AgentFa # Reset the LLM temperature to its original value llm.temperature = original_temperature + def test_run_agent_with_expected_output(): from pydantic import BaseModel from typing import Optional, List @@ -753,4 +759,3 @@ def test_agent_with_action_tool(): assert "helsinki" in response.data.output.lower() assert "SLACK_CHAT_POST_MESSAGE" in [step["tool"] for step in response.data.intermediate_steps[0]["tool_steps"]] connection.delete() - diff --git a/tests/functional/pipelines/create_test.py b/tests/functional/pipelines/create_test.py index f1dac2c4..076a637f 100644 --- a/tests/functional/pipelines/create_test.py +++ b/tests/functional/pipelines/create_test.py @@ -16,7 +16,6 @@ limitations under the License. """ import os -from aixplain.utils.cache_utils import CACHE_FOLDER from aixplain.modules.pipeline import Pipeline import json import pytest @@ -79,25 +78,3 @@ def test_create_pipeline_wrong_path(PipelineFactory): with pytest.raises(Exception): PipelineFactory.create(name=pipeline_name, pipeline="/") - - -@pytest.mark.parametrize("PipelineFactory", [PipelineFactory]) -def test_pipeline_cache_creation(PipelineFactory): - cache_file = os.path.join(CACHE_FOLDER, "pipelines.json") - if os.path.exists(cache_file): - os.remove(cache_file) - - pipeline_json = "tests/functional/pipelines/data/pipeline.json" - pipeline_name = str(uuid4()) - pipeline = PipelineFactory.create(name=pipeline_name, pipeline=pipeline_json) - - assert os.path.exists(cache_file), "Pipeline cache file was not created!" - - with open(cache_file, "r") as f: - cache_data = json.load(f) - - assert "data" in cache_data, "Cache format invalid, missing 'data'." - - pipeline.delete() - if os.path.exists(cache_file): - os.remove(cache_file) \ No newline at end of file