fix: add external tools

ai21/clients/common/maestro/run.py

@@ -11,8 +11,8 @@
     OutputOptions,
     Requirement,
     RunResponse,
-    Tool,
     ToolResources,
+    ToolDefinition,
 )
 from ai21.types import NOT_GIVEN, NotGiven
 from ai21.utils.typing import remove_not_given
@@ -26,7 +26,7 @@ def _create_body(
         *,
         input: str | List[MaestroMessage],
         models: List[str] | NotGiven,
-        tools: List[Tool] | NotGiven,
+        tools: List[ToolDefinition] | NotGiven,
         tool_resources: ToolResources | NotGiven,
         requirements: List[Requirement] | NotGiven,
         budget: Budget | NotGiven,
@@ -54,7 +54,7 @@ def create(
         *,
         input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[ToolDefinition] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
@@ -78,7 +78,7 @@ def create_and_poll(
         *,
         input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[ToolDefinition] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,

ai21/clients/studio/resources/maestro/run.py

@@ -15,8 +15,8 @@
     Requirement,
     RunResponse,
     TERMINATED_RUN_STATUSES,
-    Tool,
     ToolResources,
+    ToolDefinition,
 )
 from ai21.types import NotGiven, NOT_GIVEN
 
@@ -27,7 +27,7 @@ def create(
         *,
         input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[ToolDefinition] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
@@ -74,7 +74,7 @@ def create_and_poll(
         *,
         input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[ToolDefinition] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
@@ -107,7 +107,7 @@ async def create(
         *,
         input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[ToolDefinition] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
@@ -154,7 +154,7 @@ async def create_and_poll(
         *,
         input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
-        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tools: List[ToolDefinition] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,

ai21/models/maestro/__init__.py

@@ -5,9 +5,14 @@
     MaestroMessage,
     OutputOptions,
     Requirement,
-    Tool,
+    ToolDefinition,
     ToolResources,
     WebSearchResult,
+    HttpTool,
+    McpTool,
+    Function,
+    Endpoint,
+    Parameters,
 )
 
 __all__ = [
@@ -17,7 +22,12 @@
     "MaestroMessage",
     "OutputOptions",
     "Requirement",
-    "Tool",
+    "ToolDefinition",
     "ToolResources",
     "WebSearchResult",
+    "HttpTool",
+    "McpTool",
+    "Function",
+    "Endpoint",
+    "Parameters",
 ]

ai21/models/maestro/run.py

@@ -1,13 +1,12 @@
 from typing import Literal, List, Optional, Any, Set, Dict, Type, Union
-from typing_extensions import TypedDict
-from pydantic import BaseModel
+from typing_extensions import TypedDict, Required, Annotated
+from pydantic import BaseModel, Field
 
 from ai21.models.ai21_base_model import AI21BaseModel
 
 Budget = Literal["low", "medium", "high"]
 Role = Literal["user", "assistant"]
 RunStatus = Literal["completed", "failed", "in_progress", "requires_action"]
-ToolType = Literal["file_search", "web_search"]
 OutputOptions = Literal["data_sources", "requirements_result"]
 PrimitiveTypes = Union[Type[str], Type[int], Type[float], Type[bool]]
 PrimitiveLists = Type[List[PrimitiveTypes]]
@@ -23,11 +22,45 @@ class MaestroMessage(TypedDict):
     content: str
 
 
-class Tool(TypedDict):
-    type: ToolType
+class HTTPToolFunctionParamProperties(TypedDict):
+    type: str
+    description: str
+
+
+class Parameters(TypedDict, total=False):
+    type: Literal["object"]
+    properties: Dict[str, HTTPToolFunctionParamProperties]
+    required: List[str]
+    additional_properties: Optional[bool]
+
+
+class Function(TypedDict):
+    name: str
+    description: str
+    parameters: Parameters
+
+
+class Endpoint(TypedDict, total=False):
+    url: str
+    headers: Optional[dict]
+
+
+class HttpTool(TypedDict):
+    type: Required[Literal["http"]]
+    function: Function
+    endpoint: Endpoint
 
 
-class FileSearchToolResource(TypedDict, total=False):
+class McpTool(TypedDict, total=False):
+    type: Required[Literal["mcp"]]
+    server_label: str
+    server_url: str
+    headers: Optional[dict]
+    allowed_tools: Optional[List[str]]
+
+
+class FileSearch(TypedDict, total=False):
+    type: Literal["file_search"]
     retrieval_similarity_threshold: Optional[float]
     labels: Optional[List[str]]
     labels_filter_mode: Optional[Literal["AND", "OR"]]
@@ -37,24 +70,36 @@ class FileSearchToolResource(TypedDict, total=False):
     max_neighbors: Optional[int]
 
 
-class WebSearchToolResource(TypedDict, total=False):
+class WebSearch(TypedDict, total=False):
+    type: Literal["web_search"]
     urls: Optional[List[str]]
 
 
 class ToolResources(TypedDict, total=False):
-    file_search: Optional[FileSearchToolResource]
-    web_search: Optional[WebSearchToolResource]
+    file_search: Optional[FileSearch]
+    web_search: Optional[WebSearch]
+
+
+ToolDefinition = Annotated[
+    Union[
+        HttpTool,
+        McpTool,
+        FileSearch,
+        WebSearch,
+    ],
+    Field(discriminator="type"),
+]
 
 
 class Requirement(TypedDict, total=False):
     name: str
     description: str
-    is_mandatory: bool = False
+    is_mandatory: bool
 
 
 class RequirementResultItem(Requirement, total=False):
     score: float
-    reason: Optional[str] = None
+    reason: Optional[str]
 
 
 class RequirementsResult(TypedDict, total=False):