fix: 🔧 introduce MaestroMessage (#298)

* refactor: reorganize imports by moving maestro models to their own module

* refactor: update input type in maestro run methods to use MaestroMessage instead of ChatMessage

* test: add parameterized tests for Maestro input formats in maestro.py

* feat: add async example for multi-message processing

* feat: update async example to handle multi-message input for business proposal

* feat: include requirements result in README example for Maestro runs

* fix: update content formatting in async_run_multi_messages.py for improved readability

---------

Co-authored-by: benshuk <[email protected]>
Co-authored-by: Asaf Joseph Gardin <[email protected]>

Author: 3 people

README.md

@@ -269,6 +269,7 @@ run_result = client.beta.maestro.runs.create_and_poll(
             "description": "The poem should rhyme",
         },
     ],
+    include=["requirements_result"]
 )
 ```
 

ai21/clients/common/maestro/run.py

@@ -3,16 +3,16 @@
 from abc import ABC, abstractmethod
 from typing import List
 
-from ai21.models.chat import ChatMessage
 from ai21.models.maestro.run import (
-    Tool,
-    ToolResources,
-    RunResponse,
+    Budget,
     DEFAULT_RUN_POLL_INTERVAL,
     DEFAULT_RUN_POLL_TIMEOUT,
-    Requirement,
-    Budget,
+    MaestroMessage,
     OutputOptions,
+    Requirement,
+    RunResponse,
+    Tool,
+    ToolResources,
 )
 from ai21.types import NOT_GIVEN, NotGiven
 from ai21.utils.typing import remove_not_given
@@ -24,7 +24,7 @@ class BaseMaestroRun(ABC):
     def _create_body(
         self,
         *,
-        input: str | List[ChatMessage],
+        input: str | List[MaestroMessage],
         models: List[str] | NotGiven,
         tools: List[Tool] | NotGiven,
         tool_resources: ToolResources | NotGiven,
@@ -33,14 +33,9 @@ def _create_body(
         include: List[OutputOptions] | NotGiven,
         **kwargs,
     ) -> dict:
-        if isinstance(input, list):
-            _input = [{"role": message.role, "content": message.content} for message in input]
-        else:
-            _input = input
-
         return remove_not_given(
             {
-                "input": _input,
+                "input": input,
                 "models": models,
                 "tools": tools,
                 "tool_resources": tool_resources,
@@ -55,7 +50,7 @@ def _create_body(
     def create(
         self,
         *,
-        input: str | List[ChatMessage],
+        input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
         tools: List[Tool] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
@@ -78,7 +73,7 @@ def poll_for_status(self, *, run_id: str, poll_interval_sec: float, poll_timeout
     def create_and_poll(
         self,
         *,
-        input: str | List[ChatMessage],
+        input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
         tools: List[Tool] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,

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

@@ -6,17 +6,17 @@
 
 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.run import (
-    Tool,
-    ToolResources,
-    RunResponse,
-    TERMINATED_RUN_STATUSES,
+    Budget,
     DEFAULT_RUN_POLL_INTERVAL,
     DEFAULT_RUN_POLL_TIMEOUT,
-    Requirement,
-    Budget,
+    MaestroMessage,
     OutputOptions,
+    Requirement,
+    RunResponse,
+    TERMINATED_RUN_STATUSES,
+    Tool,
+    ToolResources,
 )
 from ai21.types import NotGiven, NOT_GIVEN
 
@@ -25,7 +25,7 @@ class MaestroRun(StudioResource, BaseMaestroRun):
     def create(
         self,
         *,
-        input: str | List[ChatMessage],
+        input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
         tools: List[Tool] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
@@ -70,7 +70,7 @@ def poll_for_status(self, *, run_id: str, poll_interval_sec: float, poll_timeout
     def create_and_poll(
         self,
         *,
-        input: str | List[ChatMessage],
+        input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
         tools: List[Tool] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
@@ -101,7 +101,7 @@ class AsyncMaestroRun(AsyncStudioResource, BaseMaestroRun):
     async def create(
         self,
         *,
-        input: str | List[ChatMessage],
+        input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
         tools: List[Tool] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
@@ -146,7 +146,7 @@ async def poll_for_status(self, *, run_id: str, poll_interval_sec: float, poll_t
     async def create_and_poll(
         self,
         *,
-        input: str | List[ChatMessage],
+        input: str | List[MaestroMessage],
         models: List[str] | NotGiven = NOT_GIVEN,
         tools: List[Tool] | NotGiven = NOT_GIVEN,
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,

ai21/models/__init__.py

@@ -8,16 +8,6 @@
     ConversationalRagSource,
 )
 from ai21.models.responses.file_response import FileResponse
-from ai21.models.maestro.run import (
-    Requirement,
-    Budget,
-    Tool,
-    ToolResources,
-    DataSources,
-    FileSearchResult,
-    WebSearchResult,
-    OutputOptions,
-)
 
 __all__ = [
     "ChatMessage",
@@ -30,12 +20,4 @@
     "FileResponse",
     "ConversationalRagResponse",
     "ConversationalRagSource",
-    "Requirement",
-    "Budget",
-    "Tool",
-    "ToolResources",
-    "DataSources",
-    "FileSearchResult",
-    "WebSearchResult",
-    "OutputOptions",
 ]

ai21/models/maestro/__init__.py

@@ -0,0 +1,23 @@
+from ai21.models.maestro.run import (
+    Budget,
+    DataSources,
+    FileSearchResult,
+    MaestroMessage,
+    OutputOptions,
+    Requirement,
+    Tool,
+    ToolResources,
+    WebSearchResult,
+)
+
+__all__ = [
+    "Budget",
+    "DataSources",
+    "FileSearchResult",
+    "MaestroMessage",
+    "OutputOptions",
+    "Requirement",
+    "Tool",
+    "ToolResources",
+    "WebSearchResult",
+]

ai21/models/maestro/run.py

@@ -18,6 +18,11 @@
 TERMINATED_RUN_STATUSES: Set[RunStatus] = {"completed", "failed", "requires_action"}
 
 
+class MaestroMessage(TypedDict):
+    role: Role
+    content: str
+
+
 class Tool(TypedDict):
     type: ToolType
 

examples/studio/maestro/async_run_multi_messages.py

@@ -0,0 +1,43 @@
+import asyncio
+
+from ai21 import AsyncAI21Client
+
+client = AsyncAI21Client()
+
+
+async def main():
+    try:
+        run_result = await client.beta.maestro.runs.create_and_poll(
+            input=[
+                {
+                    "role": "user",
+                    "content": "Write me a memo to my boss suggesting we should pivot to selling parrots. "
+                    "We now have a window cleaning business for high rises",
+                },
+                {
+                    "role": "assistant",
+                    "content": "**Subject: Proposal to Pivot Business Focus**\n"
+                    "**Section 1: Current Market Challenges**\n"
+                    "High-rise window cleaning faces increasing competition.\n"
+                    "Profit margins are shrinking due to rising operational costs.\n"
+                    "Weather conditions often disrupt our service schedule.\n"
+                    "**Section 2: New Opportunity in Parrot Sales**\n"
+                    "Parrot sales are booming with pet ownership on the rise.\n"
+                    "Parrots require minimal space and resources compared to current operations.\n"
+                    "Our existing logistics can support this new venture.\n"
+                    "**Section 3: Strategic Advantages**\n"
+                    "Parrots offer higher profit margins than window cleaning.\n"
+                    "Selling parrots expands our customer base.\n"
+                    "This pivot leverages our customer service expertise",
+                },
+                {"role": "user", "content": "great, now beef up each section with a few more sentences"},
+            ]
+        )
+
+        print(run_result)
+    except TimeoutError:
+        print("The run timed out")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

tests/integration_tests/clients/studio/test_maestro.py

@@ -1,16 +1,49 @@
 import pytest
 
 from ai21 import AsyncAI21Client
+from ai21.models.maestro import MaestroMessage
 
 
 @pytest.mark.asyncio
 async def test_maestro__when_upload__should_return_data_sources():  # file_in_library: str):
     client = AsyncAI21Client()
-    result = await client.beta.maestro.runs.create_and_poll(
-        input="When did Einstein receive a Nobel Prize?", tools=[{"type": "file_search"}], include=["data_sources"]
+    run = await client.beta.maestro.runs.create_and_poll(
+        input="When did Einstein receive a Nobel Prize?",
+        tools=[{"type": "file_search"}],
+        include=["data_sources"],
+        poll_timeout_sec=200,
     )
-    assert result.status == "completed", "Expected 'completed' status"
-    assert result.result, "Expected a non-empty answer"
-    assert result.data_sources, "Expected data sources"
-    assert len(result.data_sources["file_search"]) > 0, "Expected at least one file search data source"
-    assert result.data_sources.get("web_search") is None, "Expected no web search data sources"
+    assert run.status == "completed", f"[RUN {run.id}] Expected 'completed' status"
+    assert run.result, f"[RUN {run.id}] Expected a non-empty answer"
+    assert run.data_sources, f"[RUN {run.id}] Expected data sources"
+    assert len(run.data_sources["file_search"]) > 0, f"[RUN {run.id}] Expected at least one file search data source"
+    assert run.data_sources.get("web_search") is None, f"[RUN {run.id}] Expected no web search data sources"
+
+
+@pytest.mark.parametrize(
+    "input_data,test_description",
+    [
+        ("What is the capital of France?", "string input"),
+        (
+            [MaestroMessage(role="user", content="What is the capital of France?")],
+            "Use MaestroMessage format in a list",
+        ),
+        (
+            [
+                {"role": "user", "content": "I need help with geography."},
+                {"role": "assistant", "content": "I'd be happy to help with geography questions."},
+                {"role": "user", "content": "What is the capital of France?"},
+            ],
+            "multi-message conversation input",
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_maestro__input_formats__should_accept_string_and_list(input_data, test_description):
+    """Test that input can be passed as both string and list of dictionaries."""
+    client = AsyncAI21Client()
+
+    run = await client.beta.maestro.runs.create_and_poll(input=input_data, poll_timeout_sec=200)
+
+    assert run.status == "completed", f"[RUN {run.id}] Expected 'completed' status for {test_description}"
+    assert run.result, f"[RUN {run.id}] Expected a non-empty answer for {test_description}"