feat:Expand AI21 OpenAPI: Assistants/MCP, RAG, tool refactor, retrieval #193
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
WalkthroughExpanded Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Client
participant API as Studio API
participant AssistSvc as Assistants Service
participant Exec as Execution Engine
rect rgba(230,245,255,0.5)
note over Client,API: Create Assistant
Client->>API: POST /studio/v1/assistants (CreateAssistantRequest)
API->>AssistSvc: CreateAssistantRequest
AssistSvc-->>API: Assistant
API-->>Client: Assistant
end
rect rgba(240,255,230,0.5)
note over Client,API: Modify Assistant
Client->>API: PATCH /studio/v1/assistants/{id} (ModifyAssistantRequest)
API->>AssistSvc: ModifyAssistantRequest
AssistSvc-->>API: Assistant
API-->>Client: Assistant
end
rect rgba(255,245,230,0.5)
note over Client,Exec: Run Assistant
Client->>API: POST /studio/v1/assistants/{id}/run (RunAssistantRequest)
API->>Exec: RunAssistantRequest (tools, optimization, resources)
Exec-->>API: Run result
API-->>Client: Run result
end
sequenceDiagram
autonumber
actor Client
participant API as Studio API
participant MCP as MCP Discovery
rect rgba(230,245,255,0.5)
note over Client,MCP: MCP Tool Discovery
Client->>API: POST /studio/v1/mcp/discover (MCPDefinition)
API->>MCP: Discover(MCPDefinition)
MCP-->>API: MCPDiscoveryResponse (tools/resources)
API-->>Client: MCPDiscoveryResponse
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 inconclusive)
✅ Passed checks (2 passed)
✨ Finishing touches🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
coderabbitai
bot
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/libs/AI21/openapi.yaml (2)
2216-2237: Enum/resource mismatch:AssistantTooladdsragbut no matching resource.AssistantToolResource lacks a rag field; either add a RagToolResource (and wire it), or remove rag until supported.
AssistantToolResource: title: AssistantToolResource type: object properties: file_search: title: File Search type: object web_search: title: Web Search type: object plan_approval: title: Plan Approval type: object + rag: + title: RAG + type: object + description: 'RAG tool configuration (define RagToolResource schema and $ref it here).'
754-791: Duplicate operationId values detected — fix requiredThe operationId values "v1_library_management" and "v1_secret_storage" appear multiple times in src/libs/AI21/openapi.yaml; duplicate operationIds will cause client codegen collisions.
- Rename these operationId entries to globally unique identifiers (recommend including resource/path and verb, e.g. library.list_demos, demos.create).
- Re-run a duplicate check to confirm no remaining collisions.
🧹 Nitpick comments (11)
src/libs/AI21/openapi.yaml (11)
1018-1043: Add 404 and stronger ID typing for vector store GET.The path returns 200/422 only and treats vector_store_id as an untyped string. Recommend adding 404 for missing stores and format: uuid (if applicable).
'/studio/v1/demos/regulations/vector-store/{vector_store_id}': get: summary: Get Vector Store description: Get a vector store by ID. operationId: get_vector_store_studio_v1_demos_regulations_vector_store__vector_store_id__get parameters: - name: vector_store_id in: path required: true schema: - title: Vector Store Id - type: string + title: Vector Store Id + type: string + format: uuid responses: '200': description: Successful Response content: application/json: schema: title: Response Get Vector Store Studio V1 Demos Regulations Vector Store Vector Store Id Get type: object + '404': + description: Vector store not found '422': description: Validation Error
1520-1541: Use 201 Created for resource creation and consider 409.POST /assistants should return 201 with Location; add 409 for duplicate name (if enforced).
/studio/v1/assistants: post: summary: Create Assistant operationId: v1_create_assistant requestBody: content: application/json: schema: $ref: '#/components/schemas/CreateAssistantRequest' required: true responses: - '200': - description: Successful Response + '201': + description: Assistant created content: application/json: schema: $ref: '#/components/schemas/Assistant' + '409': + description: Conflict (e.g., name already exists) '422': description: Validation Error
1590-1617: Align read/write schema foroptimization.ModifyAssistantRequest.optimization uses RunOptimization enum, but Assistant.properties.optimization is a plain string. Make Assistant.optimization also reference RunOptimization (or vice versa).
Outside this hunk, update Assistant.properties:
- optimization: - title: Optimization - type: string + optimization: + title: Optimization + allOf: + - $ref: '#/components/schemas/RunOptimization'
1619-1647: Consider 202 Accepted for async runs.If assistant runs are asynchronous (MaestroRunResult with in_progress status), prefer 202 and optionally expose a polling URL via Location.
Please confirm whether runs start async; if so, we can update the response codes accordingly.
2301-2304: Description copy-paste error.
vector_store_idin Upload Check Compliance mentions “RFI processing”. Update to “regulations compliance check”.vector_store_id: title: Vector Store Id type: string - description: Vector store ID to use for RFI processing + description: Vector store ID to use for regulations compliance checks
2719-2725: Boundaries formax_documentsand RetrievalStrategy default.Add reasonable min/max to prevent heavy queries; keep default segments as set.
retrieval_strategy: allOf: - $ref: '#/components/schemas/RetrievalStrategy' default: segments max_documents: title: Max Documents type: integer + minimum: 1 + maximum: 100Also applies to: 2723-2725
2826-2870: CreateAssistantRequest parity with ModifyAssistantRequest.
- tools has the same issue as Assistant (array of map). See suggestion above.
- tool_resources is a free-form object here but a typed schema in ModifyAssistantRequest; use the same $ref for consistency.
tools: title: Tools type: array - items: - type: object - additionalProperties: - $ref: '#/components/schemas/AssistantTool' + items: + $ref: '#/components/schemas/AssistantTool' tool_resources: - title: Tool Resources - type: object + $ref: '#/components/schemas/AssistantToolResource'
3305-3361: FileSearchToolResource: add limits and optionalmax_documents.Guardrails improve stability and parity with ConversationalRagConfig.
FileSearchToolResource: properties: retrieval_similarity_threshold: title: Retrieval Similarity Threshold type: number + minimum: 0 + maximum: 1 labels: title: Labels type: array items: type: string ... retrieval_strategy: allOf: - $ref: '#/components/schemas/RetrievalStrategy' default: segments + max_documents: + title: Max Documents + type: integer + minimum: 1 + maximum: 100 max_neighbors: title: Max Neighbors type: integer + minimum: 1 + maximum: 50
4093-4142: Constrainincludeand typeoutput_type.Define an enum for allowed include keys and a schema for output_type to avoid free-form strings.
RunAssistantRequest: properties: include: title: Include - type: array - items: - type: string + type: array + items: + $ref: '#/components/schemas/RunInclude' default: - data_sources - requirements_result - requirements_result.metadata output_type: - title: Output Type - type: object + $ref: '#/components/schemas/MaestroPayloadType' + + RunInclude: + title: RunInclude + type: string + enum: [data_sources, requirements_result, requirements_result.metadata]
4323-4330: ToolResource wiring: looks correct.Matches execution_engine ToolDefinition variants. Consider adding brief descriptions to each field for better docs.
150-170: Large defaultlimitfor Get Workspace Files.Default 500000 may be excessive; paginate with a sensible default (e.g., 100) to avoid heavy responses.
- name: limit in: query schema: title: Limit type: integer - default: 500000 + default: 100 + minimum: 1 + maximum: 1000
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (132)
src/libs/AI21/Generated/AI21..JsonSerializerContext.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Ai21Api.GetVectorStoreStudioV1DemosRegulationsVectorStoreVectorStoreIdGet.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Ai21Api.McpToolDiscoveryStudioV1McpDiscoverPost.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Ai21Api.UploadCheckComplianceStudioV1DemosRegulationsUploadCheckCompliancePost.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Ai21Api.V1ConversationalRag.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Ai21Api.V1CreateAssistant.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Ai21Api.V1MaestroRun.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Ai21Api.V1ModifyAssistant.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Ai21Api.V1RunAssistant.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.IAi21Api.GetVectorStoreStudioV1DemosRegulationsVectorStoreVectorStoreIdGet.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.IAi21Api.McpToolDiscoveryStudioV1McpDiscoverPost.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.IAi21Api.UploadCheckComplianceStudioV1DemosRegulationsUploadCheckCompliancePost.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.IAi21Api.V1ConversationalRag.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.IAi21Api.V1CreateAssistant.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.IAi21Api.V1MaestroRun.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.IAi21Api.V1ModifyAssistant.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.IAi21Api.V1RunAssistant.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.IJambaCompleteClient.V1ChatComplete.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JambaCompleteClient.V1ChatComplete.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.AssistantType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.AssistantTypeNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.BudgetLevel.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.BudgetLevelNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.FileSearchToolResourceLabelsFilterMode.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.FileSearchToolResourceLabelsFilterModeNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.FileSearchToolResourceResponseLanguage.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.FileSearchToolResourceResponseLanguageNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.FileSearchToolResourceType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.FileSearchToolResourceTypeNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.HTTPToolFunctionParametersType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.HTTPToolFunctionParametersTypeNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.HTTPToolResourceType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.HTTPToolResourceTypeNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.LanguageStudioApiServerDataTypesChatToolDefinitionType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.LanguageStudioApiServerDataTypesChatToolDefinitionTypeNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.LanguageStudioApiServerDataTypesExecutionEngineToolDefinition.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.LanguageStudioApiServerDataTypesExecutionEngineToolDefinitionDiscriminatorType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.LanguageStudioApiServerDataTypesExecutionEngineToolDefinitionDiscriminatorTypeNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.MCPToolResourceType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.MCPToolResourceTypeNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.RunAssistantRequestResponseLanguage.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.RunAssistantRequestResponseLanguageNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.RunOptimization.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.RunOptimizationNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.Visibility.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.VisibilityNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.WebSearchToolResourceType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonConverters.WebSearchToolResourceTypeNullable.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.JsonSerializerContextTypes.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.Assistant.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.AssistantTool2.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.AssistantType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.BodyUploadCheckComplianceStudioV1DemosRegulationsUploadCheckCompliancePost.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.BudgetLevel.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ChatRequest.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ConversationalRagConfig.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ConversationalRagConfigResponseLanguage.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ConversationalRagConfigRetrievalStrategy.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.CreateAssistantRequest.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.CreateAssistantRequest.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.CreateAssistantRequestTool.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.CreateAssistantRequestTool.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.CreateAssistantRequestToolResources.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.CreateAssistantRequestToolResources.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.CreateMaestroRunsPayload.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.CreateMaestroRunsPayloadResponseLanguage.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.FileSearchToolResource.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.FileSearchToolResource.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.FileSearchToolResourceLabelsFilterMode.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.FileSearchToolResourceResponseLanguage.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.FileSearchToolResourceType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.GetVectorStoreStudioV1DemosRegulationsVectorStoreVectorStoreIdGetResponse.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.GetVectorStoreStudioV1DemosRegulationsVectorStoreVectorStoreIdGetResponse.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolEndpoint.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolEndpoint.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolEndpointHeaders.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolEndpointHeaders.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolFunction.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolFunction.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolFunctionParamProperties.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolFunctionParamProperties.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolFunctionParameters.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolFunctionParameters.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolFunctionParametersProperties.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolFunctionParametersProperties.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolFunctionParametersType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolResource.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolResource.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.HTTPToolResourceType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesAssistantRequirement.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesAssistantRequirement.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesChatToolDefinition.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesChatToolDefinition.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesChatToolDefinitionType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesExecutionEngineRequirement.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesExecutionEngineRequirement.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesExecutionEngineToolDefinition.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesExecutionEngineToolDefinition.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesExecutionEngineToolDefinitionDiscriminator.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesExecutionEngineToolDefinitionDiscriminator.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.LanguageStudioApiServerDataTypesExecutionEngineToolDefinitionDiscriminatorType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPDefinition.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPDefinition.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPDefinitionHeaders.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPDefinitionHeaders.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPDiscoveryResponse.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPDiscoveryResponse.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPTool.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPTool.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPToolInputSchema.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPToolInputSchema.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPToolResource.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPToolResource.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPToolResourceHeaders.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPToolResourceHeaders.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.MCPToolResourceType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ModifyAssistantRequest.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ModifyAssistantRequest.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ModifyAssistantRequestTool.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ModifyAssistantRequestTool.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.RunAssistantRequest.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.RunAssistantRequest.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.RunAssistantRequestOutputType.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.RunAssistantRequestOutputType.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.RunAssistantRequestResponseLanguage.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.RunOptimization.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ToolResource.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.ToolResource.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.Visibility.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.WebSearchToolResource.Json.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.WebSearchToolResource.g.csis excluded by!**/generated/**src/libs/AI21/Generated/AI21.Models.WebSearchToolResourceType.g.csis excluded by!**/generated/**
📒 Files selected for processing (1)
src/libs/AI21/openapi.yaml(23 hunks)
🔇 Additional comments (6)
src/libs/AI21/openapi.yaml (6)
2468-2469: LGTM: ChatRequest.tools now references chat ToolDefinition.This aligns the chat surface with function tools. Looks good.
2896-2901: Tool resources on Maestro runs: good decoupling.Switching to execution_engine ToolDefinition and ToolResource is the right direction.
Also applies to: 2898-2899
3894-3935: LGTM: ModifyAssistantRequest is coherent and stricter than the read model.Good use of enums and explicit flags.
4447-4451: LGTM: Visibility enum addition.Defaults to public; confirm this aligns with product defaults.
4613-4673: LGTM: Discriminated union for execution-engine tools.Good use of discriminator with explicit mapping.
2305-2339: Summary/code mismatch: no “budget” on Body_v1_library_upload.AI summary claims budgets under Body_v1_library_upload now use BudgetLevel, but this schema doesn’t expose a budget field. Either update the schema or fix the PR description.
Comment on lines
+1948
to
+1969
| post: | ||
| summary: Mcp Tool Discovery | ||
| operationId: mcp_tool_discovery_studio_v1_mcp_discover_post | ||
| requestBody: | ||
| content: | ||
| application/json: | ||
| schema: | ||
| $ref: '#/components/schemas/MCPDefinition' | ||
| required: true | ||
| responses: | ||
| '200': | ||
| description: Successful Response | ||
| content: | ||
| application/json: | ||
| schema: | ||
| $ref: '#/components/schemas/MCPDiscoveryResponse' | ||
| '422': | ||
| description: Validation Error | ||
| content: | ||
| application/json: | ||
| schema: | ||
| $ref: '#/components/schemas/HTTPValidationError' |
There was a problem hiding this comment.
SSRF risk in MCP discovery — restrict server_url.
Accepting arbitrary URLs enables SSRF. Add validation (allowlist domains, disallow private IPs), explicit pattern, and clarify errors (400/502). Consider timeouts.
/studio/v1/mcp/discover:
post:
summary: Mcp Tool Discovery
operationId: mcp_tool_discovery_studio_v1_mcp_discover_post
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/MCPDefinition'
required: true
responses:
- '200':
+ '200':
description: Successful Response
content:
application/json:
schema:
$ref: '#/components/schemas/MCPDiscoveryResponse'
+ '400':
+ description: Invalid server_url (disallowed host or scheme)
+ '502':
+ description: Upstream MCP server not reachable
'422':
description: Validation ErrorAnd in MCPDefinition:
MCPDefinition:
properties:
server_url:
title: Server Url
- type: string
+ type: string
+ pattern: '^https?://'
+ description: 'Public HTTP(S) endpoint. Private networks and link-local ranges are rejected.'📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| post: | |
| summary: Mcp Tool Discovery | |
| operationId: mcp_tool_discovery_studio_v1_mcp_discover_post | |
| requestBody: | |
| content: | |
| application/json: | |
| schema: | |
| $ref: '#/components/schemas/MCPDefinition' | |
| required: true | |
| responses: | |
| '200': | |
| description: Successful Response | |
| content: | |
| application/json: | |
| schema: | |
| $ref: '#/components/schemas/MCPDiscoveryResponse' | |
| '422': | |
| description: Validation Error | |
| content: | |
| application/json: | |
| schema: | |
| $ref: '#/components/schemas/HTTPValidationError' | |
| /studio/v1/mcp/discover: | |
| post: | |
| summary: Mcp Tool Discovery | |
| operationId: mcp_tool_discovery_studio_v1_mcp_discover_post | |
| requestBody: | |
| content: | |
| application/json: | |
| schema: | |
| $ref: '#/components/schemas/MCPDefinition' | |
| required: true | |
| responses: | |
| '200': | |
| description: Successful Response | |
| content: | |
| application/json: | |
| schema: | |
| $ref: '#/components/schemas/MCPDiscoveryResponse' | |
| '400': | |
| description: Invalid server_url (disallowed host or scheme) | |
| '502': | |
| description: Upstream MCP server not reachable | |
| '422': | |
| description: Validation Error | |
| content: | |
| application/json: | |
| schema: | |
| $ref: '#/components/schemas/HTTPValidationError' | |
| components: | |
| schemas: | |
| MCPDefinition: | |
| properties: | |
| server_url: | |
| title: Server Url | |
| type: string | |
| pattern: '^https?://' | |
| description: 'Public HTTP(S) endpoint. Private networks and link-local ranges are rejected.' |
🤖 Prompt for AI Agents
In src/libs/AI21/openapi.yaml around lines 1948 to 1969, the MCP discovery
endpoint accepts arbitrary server URLs which creates an SSRF risk; update the
operation and the referenced MCPDefinition schema to enforce strict validation:
change the server_url property to a constrained string (explicit pattern or
hostname-only format) and/or enum of allowlisted domains, validate and reject
inputs that are IP addresses in private/reserved ranges, and add server-side
checks for DNS resolution to public IPs only; return clear 400 responses for
invalid client input and 502 for upstream connection failures, and
document/implement connection timeouts and max redirect limits in the
description so callers and implementers know expected behavior.
Comment on lines
2136
to
+2153
| $ref: '#/components/schemas/AssistantToolResource' | ||
| requirements: | ||
| title: Requirements | ||
| type: array | ||
| items: | ||
| $ref: '#/components/schemas/language_studio_api_server__data_types__assistant__Requirement' | ||
| budget: | ||
| allOf: | ||
| - $ref: '#/components/schemas/BudgetLevel' | ||
| default: medium | ||
| visibility: | ||
| allOf: | ||
| - $ref: '#/components/schemas/Visibility' | ||
| default: public | ||
| assistant_type: | ||
| allOf: | ||
| - $ref: '#/components/schemas/AssistantType' | ||
| default: default |
There was a problem hiding this comment.
Assistant.tools item shape is odd (array of objects with additionalProperties).
Today this allows arbitrary object maps per item. If the intent is “array of enum AssistantTool”, make it a straight enum array. If you need keyed tool configs, define a typed object.
- tools:
- title: Tools
- type: array
- items:
- type: object
- additionalProperties:
- $ref: '#/components/schemas/AssistantTool'
+ tools:
+ title: Tools
+ type: array
+ items:
+ $ref: '#/components/schemas/AssistantTool'📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| $ref: '#/components/schemas/AssistantToolResource' | |
| requirements: | |
| title: Requirements | |
| type: array | |
| items: | |
| $ref: '#/components/schemas/language_studio_api_server__data_types__assistant__Requirement' | |
| budget: | |
| allOf: | |
| - $ref: '#/components/schemas/BudgetLevel' | |
| default: medium | |
| visibility: | |
| allOf: | |
| - $ref: '#/components/schemas/Visibility' | |
| default: public | |
| assistant_type: | |
| allOf: | |
| - $ref: '#/components/schemas/AssistantType' | |
| default: default | |
| $ref: '#/components/schemas/AssistantToolResource' | |
| tools: | |
| title: Tools | |
| type: array | |
| items: | |
| $ref: '#/components/schemas/AssistantTool' | |
| requirements: | |
| title: Requirements | |
| type: array | |
| items: | |
| $ref: '#/components/schemas/language_studio_api_server__data_types__assistant__Requirement' | |
| budget: | |
| allOf: | |
| - $ref: '#/components/schemas/BudgetLevel' | |
| default: medium | |
| visibility: | |
| allOf: | |
| - $ref: '#/components/schemas/Visibility' | |
| default: public | |
| assistant_type: | |
| allOf: | |
| - $ref: '#/components/schemas/AssistantType' | |
| default: default |
Comment on lines
+2340
to
2346
| title: BudgetLevel | ||
| enum: | ||
| - low | ||
| - medium | ||
| - high | ||
| description: An enumeration. | ||
| ChatCompletion: |
There was a problem hiding this comment.
Budget enum casing inconsistent across payloads.
You introduce BudgetLevel with lowercase values, but other payloads (e.g., Body_process_rfi_document and Body_upload_check_compliance) still use string budgets with uppercase defaults (MEDIUM). Standardize to BudgetLevel and lowercase defaults.
Proposed updates (outside this hunk):
Body_process_rfi_document_studio_v1_demos_rfi_process_rfi_post:
properties:
budget:
- title: Budget
- type: string
- description: 'Budget level: LOW, MEDIUM, or HIGH'
- default: MEDIUM
+ title: Budget
+ allOf:
+ - $ref: '#/components/schemas/BudgetLevel'
+ default: medium
Body_upload_check_compliance_studio_v1_demos_regulations_upload_check_compliance_post:
properties:
budget:
- title: Budget
- type: string
- default: MEDIUM
+ title: Budget
+ allOf:
+ - $ref: '#/components/schemas/BudgetLevel'
+ default: mediumCommittable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In src/libs/AI21/openapi.yaml around lines 2340-2346, the BudgetLevel enum is
defined with lowercase values (low, medium, high) but other schemas still use
uppercase string defaults (e.g., MEDIUM); update those payload schemas (e.g.,
Body_process_rfi_document, Body_upload_check_compliance) to reference the
BudgetLevel enum and change any hard-coded defaults or examples from uppercase
to the matching lowercase values (low|medium|high) so casing is consistent
across the spec.
Comment on lines
+3400
to
3484
| title: HTTPToolEndpoint | ||
| required: | ||
| - url | ||
| type: object | ||
| properties: | ||
| url: | ||
| title: Url | ||
| type: string | ||
| headers: | ||
| title: Headers | ||
| type: object | ||
| HTTPToolFunction: | ||
| title: HTTPToolFunction | ||
| required: | ||
| - name | ||
| - description | ||
| - parameters | ||
| type: object | ||
| properties: | ||
| name: | ||
| title: Name | ||
| type: string | ||
| description: | ||
| title: Description | ||
| type: string | ||
| parameters: | ||
| $ref: '#/components/schemas/HTTPToolFunctionParameters' | ||
| HTTPToolFunctionParamProperties: | ||
| title: HTTPToolFunctionParamProperties | ||
| required: | ||
| - type | ||
| - description | ||
| type: object | ||
| properties: | ||
| type: | ||
| title: Type | ||
| type: string | ||
| description: | ||
| title: Description | ||
| type: string | ||
| HTTPToolFunctionParameters: | ||
| title: HTTPToolFunctionParameters | ||
| required: | ||
| - properties | ||
| - required | ||
| type: object | ||
| properties: | ||
| type: | ||
| title: Type | ||
| enum: | ||
| - object | ||
| type: string | ||
| default: object | ||
| properties: | ||
| title: Properties | ||
| type: object | ||
| additionalProperties: | ||
| $ref: '#/components/schemas/HTTPToolFunctionParamProperties' | ||
| required: | ||
| title: Required | ||
| type: array | ||
| items: | ||
| type: string | ||
| additionalProperties: | ||
| title: Additionalproperties | ||
| type: boolean | ||
| default: false | ||
| HTTPToolResource: | ||
| title: HTTPToolResource | ||
| required: | ||
| - function | ||
| - endpoint | ||
| type: object | ||
| properties: | ||
| type: | ||
| title: Type | ||
| enum: | ||
| - http | ||
| type: string | ||
| default: http | ||
| function: | ||
| $ref: '#/components/schemas/HTTPToolFunction' | ||
| endpoint: | ||
| $ref: '#/components/schemas/HTTPToolEndpoint' | ||
| HTTPValidationError: |
There was a problem hiding this comment.
HTTP tool schema lacks HTTP method and timeouts.
Without a method field, execution semantics are ambiguous. Add method enum, optional timeout, and forbid header value leakage in logs via description.
HTTPToolEndpoint:
required:
- - url
+ - url
+ - method
properties:
url:
title: Url
type: string
+ pattern: '^https?://'
+ method:
+ title: Method
+ type: string
+ enum: [GET, POST, PUT, PATCH, DELETE]
+ timeout_ms:
+ title: Timeout (ms)
+ type: integer
+ minimum: 1
headers:
title: Headers
type: object
+ description: 'Sensitive headers may be redacted in logs.'📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| title: HTTPToolEndpoint | |
| required: | |
| - url | |
| type: object | |
| properties: | |
| url: | |
| title: Url | |
| type: string | |
| headers: | |
| title: Headers | |
| type: object | |
| HTTPToolFunction: | |
| title: HTTPToolFunction | |
| required: | |
| - name | |
| - description | |
| - parameters | |
| type: object | |
| properties: | |
| name: | |
| title: Name | |
| type: string | |
| description: | |
| title: Description | |
| type: string | |
| parameters: | |
| $ref: '#/components/schemas/HTTPToolFunctionParameters' | |
| HTTPToolFunctionParamProperties: | |
| title: HTTPToolFunctionParamProperties | |
| required: | |
| - type | |
| - description | |
| type: object | |
| properties: | |
| type: | |
| title: Type | |
| type: string | |
| description: | |
| title: Description | |
| type: string | |
| HTTPToolFunctionParameters: | |
| title: HTTPToolFunctionParameters | |
| required: | |
| - properties | |
| - required | |
| type: object | |
| properties: | |
| type: | |
| title: Type | |
| enum: | |
| - object | |
| type: string | |
| default: object | |
| properties: | |
| title: Properties | |
| type: object | |
| additionalProperties: | |
| $ref: '#/components/schemas/HTTPToolFunctionParamProperties' | |
| required: | |
| title: Required | |
| type: array | |
| items: | |
| type: string | |
| additionalProperties: | |
| title: Additionalproperties | |
| type: boolean | |
| default: false | |
| HTTPToolResource: | |
| title: HTTPToolResource | |
| required: | |
| - function | |
| - endpoint | |
| type: object | |
| properties: | |
| type: | |
| title: Type | |
| enum: | |
| - http | |
| type: string | |
| default: http | |
| function: | |
| $ref: '#/components/schemas/HTTPToolFunction' | |
| endpoint: | |
| $ref: '#/components/schemas/HTTPToolEndpoint' | |
| HTTPValidationError: | |
| title: HTTPToolEndpoint | |
| required: | |
| - url | |
| - method | |
| type: object | |
| properties: | |
| url: | |
| title: Url | |
| type: string | |
| pattern: '^https?://' | |
| method: | |
| title: Method | |
| type: string | |
| enum: [GET, POST, PUT, PATCH, DELETE] | |
| timeout_ms: | |
| title: Timeout (ms) | |
| type: integer | |
| minimum: 1 | |
| headers: | |
| title: Headers | |
| type: object | |
| description: 'Sensitive headers may be redacted in logs.' | |
| HTTPToolFunction: | |
| title: HTTPToolFunction | |
| required: | |
| - name | |
| - description | |
| - parameters | |
| type: object | |
| properties: | |
| name: | |
| title: Name | |
| type: string | |
| description: | |
| title: Description | |
| type: string | |
| parameters: | |
| $ref: '#/components/schemas/HTTPToolFunctionParameters' | |
| HTTPToolFunctionParamProperties: | |
| title: HTTPToolFunctionParamProperties | |
| required: | |
| - type | |
| - description | |
| type: object | |
| properties: | |
| type: | |
| title: Type | |
| type: string | |
| description: | |
| title: Description | |
| type: string | |
| HTTPToolFunctionParameters: | |
| title: HTTPToolFunctionParameters | |
| required: | |
| - properties | |
| - required | |
| type: object | |
| properties: | |
| type: | |
| title: Type | |
| enum: | |
| - object | |
| type: string | |
| default: object | |
| properties: | |
| title: Properties | |
| type: object | |
| additionalProperties: | |
| $ref: '#/components/schemas/HTTPToolFunctionParamProperties' | |
| required: | |
| title: Required | |
| type: array | |
| items: | |
| type: string | |
| additionalProperties: | |
| title: Additionalproperties | |
| type: boolean | |
| default: false | |
| HTTPToolResource: | |
| title: HTTPToolResource | |
| required: | |
| - function | |
| - endpoint | |
| type: object | |
| properties: | |
| type: | |
| title: Type | |
| enum: | |
| - http | |
| type: string | |
| default: http | |
| function: | |
| $ref: '#/components/schemas/HTTPToolFunction' | |
| endpoint: | |
| $ref: '#/components/schemas/HTTPToolEndpoint' | |
| HTTPValidationError: |
Comment on lines
+3672
to
+3742
| title: MCPDefinition | ||
| required: | ||
| - server_url | ||
| type: object | ||
| properties: | ||
| server_url: | ||
| title: Server Url | ||
| type: string | ||
| server_label: | ||
| title: Server Label | ||
| type: string | ||
| headers: | ||
| title: Headers | ||
| type: object | ||
| MCPDiscoveryResponse: | ||
| title: MCPDiscoveryResponse | ||
| required: | ||
| - mcp_definition | ||
| - tools | ||
| type: object | ||
| properties: | ||
| mcp_definition: | ||
| $ref: '#/components/schemas/MCPDefinition' | ||
| tools: | ||
| title: Tools | ||
| type: array | ||
| items: | ||
| $ref: '#/components/schemas/MCPTool' | ||
| MCPTool: | ||
| title: MCPTool | ||
| required: | ||
| - name | ||
| - inputSchema | ||
| type: object | ||
| properties: | ||
| name: | ||
| title: Name | ||
| type: string | ||
| description: | ||
| title: Description | ||
| type: string | ||
| inputSchema: | ||
| title: Inputschema | ||
| type: object | ||
| MCPToolResource: | ||
| title: MCPToolResource | ||
| required: | ||
| - server_label | ||
| - server_url | ||
| type: object | ||
| properties: | ||
| type: | ||
| title: Type | ||
| enum: | ||
| - mcp | ||
| type: string | ||
| default: mcp | ||
| server_label: | ||
| title: Server Label | ||
| type: string | ||
| server_url: | ||
| title: Server Url | ||
| type: string | ||
| headers: | ||
| title: Headers | ||
| type: object | ||
| allowed_tools: | ||
| title: Allowed Tools | ||
| type: array | ||
| items: | ||
| type: string |
There was a problem hiding this comment.
🛠️ Refactor suggestion
MCP tool/resource: also enforce allowlists and tool filtering.
allowed_tools exists; document that undefined tools are rejected and enforce domain allowlist alongside the earlier SSRF guard.
MCPToolResource:
properties:
allowed_tools:
title: Allowed Tools
type: array
items:
type: string
+ description: 'Only these tool names are callable; all others are rejected.'📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| title: MCPDefinition | |
| required: | |
| - server_url | |
| type: object | |
| properties: | |
| server_url: | |
| title: Server Url | |
| type: string | |
| server_label: | |
| title: Server Label | |
| type: string | |
| headers: | |
| title: Headers | |
| type: object | |
| MCPDiscoveryResponse: | |
| title: MCPDiscoveryResponse | |
| required: | |
| - mcp_definition | |
| - tools | |
| type: object | |
| properties: | |
| mcp_definition: | |
| $ref: '#/components/schemas/MCPDefinition' | |
| tools: | |
| title: Tools | |
| type: array | |
| items: | |
| $ref: '#/components/schemas/MCPTool' | |
| MCPTool: | |
| title: MCPTool | |
| required: | |
| - name | |
| - inputSchema | |
| type: object | |
| properties: | |
| name: | |
| title: Name | |
| type: string | |
| description: | |
| title: Description | |
| type: string | |
| inputSchema: | |
| title: Inputschema | |
| type: object | |
| MCPToolResource: | |
| title: MCPToolResource | |
| required: | |
| - server_label | |
| - server_url | |
| type: object | |
| properties: | |
| type: | |
| title: Type | |
| enum: | |
| - mcp | |
| type: string | |
| default: mcp | |
| server_label: | |
| title: Server Label | |
| type: string | |
| server_url: | |
| title: Server Url | |
| type: string | |
| headers: | |
| title: Headers | |
| type: object | |
| allowed_tools: | |
| title: Allowed Tools | |
| type: array | |
| items: | |
| type: string | |
| title: MCPDefinition | |
| required: | |
| - server_url | |
| type: object | |
| properties: | |
| server_url: | |
| title: Server Url | |
| type: string | |
| server_label: | |
| title: Server Label | |
| type: string | |
| headers: | |
| title: Headers | |
| type: object | |
| MCPDiscoveryResponse: | |
| title: MCPDiscoveryResponse | |
| required: | |
| - mcp_definition | |
| - tools | |
| type: object | |
| properties: | |
| mcp_definition: | |
| $ref: '#/components/schemas/MCPDefinition' | |
| tools: | |
| title: Tools | |
| type: array | |
| items: | |
| $ref: '#/components/schemas/MCPTool' | |
| MCPTool: | |
| title: MCPTool | |
| required: | |
| - name | |
| - inputSchema | |
| type: object | |
| properties: | |
| name: | |
| title: Name | |
| type: string | |
| description: | |
| title: Description | |
| type: string | |
| inputSchema: | |
| title: Inputschema | |
| type: object | |
| MCPToolResource: | |
| title: MCPToolResource | |
| required: | |
| - server_label | |
| - server_url | |
| type: object | |
| properties: | |
| type: | |
| title: Type | |
| enum: | |
| - mcp | |
| type: string | |
| default: mcp | |
| server_label: | |
| title: Server Label | |
| type: string | |
| server_url: | |
| title: Server Url | |
| type: string | |
| headers: | |
| title: Headers | |
| type: object | |
| allowed_tools: | |
| title: Allowed Tools | |
| type: array | |
| items: | |
| type: string | |
| description: 'Only these tool names are callable; all others are rejected.' |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Refactor