From 6efbca7fd236df7ff90acda8ef0ed41c3dddd93b Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 11 Mar 2025 13:07:40 +0200 Subject: [PATCH 01/21] feat: :sparkles: introduce Maestro --- ai21/clients/common/maestro/__init__.py | 0 ai21/clients/common/maestro/maestro.py | 8 + ai21/clients/common/maestro/runs.py | 103 ++++++++++ ai21/clients/studio/ai21_client.py | 2 + ai21/clients/studio/async_ai21_client.py | 2 + .../studio/resources/maestro/__init__.py | 0 .../studio/resources/maestro/maestro.py | 19 ++ ai21/clients/studio/resources/maestro/runs.py | 186 ++++++++++++++++++ ai21/models/maestro/__init__.py | 0 ai21/models/maestro/runs.py | 53 +++++ examples/studio/maestro/__init__.py | 0 11 files changed, 373 insertions(+) create mode 100644 ai21/clients/common/maestro/__init__.py create mode 100644 ai21/clients/common/maestro/maestro.py create mode 100644 ai21/clients/common/maestro/runs.py create mode 100644 ai21/clients/studio/resources/maestro/__init__.py create mode 100644 ai21/clients/studio/resources/maestro/maestro.py create mode 100644 ai21/clients/studio/resources/maestro/runs.py create mode 100644 ai21/models/maestro/__init__.py create mode 100644 ai21/models/maestro/runs.py create mode 100644 examples/studio/maestro/__init__.py diff --git a/ai21/clients/common/maestro/__init__.py b/ai21/clients/common/maestro/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ai21/clients/common/maestro/maestro.py b/ai21/clients/common/maestro/maestro.py new file mode 100644 index 00000000..52442c0c --- /dev/null +++ b/ai21/clients/common/maestro/maestro.py @@ -0,0 +1,8 @@ +from abc import ABC + +from ai21.clients.common.maestro.runs import BaseMaestroRun + + +class BaseMaestro(ABC): + _module_name = "maestro" + runs: BaseMaestroRun diff --git a/ai21/clients/common/maestro/runs.py b/ai21/clients/common/maestro/runs.py new file mode 100644 index 00000000..5c965769 --- /dev/null +++ b/ai21/clients/common/maestro/runs.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import List, Dict, Any + +from ai21.models.maestro.runs import ( + Message, + Tool, + ToolResources, + Budget, + RunResponse, + DEFAULT_RUN_POLL_INTERVAL, + DEFAULT_RUN_POLL_TIMEOUT, +) +from ai21.types import NOT_GIVEN, NotGiven +from ai21.utils.typing import remove_not_given + + +class BaseMaestroRun(ABC): + _module_name = "maestro/runs" + + def _create_body( + self, + *, + messages: List[Message] | NotGiven, + instruction: str | NotGiven, + output_type: Dict[str, Any] | NotGiven, + models: List[str] | NotGiven, + tools: List[Tool] | NotGiven, + tool_resources: ToolResources | NotGiven, + context: Dict[str, Any] | NotGiven, + requirements: List[str] | NotGiven, + budget: Budget | NotGiven, + **kwargs, + ) -> dict: + messages_payload = remove_not_given({"messages": messages, "instruction": instruction}) + if not messages_payload: + # not messages nor instruction were given + raise ValueError("Must provide either `messages` or `instruction`") + elif len(messages_payload.keys()) > 1: + # both messages and instruction were given + raise ValueError("Must provide only one of `messages` or `instruction`") + elif "instruction" in messages_payload: + # instruction was given and should be modified accordingly + messages = [{"role": "user", "content": instruction}] + + return remove_not_given( + { + "messages": messages, + "output_type": output_type, + "models": models, + "tools": tools, + "tool_resources": tool_resources, + "context": context, + "requirements": requirements, + "budget": budget, + **kwargs, + } + ) + + @abstractmethod + def create( + self, + *, + messages: List[Message] | NotGiven = NOT_GIVEN, + instruction: str | NotGiven = NOT_GIVEN, + output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + models: List[str] | NotGiven = NOT_GIVEN, + tools: List[Tool] | NotGiven = NOT_GIVEN, + tool_resources: ToolResources | NotGiven = NOT_GIVEN, + context: Dict[str, Any] | NotGiven = NOT_GIVEN, + requirements: List[str] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, + **kwargs, + ) -> RunResponse: + pass + + @abstractmethod + def retrieve(self, run_id: str) -> RunResponse: + pass + + @abstractmethod + def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_timeout: float) -> RunResponse: + pass + + @abstractmethod + def create_and_poll( + self, + *, + instruction: str | NotGiven = NOT_GIVEN, + messages: List[Message] | NotGiven = NOT_GIVEN, + output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + models: List[str] | NotGiven = NOT_GIVEN, + tools: List[Tool] | NotGiven = NOT_GIVEN, + tool_resources: ToolResources | NotGiven = NOT_GIVEN, + context: Dict[str, Any] | NotGiven = NOT_GIVEN, + requirements: List[str] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, + poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL, + poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT, + **kwargs, + ) -> RunResponse: + pass diff --git a/ai21/clients/studio/ai21_client.py b/ai21/clients/studio/ai21_client.py index 34cdb491..48ee933f 100644 --- a/ai21/clients/studio/ai21_client.py +++ b/ai21/clients/studio/ai21_client.py @@ -5,6 +5,7 @@ from ai21.ai21_env_config import _AI21EnvConfig, AI21EnvConfig from ai21.clients.studio.client_url_parser import create_client_url from ai21.clients.studio.resources.beta.beta import Beta +from ai21.clients.studio.resources.maestro.maestro import Maestro from ai21.clients.studio.resources.studio_chat import StudioChat from ai21.clients.studio.resources.studio_library import StudioLibrary from ai21.http_client.http_client import AI21HTTPClient @@ -42,4 +43,5 @@ def __init__( ) self.chat: StudioChat = StudioChat(self) self.library = StudioLibrary(self) + self.maestro = Maestro(self) self.beta = Beta(self) diff --git a/ai21/clients/studio/async_ai21_client.py b/ai21/clients/studio/async_ai21_client.py index 4f7ed322..0daf9563 100644 --- a/ai21/clients/studio/async_ai21_client.py +++ b/ai21/clients/studio/async_ai21_client.py @@ -5,6 +5,7 @@ from ai21.ai21_env_config import _AI21EnvConfig, AI21EnvConfig from ai21.clients.studio.client_url_parser import create_client_url from ai21.clients.studio.resources.beta.async_beta import AsyncBeta +from ai21.clients.studio.resources.maestro.maestro import AsyncMaestro from ai21.clients.studio.resources.studio_chat import AsyncStudioChat from ai21.clients.studio.resources.studio_library import AsyncStudioLibrary from ai21.http_client.async_http_client import AsyncAI21HTTPClient @@ -41,4 +42,5 @@ def __init__( self.chat: AsyncStudioChat = AsyncStudioChat(self) self.library = AsyncStudioLibrary(self) + self.maestro = AsyncMaestro(self) self.beta = AsyncBeta(self) diff --git a/ai21/clients/studio/resources/maestro/__init__.py b/ai21/clients/studio/resources/maestro/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ai21/clients/studio/resources/maestro/maestro.py b/ai21/clients/studio/resources/maestro/maestro.py new file mode 100644 index 00000000..43a2eeee --- /dev/null +++ b/ai21/clients/studio/resources/maestro/maestro.py @@ -0,0 +1,19 @@ +from ai21.clients.common.maestro.maestro import BaseMaestro +from ai21.clients.studio.resources.maestro.runs import MaestroRun, AsyncMaestroRun +from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource +from ai21.http_client.async_http_client import AsyncAI21HTTPClient +from ai21.http_client.http_client import AI21HTTPClient + + +class Maestro(StudioResource, BaseMaestro): + def __init__(self, client: AI21HTTPClient): + super().__init__(client) + + self.runs = MaestroRun(client) + + +class AsyncMaestro(AsyncStudioResource, BaseMaestro): + def __init__(self, client: AsyncAI21HTTPClient): + super().__init__(client) + + self.runs = AsyncMaestroRun(client) diff --git a/ai21/clients/studio/resources/maestro/runs.py b/ai21/clients/studio/resources/maestro/runs.py new file mode 100644 index 00000000..22a41e87 --- /dev/null +++ b/ai21/clients/studio/resources/maestro/runs.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +import asyncio +import time +from typing import Any, List, Dict + +from ai21.clients.common.maestro.runs import BaseMaestroRun +from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource +from ai21.models.maestro.runs import ( + Message, + Tool, + ToolResources, + Budget, + RunResponse, + TERMINATED_RUN_STATUSES, + DEFAULT_RUN_POLL_INTERVAL, + DEFAULT_RUN_POLL_TIMEOUT, +) + +from ai21.types import NotGiven, NOT_GIVEN + + +class MaestroRun(StudioResource, BaseMaestroRun): + def create( + self, + *, + messages: List[Message] | NotGiven = NOT_GIVEN, + instruction: str | NotGiven = NOT_GIVEN, + output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + models: List[str] | NotGiven = NOT_GIVEN, + tools: List[Tool] | NotGiven = NOT_GIVEN, + tool_resources: ToolResources | NotGiven = NOT_GIVEN, + context: Dict[str, Any] | NotGiven = NOT_GIVEN, + requirements: List[str] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, + **kwargs, + ) -> RunResponse: + body = self._create_body( + messages=messages, + instruction=instruction, + output_type=output_type, + models=models, + tools=tools, + tool_resources=tool_resources, + context=context, + requirements=requirements, + budget=budget, + **kwargs, + ) + + return self._post(path=f"/{self._module_name}", body=body, response_cls=RunResponse) + + def retrieve( + self, + run_id: str, + ) -> RunResponse: + return self._get(path=f"/{self._module_name}/{run_id}", response_cls=RunResponse) + + def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_timeout: float) -> RunResponse: + start_time = time.time() + + while True: + run = self.retrieve(run_id) + + if run.status in TERMINATED_RUN_STATUSES: + return run + + if (time.time() - start_time) >= poll_timeout: + return run + + time.sleep(poll_interval) + + def create_and_poll( + self, + *, + messages: List[Message] | NotGiven = NOT_GIVEN, + instruction: str | NotGiven = NOT_GIVEN, + output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + models: List[str] | NotGiven = NOT_GIVEN, + tools: List[Tool] | NotGiven = NOT_GIVEN, + tool_resources: ToolResources | NotGiven = NOT_GIVEN, + context: Dict[str, Any] | NotGiven = NOT_GIVEN, + requirements: List[str] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, + poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL, + poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT, + **kwargs, + ) -> RunResponse: + run = self.create( + messages=messages, + instruction=instruction, + output_type=output_type, + models=models, + tools=tools, + tool_resources=tool_resources, + context=context, + requirements=requirements, + budget=budget, + **kwargs, + ) + + return self._poll_for_status(run_id=run.id, poll_interval=poll_interval_sec, poll_timeout=poll_timeout_sec) + + +class AsyncMaestroRun(AsyncStudioResource, BaseMaestroRun): + async def create( + self, + *, + messages: List[Message] | NotGiven = NOT_GIVEN, + instruction: str | NotGiven = NOT_GIVEN, + output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + models: List[str] | NotGiven = NOT_GIVEN, + tools: List[Tool] | NotGiven = NOT_GIVEN, + tool_resources: ToolResources | NotGiven = NOT_GIVEN, + context: Dict[str, Any] | NotGiven = NOT_GIVEN, + requirements: List[str] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, + **kwargs, + ) -> RunResponse: + body = self._create_body( + messages=messages, + instruction=instruction, + output_type=output_type, + models=models, + tools=tools, + tool_resources=tool_resources, + context=context, + requirements=requirements, + budget=budget, + **kwargs, + ) + + return await self._post(path=f"/{self._module_name}", body=body, response_cls=RunResponse) + + async def retrieve( + self, + run_id: str, + ) -> RunResponse: + return await self._get(path=f"/{self._module_name}/{run_id}", response_cls=RunResponse) + + async def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_timeout: float) -> RunResponse: + start_time = time.time() + + while True: + run = await self.retrieve(run_id) + + if run.status in TERMINATED_RUN_STATUSES: + return run + + if (time.time() - start_time) >= poll_timeout: + return run + + await asyncio.sleep(poll_interval) + + async def create_and_poll( + self, + *, + messages: List[Message] | NotGiven = NOT_GIVEN, + instruction: str | NotGiven = NOT_GIVEN, + output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + models: List[str] | NotGiven = NOT_GIVEN, + tools: List[Tool] | NotGiven = NOT_GIVEN, + tool_resources: ToolResources | NotGiven = NOT_GIVEN, + context: Dict[str, Any] | NotGiven = NOT_GIVEN, + requirements: List[str] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, + poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL, + poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT, + **kwargs, + ) -> RunResponse: + run = await self.create( + messages=messages, + instruction=instruction, + output_type=output_type, + models=models, + tools=tools, + tool_resources=tool_resources, + context=context, + requirements=requirements, + budget=budget, + **kwargs, + ) + + return await self._poll_for_status( + run_id=run.id, poll_interval=poll_interval_sec, poll_timeout=poll_timeout_sec + ) diff --git a/ai21/models/maestro/__init__.py b/ai21/models/maestro/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ai21/models/maestro/runs.py b/ai21/models/maestro/runs.py new file mode 100644 index 00000000..4975296a --- /dev/null +++ b/ai21/models/maestro/runs.py @@ -0,0 +1,53 @@ +from typing import TypedDict, Literal, List, Optional, Any, Set + +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"] + +DEFAULT_RUN_POLL_INTERVAL: float = 1 # seconds +DEFAULT_RUN_POLL_TIMEOUT: float = 120 # seconds +TERMINATED_RUN_STATUSES: Set[RunStatus] = {"completed", "failed", "requires_action"} + + +class Tool(TypedDict): + type: ToolType + + +class FileSearchToolResource(TypedDict, total=False): + retrieval_similarity_threshold: Optional[float] + labels: Optional[List[str]] + labels_filter_mode: Optional[Literal["AND", "OR"]] + labels_filter: Optional[dict] + file_ids: Optional[List[str]] + retrieval_strategy: Optional[str] + max_neighbors: Optional[int] + + +class WebSearchToolResource(TypedDict, total=False): + urls: Optional[List[str]] + fallback_to_web: Optional[bool] + + +class ToolResources(TypedDict): + file_search: FileSearchToolResource + web_search: WebSearchToolResource + + +class Message(TypedDict): + role: Role + content: str + + +class Constraint(TypedDict): + name: str + description: str + is_mandatory: bool + + +class RunResponse(AI21BaseModel): + id: str + status: RunStatus + result: Any diff --git a/examples/studio/maestro/__init__.py b/examples/studio/maestro/__init__.py new file mode 100644 index 00000000..e69de29b From f72286e027e54ccaf4d19d3ab1a21fbd3575e73c Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Tue, 11 Mar 2025 13:29:25 +0200 Subject: [PATCH 02/21] ci: Added remove cache --- .github/workflows/integration-tests.yaml | 1 + .github/workflows/test.yaml | 6 ++++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/workflows/integration-tests.yaml b/.github/workflows/integration-tests.yaml index 6e9c89de..3d71fb48 100644 --- a/.github/workflows/integration-tests.yaml +++ b/.github/workflows/integration-tests.yaml @@ -53,6 +53,7 @@ jobs: - name: Set Poetry environment run: | poetry env use ${{ matrix.python-version }} + poetry cache clear --all pypi - name: Install dependencies run: | poetry install --all-extras diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 3e22653f..8c4ffb26 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -25,9 +25,10 @@ jobs: python-version: ${{ matrix.python-version }} cache: poetry cache-dependency-path: poetry.lock - - name: Set Poetry environment + - name: Set Poetry environment and clear cache run: | poetry env use ${{ matrix.python-version }} + poetry cache clear --all pypi - name: Install dependencies run: | poetry install --no-root --only dev --all-extras @@ -58,9 +59,10 @@ jobs: python-version: ${{ matrix.python-version }} cache: poetry cache-dependency-path: poetry.lock - - name: Set Poetry environment + - name: Set Poetry environment and clear cache run: | poetry env use ${{ matrix.python-version }} + poetry cache clear --all pypi - name: Override Pydantic version run: | if [[ "${{ matrix.pydantic-version }}" == ^1.* ]]; then From a267dbcf52a953623a6dac2e219d40dd3d28c36b Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 11 Mar 2025 14:07:56 +0200 Subject: [PATCH 03/21] chore: :refactor: use `ChatMessage` instead of `Message` pretty TypedDict :( --- README.md | 19 +++++++++++++++++++ ai21/clients/common/maestro/runs.py | 13 +++++++------ ai21/clients/studio/resources/maestro/runs.py | 10 +++++----- ai21/models/maestro/runs.py | 5 ----- examples/studio/maestro/async_runs.py | 19 +++++++++++++++++++ 5 files changed, 50 insertions(+), 16 deletions(-) create mode 100644 examples/studio/maestro/async_runs.py diff --git a/README.md b/README.md index f61214e4..0ba94912 100644 --- a/README.md +++ b/README.md @@ -359,6 +359,25 @@ asyncio.run(main()) --- +### Maestro + +AI Planning & Orchestration System built for the enterprise. Read more [here](https://www.ai21.com/maestro/). + +```python +from ai21 import AI21Client + +client = AI21Client() + +run_result = client.maestro.runs.create_and_poll( + instruction="Who was the Maestro in the movie 'The Maestro'?", + tools=[{"type": "web_search"}], +) +``` + +For a more detailed example, see maestro [sync](examples/studio/maestro/runs.py) and [async](examples/studio/maestro/async_runs.py) examples. + +--- + ### Conversational RAG (Beta) Like chat, but with the ability to retrieve information from your Studio library. diff --git a/ai21/clients/common/maestro/runs.py b/ai21/clients/common/maestro/runs.py index 5c965769..006ff7e2 100644 --- a/ai21/clients/common/maestro/runs.py +++ b/ai21/clients/common/maestro/runs.py @@ -3,8 +3,9 @@ from abc import ABC, abstractmethod from typing import List, Dict, Any +from ai21.models._pydantic_compatibility import _to_dict +from ai21.models.chat import ChatMessage from ai21.models.maestro.runs import ( - Message, Tool, ToolResources, Budget, @@ -22,8 +23,8 @@ class BaseMaestroRun(ABC): def _create_body( self, *, - messages: List[Message] | NotGiven, instruction: str | NotGiven, + messages: List[ChatMessage] | NotGiven, output_type: Dict[str, Any] | NotGiven, models: List[str] | NotGiven, tools: List[Tool] | NotGiven, @@ -42,11 +43,11 @@ def _create_body( raise ValueError("Must provide only one of `messages` or `instruction`") elif "instruction" in messages_payload: # instruction was given and should be modified accordingly - messages = [{"role": "user", "content": instruction}] + messages = [ChatMessage(role="user", content=instruction)] return remove_not_given( { - "messages": messages, + "messages": [_to_dict(message) for message in messages], "output_type": output_type, "models": models, "tools": tools, @@ -62,8 +63,8 @@ def _create_body( def create( self, *, - messages: List[Message] | NotGiven = NOT_GIVEN, instruction: str | NotGiven = NOT_GIVEN, + messages: List[ChatMessage] | NotGiven = NOT_GIVEN, output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, @@ -88,7 +89,7 @@ def create_and_poll( self, *, instruction: str | NotGiven = NOT_GIVEN, - messages: List[Message] | NotGiven = NOT_GIVEN, + messages: List[ChatMessage] | NotGiven = NOT_GIVEN, output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, diff --git a/ai21/clients/studio/resources/maestro/runs.py b/ai21/clients/studio/resources/maestro/runs.py index 22a41e87..823ec467 100644 --- a/ai21/clients/studio/resources/maestro/runs.py +++ b/ai21/clients/studio/resources/maestro/runs.py @@ -6,8 +6,8 @@ from ai21.clients.common.maestro.runs import BaseMaestroRun from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource +from ai21.models.chat import ChatMessage from ai21.models.maestro.runs import ( - Message, Tool, ToolResources, Budget, @@ -24,8 +24,8 @@ class MaestroRun(StudioResource, BaseMaestroRun): def create( self, *, - messages: List[Message] | NotGiven = NOT_GIVEN, instruction: str | NotGiven = NOT_GIVEN, + messages: List[ChatMessage] | NotGiven = NOT_GIVEN, output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, @@ -73,8 +73,8 @@ def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_timeout: f def create_and_poll( self, *, - messages: List[Message] | NotGiven = NOT_GIVEN, instruction: str | NotGiven = NOT_GIVEN, + messages: List[ChatMessage] | NotGiven = NOT_GIVEN, output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, @@ -106,8 +106,8 @@ class AsyncMaestroRun(AsyncStudioResource, BaseMaestroRun): async def create( self, *, - messages: List[Message] | NotGiven = NOT_GIVEN, instruction: str | NotGiven = NOT_GIVEN, + messages: List[ChatMessage] | NotGiven = NOT_GIVEN, output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, @@ -155,8 +155,8 @@ async def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_time async def create_and_poll( self, *, - messages: List[Message] | NotGiven = NOT_GIVEN, instruction: str | NotGiven = NOT_GIVEN, + messages: List[ChatMessage] | NotGiven = NOT_GIVEN, output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, diff --git a/ai21/models/maestro/runs.py b/ai21/models/maestro/runs.py index 4975296a..99874de0 100644 --- a/ai21/models/maestro/runs.py +++ b/ai21/models/maestro/runs.py @@ -36,11 +36,6 @@ class ToolResources(TypedDict): web_search: WebSearchToolResource -class Message(TypedDict): - role: Role - content: str - - class Constraint(TypedDict): name: str description: str diff --git a/examples/studio/maestro/async_runs.py b/examples/studio/maestro/async_runs.py new file mode 100644 index 00000000..703f2bc0 --- /dev/null +++ b/examples/studio/maestro/async_runs.py @@ -0,0 +1,19 @@ +import asyncio + +from ai21 import AsyncAI21Client + +client = AsyncAI21Client() + + +async def main(): + run_result = await client.maestro.runs.create_and_poll( + instruction="Who was the Maestro in the movie 'Maestro'?", + tools=[{"type": "web_search"}], + tool_resources={"web_search": {"urls": ["google.com"]}}, + ) + + print(run_result) + + +if __name__ == "__main__": + asyncio.run(main()) From 548d81ecf50118466d7d0801f20aa0495b65deba Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 11 Mar 2025 14:08:35 +0200 Subject: [PATCH 04/21] chore: :refactor: use `ChatMessage` instead of `Message` pretty TypedDict :( --- examples/studio/maestro/runs.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 examples/studio/maestro/runs.py diff --git a/examples/studio/maestro/runs.py b/examples/studio/maestro/runs.py new file mode 100644 index 00000000..eb80e152 --- /dev/null +++ b/examples/studio/maestro/runs.py @@ -0,0 +1,17 @@ +from ai21 import AI21Client + +client = AI21Client() + + +def main(): + run_result = client.maestro.runs.create_and_poll( + instruction="Who was the Maestro in the movie 'Maestro'?", + tools=[{"type": "web_search"}], + tool_resources={"web_search": {"urls": ["google.com"]}}, + ) + + print(run_result) + + +if __name__ == "__main__": + main() From e76bba634201d47a4884b4fc05c478b820a16ffe Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 11 Mar 2025 14:50:41 +0200 Subject: [PATCH 05/21] chore: :truck: rename `maestro/runs` files to `maestro/run` --- ai21/clients/common/maestro/maestro.py | 2 +- ai21/clients/common/maestro/{runs.py => run.py} | 2 +- ai21/clients/studio/resources/maestro/maestro.py | 2 +- ai21/clients/studio/resources/maestro/{runs.py => run.py} | 4 ++-- ai21/models/maestro/{runs.py => run.py} | 0 5 files changed, 5 insertions(+), 5 deletions(-) rename ai21/clients/common/maestro/{runs.py => run.py} (98%) rename ai21/clients/studio/resources/maestro/{runs.py => run.py} (98%) rename ai21/models/maestro/{runs.py => run.py} (100%) diff --git a/ai21/clients/common/maestro/maestro.py b/ai21/clients/common/maestro/maestro.py index 52442c0c..393e39c9 100644 --- a/ai21/clients/common/maestro/maestro.py +++ b/ai21/clients/common/maestro/maestro.py @@ -1,6 +1,6 @@ from abc import ABC -from ai21.clients.common.maestro.runs import BaseMaestroRun +from ai21.clients.common.maestro.run import BaseMaestroRun class BaseMaestro(ABC): diff --git a/ai21/clients/common/maestro/runs.py b/ai21/clients/common/maestro/run.py similarity index 98% rename from ai21/clients/common/maestro/runs.py rename to ai21/clients/common/maestro/run.py index 006ff7e2..6e85927b 100644 --- a/ai21/clients/common/maestro/runs.py +++ b/ai21/clients/common/maestro/run.py @@ -5,7 +5,7 @@ from ai21.models._pydantic_compatibility import _to_dict from ai21.models.chat import ChatMessage -from ai21.models.maestro.runs import ( +from ai21.models.maestro.run import ( Tool, ToolResources, Budget, diff --git a/ai21/clients/studio/resources/maestro/maestro.py b/ai21/clients/studio/resources/maestro/maestro.py index 43a2eeee..346b0cc0 100644 --- a/ai21/clients/studio/resources/maestro/maestro.py +++ b/ai21/clients/studio/resources/maestro/maestro.py @@ -1,5 +1,5 @@ from ai21.clients.common.maestro.maestro import BaseMaestro -from ai21.clients.studio.resources.maestro.runs import MaestroRun, AsyncMaestroRun +from ai21.clients.studio.resources.maestro.run import MaestroRun, AsyncMaestroRun from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource from ai21.http_client.async_http_client import AsyncAI21HTTPClient from ai21.http_client.http_client import AI21HTTPClient diff --git a/ai21/clients/studio/resources/maestro/runs.py b/ai21/clients/studio/resources/maestro/run.py similarity index 98% rename from ai21/clients/studio/resources/maestro/runs.py rename to ai21/clients/studio/resources/maestro/run.py index 823ec467..81032249 100644 --- a/ai21/clients/studio/resources/maestro/runs.py +++ b/ai21/clients/studio/resources/maestro/run.py @@ -4,10 +4,10 @@ import time from typing import Any, List, Dict -from ai21.clients.common.maestro.runs import BaseMaestroRun +from ai21.clients.common.maestro.run import BaseMaestroRun from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource from ai21.models.chat import ChatMessage -from ai21.models.maestro.runs import ( +from ai21.models.maestro.run import ( Tool, ToolResources, Budget, diff --git a/ai21/models/maestro/runs.py b/ai21/models/maestro/run.py similarity index 100% rename from ai21/models/maestro/runs.py rename to ai21/models/maestro/run.py From c885a72e2bf75e9437602d070493d861c7a59997 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 11 Mar 2025 19:20:58 +0200 Subject: [PATCH 06/21] fix: :recycle: bug fixes and support for more params - update examples (async & sync) - support passing types in `output_type` - deprecate `instruction` --- README.md | 2 +- ai21/clients/common/maestro/run.py | 71 +++++++++++++------- ai21/clients/studio/resources/maestro/run.py | 25 +++---- ai21/models/_pydantic_compatibility.py | 9 ++- ai21/models/maestro/run.py | 12 ++-- examples/studio/maestro/async_runs.py | 6 +- examples/studio/maestro/runs.py | 6 +- 7 files changed, 80 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index 0ba94912..5f18086b 100644 --- a/README.md +++ b/README.md @@ -369,7 +369,7 @@ from ai21 import AI21Client client = AI21Client() run_result = client.maestro.runs.create_and_poll( - instruction="Who was the Maestro in the movie 'The Maestro'?", + messages=[{"role": "user", "content": "Who was the Maestro in the movie 'The Maestro'?"}], tools=[{"type": "web_search"}], ) ``` diff --git a/ai21/clients/common/maestro/run.py b/ai21/clients/common/maestro/run.py index 6e85927b..17b1b6f2 100644 --- a/ai21/clients/common/maestro/run.py +++ b/ai21/clients/common/maestro/run.py @@ -1,9 +1,12 @@ from __future__ import annotations +import inspect from abc import ABC, abstractmethod -from typing import List, Dict, Any +from typing import List, Dict, Any, get_origin, get_args, Optional -from ai21.models._pydantic_compatibility import _to_dict +from pydantic import BaseModel + +from ai21.models._pydantic_compatibility import _to_dict, _to_schema from ai21.models.chat import ChatMessage from ai21.models.maestro.run import ( Tool, @@ -12,20 +15,55 @@ RunResponse, DEFAULT_RUN_POLL_INTERVAL, DEFAULT_RUN_POLL_TIMEOUT, + OutputType, ) from ai21.types import NOT_GIVEN, NotGiven from ai21.utils.typing import remove_not_given +def _primitive_to_schema(output_type): + type_mapping = { + int: "integer", + float: "number", + str: "string", + bool: "boolean", + } + + if output_type in type_mapping: + return {"type": type_mapping[output_type]} + + origin = get_origin(output_type) + args = get_args(output_type) + + if origin is list and len(args) == 1 and args[0] in type_mapping: + return {"type": "array", "items": {"type": type_mapping[args[0]]}} + + class BaseMaestroRun(ABC): _module_name = "maestro/runs" + def _output_type_to_json_schema(self, output_type: OutputType | NotGiven) -> Optional[Dict[str, Any]]: + if not output_type or isinstance(output_type, NotGiven): + return NOT_GIVEN + + if inspect.isclass(output_type) and issubclass(output_type, BaseModel): + result = _to_schema(output_type) + else: + result = _primitive_to_schema(output_type) + + if not result: + raise ValueError( + "Unsupported type. Supported types are: primitives types, List of primitive types, json schema dict, " + "and pydantic models." + ) + + return result + def _create_body( self, *, - instruction: str | NotGiven, - messages: List[ChatMessage] | NotGiven, - output_type: Dict[str, Any] | NotGiven, + messages: List[ChatMessage], + output_type: OutputType | NotGiven, models: List[str] | NotGiven, tools: List[Tool] | NotGiven, tool_resources: ToolResources | NotGiven, @@ -34,21 +72,10 @@ def _create_body( budget: Budget | NotGiven, **kwargs, ) -> dict: - messages_payload = remove_not_given({"messages": messages, "instruction": instruction}) - if not messages_payload: - # not messages nor instruction were given - raise ValueError("Must provide either `messages` or `instruction`") - elif len(messages_payload.keys()) > 1: - # both messages and instruction were given - raise ValueError("Must provide only one of `messages` or `instruction`") - elif "instruction" in messages_payload: - # instruction was given and should be modified accordingly - messages = [ChatMessage(role="user", content=instruction)] - return remove_not_given( { "messages": [_to_dict(message) for message in messages], - "output_type": output_type, + "output_type": self._output_type_to_json_schema(output_type), "models": models, "tools": tools, "tool_resources": tool_resources, @@ -63,9 +90,8 @@ def _create_body( def create( self, *, - instruction: str | NotGiven = NOT_GIVEN, - messages: List[ChatMessage] | NotGiven = NOT_GIVEN, - output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + messages: List[ChatMessage], + output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -88,9 +114,8 @@ def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_timeout: f def create_and_poll( self, *, - instruction: str | NotGiven = NOT_GIVEN, - messages: List[ChatMessage] | NotGiven = NOT_GIVEN, - output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + messages: List[ChatMessage], + output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, diff --git a/ai21/clients/studio/resources/maestro/run.py b/ai21/clients/studio/resources/maestro/run.py index 81032249..6a5638c8 100644 --- a/ai21/clients/studio/resources/maestro/run.py +++ b/ai21/clients/studio/resources/maestro/run.py @@ -15,6 +15,7 @@ TERMINATED_RUN_STATUSES, DEFAULT_RUN_POLL_INTERVAL, DEFAULT_RUN_POLL_TIMEOUT, + OutputType, ) from ai21.types import NotGiven, NOT_GIVEN @@ -24,9 +25,8 @@ class MaestroRun(StudioResource, BaseMaestroRun): def create( self, *, - instruction: str | NotGiven = NOT_GIVEN, - messages: List[ChatMessage] | NotGiven = NOT_GIVEN, - output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + messages: List[ChatMessage], + output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -37,7 +37,6 @@ def create( ) -> RunResponse: body = self._create_body( messages=messages, - instruction=instruction, output_type=output_type, models=models, tools=tools, @@ -73,9 +72,8 @@ def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_timeout: f def create_and_poll( self, *, - instruction: str | NotGiven = NOT_GIVEN, - messages: List[ChatMessage] | NotGiven = NOT_GIVEN, - output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + messages: List[ChatMessage], + output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -88,7 +86,6 @@ def create_and_poll( ) -> RunResponse: run = self.create( messages=messages, - instruction=instruction, output_type=output_type, models=models, tools=tools, @@ -106,9 +103,8 @@ class AsyncMaestroRun(AsyncStudioResource, BaseMaestroRun): async def create( self, *, - instruction: str | NotGiven = NOT_GIVEN, - messages: List[ChatMessage] | NotGiven = NOT_GIVEN, - output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + messages: List[ChatMessage], + output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -119,7 +115,6 @@ async def create( ) -> RunResponse: body = self._create_body( messages=messages, - instruction=instruction, output_type=output_type, models=models, tools=tools, @@ -155,9 +150,8 @@ async def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_time async def create_and_poll( self, *, - instruction: str | NotGiven = NOT_GIVEN, - messages: List[ChatMessage] | NotGiven = NOT_GIVEN, - output_type: Dict[str, Any] | NotGiven = NOT_GIVEN, + messages: List[ChatMessage], + output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -170,7 +164,6 @@ async def create_and_poll( ) -> RunResponse: run = await self.create( messages=messages, - instruction=instruction, output_type=output_type, models=models, tools=tools, diff --git a/ai21/models/_pydantic_compatibility.py b/ai21/models/_pydantic_compatibility.py index 5f58c0de..41d08509 100644 --- a/ai21/models/_pydantic_compatibility.py +++ b/ai21/models/_pydantic_compatibility.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Dict, Any +from typing import Dict, Any, Type from pydantic import VERSION, BaseModel @@ -33,3 +33,10 @@ def _from_json(obj: "AI21BaseModel", json_str: str, **kwargs) -> BaseModel: # n return obj.model_validate_json(json_str, **kwargs) return obj.parse_raw(json_str, **kwargs) + + +def _to_schema(model_object: Type[BaseModel], **kwargs) -> Dict[str, Any]: + if IS_PYDANTIC_V2: + return model_object.model_json_schema(**kwargs) + + return model_object.schema(**kwargs) diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py index 99874de0..eb0b7476 100644 --- a/ai21/models/maestro/run.py +++ b/ai21/models/maestro/run.py @@ -1,4 +1,6 @@ -from typing import TypedDict, Literal, List, Optional, Any, Set +from typing import TypedDict, Literal, List, Optional, Any, Set, Dict, Type + +from pydantic import BaseModel from ai21.models.ai21_base_model import AI21BaseModel @@ -6,6 +8,8 @@ Role = Literal["user", "assistant"] RunStatus = Literal["completed", "failed", "in_progress", "requires_action"] ToolType = Literal["file_search", "web_search"] +PrimitiveType = Type[str] | Type[int] | Type[float] | Type[bool] | Type[List[str]] | Type[List[int]] | Type[List[float]] +OutputType = PrimitiveType | Type[BaseModel] | Dict[str, Any] DEFAULT_RUN_POLL_INTERVAL: float = 1 # seconds DEFAULT_RUN_POLL_TIMEOUT: float = 120 # seconds @@ -31,9 +35,9 @@ class WebSearchToolResource(TypedDict, total=False): fallback_to_web: Optional[bool] -class ToolResources(TypedDict): - file_search: FileSearchToolResource - web_search: WebSearchToolResource +class ToolResources(TypedDict, total=False): + file_search: Optional[FileSearchToolResource] + web_search: Optional[WebSearchToolResource] class Constraint(TypedDict): diff --git a/examples/studio/maestro/async_runs.py b/examples/studio/maestro/async_runs.py index 703f2bc0..48fc6d35 100644 --- a/examples/studio/maestro/async_runs.py +++ b/examples/studio/maestro/async_runs.py @@ -1,15 +1,15 @@ import asyncio from ai21 import AsyncAI21Client +from ai21.models.chat import ChatMessage client = AsyncAI21Client() async def main(): run_result = await client.maestro.runs.create_and_poll( - instruction="Who was the Maestro in the movie 'Maestro'?", - tools=[{"type": "web_search"}], - tool_resources={"web_search": {"urls": ["google.com"]}}, + messages=[ChatMessage(role="user", content="Analyze the text below and determine who's the best pokemon ever")], + context={"text": "Psyduck is the best pokemon."}, ) print(run_result) diff --git a/examples/studio/maestro/runs.py b/examples/studio/maestro/runs.py index eb80e152..f12bd6fc 100644 --- a/examples/studio/maestro/runs.py +++ b/examples/studio/maestro/runs.py @@ -1,13 +1,13 @@ from ai21 import AI21Client +from ai21.models.chat import ChatMessage client = AI21Client() def main(): run_result = client.maestro.runs.create_and_poll( - instruction="Who was the Maestro in the movie 'Maestro'?", - tools=[{"type": "web_search"}], - tool_resources={"web_search": {"urls": ["google.com"]}}, + messages=[ChatMessage(role="user", content="Analyze the text below and determine who's the best pokemon ever")], + context={"text": "Psyduck is the best pokemon."}, ) print(run_result) From ca3424a76e3dfd5487b40148a598b1f19027ace4 Mon Sep 17 00:00:00 2001 From: benshuk Date: Wed, 12 Mar 2025 15:42:32 +0200 Subject: [PATCH 07/21] fix: :fire: remove unsupported parameters --- ai21/clients/common/maestro/run.py | 59 ++++++++++---------- ai21/clients/studio/resources/maestro/run.py | 28 ++-------- ai21/models/maestro/run.py | 3 +- 3 files changed, 36 insertions(+), 54 deletions(-) diff --git a/ai21/clients/common/maestro/run.py b/ai21/clients/common/maestro/run.py index 17b1b6f2..c08870f0 100644 --- a/ai21/clients/common/maestro/run.py +++ b/ai21/clients/common/maestro/run.py @@ -11,17 +11,17 @@ from ai21.models.maestro.run import ( Tool, ToolResources, - Budget, RunResponse, DEFAULT_RUN_POLL_INTERVAL, DEFAULT_RUN_POLL_TIMEOUT, OutputType, + Requirement, ) from ai21.types import NOT_GIVEN, NotGiven from ai21.utils.typing import remove_not_given -def _primitive_to_schema(output_type): +def __primitive_to_schema(output_type): type_mapping = { int: "integer", float: "number", @@ -39,49 +39,54 @@ def _primitive_to_schema(output_type): return {"type": "array", "items": {"type": type_mapping[args[0]]}} -class BaseMaestroRun(ABC): - _module_name = "maestro/runs" - - def _output_type_to_json_schema(self, output_type: OutputType | NotGiven) -> Optional[Dict[str, Any]]: - if not output_type or isinstance(output_type, NotGiven): - return NOT_GIVEN +def __output_type_to_json_schema(output_type: OutputType | NotGiven) -> Optional[Dict[str, Any]]: + if not output_type or isinstance(output_type, NotGiven): + return NOT_GIVEN + + if inspect.isclass(output_type) and issubclass(output_type, BaseModel): + result = _to_schema(output_type) + else: + result = __primitive_to_schema(output_type) + + if not result: + raise ValueError( + "Unsupported output type. Supported types are:\n" + "- Pydantic model\n" + "- str\n" + "- int\n" + "- float\n" + "- bool\n" + "- List[str]\n" + "- List[int]\n" + "- List[float]\n" + "- Valid JSON schema dict" + ) - if inspect.isclass(output_type) and issubclass(output_type, BaseModel): - result = _to_schema(output_type) - else: - result = _primitive_to_schema(output_type) + return result - if not result: - raise ValueError( - "Unsupported type. Supported types are: primitives types, List of primitive types, json schema dict, " - "and pydantic models." - ) - return result +class BaseMaestroRun(ABC): + _module_name = "maestro/runs" def _create_body( self, *, messages: List[ChatMessage], - output_type: OutputType | NotGiven, models: List[str] | NotGiven, tools: List[Tool] | NotGiven, tool_resources: ToolResources | NotGiven, context: Dict[str, Any] | NotGiven, - requirements: List[str] | NotGiven, - budget: Budget | NotGiven, + requirements: List[Requirement] | NotGiven, **kwargs, ) -> dict: return remove_not_given( { "messages": [_to_dict(message) for message in messages], - "output_type": self._output_type_to_json_schema(output_type), "models": models, "tools": tools, "tool_resources": tool_resources, "context": context, "requirements": requirements, - "budget": budget, **kwargs, } ) @@ -91,13 +96,11 @@ def create( self, *, messages: List[ChatMessage], - output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, - requirements: List[str] | NotGiven = NOT_GIVEN, - budget: Budget | NotGiven = NOT_GIVEN, + requirements: List[Requirement] | NotGiven = NOT_GIVEN, **kwargs, ) -> RunResponse: pass @@ -115,13 +118,11 @@ def create_and_poll( self, *, messages: List[ChatMessage], - output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, - requirements: List[str] | NotGiven = NOT_GIVEN, - budget: Budget | NotGiven = NOT_GIVEN, + requirements: List[Requirement] | NotGiven = NOT_GIVEN, poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL, poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT, **kwargs, diff --git a/ai21/clients/studio/resources/maestro/run.py b/ai21/clients/studio/resources/maestro/run.py index 6a5638c8..860cc78e 100644 --- a/ai21/clients/studio/resources/maestro/run.py +++ b/ai21/clients/studio/resources/maestro/run.py @@ -10,14 +10,12 @@ from ai21.models.maestro.run import ( Tool, ToolResources, - Budget, RunResponse, TERMINATED_RUN_STATUSES, DEFAULT_RUN_POLL_INTERVAL, DEFAULT_RUN_POLL_TIMEOUT, - OutputType, + Requirement, ) - from ai21.types import NotGiven, NOT_GIVEN @@ -26,24 +24,20 @@ def create( self, *, messages: List[ChatMessage], - output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, - requirements: List[str] | NotGiven = NOT_GIVEN, - budget: Budget | NotGiven = NOT_GIVEN, + requirements: List[Requirement] | NotGiven = NOT_GIVEN, **kwargs, ) -> RunResponse: body = self._create_body( messages=messages, - output_type=output_type, models=models, tools=tools, tool_resources=tool_resources, context=context, requirements=requirements, - budget=budget, **kwargs, ) @@ -73,26 +67,22 @@ def create_and_poll( self, *, messages: List[ChatMessage], - output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, - requirements: List[str] | NotGiven = NOT_GIVEN, - budget: Budget | NotGiven = NOT_GIVEN, + requirements: List[Requirement] | NotGiven = NOT_GIVEN, poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL, poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT, **kwargs, ) -> RunResponse: run = self.create( messages=messages, - output_type=output_type, models=models, tools=tools, tool_resources=tool_resources, context=context, requirements=requirements, - budget=budget, **kwargs, ) @@ -104,24 +94,20 @@ async def create( self, *, messages: List[ChatMessage], - output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, - requirements: List[str] | NotGiven = NOT_GIVEN, - budget: Budget | NotGiven = NOT_GIVEN, + requirements: List[Requirement] | NotGiven = NOT_GIVEN, **kwargs, ) -> RunResponse: body = self._create_body( messages=messages, - output_type=output_type, models=models, tools=tools, tool_resources=tool_resources, context=context, requirements=requirements, - budget=budget, **kwargs, ) @@ -151,26 +137,22 @@ async def create_and_poll( self, *, messages: List[ChatMessage], - output_type: OutputType | NotGiven = NOT_GIVEN, models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, - requirements: List[str] | NotGiven = NOT_GIVEN, - budget: Budget | NotGiven = NOT_GIVEN, + requirements: List[Requirement] | NotGiven = NOT_GIVEN, poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL, poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT, **kwargs, ) -> RunResponse: run = await self.create( messages=messages, - output_type=output_type, models=models, tools=tools, tool_resources=tool_resources, context=context, requirements=requirements, - budget=budget, **kwargs, ) diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py index eb0b7476..cac7752c 100644 --- a/ai21/models/maestro/run.py +++ b/ai21/models/maestro/run.py @@ -40,10 +40,9 @@ class ToolResources(TypedDict, total=False): web_search: Optional[WebSearchToolResource] -class Constraint(TypedDict): +class Requirement(TypedDict): name: str description: str - is_mandatory: bool class RunResponse(AI21BaseModel): From 74bc78e2ce9b355a549f3f8ed931415d50f667ce Mon Sep 17 00:00:00 2001 From: benshuk Date: Wed, 12 Mar 2025 16:05:37 +0200 Subject: [PATCH 08/21] fix: :fire: remove unsupported parameters --- ai21/models/maestro/run.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py index cac7752c..e8defd34 100644 --- a/ai21/models/maestro/run.py +++ b/ai21/models/maestro/run.py @@ -32,7 +32,6 @@ class FileSearchToolResource(TypedDict, total=False): class WebSearchToolResource(TypedDict, total=False): urls: Optional[List[str]] - fallback_to_web: Optional[bool] class ToolResources(TypedDict, total=False): From a1c2f0ad81fb87c688633787c0482c48b2d26a18 Mon Sep 17 00:00:00 2001 From: benshuk Date: Wed, 12 Mar 2025 16:06:45 +0200 Subject: [PATCH 09/21] fix: :bug: try saving the day --- ai21/models/maestro/run.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py index e8defd34..8b158342 100644 --- a/ai21/models/maestro/run.py +++ b/ai21/models/maestro/run.py @@ -8,8 +8,7 @@ Role = Literal["user", "assistant"] RunStatus = Literal["completed", "failed", "in_progress", "requires_action"] ToolType = Literal["file_search", "web_search"] -PrimitiveType = Type[str] | Type[int] | Type[float] | Type[bool] | Type[List[str]] | Type[List[int]] | Type[List[float]] -OutputType = PrimitiveType | Type[BaseModel] | Dict[str, Any] +OutputType = Type[BaseModel] | Type[int] | Dict[str, Any] DEFAULT_RUN_POLL_INTERVAL: float = 1 # seconds DEFAULT_RUN_POLL_TIMEOUT: float = 120 # seconds From 8879087c794b3fec35559f739765055a7d943429 Mon Sep 17 00:00:00 2001 From: benshuk Date: Wed, 12 Mar 2025 16:09:12 +0200 Subject: [PATCH 10/21] fix: :bug: try saving the day --- ai21/models/maestro/run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py index 8b158342..133badb7 100644 --- a/ai21/models/maestro/run.py +++ b/ai21/models/maestro/run.py @@ -8,7 +8,7 @@ Role = Literal["user", "assistant"] RunStatus = Literal["completed", "failed", "in_progress", "requires_action"] ToolType = Literal["file_search", "web_search"] -OutputType = Type[BaseModel] | Type[int] | Dict[str, Any] +OutputType = Type[BaseModel] | Dict[str, Any] DEFAULT_RUN_POLL_INTERVAL: float = 1 # seconds DEFAULT_RUN_POLL_TIMEOUT: float = 120 # seconds From 50b09aa833b0c64a21a2a44555fd7919430cc178 Mon Sep 17 00:00:00 2001 From: benshuk Date: Wed, 12 Mar 2025 16:11:15 +0200 Subject: [PATCH 11/21] fix: :bug: try saving the day --- ai21/models/maestro/run.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py index 133badb7..72fa9599 100644 --- a/ai21/models/maestro/run.py +++ b/ai21/models/maestro/run.py @@ -1,4 +1,4 @@ -from typing import TypedDict, Literal, List, Optional, Any, Set, Dict, Type +from typing import TypedDict, Literal, List, Optional, Any, Set, Dict, Type, Union from pydantic import BaseModel @@ -8,7 +8,8 @@ Role = Literal["user", "assistant"] RunStatus = Literal["completed", "failed", "in_progress", "requires_action"] ToolType = Literal["file_search", "web_search"] -OutputType = Type[BaseModel] | Dict[str, Any] +PrimitiveTypes = Union[Type[str], Type[int], Type[float], Type[bool]] +OutputType = Union[Type[BaseModel], PrimitiveTypes, Dict[str, Any]] DEFAULT_RUN_POLL_INTERVAL: float = 1 # seconds DEFAULT_RUN_POLL_TIMEOUT: float = 120 # seconds From 87953d13337d88211b5e27c5f7afb1cb5c294a37 Mon Sep 17 00:00:00 2001 From: benshuk Date: Wed, 12 Mar 2025 16:13:52 +0200 Subject: [PATCH 12/21] fix: :bug: let's give it another go shall we --- ai21/models/maestro/run.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py index 72fa9599..110c90f1 100644 --- a/ai21/models/maestro/run.py +++ b/ai21/models/maestro/run.py @@ -9,6 +9,7 @@ RunStatus = Literal["completed", "failed", "in_progress", "requires_action"] ToolType = Literal["file_search", "web_search"] PrimitiveTypes = Union[Type[str], Type[int], Type[float], Type[bool]] +PrimitiveLists = Type[List[PrimitiveTypes]] OutputType = Union[Type[BaseModel], PrimitiveTypes, Dict[str, Any]] DEFAULT_RUN_POLL_INTERVAL: float = 1 # seconds From b0104f017653d136802068272d6ab6547b9f5341 Mon Sep 17 00:00:00 2001 From: benshuk Date: Thu, 13 Mar 2025 11:15:58 +0200 Subject: [PATCH 13/21] fix: :fire: remove unused functions --- ai21/clients/common/maestro/run.py | 52 ++---------------------------- 1 file changed, 2 insertions(+), 50 deletions(-) diff --git a/ai21/clients/common/maestro/run.py b/ai21/clients/common/maestro/run.py index c08870f0..4beec466 100644 --- a/ai21/clients/common/maestro/run.py +++ b/ai21/clients/common/maestro/run.py @@ -1,12 +1,9 @@ from __future__ import annotations -import inspect from abc import ABC, abstractmethod -from typing import List, Dict, Any, get_origin, get_args, Optional +from typing import List, Dict, Any -from pydantic import BaseModel - -from ai21.models._pydantic_compatibility import _to_dict, _to_schema +from ai21.models._pydantic_compatibility import _to_dict from ai21.models.chat import ChatMessage from ai21.models.maestro.run import ( Tool, @@ -14,57 +11,12 @@ RunResponse, DEFAULT_RUN_POLL_INTERVAL, DEFAULT_RUN_POLL_TIMEOUT, - OutputType, Requirement, ) from ai21.types import NOT_GIVEN, NotGiven from ai21.utils.typing import remove_not_given -def __primitive_to_schema(output_type): - type_mapping = { - int: "integer", - float: "number", - str: "string", - bool: "boolean", - } - - if output_type in type_mapping: - return {"type": type_mapping[output_type]} - - origin = get_origin(output_type) - args = get_args(output_type) - - if origin is list and len(args) == 1 and args[0] in type_mapping: - return {"type": "array", "items": {"type": type_mapping[args[0]]}} - - -def __output_type_to_json_schema(output_type: OutputType | NotGiven) -> Optional[Dict[str, Any]]: - if not output_type or isinstance(output_type, NotGiven): - return NOT_GIVEN - - if inspect.isclass(output_type) and issubclass(output_type, BaseModel): - result = _to_schema(output_type) - else: - result = __primitive_to_schema(output_type) - - if not result: - raise ValueError( - "Unsupported output type. Supported types are:\n" - "- Pydantic model\n" - "- str\n" - "- int\n" - "- float\n" - "- bool\n" - "- List[str]\n" - "- List[int]\n" - "- List[float]\n" - "- Valid JSON schema dict" - ) - - return result - - class BaseMaestroRun(ABC): _module_name = "maestro/runs" From 12151a6ef1d3cdd43e7152a823987c15b3dfe000 Mon Sep 17 00:00:00 2001 From: benshuk Date: Thu, 13 Mar 2025 13:27:42 +0200 Subject: [PATCH 14/21] test: :white_check_mark: tests --- tests/integration_tests/clients/test_studio.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/integration_tests/clients/test_studio.py b/tests/integration_tests/clients/test_studio.py index 1e10e10e..5dea6331 100644 --- a/tests/integration_tests/clients/test_studio.py +++ b/tests/integration_tests/clients/test_studio.py @@ -56,12 +56,16 @@ def test_studio(test_file_name: str): ("chat/async_stream_chat_completions.py",), ("conversational_rag/conversational_rag.py",), ("conversational_rag/async_conversational_rag.py",), + ("maestro/runs.py",), + ("maestro/async_runs.py",), ], ids=[ "when_chat_completions__should_return_ok", "when_stream_chat_completions__should_return_ok", "when_conversational_rag__should_return_ok", "when_async_conversational_rag__should_return_ok", + "when_maestro_runs__should_return_ok", + "when_maestro_async_runs__should_return_ok", ], ) async def test_async_studio(test_file_name: str): From 495ce5b44e065774eaf385c72c17774e3617d966 Mon Sep 17 00:00:00 2001 From: benshuk Date: Sun, 16 Mar 2025 11:56:48 +0200 Subject: [PATCH 15/21] refactor: :truck: rename messages and constraints --- ai21/clients/common/maestro/run.py | 9 ++++----- ai21/clients/studio/resources/maestro/run.py | 16 ++++++++-------- examples/studio/maestro/async_runs.py | 3 +-- examples/studio/maestro/runs.py | 3 +-- 4 files changed, 14 insertions(+), 17 deletions(-) diff --git a/ai21/clients/common/maestro/run.py b/ai21/clients/common/maestro/run.py index 4beec466..c3ef5783 100644 --- a/ai21/clients/common/maestro/run.py +++ b/ai21/clients/common/maestro/run.py @@ -3,7 +3,6 @@ from abc import ABC, abstractmethod from typing import List, Dict, Any -from ai21.models._pydantic_compatibility import _to_dict from ai21.models.chat import ChatMessage from ai21.models.maestro.run import ( Tool, @@ -23,7 +22,7 @@ class BaseMaestroRun(ABC): def _create_body( self, *, - messages: List[ChatMessage], + input: str | List[ChatMessage], models: List[str] | NotGiven, tools: List[Tool] | NotGiven, tool_resources: ToolResources | NotGiven, @@ -33,7 +32,7 @@ def _create_body( ) -> dict: return remove_not_given( { - "messages": [_to_dict(message) for message in messages], + "input": input, "models": models, "tools": tools, "tool_resources": tool_resources, @@ -47,7 +46,7 @@ def _create_body( def create( self, *, - messages: List[ChatMessage], + input: str | List[ChatMessage], models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -69,7 +68,7 @@ def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_timeout: f def create_and_poll( self, *, - messages: List[ChatMessage], + input: str | List[ChatMessage], models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, diff --git a/ai21/clients/studio/resources/maestro/run.py b/ai21/clients/studio/resources/maestro/run.py index 860cc78e..4e83f757 100644 --- a/ai21/clients/studio/resources/maestro/run.py +++ b/ai21/clients/studio/resources/maestro/run.py @@ -23,7 +23,7 @@ class MaestroRun(StudioResource, BaseMaestroRun): def create( self, *, - messages: List[ChatMessage], + input: str | List[ChatMessage], models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -32,7 +32,7 @@ def create( **kwargs, ) -> RunResponse: body = self._create_body( - messages=messages, + input=input, models=models, tools=tools, tool_resources=tool_resources, @@ -66,7 +66,7 @@ def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_timeout: f def create_and_poll( self, *, - messages: List[ChatMessage], + input: str | List[ChatMessage], models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -77,7 +77,7 @@ def create_and_poll( **kwargs, ) -> RunResponse: run = self.create( - messages=messages, + input=input, models=models, tools=tools, tool_resources=tool_resources, @@ -93,7 +93,7 @@ class AsyncMaestroRun(AsyncStudioResource, BaseMaestroRun): async def create( self, *, - messages: List[ChatMessage], + input: str | List[ChatMessage], models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -102,7 +102,7 @@ async def create( **kwargs, ) -> RunResponse: body = self._create_body( - messages=messages, + input=input, models=models, tools=tools, tool_resources=tool_resources, @@ -136,7 +136,7 @@ async def _poll_for_status(self, *, run_id: str, poll_interval: float, poll_time async def create_and_poll( self, *, - messages: List[ChatMessage], + input: str | List[ChatMessage], models: List[str] | NotGiven = NOT_GIVEN, tools: List[Tool] | NotGiven = NOT_GIVEN, tool_resources: ToolResources | NotGiven = NOT_GIVEN, @@ -147,7 +147,7 @@ async def create_and_poll( **kwargs, ) -> RunResponse: run = await self.create( - messages=messages, + input=input, models=models, tools=tools, tool_resources=tool_resources, diff --git a/examples/studio/maestro/async_runs.py b/examples/studio/maestro/async_runs.py index 48fc6d35..6104b8e1 100644 --- a/examples/studio/maestro/async_runs.py +++ b/examples/studio/maestro/async_runs.py @@ -1,14 +1,13 @@ import asyncio from ai21 import AsyncAI21Client -from ai21.models.chat import ChatMessage client = AsyncAI21Client() async def main(): run_result = await client.maestro.runs.create_and_poll( - messages=[ChatMessage(role="user", content="Analyze the text below and determine who's the best pokemon ever")], + input="Analyze the text below and determine who's the best pokemon ever", context={"text": "Psyduck is the best pokemon."}, ) diff --git a/examples/studio/maestro/runs.py b/examples/studio/maestro/runs.py index f12bd6fc..19a8051e 100644 --- a/examples/studio/maestro/runs.py +++ b/examples/studio/maestro/runs.py @@ -1,12 +1,11 @@ from ai21 import AI21Client -from ai21.models.chat import ChatMessage client = AI21Client() def main(): run_result = client.maestro.runs.create_and_poll( - messages=[ChatMessage(role="user", content="Analyze the text below and determine who's the best pokemon ever")], + input="Analyze the text below and determine who's the best pokemon ever", context={"text": "Psyduck is the best pokemon."}, ) From 02efd99d43cd22a52666fb22065ec39f709d4417 Mon Sep 17 00:00:00 2001 From: benshuk Date: Sun, 16 Mar 2025 17:29:42 +0200 Subject: [PATCH 16/21] ci: :technologist: add git hooks to check for in commit content --- .git-hooks/check_api_key.sh | 9 +++++++++ .pre-commit-config.yaml | 7 +++++++ examples/studio/maestro/async_runs.py | 2 +- 3 files changed, 17 insertions(+), 1 deletion(-) create mode 100755 .git-hooks/check_api_key.sh diff --git a/.git-hooks/check_api_key.sh b/.git-hooks/check_api_key.sh new file mode 100755 index 00000000..6bb06596 --- /dev/null +++ b/.git-hooks/check_api_key.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +# Check for `api_key=` in staged changes +if git diff --cached | grep -q "api_key="; then + echo "❌ Commit blocked: Found 'api_key=' in staged changes." + exit 1 # Prevent commit +fi + +exit 0 # Allow commit diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9c210e34..eece3b24 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -118,3 +118,10 @@ repos: entry: hadolint/hadolint:v2.10.0 hadolint types: - dockerfile + - repo: local + hooks: + - id: check-api-key + name: Check for API keys + entry: .git-hooks/check_api_key.sh + language: system + stages: [commit] diff --git a/examples/studio/maestro/async_runs.py b/examples/studio/maestro/async_runs.py index 6104b8e1..3d401608 100644 --- a/examples/studio/maestro/async_runs.py +++ b/examples/studio/maestro/async_runs.py @@ -6,7 +6,7 @@ async def main(): - run_result = await client.maestro.runs.create_and_poll( + run_result = await client.beta.maestro.runs.create_and_poll( input="Analyze the text below and determine who's the best pokemon ever", context={"text": "Psyduck is the best pokemon."}, ) From 2e7f9f3906bbfa331763738f4833cccf5b5c5278 Mon Sep 17 00:00:00 2001 From: benshuk Date: Sun, 16 Mar 2025 17:32:06 +0200 Subject: [PATCH 17/21] chore: :truck: move `maestro` under beta --- ai21/clients/studio/ai21_client.py | 2 -- ai21/clients/studio/async_ai21_client.py | 2 -- ai21/clients/studio/resources/beta/async_beta.py | 2 ++ ai21/clients/studio/resources/beta/beta.py | 2 ++ examples/studio/maestro/runs.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/ai21/clients/studio/ai21_client.py b/ai21/clients/studio/ai21_client.py index 48ee933f..34cdb491 100644 --- a/ai21/clients/studio/ai21_client.py +++ b/ai21/clients/studio/ai21_client.py @@ -5,7 +5,6 @@ from ai21.ai21_env_config import _AI21EnvConfig, AI21EnvConfig from ai21.clients.studio.client_url_parser import create_client_url from ai21.clients.studio.resources.beta.beta import Beta -from ai21.clients.studio.resources.maestro.maestro import Maestro from ai21.clients.studio.resources.studio_chat import StudioChat from ai21.clients.studio.resources.studio_library import StudioLibrary from ai21.http_client.http_client import AI21HTTPClient @@ -43,5 +42,4 @@ def __init__( ) self.chat: StudioChat = StudioChat(self) self.library = StudioLibrary(self) - self.maestro = Maestro(self) self.beta = Beta(self) diff --git a/ai21/clients/studio/async_ai21_client.py b/ai21/clients/studio/async_ai21_client.py index 0daf9563..4f7ed322 100644 --- a/ai21/clients/studio/async_ai21_client.py +++ b/ai21/clients/studio/async_ai21_client.py @@ -5,7 +5,6 @@ from ai21.ai21_env_config import _AI21EnvConfig, AI21EnvConfig from ai21.clients.studio.client_url_parser import create_client_url from ai21.clients.studio.resources.beta.async_beta import AsyncBeta -from ai21.clients.studio.resources.maestro.maestro import AsyncMaestro from ai21.clients.studio.resources.studio_chat import AsyncStudioChat from ai21.clients.studio.resources.studio_library import AsyncStudioLibrary from ai21.http_client.async_http_client import AsyncAI21HTTPClient @@ -42,5 +41,4 @@ def __init__( self.chat: AsyncStudioChat = AsyncStudioChat(self) self.library = AsyncStudioLibrary(self) - self.maestro = AsyncMaestro(self) self.beta = AsyncBeta(self) diff --git a/ai21/clients/studio/resources/beta/async_beta.py b/ai21/clients/studio/resources/beta/async_beta.py index a94ddc95..8c5bb35a 100644 --- a/ai21/clients/studio/resources/beta/async_beta.py +++ b/ai21/clients/studio/resources/beta/async_beta.py @@ -1,3 +1,4 @@ +from ai21.clients.studio.resources.maestro.maestro import AsyncMaestro from ai21.clients.studio.resources.studio_conversational_rag import AsyncStudioConversationalRag from ai21.clients.studio.resources.studio_resource import AsyncStudioResource from ai21.http_client.async_http_client import AsyncAI21HTTPClient @@ -8,3 +9,4 @@ def __init__(self, client: AsyncAI21HTTPClient): super().__init__(client) self.conversational_rag = AsyncStudioConversationalRag(client) + self.maestro = AsyncMaestro(client) diff --git a/ai21/clients/studio/resources/beta/beta.py b/ai21/clients/studio/resources/beta/beta.py index 1269f970..c33d3bc0 100644 --- a/ai21/clients/studio/resources/beta/beta.py +++ b/ai21/clients/studio/resources/beta/beta.py @@ -1,3 +1,4 @@ +from ai21.clients.studio.resources.maestro.maestro import Maestro from ai21.clients.studio.resources.studio_conversational_rag import StudioConversationalRag from ai21.clients.studio.resources.studio_resource import StudioResource from ai21.http_client.http_client import AI21HTTPClient @@ -8,3 +9,4 @@ def __init__(self, client: AI21HTTPClient): super().__init__(client) self.conversational_rag = StudioConversationalRag(client) + self.maestro = Maestro(client) diff --git a/examples/studio/maestro/runs.py b/examples/studio/maestro/runs.py index 19a8051e..c5f87006 100644 --- a/examples/studio/maestro/runs.py +++ b/examples/studio/maestro/runs.py @@ -4,7 +4,7 @@ def main(): - run_result = client.maestro.runs.create_and_poll( + run_result = client.beta.maestro.runs.create_and_poll( input="Analyze the text below and determine who's the best pokemon ever", context={"text": "Psyduck is the best pokemon."}, ) From df840314f332ef18dd571602b4e0fcbb7ce33c31 Mon Sep 17 00:00:00 2001 From: benshuk Date: Tue, 18 Mar 2025 16:52:48 +0200 Subject: [PATCH 18/21] docs: :memo: update examples --- examples/studio/maestro/async_runs.py | 13 +++++++++++-- examples/studio/maestro/runs.py | 13 +++++++++++-- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/examples/studio/maestro/async_runs.py b/examples/studio/maestro/async_runs.py index 3d401608..f1f00180 100644 --- a/examples/studio/maestro/async_runs.py +++ b/examples/studio/maestro/async_runs.py @@ -7,8 +7,17 @@ async def main(): run_result = await client.beta.maestro.runs.create_and_poll( - input="Analyze the text below and determine who's the best pokemon ever", - context={"text": "Psyduck is the best pokemon."}, + input="Write a poem about the ocean", + requirements=[ + { + "name": "length requirement", + "description": "The length of the poem should be less than 1000 characters", + }, + { + "name": "rhyme requirement", + "description": "The poem should rhyme", + }, + ], ) print(run_result) diff --git a/examples/studio/maestro/runs.py b/examples/studio/maestro/runs.py index c5f87006..72024f33 100644 --- a/examples/studio/maestro/runs.py +++ b/examples/studio/maestro/runs.py @@ -5,8 +5,17 @@ def main(): run_result = client.beta.maestro.runs.create_and_poll( - input="Analyze the text below and determine who's the best pokemon ever", - context={"text": "Psyduck is the best pokemon."}, + input="Write a poem about the ocean", + requirements=[ + { + "name": "length requirement", + "description": "The length of the poem should be less than 1000 characters", + }, + { + "name": "rhyme requirement", + "description": "The poem should rhyme", + }, + ], ) print(run_result) From c124fa0f5fc15a049b3aa3a51c8fcb592baca046 Mon Sep 17 00:00:00 2001 From: benshuk Date: Wed, 19 Mar 2025 11:00:17 +0200 Subject: [PATCH 19/21] docs: :memo: update README --- README.md | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 382adad3..8988764e 100644 --- a/README.md +++ b/README.md @@ -257,9 +257,18 @@ from ai21 import AI21Client client = AI21Client() -run_result = client.maestro.runs.create_and_poll( - messages=[{"role": "user", "content": "Who was the Maestro in the movie 'The Maestro'?"}], - tools=[{"type": "web_search"}], +run_result = client.beta.maestro.runs.create_and_poll( + input="Write a poem about the ocean", + requirements=[ + { + "name": "length requirement", + "description": "The length of the poem should be less than 1000 characters", + }, + { + "name": "rhyme requirement", + "description": "The poem should rhyme", + }, + ], ) ``` From ccacd97605949a5ffbb0678b589e105fb077e7b1 Mon Sep 17 00:00:00 2001 From: benshuk Date: Wed, 19 Mar 2025 11:01:33 +0200 Subject: [PATCH 20/21] refactor: :truck: rename maestro runs examples --- README.md | 2 +- examples/studio/maestro/{async_runs.py => async_run.py} | 0 examples/studio/maestro/{runs.py => run.py} | 0 tests/integration_tests/clients/test_studio.py | 4 ++-- 4 files changed, 3 insertions(+), 3 deletions(-) rename examples/studio/maestro/{async_runs.py => async_run.py} (100%) rename examples/studio/maestro/{runs.py => run.py} (100%) diff --git a/README.md b/README.md index 8988764e..646fc3b5 100644 --- a/README.md +++ b/README.md @@ -272,7 +272,7 @@ run_result = client.beta.maestro.runs.create_and_poll( ) ``` -For a more detailed example, see maestro [sync](examples/studio/maestro/runs.py) and [async](examples/studio/maestro/async_runs.py) examples. +For a more detailed example, see maestro [sync](examples/studio/maestro/run.py) and [async](examples/studio/maestro/async_run.py) examples. --- diff --git a/examples/studio/maestro/async_runs.py b/examples/studio/maestro/async_run.py similarity index 100% rename from examples/studio/maestro/async_runs.py rename to examples/studio/maestro/async_run.py diff --git a/examples/studio/maestro/runs.py b/examples/studio/maestro/run.py similarity index 100% rename from examples/studio/maestro/runs.py rename to examples/studio/maestro/run.py diff --git a/tests/integration_tests/clients/test_studio.py b/tests/integration_tests/clients/test_studio.py index 5dea6331..a23d6e07 100644 --- a/tests/integration_tests/clients/test_studio.py +++ b/tests/integration_tests/clients/test_studio.py @@ -56,8 +56,8 @@ def test_studio(test_file_name: str): ("chat/async_stream_chat_completions.py",), ("conversational_rag/conversational_rag.py",), ("conversational_rag/async_conversational_rag.py",), - ("maestro/runs.py",), - ("maestro/async_runs.py",), + ("maestro/run.py",), + ("maestro/async_run.py",), ], ids=[ "when_chat_completions__should_return_ok", From f3604fd69670444247fc6e61bd1aea087698ab4b Mon Sep 17 00:00:00 2001 From: benshuk Date: Wed, 19 Mar 2025 11:04:25 +0200 Subject: [PATCH 21/21] chore: :wrench: add budget support --- ai21/clients/common/maestro/run.py | 5 +++++ ai21/clients/studio/resources/maestro/run.py | 9 +++++++++ 2 files changed, 14 insertions(+) diff --git a/ai21/clients/common/maestro/run.py b/ai21/clients/common/maestro/run.py index c3ef5783..43851dec 100644 --- a/ai21/clients/common/maestro/run.py +++ b/ai21/clients/common/maestro/run.py @@ -11,6 +11,7 @@ DEFAULT_RUN_POLL_INTERVAL, DEFAULT_RUN_POLL_TIMEOUT, Requirement, + Budget, ) from ai21.types import NOT_GIVEN, NotGiven from ai21.utils.typing import remove_not_given @@ -28,6 +29,7 @@ def _create_body( tool_resources: ToolResources | NotGiven, context: Dict[str, Any] | NotGiven, requirements: List[Requirement] | NotGiven, + budget: Budget | NotGiven, **kwargs, ) -> dict: return remove_not_given( @@ -38,6 +40,7 @@ def _create_body( "tool_resources": tool_resources, "context": context, "requirements": requirements, + "budget": budget, **kwargs, } ) @@ -52,6 +55,7 @@ def create( tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, requirements: List[Requirement] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, **kwargs, ) -> RunResponse: pass @@ -74,6 +78,7 @@ def create_and_poll( tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, requirements: List[Requirement] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL, poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT, **kwargs, diff --git a/ai21/clients/studio/resources/maestro/run.py b/ai21/clients/studio/resources/maestro/run.py index 4e83f757..5166b714 100644 --- a/ai21/clients/studio/resources/maestro/run.py +++ b/ai21/clients/studio/resources/maestro/run.py @@ -15,6 +15,7 @@ DEFAULT_RUN_POLL_INTERVAL, DEFAULT_RUN_POLL_TIMEOUT, Requirement, + Budget, ) from ai21.types import NotGiven, NOT_GIVEN @@ -29,6 +30,7 @@ def create( tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, requirements: List[Requirement] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, **kwargs, ) -> RunResponse: body = self._create_body( @@ -38,6 +40,7 @@ def create( tool_resources=tool_resources, context=context, requirements=requirements, + budget=budget, **kwargs, ) @@ -72,6 +75,7 @@ def create_and_poll( tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, requirements: List[Requirement] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL, poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT, **kwargs, @@ -83,6 +87,7 @@ def create_and_poll( tool_resources=tool_resources, context=context, requirements=requirements, + budget=budget, **kwargs, ) @@ -99,6 +104,7 @@ async def create( tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, requirements: List[Requirement] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, **kwargs, ) -> RunResponse: body = self._create_body( @@ -108,6 +114,7 @@ async def create( tool_resources=tool_resources, context=context, requirements=requirements, + budget=budget, **kwargs, ) @@ -142,6 +149,7 @@ async def create_and_poll( tool_resources: ToolResources | NotGiven = NOT_GIVEN, context: Dict[str, Any] | NotGiven = NOT_GIVEN, requirements: List[Requirement] | NotGiven = NOT_GIVEN, + budget: Budget | NotGiven = NOT_GIVEN, poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL, poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT, **kwargs, @@ -153,6 +161,7 @@ async def create_and_poll( tool_resources=tool_resources, context=context, requirements=requirements, + budget=budget, **kwargs, )