feat: review

Author: MikeSquall

src/resources/maestro/runs.ts

@@ -4,6 +4,7 @@ import { MaestroRunRequest, MaestroRunResponse, RequestOptions, MaestroRunReques
 import { DEFAULT_INTERVAL, DEFAULT_TIMEOUT } from '../../Constants';
 
 const MAESTRO_PATH = '/maestro/runs';
+const MAESTRO_TERMINATED_RUN_STATUSES = ['completed', 'failed', 'requires_action'];
 
 export class Runs extends APIResource {
   async create(body: MaestroRunRequest): Promise<MaestroRunResponse> {
@@ -31,7 +32,7 @@ export class Runs extends APIResource {
 
     while (Date.now() - startTime < timeout) {
       const response = await this.get(runId);
-      if (response.status === 'completed') {
+      if (MAESTRO_TERMINATED_RUN_STATUSES.includes(response.status)) {
         return response;
       }
 

src/types/maestro/MaestroRunRequest.ts

@@ -15,31 +15,6 @@ type MaestroRunRequirement = {
 
 type MaestroRunTool = HttpTool | MCPTool | FileSearchTool | WebSearchTool;
 
-export type MaestroToolResources = {
-  /* 
-  When provided, this object defines filters that AI21 Maestro will apply whenever it performs a file search.
-  */
-  file_search?: {
-    /* 
-    Restrict file search to these file IDs.
-    */
-    file_ids?: string[];
-    /* 
-    Restrict file search to files with these labels.
-    */
-    labels?: string[];
-  };
-  /* 
-  When provided, this object defines filters that AI21 Maestro will apply whenever it performs a web search.
-  */
-  web_search?: {
-    /* 
-    Restrict web search to the specified URL prefixes.
-    */
-    urls: string[];
-  };
-};
-
 type MaestroRunFirstPartyModel = 'jamba-mini' | 'jamba-large';
 
 type MaestroRunManagedThirdPartyModel = string;
@@ -50,7 +25,7 @@ type MaestroRunBudget = 'low' | 'medium' | 'high';
 
 type MaestroRunIncludeFields = 'data_sources' | 'requirements_result' | string;
 
-type MaestroRunResponseLanguage =
+export type MaestroRunResponseLanguage =
   | 'arabic'
   | 'dutch'
   | 'english'
@@ -60,6 +35,7 @@ type MaestroRunResponseLanguage =
   | 'italian'
   | 'portuguese'
   | 'spanish'
+  | 'unset'
   | string;
 
 export interface MaestroRunRequest {
@@ -76,10 +52,6 @@ export interface MaestroRunRequest {
   */
   tools?: MaestroRunTool[];
   /* 
-  A set of resources used by AI21 Maestro’s tools.
-  */
-  tool_resources?: MaestroToolResources;
-  /* 
   The models to use for the maestro run.
   */
   models?: MaestroRunModel[];

src/types/maestro/MaestroRunResponse.ts

@@ -1,6 +1,6 @@
-import { MaestroToolResources } from './MaestroRunRequest';
+import { FileSearchTool, WebSearchTool } from './MaestroTools';
 
-type MaestroRunResponseStatus = 'completed' | 'failed' | 'in_progress';
+type MaestroRunResponseStatus = 'completed' | 'failed' | 'in_progress' | 'requires_action';
 
 type MaestroRunRequirementResult = {
   score: number;
@@ -13,6 +13,11 @@ type MaestroRunRequirementResult = {
   };
 };
 
+type DataSources = {
+  file_search?: FileSearchTool;
+  web_search?: WebSearchTool;
+};
+
 export interface MaestroRunResponse {
   /* 
   The ID of the maestro run for polling the status.
@@ -26,7 +31,7 @@ export interface MaestroRunResponse {
   /* 
   Specifies the data sources used to retrieve contextual information for a run.
   */
-  data_sources: MaestroToolResources;
+  data_sources: DataSources;
   /* 
   Detailed results for each requirement.
   */

src/types/maestro/MaestroTools.ts

@@ -37,8 +37,13 @@ export interface MCPTool {
 
 export interface FileSearchTool {
   type: 'file_search';
-  file_ids?: string[];
+  retrieval_similarity_threshold?: number;
   labels?: string[];
+  labels_filter_mode?: 'AND' | 'OR';
+  labels_filter?: Record<string, unknown>;
+  file_ids?: string[];
+  retrieval_strategy?: string;
+  max_neighbors?: number;
 }
 
 export interface WebSearchTool {

tests/unittests/resources/maestro/runs.test.ts

@@ -94,15 +94,6 @@ describe('Maestro Runs', () => {
             allowed_tools: ['get_weather'],
           },
         ],
-        tool_resources: {
-          file_search: {
-            file_ids: ['file_123', 'file_456'],
-            labels: ['ai', 'ml', 'nlp'],
-          },
-          web_search: {
-            urls: ['https://arxiv.org', 'https://openai.com'],
-          },
-        },
         models: ['jamba-large', 'gpt-4o'],
         budget: 'high',
         include: ['data_sources', 'requirements_result'],
@@ -115,10 +106,12 @@ describe('Maestro Runs', () => {
         result: null,
         data_sources: {
           file_search: {
+            type: 'file_search',
             file_ids: ['file_123', 'file_456'],
             labels: ['ai', 'ml', 'nlp'],
           },
           web_search: {
+            type: 'web_search',
             urls: ['https://arxiv.org', 'https://openai.com'],
           },
         },
@@ -165,6 +158,7 @@ describe('Maestro Runs', () => {
         result: 'Here is a comprehensive summary of the latest AI developments...',
         data_sources: {
           web_search: {
+            type: 'web_search',
             urls: ['https://arxiv.org', 'https://openai.com'],
           },
         },