mcp-oci-integration/new_mcp_consumption.py at main · luigisaetta/mcp-oci-integration · GitHub

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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
"""
MCP Server to query data regarding OCI usage and consumption (refactored).

To test a new organization of the MCP server using a class-based approach.
"""

from typing import Any, Dict

from mcp_base import BaseMCPServer, expose_tool
from utils import get_console_logger
from config import DEBUG

# Your existing domain logic (unchanged)
from consumption_utils import (
    usage_summary_by_service_structured,
    usage_summary_by_compartment_structured,
    fetch_consumption_by_compartment,
    usage_summary_by_service_for_compartment,
)


class OciConsumptionServer(BaseMCPServer):
    """
    Class-based MCP server for OCI Consumption queries.
    """

    def __init__(self, *, debug: bool = False):
        super().__init__("OCI Consumption MCP server", debug=debug)
        self.logger = get_console_logger()

    @expose_tool()
    def usage_summary_by_service(
        self, start_date: str, end_date: str
    ) -> Dict[str, Any]:
        """
        Return the total consumption aggregated by service within a specified time period.

        Args:
            start_date (str): Start date of the period, in ISO format (YYYY-MM-DD).
            end_date (str): End date of the period, in ISO format (YYYY-MM-DD).
                The time window between start_date and end_date must not exceed 93 days.

        Returns:
            dict: A structured dictionary containing consumption details aggregated by service.

        Raises:
            Error: If the time period exceeds 93 days, or any other errors occurs.
        """
        if self.debug:
            self.logger.info("Called usage_summary_by_service...")
        return usage_summary_by_service_structured(start_date, end_date)

    @expose_tool()
    def usage_summary_by_compartment(
        self, start_date: str, end_date: str
    ) -> Dict[str, Any]:
        """
        Return the total consumption aggregated by compartment within a specified period.

        Args:
            start_date (str): Start date of the period, in ISO format (YYYY-MM-DD).
            end_date (str): End date of the period, in ISO format (YYYY-MM-DD).
                The time window between start_date and end_date must not exceed 93 days.

        Returns:
            dict: A structured dictionary containing consumption details aggregated by compartment.

        Raises:
            Error: If the time period exceeds 93 days, or any other errors occurs.
        """
        if self.debug:
            self.logger.info("Called usage_summary_by_compartment...")
        return usage_summary_by_compartment_structured(start_date, end_date)

    @expose_tool(name="usage_breakdown_for_service_by_compartment")
    def breakdown_service_by_compartment(
        self, start_date: str, end_date: str, service_name: str
    ) -> Dict[str, Any]:
        """
        Return the consumption for a specific service within a given time period,
        broken down by compartment.

        Args:
            start_date (str): Start date of the period, in ISO format (YYYY-MM-DD).
            end_date (str): End date of the period, in ISO format (YYYY-MM-DD).
                The time window between start_date and end_date must not exceed 93 days.
            service_name (str): Name of the service to filter by.
                Case-insensitive and substring matches are allowed.

        Returns:
            dict: A structured dictionary containing consumption details by compartment.
                Each entry corresponds to one compartment.

        Raises:
            Error: If the time period exceeds 93 days, or any other errors occurs.
        """
        return fetch_consumption_by_compartment(start_date, end_date, service_name)

    @expose_tool(name="usage_breakdown_for_compartment_by_service")
    def breakdown_compartment_by_service(
        self, start_date: str, end_date: str, compartment_name: str
    ) -> Dict[str, Any]:
        """
        Return the consumption for a specific compartment within a given time period,
        broken down by service.

        Args:
            start_date (str): Start date of the period, in ISO format (YYYY-MM-DD).
            end_date (str): End date of the period, in ISO format (YYYY-MM-DD).
                The time window between start_date and end_date must not exceed 93 days.
            compartment_name (str): Name of the compartment to filter by.
                Case-insensitive and substring matches are allowed.

        Returns:
            dict: A structured dictionary containing consumption details by service.
                Each entry corresponds to one service.
        Raises:
            Error: If the time period exceeds 93 days, or any other errors occurs.
        """
        return usage_summary_by_service_for_compartment(
            start_date, end_date, compartment_name
        )


if __name__ == "__main__":
    OciConsumptionServer(debug=DEBUG).run()