wip

Author: harel-tussi

ai21/clients/common/agents/agents.py

@@ -30,11 +30,10 @@ def _create_body(
         tool_resources: Dict[str, Any] | NotGiven,
         requirements: List[Requirement] | NotGiven,
         budget: BudgetLevel | NotGiven,
-        agent_type: AgentType | NotGiven,
         response_language: ResponseLanguage | NotGiven,
         **kwargs,
     ) -> dict:
-        body = remove_not_given(
+        return remove_not_given(
             {
                 "name": name,
                 "description": description,
@@ -43,15 +42,10 @@ def _create_body(
                 "tool_resources": tool_resources,
                 "requirements": requirements,
                 "budget": budget,
-                "agent_type": agent_type,
                 "response_language": response_language,
                 **kwargs,
             }
         )
-        # Map agent_type to assistant_type for API compatibility
-        if "agent_type" in body:
-            body["assistant_type"] = body.pop("agent_type")
-        return body
 
     def _modify_body(
         self,
@@ -115,9 +109,6 @@ def modify(
         *,
         name: str | NotGiven = NOT_GIVEN,
         description: str | NotGiven = NOT_GIVEN,
-        optimization: str | NotGiven = NOT_GIVEN,
-        avatar: str | NotGiven = NOT_GIVEN,
-        is_archived: bool | NotGiven = NOT_GIVEN,
         models: List[str] | NotGiven = NOT_GIVEN,
         tools: List[Dict[str, Any]] | NotGiven = NOT_GIVEN,
         tool_resources: Dict[str, Any] | NotGiven = NOT_GIVEN,

ai21/clients/common/agents/run.py

@@ -1,10 +1,29 @@
 from abc import ABC
 
 from ai21.models.agents import Agent
+from ai21.models.maestro.run import Requirement
 from ai21.utils.typing import remove_not_given
 
 
 class BaseAgentRun(ABC):
+    def _convert_requirements(self, requirements):
+        """Convert agent requirements to maestro requirements, filtering out invalid ones."""
+        if not requirements:
+            return None
+
+        converted_requirements = []
+        for req in requirements:
+            if req.title and req.description and req.type:
+                converted_requirements.append(
+                    Requirement(
+                        name=req.title,
+                        description=req.description,
+                        is_mandatory=req.type == "mandatory",
+                    )
+                )
+
+        return converted_requirements if converted_requirements else None
+
     def convert_agent_to_maestro_run_payload(
         self,
         agent: Agent,
@@ -15,7 +34,7 @@ def convert_agent_to_maestro_run_payload(
                 models=agent.models,
                 tools=agent.tools,
                 tool_resources=agent.tool_resources,
-                requirements=agent.requirements,
+                requirements=self._convert_requirements(agent.requirements),
                 budget=agent.budget,
                 response_language=agent.response_language,
                 **kwargs,

ai21/models/agents/agent.py

@@ -49,19 +49,19 @@ class Requirement(AI21BaseModel):
 class Agent(AI21BaseModel):
     id: str
     name: str
-    description: Union[str, None] = None
+    description: str | None = None
     organization_id: str
     user_id: str
-    models: Union[List[str], None] = None
-    tools: Union[List[ToolDefinition], None] = None
-    tool_resources: Union[ToolResources, None] = None
-    requirements: Union[List[Requirement], None] = None
-    budget: Union[Budget, None] = None
-    visibility: Union[Visibility, None] = None
-    assistant_type: Union[AgentType, None] = None
+    models: list[str] | None = None
+    tools: list[ToolDefinition] | None = None
+    tool_resources: ToolResources | None = None
+    requirements: list[Requirement] | None = None
+    budget: Budget | None = None
+    visibility: Visibility | None = None
+    assistant_type: AgentType | None = None
     created_at: datetime
     updated_at: datetime
-    response_language: Union[ResponseLanguage, None] = None
+    response_language: ResponseLanguage | None = None
 
 
 class ListAgentsResponse(AI21BaseModel):