diff --git a/src/utils/models_dumper.py b/src/utils/models_dumper.py index 05acf0a37..b254b6828 100644 --- a/src/utils/models_dumper.py +++ b/src/utils/models_dumper.py @@ -1,5 +1,9 @@ """Function to dump the schema of all data models into OpenAPI-compatible format.""" +from typing import Optional + +from pydantic import BaseModel + import models.api.requests as r import models.api.responses.error as e import models.api.responses.successful as s @@ -9,6 +13,137 @@ import models.compaction as models_compaction from utils.openapi_schema_dumper import dump_openapi_schema +conversation_summary_models: list[type[BaseModel]] = [ + models_compaction.ConversationSummary +] + +requests_models: list[type[BaseModel]] = [ + r.ConversationUpdateRequest, + r.FeedbackRequest, + r.FeedbackStatusUpdateRequest, + r.MCPServerRegistrationRequest, + r.ModelFilter, + r.PromptCreateRequest, + r.PromptUpdateRequest, + r.QueryRequest, + r.ResponsesRequest, + r.RlsapiV1Attachment, + r.RlsapiV1CLA, + r.RlsapiV1Context, + r.RlsapiV1InferRequest, + r.RlsapiV1SystemInfo, + r.RlsapiV1Terminal, + r.StreamingInterruptRequest, + r.VectorStoreCreateRequest, + r.VectorStoreFileCreateRequest, + r.VectorStoreUpdateRequest, +] + +successful_responses_models: list[type[BaseModel]] = [ + s.AuthorizedResponse, + s.ConfigurationResponse, + s.ConversationDeleteResponse, + s.ConversationResponse, + s.ConversationUpdateResponse, + s.ConversationsListResponse, + s.ConversationsListResponseV2, + s.FeedbackResponse, + s.FeedbackStatusUpdateResponse, + s.FileResponse, + s.InfoResponse, + s.LivenessResponse, + s.MCPClientAuthOptionsResponse, + s.MCPServerDeleteResponse, + s.MCPServerListResponse, + s.MCPServerRegistrationResponse, + s.ModelsResponse, + s.PromptDeleteResponse, + s.PromptResourceResponse, + s.PromptsListResponse, + s.ProviderResponse, + s.ProvidersListResponse, + s.QueryResponse, + s.RAGInfoResponse, + s.RAGListResponse, + s.ReadinessResponse, + s.ResponsesResponse, + s.RlsapiV1InferData, + s.RlsapiV1InferResponse, + s.ShieldsResponse, + s.StatusResponse, + s.StreamingInterruptResponse, + s.StreamingQueryResponse, + s.ToolsResponse, + s.VectorStoreDeleteResponse, + s.VectorStoreFileDeleteResponse, + s.VectorStoreFileResponse, + s.VectorStoreFilesListResponse, + s.VectorStoreResponse, + s.VectorStoresListResponse, +] + +error_responses_models: list[type[BaseModel]] = [ + e.AbstractErrorResponse, + e.BadRequestResponse, + e.ConflictResponse, + e.DetailModel, + e.FileTooLargeResponse, + e.ForbiddenResponse, + e.InternalServerErrorResponse, + e.NotFoundResponse, + e.PromptTooLongResponse, + e.QuotaExceededResponse, + e.ServiceUnavailableResponse, + e.UnauthorizedResponse, + e.UnprocessableEntityResponse, +] + +common_models: list[type[BaseModel]] = [ + c.Attachment, + c.ConversationData, + c.ConversationDetails, + c.ConversationTurn, + c.MCPListToolsSummary, + c.MCPServerAuthInfo, + c.MCPServerInfo, + c.Message, + c.ProviderHealthStatus, + c.RAGChunk, + c.RAGContext, + c.ReferencedDocument, + c.ShieldModerationBlocked, + c.ShieldModerationPassed, + c.SolrVectorSearchRequest, + c.ToolCallSummary, + c.ToolInfoSummary, + c.ToolResultSummary, + c.Transcript, + c.TranscriptMetadata, + c.TurnSummary, +] + +agents_models: list[type[BaseModel]] = [ + a.EndEventData, + a.EndStreamPayload, + a.ErrorEventData, + a.ErrorStreamPayload, + a.InterruptedEventData, + a.InterruptedStreamPayload, + a.StartEventData, + a.StartStreamPayload, + a.StreamPayloadBase, + a.TokenChunkData, + a.TokenStreamPayload, + a.ToolCallStreamPayload, + a.ToolResultStreamPayload, + a.TurnCompleteStreamPayload, +] + +common_responses_models: list[type[BaseModel]] = [ + cr.InputToolMCP, + cr.ResponsesApiParams, +] + def dump_models(filename: str) -> None: """Dump the schema of all models into OpenAPI-compatible JSON file. @@ -25,119 +160,80 @@ def dump_models(filename: str) -> None: ------ IOError: If the file cannot be written. """ - models = [models_compaction.ConversationSummary] - - # add all requests data models - for model in [ - r.ConversationUpdateRequest, - r.FeedbackRequest, - r.FeedbackStatusUpdateRequest, - r.MCPServerRegistrationRequest, - r.ModelFilter, - r.PromptCreateRequest, - r.PromptUpdateRequest, - r.QueryRequest, - r.ResponsesRequest, - r.RlsapiV1Attachment, - r.RlsapiV1CLA, - r.RlsapiV1Context, - r.RlsapiV1InferRequest, - r.RlsapiV1SystemInfo, - r.RlsapiV1Terminal, - r.StreamingInterruptRequest, - r.VectorStoreCreateRequest, - r.VectorStoreFileCreateRequest, - r.VectorStoreUpdateRequest, - s.AuthorizedResponse, - s.ConfigurationResponse, - s.ConversationDeleteResponse, - s.ConversationResponse, - s.ConversationUpdateResponse, - s.ConversationsListResponse, - s.ConversationsListResponseV2, - s.FeedbackResponse, - s.FeedbackStatusUpdateResponse, - s.FileResponse, - s.InfoResponse, - s.LivenessResponse, - s.MCPClientAuthOptionsResponse, - s.MCPServerDeleteResponse, - s.MCPServerListResponse, - s.MCPServerRegistrationResponse, - s.ModelsResponse, - s.PromptDeleteResponse, - s.PromptResourceResponse, - s.PromptsListResponse, - s.ProviderResponse, - s.ProvidersListResponse, - s.QueryResponse, - s.RAGInfoResponse, - s.RAGListResponse, - s.ReadinessResponse, - s.ResponsesResponse, - s.RlsapiV1InferData, - s.RlsapiV1InferResponse, - s.ShieldsResponse, - s.StatusResponse, - s.StreamingInterruptResponse, - s.StreamingQueryResponse, - s.ToolsResponse, - s.VectorStoreDeleteResponse, - s.VectorStoreFileDeleteResponse, - s.VectorStoreFileResponse, - s.VectorStoreFilesListResponse, - s.VectorStoreResponse, - s.VectorStoresListResponse, - e.AbstractErrorResponse, - e.BadRequestResponse, - e.ConflictResponse, - e.DetailModel, - e.FileTooLargeResponse, - e.ForbiddenResponse, - e.InternalServerErrorResponse, - e.NotFoundResponse, - e.PromptTooLongResponse, - e.QuotaExceededResponse, - e.ServiceUnavailableResponse, - e.UnauthorizedResponse, - e.UnprocessableEntityResponse, - c.Attachment, - c.ConversationData, - c.ConversationDetails, - c.ConversationTurn, - c.MCPListToolsSummary, - c.MCPServerAuthInfo, - c.MCPServerInfo, - c.Message, - c.ProviderHealthStatus, - c.RAGChunk, - c.RAGContext, - c.ReferencedDocument, - c.ShieldModerationBlocked, - c.ShieldModerationPassed, - c.SolrVectorSearchRequest, - c.ToolCallSummary, - c.ToolInfoSummary, - c.ToolResultSummary, - c.Transcript, - c.TranscriptMetadata, - c.TurnSummary, - a.EndEventData, - a.EndStreamPayload, - a.ErrorEventData, - a.ErrorStreamPayload, - a.InterruptedEventData, - a.InterruptedStreamPayload, - a.StartEventData, - a.StartStreamPayload, - a.StreamPayloadBase, - a.TokenChunkData, - a.TokenStreamPayload, - a.ToolCallStreamPayload, - a.ToolResultStreamPayload, - a.TurnCompleteStreamPayload, - cr.InputToolMCP, - cr.ResponsesApiParams, - ]: - models.append(model) + # construct a list with all models + models = ( + conversation_summary_models + + requests_models + + successful_responses_models + + error_responses_models + + common_models + + agents_models + + common_responses_models + ) + + # dump all the models into one OpenAPI-compatible JSON file + dump_openapi_schema(models, filename) + + +def get_models_for_group(model_group: str) -> list[type[BaseModel]]: + """Return the list of Pydantic model classes for the given model group. + + Supported groups: + - "requests" + - "successful_responses" + - "error_responses" + - "common" + - "agents" + - "common_responses" + + Parameters: + ---------- + model_group: The name of the model group to look up. + + Returns: + ------- + A list of Pydantic model classes belonging to the requested group. + + Raises: + ------ + Exception: If model_group is not a recognized group name. + """ + match model_group: + case "requests": + return requests_models + case "successful_responses": + return successful_responses_models + case "error_responses": + return error_responses_models + case "common": + return common_models + case "agents": + return agents_models + case "common_responses": + return common_responses_models + case _: + raise ValueError(f"Unknown model group provided: {model_group}") + + +def dump_models_group(model_group: str, filename: Optional[str] = None) -> None: + """Dump the schema of selected models group into OpenAPI-compatible JSON file. + + Parameters: + ---------- + - model_group: str - name of model group to export the schema to + + Returns: + ------- + - None + + Raises: + ------ + IOError: If the file cannot be written. + """ + models = get_models_for_group(model_group) + + if filename is None: + filename = f"{model_group}.json" + + # dump all selected models into one OpenAPI-compatible JSON file dump_openapi_schema(models, filename) diff --git a/tests/unit/utils/test_models_dumper.py b/tests/unit/utils/test_models_dumper.py index 2902fa25f..6189453e7 100644 --- a/tests/unit/utils/test_models_dumper.py +++ b/tests/unit/utils/test_models_dumper.py @@ -5,7 +5,9 @@ from json import load from pathlib import Path -from utils.models_dumper import dump_models +import pytest + +from utils.models_dumper import dump_models, dump_models_group def test_dump_models(tmpdir: Path) -> None: @@ -9326,3 +9328,215 @@ def test_dump_models(tmpdir: Path) -> None: ) for expected_schema in expected_schemas: assert expected_schema in schemas + + +def check_json_file_content(filename: str, expected_schemas: list[str]) -> None: + """Check the content of provided JSON file with OpenAPI-compatible schema.""" + with open(filename, "r", encoding="utf-8") as fin: + # schema should be stored in JSON format + content = load(fin) + assert content is not None + + # top-level keys test + keys = ("openapi", "info", "components", "paths") + for key in keys: + assert key in content + + # components should be top-level node + components = content["components"] + assert components is not None + + # schemas should be a node stored inside components node + assert "schemas" in components + schemas = components["schemas"] + assert schemas is not None + + for expected_schema in expected_schemas: + assert expected_schema in schemas + + +def test_dump_models_group_requests(tmpdir: Path) -> None: + """Test that selected models can be dump into a JSON file.""" + group = "requests" + filename = tmpdir / "foo.json" + dump_models_group(group, filename) + + # list of schemas expected in a dump + expected_schemas = ( + "ConversationUpdateRequest", + "FeedbackRequest", + "FeedbackStatusUpdateRequest", + "MCPServerRegistrationRequest", + "ModelFilter", + "PromptCreateRequest", + "PromptUpdateRequest", + "QueryRequest", + "ResponsesRequest", + "RlsapiV1Attachment", + "RlsapiV1CLA", + "RlsapiV1Context", + "RlsapiV1InferRequest", + "RlsapiV1SystemInfo", + "RlsapiV1Terminal", + "StreamingInterruptRequest", + "VectorStoreCreateRequest", + "VectorStoreFileCreateRequest", + "VectorStoreUpdateRequest", + ) + check_json_file_content(filename, expected_schemas) + + +def test_dump_models_group_successful_responses(tmpdir: Path) -> None: + """Test that selected models can be dump into a JSON file.""" + group = "successful_responses" + filename = tmpdir / "foo.json" + dump_models_group(group, filename) + + # list of schemas expected in a dump + expected_schemas = ( + "AuthorizedResponse", + "ConfigurationResponse", + "ConversationDeleteResponse", + "ConversationResponse", + "ConversationUpdateResponse", + "ConversationsListResponse", + "ConversationsListResponseV2", + "FeedbackResponse", + "FeedbackStatusUpdateResponse", + "FileResponse", + "InfoResponse", + "LivenessResponse", + "MCPClientAuthOptionsResponse", + "MCPServerDeleteResponse", + "MCPServerListResponse", + "MCPServerRegistrationResponse", + "ModelsResponse", + "PromptDeleteResponse", + "PromptResourceResponse", + "PromptsListResponse", + "ProviderResponse", + "ProvidersListResponse", + "QueryResponse", + "RAGInfoResponse", + "RAGListResponse", + "ReadinessResponse", + "ResponsesResponse", + "RlsapiV1InferData", + "RlsapiV1InferResponse", + "ShieldsResponse", + "StatusResponse", + "StreamingInterruptResponse", + "StreamingQueryResponse", + "ToolsResponse", + "VectorStoreDeleteResponse", + "VectorStoreFileDeleteResponse", + "VectorStoreFileResponse", + "VectorStoreFilesListResponse", + "VectorStoreResponse", + "VectorStoresListResponse", + ) + check_json_file_content(filename, expected_schemas) + + +def test_dump_models_group_error_responses(tmpdir: Path) -> None: + """Test that selected models can be dump into a JSON file.""" + group = "error_responses" + filename = tmpdir / "foo.json" + dump_models_group(group, filename) + + # list of schemas expected in a dump + expected_schemas = ( + "AbstractErrorResponse", + "BadRequestResponse", + "ConflictResponse", + "DetailModel", + "FileTooLargeResponse", + "ForbiddenResponse", + "InternalServerErrorResponse", + "NotFoundResponse", + "PromptTooLongResponse", + "QuotaExceededResponse", + "ServiceUnavailableResponse", + "UnauthorizedResponse", + "UnprocessableEntityResponse", + ) + check_json_file_content(filename, expected_schemas) + + +def test_dump_models_group_common_models(tmpdir: Path) -> None: + """Test that selected models can be dump into a JSON file.""" + group = "common" + filename = tmpdir / "foo.json" + dump_models_group(group, filename) + + # list of schemas expected in a dump + expected_schemas = ( + "Attachment", + "ConversationData", + "ConversationDetails", + "ConversationTurn", + "MCPListToolsSummary", + "MCPServerAuthInfo", + "MCPServerInfo", + "Message", + "ProviderHealthStatus", + "RAGChunk", + "RAGContext", + "ReferencedDocument", + "ShieldModerationBlocked", + "ShieldModerationPassed", + "SolrVectorSearchRequest", + "ToolCallSummary", + "ToolInfoSummary", + "ToolResultSummary", + "Transcript", + "TranscriptMetadata", + "TurnSummary", + ) + check_json_file_content(filename, expected_schemas) + + +def test_dump_models_group_agent_models(tmpdir: Path) -> None: + """Test that selected models can be dump into a JSON file.""" + group = "agents" + filename = tmpdir / "foo.json" + dump_models_group(group, filename) + + # list of schemas expected in a dump + expected_schemas = ( + "EndEventData", + "EndStreamPayload", + "ErrorEventData", + "ErrorStreamPayload", + "InterruptedEventData", + "InterruptedStreamPayload", + "StartEventData", + "StartStreamPayload", + "StreamPayloadBase", + "TokenChunkData", + "TokenStreamPayload", + "ToolCallStreamPayload", + "ToolResultStreamPayload", + "TurnCompleteStreamPayload", + ) + check_json_file_content(filename, expected_schemas) + + +def test_dump_models_common_responses_models(tmpdir: Path) -> None: + """Test that selected models can be dump into a JSON file.""" + group = "common_responses" + filename = tmpdir / "foo.json" + dump_models_group(group, filename) + + # list of schemas expected in a dump + expected_schemas = ( + "InputToolMCP", + "ResponsesApiParams", + ) + check_json_file_content(filename, expected_schemas) + + +def test_dump_models_unknown_group() -> None: + """Test that exception is raised for unknown model group.""" + with pytest.raises(ValueError, match="Unknown model group provided: unknown"): + dump_models_group("unknown")