From 562fa80fbde2be57b3cb2bb25013534081f42a3c Mon Sep 17 00:00:00 2001
From: benshuk <bens@ai21.com>
Date: Sun, 25 May 2025 12:02:58 +0300
Subject: [PATCH 1/3] feat: :sparkles: Enhance Requirement and Output Options
 in Maestro

---
 ai21/models/maestro/run.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py
index c60b1641..15a8c9e5 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"]
-OutputOptions = Literal["data_sources"]
+OutputOptions = Literal["data_sources", "requirements_result"]
 PrimitiveTypes = Union[Type[str], Type[int], Type[float], Type[bool]]
 PrimitiveLists = Type[List[PrimitiveTypes]]
 OutputType = Union[Type[BaseModel], PrimitiveTypes, Dict[str, Any]]
@@ -44,6 +44,18 @@ class ToolResources(TypedDict, total=False):
 class Requirement(TypedDict, total=False):
     name: str
     description: str
+    is_mandatory: Optional[bool] = False
+
+
+class RequirementResultItem(Requirement, total=False):
+    score: float
+    reason: Optional[str] = None
+
+
+class RequirementsResult(TypedDict, total=False):
+    score: float
+    finish_reason: str
+    requirements: list[RequirementResultItem]
 
 
 class FileSearchResult(TypedDict, total=False):
@@ -70,3 +82,4 @@ class RunResponse(AI21BaseModel):
     status: RunStatus
     result: Any
     data_sources: Optional[DataSources] = None
+    requirements_result: Optional[RequirementsResult] = None

From f22acd3dc62bf5c1c693fc3e74d1f37ad4a1bbe7 Mon Sep 17 00:00:00 2001
From: benshuk <bens@ai21.com>
Date: Sun, 25 May 2025 12:05:55 +0300
Subject: [PATCH 2/3] chore: :wrench: Include Requirements Result in Maestro
 Async and Sync Runs examples

---
 examples/studio/maestro/async_run.py | 1 +
 examples/studio/maestro/run.py       | 1 +
 2 files changed, 2 insertions(+)

diff --git a/examples/studio/maestro/async_run.py b/examples/studio/maestro/async_run.py
index f1f00180..f96f220a 100644
--- a/examples/studio/maestro/async_run.py
+++ b/examples/studio/maestro/async_run.py
@@ -18,6 +18,7 @@ async def main():
                 "description": "The poem should rhyme",
             },
         ],
+        include=["requirements_result"],
     )
 
     print(run_result)
diff --git a/examples/studio/maestro/run.py b/examples/studio/maestro/run.py
index 72024f33..7d6fa5b5 100644
--- a/examples/studio/maestro/run.py
+++ b/examples/studio/maestro/run.py
@@ -16,6 +16,7 @@ def main():
                 "description": "The poem should rhyme",
             },
         ],
+        include=["requirements_result"],
     )
 
     print(run_result)

From 816be14488061bd3cba59b8167938ca026604712 Mon Sep 17 00:00:00 2001
From: benshuk <bens@ai21.com>
Date: Sun, 25 May 2025 13:06:22 +0300
Subject: [PATCH 3/3] fix: :label: requirements typings

---
 ai21/models/maestro/run.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py
index 15a8c9e5..4ef078dd 100644
--- a/ai21/models/maestro/run.py
+++ b/ai21/models/maestro/run.py
@@ -44,7 +44,7 @@ class ToolResources(TypedDict, total=False):
 class Requirement(TypedDict, total=False):
     name: str
     description: str
-    is_mandatory: Optional[bool] = False
+    is_mandatory: bool = False
 
 
 class RequirementResultItem(Requirement, total=False):
@@ -55,7 +55,7 @@ class RequirementResultItem(Requirement, total=False):
 class RequirementsResult(TypedDict, total=False):
     score: float
     finish_reason: str
-    requirements: list[RequirementResultItem]
+    requirements: List[RequirementResultItem]
 
 
 class FileSearchResult(TypedDict, total=False):