refactor location of proxy

litellm_proxy/common_utils/proxy_state.py

@@ -0,0 +1,36 @@
+"""
+This file is used to store the state variables of the proxy server.
+
+Example: `spend_logs_row_count` is used to store the number of rows in the `LiteLLM_SpendLogs` table.
+"""
+
+from typing import Any, Literal
+
+from litellm_proxy._types import ProxyStateVariables
+
+
+class ProxyState:
+    """
+    Proxy state class has get/set methods for Proxy state variables.
+    """
+
+    # Note: mypy does not recognize when we fetch ProxyStateVariables.annotations.keys(), so we also need to add the valid keys here
+    valid_keys_literal = Literal["spend_logs_row_count"]
+
+    def __init__(self) -> None:
+        self.proxy_state_variables: ProxyStateVariables = ProxyStateVariables(
+            spend_logs_row_count=0,
+        )
+
+    def get_proxy_state_variable(
+        self,
+        variable_name: valid_keys_literal,
+    ) -> Any:
+        return self.proxy_state_variables.get(variable_name, None)
+
+    def set_proxy_state_variable(
+        self,
+        variable_name: valid_keys_literal,
+        value: Any,
+    ) -> None:
+        self.proxy_state_variables[variable_name] = value