wip

Author: harel-tussi

ai21/clients/studio/resources/agents/utils.py

@@ -5,7 +5,7 @@
 from ai21.models.ai21_base_model import AI21BaseModel
 from ai21.models.agents import Agent
 from ai21.models.maestro.run import MaestroMessage
-from ai21.types import NOT_GIVEN, NotGiven
+from typing import Optional
 
 
 class AgentToMaestroRunConverter(AI21BaseModel):
@@ -15,27 +15,27 @@ class AgentToMaestroRunConverter(AI21BaseModel):
     agent_id: str
     input: List[Dict[str, str]]
     verbose: bool = False
-    output_type: Dict[str, Any] | NotGiven = NOT_GIVEN
-    include: List[str] | NotGiven = NOT_GIVEN
+    output_type: Optional[Dict[str, Any]] = None
+    include: Optional[List[str]] = None
     structured_rag_enabled: bool = False
     dynamic_planning_enabled: bool = False
     response_language: str = "english"
 
     # Agent configuration (will be populated from agent fetch)
-    agent_models: List[str] | NotGiven = NOT_GIVEN
-    agent_tools: List[Dict[str, Any]] | NotGiven = NOT_GIVEN
-    agent_tool_resources: Dict[str, Any] | NotGiven = NOT_GIVEN
-    agent_requirements: List[Dict[str, Any]] | NotGiven = NOT_GIVEN
-    agent_budget: str | NotGiven = NOT_GIVEN
+    agent_models: Optional[List[str]] = None
+    agent_tools: Optional[List[Dict[str, Any]]] = None
+    agent_tool_resources: Optional[Dict[str, Any]] = None
+    agent_requirements: Optional[List[Dict[str, Any]]] = None
+    agent_budget: Optional[str] = None
 
     # Computed maestro parameters
     maestro_input: List[MaestroMessage] = Field(default_factory=list)
-    maestro_models: List[str] | NotGiven = NOT_GIVEN
-    maestro_tools: List[Dict[str, Any]] | NotGiven = NOT_GIVEN
-    maestro_tool_resources: Dict[str, Any] | NotGiven = NOT_GIVEN
-    maestro_requirements: List[Dict[str, Any]] | NotGiven = NOT_GIVEN
-    maestro_budget: str | NotGiven = NOT_GIVEN
-    maestro_include: List[str] | NotGiven = NOT_GIVEN
+    maestro_models: Optional[List[str]] = None
+    maestro_tools: Optional[List[Dict[str, Any]]] = None
+    maestro_tool_resources: Optional[Dict[str, Any]] = None
+    maestro_requirements: Optional[List[Dict[str, Any]]] = None
+    maestro_budget: Optional[str] = None
+    maestro_include: Optional[List[str]] = None
 
     def model_post_init(self, __context) -> None:
         """Convert agent parameters to maestro parameters after initialization"""
@@ -48,7 +48,7 @@ def model_post_init(self, __context) -> None:
         self.maestro_tool_resources = self.agent_tool_resources
         self.maestro_requirements = self.agent_requirements
         self.maestro_budget = self.agent_budget
-        self.maestro_include = self.include if self.include is not NOT_GIVEN else NOT_GIVEN
+        self.maestro_include = self.include
 
     @classmethod
     def from_agent_and_params(
@@ -57,15 +57,15 @@ def from_agent_and_params(
         agent_id: str,
         input: List[Dict[str, str]],
         verbose: bool = False,
-        output_type: Dict[str, Any] | NotGiven = NOT_GIVEN,
-        include: List[str] | NotGiven = NOT_GIVEN,
+        output_type: Optional[Dict[str, Any]] = None,
+        include: Optional[List[str]] = None,
         structured_rag_enabled: bool = False,
         dynamic_planning_enabled: bool = False,
         response_language: str = "english",
         **kwargs,
     ) -> "AgentToMaestroRunConverter":
         """Create converter from agent configuration and run parameters"""
-        agent_requirements = NOT_GIVEN
+        agent_requirements = None
         if agent.requirements:
             agent_requirements = [req.model_dump() for req in agent.requirements]
 
@@ -78,11 +78,11 @@ def from_agent_and_params(
             structured_rag_enabled=structured_rag_enabled,
             dynamic_planning_enabled=dynamic_planning_enabled,
             response_language=response_language,
-            agent_models=agent.models or NOT_GIVEN,
-            agent_tools=agent.tools or NOT_GIVEN,
-            agent_tool_resources=agent.tool_resources or NOT_GIVEN,
+            agent_models=agent.models,
+            agent_tools=agent.tools,
+            agent_tool_resources=agent.tool_resources,
             agent_requirements=agent_requirements,
-            agent_budget=agent.budget.value if agent.budget else NOT_GIVEN,
+            agent_budget=agent.budget.value if agent.budget else None,
             **kwargs,
         )
 
@@ -98,5 +98,5 @@ def to_maestro_create_params(self) -> dict:
             "include": self.maestro_include,
             "response_language": self.response_language,
         }
-        # Remove NOT_GIVEN values
-        return {k: v for k, v in params.items() if v is not NOT_GIVEN}
+        # Remove None values
+        return {k: v for k, v in params.items() if v is not None}