digitalkin 0.2.25rc0__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (122) hide show
  1. base_server/server_async_insecure.py +6 -5
  2. base_server/server_async_secure.py +6 -5
  3. base_server/server_sync_insecure.py +5 -4
  4. base_server/server_sync_secure.py +5 -4
  5. digitalkin/__version__.py +1 -1
  6. digitalkin/core/__init__.py +1 -0
  7. digitalkin/core/common/__init__.py +9 -0
  8. digitalkin/core/common/factories.py +156 -0
  9. digitalkin/core/job_manager/__init__.py +1 -0
  10. digitalkin/{modules → core}/job_manager/base_job_manager.py +138 -32
  11. digitalkin/core/job_manager/single_job_manager.py +373 -0
  12. digitalkin/{modules → core}/job_manager/taskiq_broker.py +121 -26
  13. digitalkin/core/job_manager/taskiq_job_manager.py +541 -0
  14. digitalkin/core/task_manager/__init__.py +1 -0
  15. digitalkin/core/task_manager/base_task_manager.py +539 -0
  16. digitalkin/core/task_manager/local_task_manager.py +108 -0
  17. digitalkin/core/task_manager/remote_task_manager.py +87 -0
  18. digitalkin/core/task_manager/surrealdb_repository.py +266 -0
  19. digitalkin/core/task_manager/task_executor.py +249 -0
  20. digitalkin/core/task_manager/task_session.py +368 -0
  21. digitalkin/grpc_servers/__init__.py +1 -19
  22. digitalkin/grpc_servers/_base_server.py +3 -3
  23. digitalkin/grpc_servers/module_server.py +120 -195
  24. digitalkin/grpc_servers/module_servicer.py +81 -44
  25. digitalkin/grpc_servers/utils/__init__.py +1 -0
  26. digitalkin/grpc_servers/utils/exceptions.py +0 -8
  27. digitalkin/grpc_servers/utils/grpc_client_wrapper.py +25 -9
  28. digitalkin/grpc_servers/utils/grpc_error_handler.py +53 -0
  29. digitalkin/grpc_servers/utils/utility_schema_extender.py +100 -0
  30. digitalkin/logger.py +64 -27
  31. digitalkin/mixins/__init__.py +19 -0
  32. digitalkin/mixins/base_mixin.py +10 -0
  33. digitalkin/mixins/callback_mixin.py +24 -0
  34. digitalkin/mixins/chat_history_mixin.py +110 -0
  35. digitalkin/mixins/cost_mixin.py +76 -0
  36. digitalkin/mixins/file_history_mixin.py +93 -0
  37. digitalkin/mixins/filesystem_mixin.py +46 -0
  38. digitalkin/mixins/logger_mixin.py +51 -0
  39. digitalkin/mixins/storage_mixin.py +79 -0
  40. digitalkin/models/__init__.py +1 -1
  41. digitalkin/models/core/__init__.py +1 -0
  42. digitalkin/{modules/job_manager → models/core}/job_manager_models.py +3 -11
  43. digitalkin/models/core/task_monitor.py +74 -0
  44. digitalkin/models/grpc_servers/__init__.py +1 -0
  45. digitalkin/{grpc_servers/utils → models/grpc_servers}/models.py +92 -7
  46. digitalkin/models/module/__init__.py +18 -11
  47. digitalkin/models/module/base_types.py +61 -0
  48. digitalkin/models/module/module.py +9 -1
  49. digitalkin/models/module/module_context.py +282 -6
  50. digitalkin/models/module/module_types.py +29 -105
  51. digitalkin/models/module/setup_types.py +490 -0
  52. digitalkin/models/module/tool_cache.py +68 -0
  53. digitalkin/models/module/tool_reference.py +117 -0
  54. digitalkin/models/module/utility.py +167 -0
  55. digitalkin/models/services/__init__.py +9 -0
  56. digitalkin/models/services/cost.py +1 -0
  57. digitalkin/models/services/registry.py +35 -0
  58. digitalkin/models/services/storage.py +39 -5
  59. digitalkin/modules/__init__.py +5 -1
  60. digitalkin/modules/_base_module.py +265 -167
  61. digitalkin/modules/archetype_module.py +6 -1
  62. digitalkin/modules/tool_module.py +16 -3
  63. digitalkin/modules/trigger_handler.py +7 -6
  64. digitalkin/modules/triggers/__init__.py +8 -0
  65. digitalkin/modules/triggers/healthcheck_ping_trigger.py +45 -0
  66. digitalkin/modules/triggers/healthcheck_services_trigger.py +63 -0
  67. digitalkin/modules/triggers/healthcheck_status_trigger.py +52 -0
  68. digitalkin/services/__init__.py +4 -0
  69. digitalkin/services/communication/__init__.py +7 -0
  70. digitalkin/services/communication/communication_strategy.py +76 -0
  71. digitalkin/services/communication/default_communication.py +101 -0
  72. digitalkin/services/communication/grpc_communication.py +234 -0
  73. digitalkin/services/cost/__init__.py +9 -2
  74. digitalkin/services/cost/grpc_cost.py +9 -42
  75. digitalkin/services/filesystem/default_filesystem.py +0 -2
  76. digitalkin/services/filesystem/grpc_filesystem.py +10 -39
  77. digitalkin/services/registry/__init__.py +22 -1
  78. digitalkin/services/registry/default_registry.py +135 -4
  79. digitalkin/services/registry/exceptions.py +47 -0
  80. digitalkin/services/registry/grpc_registry.py +306 -0
  81. digitalkin/services/registry/registry_models.py +15 -0
  82. digitalkin/services/registry/registry_strategy.py +88 -4
  83. digitalkin/services/services_config.py +25 -3
  84. digitalkin/services/services_models.py +5 -1
  85. digitalkin/services/setup/default_setup.py +6 -7
  86. digitalkin/services/setup/grpc_setup.py +52 -15
  87. digitalkin/services/storage/grpc_storage.py +4 -4
  88. digitalkin/services/user_profile/__init__.py +12 -0
  89. digitalkin/services/user_profile/default_user_profile.py +55 -0
  90. digitalkin/services/user_profile/grpc_user_profile.py +69 -0
  91. digitalkin/services/user_profile/user_profile_strategy.py +25 -0
  92. digitalkin/utils/__init__.py +28 -0
  93. digitalkin/utils/arg_parser.py +1 -1
  94. digitalkin/utils/development_mode_action.py +2 -2
  95. digitalkin/utils/dynamic_schema.py +483 -0
  96. digitalkin/utils/package_discover.py +1 -2
  97. digitalkin/utils/schema_splitter.py +207 -0
  98. {digitalkin-0.2.25rc0.dist-info → digitalkin-0.3.2.dev14.dist-info}/METADATA +11 -30
  99. digitalkin-0.3.2.dev14.dist-info/RECORD +143 -0
  100. {digitalkin-0.2.25rc0.dist-info → digitalkin-0.3.2.dev14.dist-info}/top_level.txt +1 -0
  101. modules/archetype_with_tools_module.py +244 -0
  102. modules/cpu_intensive_module.py +1 -1
  103. modules/dynamic_setup_module.py +338 -0
  104. modules/minimal_llm_module.py +1 -1
  105. modules/text_transform_module.py +1 -1
  106. monitoring/digitalkin_observability/__init__.py +46 -0
  107. monitoring/digitalkin_observability/http_server.py +150 -0
  108. monitoring/digitalkin_observability/interceptors.py +176 -0
  109. monitoring/digitalkin_observability/metrics.py +201 -0
  110. monitoring/digitalkin_observability/prometheus.py +137 -0
  111. monitoring/tests/test_metrics.py +172 -0
  112. services/filesystem_module.py +7 -5
  113. services/storage_module.py +4 -2
  114. digitalkin/grpc_servers/registry_server.py +0 -65
  115. digitalkin/grpc_servers/registry_servicer.py +0 -456
  116. digitalkin/grpc_servers/utils/factory.py +0 -180
  117. digitalkin/modules/job_manager/single_job_manager.py +0 -294
  118. digitalkin/modules/job_manager/taskiq_job_manager.py +0 -290
  119. digitalkin-0.2.25rc0.dist-info/RECORD +0 -89
  120. /digitalkin/{grpc_servers/utils → models/grpc_servers}/types.py +0 -0
  121. {digitalkin-0.2.25rc0.dist-info → digitalkin-0.3.2.dev14.dist-info}/WHEEL +0 -0
  122. {digitalkin-0.2.25rc0.dist-info → digitalkin-0.3.2.dev14.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,338 @@
1
+ """Example module demonstrating dynamic schema fields in SetupModel.
2
+
3
+ This example shows how to use the Dynamic metadata class with async fetchers
4
+ to populate field schemas (like enums) at runtime. This is useful when the
5
+ available options come from external sources like databases or APIs.
6
+
7
+ Usage:
8
+ # Start the module server
9
+ python examples/modules/dynamic_setup_module.py
10
+
11
+ # Or import and use in your own code
12
+ from examples.modules.dynamic_setup_module import DynamicSetupModule
13
+ """
14
+
15
+ import asyncio
16
+ import logging
17
+ from typing import Annotated, Any, ClassVar
18
+
19
+ from pydantic import BaseModel, Field
20
+
21
+ from digitalkin.models.module.module_context import ModuleContext
22
+ from digitalkin.models.module.module_types import DataModel, DataTrigger, SetupModel
23
+ from digitalkin.modules._base_module import BaseModule
24
+ from digitalkin.services.services_models import ServicesStrategy
25
+ from digitalkin.utils import Dynamic
26
+
27
+ # Configure logging
28
+ logging.basicConfig(
29
+ level=logging.DEBUG,
30
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
31
+ )
32
+ logger = logging.getLogger(__name__)
33
+
34
+
35
+ # =============================================================================
36
+ # Simulated External Services (replace with real implementations)
37
+ # =============================================================================
38
+
39
+
40
+ class MockModelRegistry:
41
+ """Simulates an external model registry service.
42
+
43
+ In a real application, this would be a connection to a database,
44
+ API service, or configuration management system.
45
+ """
46
+
47
+ _models: ClassVar[list[str]] = ["gpt-4", "gpt-4-turbo", "gpt-3.5-turbo"]
48
+ _languages: ClassVar[list[str]] = ["en", "fr", "de", "es", "it", "pt"]
49
+
50
+ @classmethod
51
+ async def fetch_available_models(cls) -> list[str]:
52
+ """Fetch available models from the registry.
53
+
54
+ Simulates an async API call with a small delay.
55
+ """
56
+ await asyncio.sleep(0.1) # Simulate network latency
57
+ logger.info("Fetched %d models from registry", len(cls._models))
58
+ return cls._models.copy()
59
+
60
+ @classmethod
61
+ async def fetch_supported_languages(cls) -> list[str]:
62
+ """Fetch supported languages from the registry."""
63
+ await asyncio.sleep(0.05) # Simulate network latency
64
+ logger.info("Fetched %d languages from registry", len(cls._languages))
65
+ return cls._languages.copy()
66
+
67
+ @classmethod
68
+ def get_default_model(cls) -> str:
69
+ """Get the default model (sync fetcher example)."""
70
+ return cls._models[0] if cls._models else "gpt-4"
71
+
72
+
73
+ # =============================================================================
74
+ # Dynamic Fetcher Functions
75
+ # =============================================================================
76
+
77
+
78
+ async def fetch_models() -> list[str]:
79
+ """Async fetcher for available model names.
80
+
81
+ This function is called when SetupModel.get_clean_model(force=True)
82
+ is invoked, typically during module initialization or schema refresh.
83
+ """
84
+ return await MockModelRegistry.fetch_available_models()
85
+
86
+
87
+ async def fetch_languages() -> list[str]:
88
+ """Async fetcher for supported languages."""
89
+ return await MockModelRegistry.fetch_supported_languages()
90
+
91
+
92
+ def get_temperature_range() -> dict[str, float]:
93
+ """Sync fetcher example returning min/max for temperature.
94
+
95
+ Demonstrates that fetchers can return any JSON-serializable value,
96
+ not just lists for enums.
97
+ """
98
+ return {"minimum": 0.0, "maximum": 2.0}
99
+
100
+
101
+ # =============================================================================
102
+ # Setup Model with Dynamic Fields
103
+ # =============================================================================
104
+
105
+
106
+ class DynamicAgentSetup(SetupModel):
107
+ """Setup model demonstrating dynamic schema fields.
108
+
109
+ Fields marked with Dynamic(...) will have their schema values
110
+ refreshed at runtime when get_clean_model(force=True) is called.
111
+
112
+ Attributes:
113
+ model_name: The LLM model to use. Enum values fetched from registry.
114
+ language: Output language. Enum values fetched dynamically.
115
+ temperature: Sampling temperature. Static field for comparison.
116
+ max_tokens: Maximum tokens to generate.
117
+ system_prompt: The system prompt for the model.
118
+ """
119
+
120
+ # Dynamic field: enum values fetched asynchronously from model registry
121
+ model_name: Annotated[str, Dynamic(enum=fetch_models)] = Field(
122
+ default="gpt-4",
123
+ title="Model Name",
124
+ description="The LLM model to use for generation.",
125
+ json_schema_extra={
126
+ "config": True, # Shown in initial configuration
127
+ "ui:widget": "select",
128
+ },
129
+ )
130
+
131
+ # Dynamic field: language options fetched asynchronously
132
+ language: Annotated[str, Dynamic(enum=fetch_languages)] = Field(
133
+ default="en",
134
+ title="Output Language",
135
+ description="The language for generated responses.",
136
+ json_schema_extra={
137
+ "config": True,
138
+ "ui:widget": "select",
139
+ },
140
+ )
141
+
142
+ # Static field: no dynamic fetcher, values defined at class time
143
+ temperature: float = Field(
144
+ default=0.7,
145
+ ge=0.0,
146
+ le=2.0,
147
+ title="Temperature",
148
+ description="Controls randomness. Higher values = more creative.",
149
+ json_schema_extra={"config": True},
150
+ )
151
+
152
+ # Static field with hidden flag (runtime-only, not in initial config)
153
+ max_tokens: int = Field(
154
+ default=1024,
155
+ ge=1,
156
+ le=4096,
157
+ title="Max Tokens",
158
+ description="Maximum tokens in the response.",
159
+ json_schema_extra={"hidden": True},
160
+ )
161
+
162
+ # Static field without any special flags
163
+ system_prompt: str = Field(
164
+ default="You are a helpful assistant.",
165
+ title="System Prompt",
166
+ description="The system prompt defining assistant behavior.",
167
+ )
168
+
169
+
170
+ # =============================================================================
171
+ # Input/Output Models (Using DataModel/DataTrigger pattern)
172
+ # =============================================================================
173
+
174
+
175
+ class MessageInputTrigger(DataTrigger):
176
+ """Message input trigger following DigitalKin DataTrigger pattern.
177
+
178
+ The protocol field determines which trigger handler processes this input.
179
+ """
180
+
181
+ protocol: str = "message"
182
+ content: str = Field(default="", description="The user message content.")
183
+
184
+
185
+ class DynamicModuleInput(DataModel[MessageInputTrigger]):
186
+ """Input model following DigitalKin DataModel pattern.
187
+
188
+ Wraps the trigger in a root field with optional annotations.
189
+ """
190
+
191
+ root: MessageInputTrigger = Field(default_factory=MessageInputTrigger)
192
+
193
+
194
+ class MessageOutputTrigger(DataTrigger):
195
+ """Message output trigger following DigitalKin DataTrigger pattern."""
196
+
197
+ protocol: str = "message"
198
+ content: str = Field(default="", description="The generated response.")
199
+ model_used: str = Field(default="", description="The model that generated this response.")
200
+ language: str = Field(default="", description="The output language.")
201
+
202
+
203
+ class DynamicModuleOutput(DataModel[MessageOutputTrigger]):
204
+ """Output model following DigitalKin DataModel pattern."""
205
+
206
+ root: MessageOutputTrigger = Field(default_factory=MessageOutputTrigger)
207
+
208
+
209
+ class DynamicModuleSecret(BaseModel):
210
+ """Secret model (empty for this example)."""
211
+
212
+
213
+ # =============================================================================
214
+ # Module Implementation
215
+ # =============================================================================
216
+
217
+
218
+ class DynamicSetupModule(
219
+ BaseModule[
220
+ DynamicModuleInput,
221
+ DynamicModuleOutput,
222
+ DynamicAgentSetup,
223
+ DynamicModuleSecret,
224
+ ]
225
+ ):
226
+ """Example module demonstrating dynamic schema in SetupModel.
227
+
228
+ This module shows how to:
229
+ 1. Define setup fields with Dynamic() metadata for runtime enum fetching
230
+ 2. Mix static and dynamic fields in the same SetupModel
231
+ 3. Use async fetchers that simulate external service calls
232
+ 4. Follow DigitalKin's DataModel/DataTrigger pattern for I/O
233
+
234
+ The key integration point is in the gRPC servicer, which calls
235
+ SetupModel.get_clean_model(force=True) to refresh dynamic values
236
+ before returning schema information to clients.
237
+ """
238
+
239
+ name = "DynamicSetupModule"
240
+ description = "Demonstrates dynamic schema fields in module setup"
241
+
242
+ # Schema format definitions
243
+ input_format = DynamicModuleInput
244
+ output_format = DynamicModuleOutput
245
+ setup_format = DynamicAgentSetup
246
+ secret_format = DynamicModuleSecret
247
+
248
+ # Module metadata
249
+ metadata: ClassVar[dict[str, Any]] = {
250
+ "name": "DynamicSetupModule",
251
+ "description": "Example module with dynamic setup schema",
252
+ "version": "1.0.0",
253
+ "tags": ["example", "dynamic-schema"],
254
+ }
255
+
256
+ # Services configuration (empty for this example)
257
+ services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]] = {}
258
+ services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]] = {}
259
+
260
+ async def initialize(self, context: ModuleContext, setup_data: DynamicAgentSetup) -> None:
261
+ """Initialize the module with setup data.
262
+
263
+ Args:
264
+ context: The module context with services and session info.
265
+ setup_data: The validated setup configuration.
266
+ """
267
+ logger.info(
268
+ "Initializing DynamicSetupModule with model=%s, language=%s",
269
+ setup_data.model_name,
270
+ setup_data.language,
271
+ )
272
+ self.setup = setup_data
273
+
274
+ async def cleanup(self) -> None:
275
+ """Clean up resources."""
276
+ logger.info("Cleaning up DynamicSetupModule")
277
+
278
+
279
+ # =============================================================================
280
+ # Demonstration Script
281
+ # =============================================================================
282
+
283
+
284
+ async def demonstrate_dynamic_schema() -> None:
285
+ """Demonstrate the dynamic schema functionality."""
286
+ # 1. Show schema WITHOUT force (dynamic fields not resolved)
287
+
288
+ model_no_force = await DynamicAgentSetup.get_clean_model(
289
+ config_fields=True,
290
+ hidden_fields=False,
291
+ force=False,
292
+ )
293
+ schema_no_force = model_no_force.model_json_schema()
294
+
295
+ # Check if enum is present
296
+ model_name_schema = schema_no_force.get("properties", {}).get("model_name", {})
297
+ if "enum" in model_name_schema:
298
+ pass
299
+
300
+ # 2. Show schema WITH force (dynamic fields resolved)
301
+
302
+ model_with_force = await DynamicAgentSetup.get_clean_model(
303
+ config_fields=True,
304
+ hidden_fields=False,
305
+ force=True,
306
+ )
307
+ schema_with_force = model_with_force.model_json_schema()
308
+
309
+ # Check enum values after force
310
+ model_name_schema = schema_with_force.get("properties", {}).get("model_name", {})
311
+ if "enum" in model_name_schema:
312
+ pass
313
+
314
+ language_schema = schema_with_force.get("properties", {}).get("language", {})
315
+ if "enum" in language_schema:
316
+ pass
317
+
318
+ # 3. Show that static json_schema_extra is preserved
319
+
320
+ # 4. Show field filtering
321
+
322
+ # Config fields only (hidden excluded)
323
+ await DynamicAgentSetup.get_clean_model(
324
+ config_fields=True,
325
+ hidden_fields=False,
326
+ force=False,
327
+ )
328
+
329
+ # All fields including hidden
330
+ await DynamicAgentSetup.get_clean_model(
331
+ config_fields=True,
332
+ hidden_fields=True,
333
+ force=False,
334
+ )
335
+
336
+
337
+ if __name__ == "__main__":
338
+ asyncio.run(demonstrate_dynamic_schema())
@@ -6,9 +6,9 @@ from collections.abc import Callable
6
6
  from typing import Any, ClassVar, Literal
