diff --git a/ai21/models/maestro/run.py b/ai21/models/maestro/run.py index c60b1641..4ef078dd 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: 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 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)