mosec 0.9.0

Model Serving made Efficient in the Cloud.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

// Copyright 2023 MOSEC Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use utoipa::openapi::{Components, OpenApi};

use crate::config::Route;

#[derive(Default, Clone)]
pub(crate) struct MosecOpenAPI {
    pub api: OpenApi,
}

impl MosecOpenAPI {
    /// Merge the route request_body/response/schemas into the OpenAPI.
    pub fn merge_route(&mut self, route: &Route) -> &mut Self {
        let reserved = match route.is_sse {
            true => "/openapi/reserved/inference",
            false => "/openapi/reserved/inference_sse",
        };
        let mut path = self.api.paths.paths.get(reserved).unwrap().clone();
        if let Some(mut op) = path.post.clone() {
            if let Some(mut user_schemas) = route.schemas.clone() {
                if self.api.components.is_none() {
                    self.api.components = Some(Components::default());
                }
                self.api
                    .components
                    .as_mut()
                    .unwrap()
                    .schemas
                    .append(&mut user_schemas);
            };
            if let Some(req) = route.request_body.clone() {
                op.request_body = Some(req);
            };

            if let Some(mut responses) = route.responses.clone() {
                op.responses.responses.append(&mut responses);
            };
            path.post = Some(op);
        }
        self.api.paths.paths.insert(route.endpoint.clone(), path);

        self
    }

    /// Removes the reserved paths from the OpenAPI spec.
    pub fn clean(&mut self) -> &mut Self {
        self.api.paths.paths.remove("/openapi/reserved/inference");
        self.api
            .paths
            .paths
            .remove("/openapi/reserved/inference_sse");
        self
    }
}