7
7
 
8
8
  import openai
9
+ from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
9
10
  from pydantic import BaseModel, Field
10
11
 
11
- from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
12
12
  from digitalkin.modules._base_module import BaseModule
13
13
  from digitalkin.services.services_models import ServicesStrategy
14
14
 
@@ -4,9 +4,9 @@ import logging
4
4
  from collections.abc import Callable
5
5
  from typing import Any, ClassVar
6
6
 
7
+ from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
7
8
  from pydantic import BaseModel
8
9
 
9
- from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
10
10
  from digitalkin.modules._base_module import BaseModule
11
11
  from digitalkin.services.setup.setup_strategy import SetupData
12
12
  from digitalkin.services.storage.storage_strategy import DataType, StorageRecord
@@ -0,0 +1,46 @@
1
+ """Standalone observability module for DigitalKin.
2
+
3
+ This module can be copied into your project and used independently.
4
+ It has no dependencies on the digitalkin package.
5
+
6
+ Usage:
7
+ from digitalkin_observability import (
8
+ MetricsCollector,
9
+ MetricsServer,
10
+ MetricsServerInterceptor,
11
+ PrometheusExporter,
12
+ get_metrics,
13
+ start_metrics_server,
14
+ stop_metrics_server,
15
+ )
16
+
17
+ # Start metrics HTTP server
18
+ start_metrics_server(port=8081)
19
+
20
+ # Track metrics
21
+ metrics = get_metrics()
22
+ metrics.inc_jobs_started("my_module")
23
+ metrics.inc_jobs_completed("my_module", duration=1.5)
24
+
25
+ # Export to Prometheus format
26
+ print(PrometheusExporter.export())
27
+ """
28
+
29
+ from digitalkin_observability.http_server import (
30
+ MetricsServer,
31
+ start_metrics_server,
32
+ stop_metrics_server,
33
+ )
34
+ from digitalkin_observability.interceptors import MetricsServerInterceptor
35
+ from digitalkin_observability.metrics import MetricsCollector, get_metrics
36
+ from digitalkin_observability.prometheus import PrometheusExporter
37
+
38
+ __all__ = [
39
+ "MetricsCollector",
40
+ "MetricsServer",
41
+ "MetricsServerInterceptor",
42
+ "PrometheusExporter",
43
+ "get_metrics",
44
+ "start_metrics_server",
45
+ "stop_metrics_server",
46
+ ]
@@ -0,0 +1,150 @@
1
+ """Simple HTTP server for exposing Prometheus metrics.
2
+
3
+ This module provides an HTTP server that exposes metrics at /metrics endpoint.
4
+ No external dependencies required beyond Python standard library.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import logging
10
+ from http.server import BaseHTTPRequestHandler, HTTPServer
11
+ from threading import Thread
12
+ from typing import TYPE_CHECKING, ClassVar
13
+
14
+ if TYPE_CHECKING:
15
+ from typing import Self
16
+
17
+ logger = logging.getLogger(__name__)
18
+
19
+
20
+ class MetricsHandler(BaseHTTPRequestHandler):
21
+ """HTTP request handler for metrics endpoint."""
22
+
23
+ def do_GET(self) -> None:
24
+ """Handle GET requests."""
25
+ if self.path == "/metrics":
26
+ self._serve_metrics()
27
+ elif self.path == "/health":
28
+ self._serve_health()
29
+ else:
30
+ self.send_error(404, "Not Found")
31
+
32
+ def _serve_metrics(self) -> None:
33
+ """Serve Prometheus metrics."""
34
+ from digitalkin_observability.prometheus import PrometheusExporter
35
+
36
+ content = PrometheusExporter.export()
37
+ self.send_response(200)
38
+ self.send_header("Content-Type", "text/plain; charset=utf-8")
39
+ self.send_header("Content-Length", str(len(content)))
40
+ self.end_headers()
41
+ self.wfile.write(content.encode("utf-8"))
42
+
43
+ def _serve_health(self) -> None:
44
+ """Serve health check."""
45
+ content = '{"status": "ok"}'
46
+ self.send_response(200)
47
+ self.send_header("Content-Type", "application/json")
48
+ self.send_header("Content-Length", str(len(content)))
49
+ self.end_headers()
50
+ self.wfile.write(content.encode("utf-8"))
51
+
52
+ def log_message(self, format: str, *args: object) -> None:
53
+ """Suppress default logging."""
54
+
55
+
56
+ class MetricsServer:
57
+ """HTTP server for exposing metrics to Prometheus.
58
+
59
+ Usage:
60
+ server = MetricsServer(port=8081)
61
+ server.start()
62
+ # ... run your application ...
63
+ server.stop()
64
+
65
+ Or as context manager:
66
+ with MetricsServer(port=8081):
67
+ # ... run your application ...
68
+
69
+ Or as async context manager:
70
+ async with MetricsServer(port=8081):
71
+ # ... run your application ...
72
+ """
73
+
74
+ instance: ClassVar["MetricsServer | None"] = None
75
+
76
+ def __init__(self, host: str = "0.0.0.0", port: int = 8081) -> None:
77
+ """Initialize the metrics server.
78
+
79
+ Args:
80
+ host: Host to bind to (default: 0.0.0.0 for all interfaces).
81
+ port: Port to listen on (default: 8081).
82
+ """
83
+ self.host = host
84
+ self.port = port
85
+ self._server: HTTPServer | None = None
86
+ self._thread: Thread | None = None
87
+
88
+ def start(self) -> None:
89
+ """Start the metrics server in a background thread."""
90
+ if self._server is not None:
91
+ logger.warning("Metrics server already running")
92
+ return
93
+
94
+ self._server = HTTPServer((self.host, self.port), MetricsHandler)
95
+ self._thread = Thread(target=self._server.serve_forever, daemon=True)
96
+ self._thread.start()
97
+ logger.info(
98
+ "Metrics server started on http://%s:%s/metrics",
99
+ self.host,
100
+ self.port,
101
+ )
102
+
103
+ def stop(self) -> None:
104
+ """Stop the metrics server."""
105
+ if self._server is not None:
106
+ self._server.shutdown()
107
+ self._server = None
108
+ self._thread = None
109
+ logger.info("Metrics server stopped")
110
+
111
+ async def __aenter__(self) -> "Self":
112
+ """Async context manager entry."""
113
+ self.start()
114
+ return self
115
+
116
+ async def __aexit__(self, *args: object) -> None:
117
+ """Async context manager exit."""
118
+ self.stop()
119
+
120
+ def __enter__(self) -> "Self":
121
+ """Context manager entry."""
122
+ self.start()
123
+ return self
124
+
125
+ def __exit__(self, *args: object) -> None:
126
+ """Context manager exit."""
127
+ self.stop()
128
+
129
+
130
+ def start_metrics_server(host: str = "0.0.0.0", port: int = 8081) -> MetricsServer:
131
+ """Start a metrics server singleton.
132
+
133
+ Args:
134
+ host: Host to bind to.
135
+ port: Port to listen on.
136
+
137
+ Returns:
138
+ The MetricsServer instance.
139
+ """
140
+ if MetricsServer.instance is None:
141
+ MetricsServer.instance = MetricsServer(host, port)
142
+ MetricsServer.instance.start()
143
+ return MetricsServer.instance
144
+
145
+
146
+ def stop_metrics_server() -> None:
147
+ """Stop the metrics server singleton."""
148
+ if MetricsServer.instance is not None:
149
+ MetricsServer.instance.stop()
150
+ MetricsServer.instance = None