ai-lib-python 0.5.0__py3-none-any.whl → 0.7.0__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 (98) hide show
  1. ai_lib_python/__init__.py +62 -43
  2. ai_lib_python/_features.py +51 -0
  3. ai_lib_python/batch/__init__.py +15 -15
  4. ai_lib_python/batch/collector.py +244 -244
  5. ai_lib_python/batch/executor.py +224 -224
  6. ai_lib_python/cache/__init__.py +26 -26
  7. ai_lib_python/cache/backends.py +380 -380
  8. ai_lib_python/cache/key.py +237 -237
  9. ai_lib_python/cache/manager.py +332 -332
  10. ai_lib_python/client/__init__.py +37 -37
  11. ai_lib_python/client/builder.py +528 -528
  12. ai_lib_python/client/cancel.py +368 -368
  13. ai_lib_python/client/core.py +434 -433
  14. ai_lib_python/client/response.py +134 -134
  15. ai_lib_python/computer_use/__init__.py +228 -0
  16. ai_lib_python/drivers/__init__.py +140 -0
  17. ai_lib_python/drivers/anthropic.py +173 -0
  18. ai_lib_python/drivers/gemini.py +177 -0
  19. ai_lib_python/drivers/openai.py +133 -0
  20. ai_lib_python/embeddings/__init__.py +36 -36
  21. ai_lib_python/embeddings/client.py +339 -339
  22. ai_lib_python/embeddings/types.py +234 -234
  23. ai_lib_python/embeddings/vectors.py +246 -246
  24. ai_lib_python/errors/__init__.py +54 -41
  25. ai_lib_python/errors/base.py +325 -316
  26. ai_lib_python/errors/classification.py +240 -210
  27. ai_lib_python/errors/standard_codes.py +280 -0
  28. ai_lib_python/guardrails/__init__.py +35 -35
  29. ai_lib_python/guardrails/base.py +336 -336
  30. ai_lib_python/guardrails/filters.py +583 -583
  31. ai_lib_python/guardrails/validators.py +475 -475
  32. ai_lib_python/mcp/__init__.py +181 -0
  33. ai_lib_python/multimodal/__init__.py +138 -0
  34. ai_lib_python/pipeline/__init__.py +55 -55
  35. ai_lib_python/pipeline/accumulate.py +248 -248
  36. ai_lib_python/pipeline/base.py +240 -240
  37. ai_lib_python/pipeline/decode.py +281 -281
  38. ai_lib_python/pipeline/event_map.py +507 -506
  39. ai_lib_python/pipeline/fan_out.py +284 -284
  40. ai_lib_python/pipeline/select.py +297 -297
  41. ai_lib_python/plugins/__init__.py +32 -32
  42. ai_lib_python/plugins/base.py +294 -294
  43. ai_lib_python/plugins/hooks.py +296 -296
  44. ai_lib_python/plugins/middleware.py +285 -285
  45. ai_lib_python/plugins/registry.py +294 -294
  46. ai_lib_python/protocol/__init__.py +71 -71
  47. ai_lib_python/protocol/loader.py +317 -317
  48. ai_lib_python/protocol/manifest.py +385 -385
  49. ai_lib_python/protocol/v2/__init__.py +22 -0
  50. ai_lib_python/protocol/v2/capabilities.py +198 -0
  51. ai_lib_python/protocol/v2/manifest.py +256 -0
  52. ai_lib_python/protocol/validator.py +460 -460
  53. ai_lib_python/py.typed +1 -1
  54. ai_lib_python/registry/__init__.py +174 -0
  55. ai_lib_python/resilience/__init__.py +102 -102
  56. ai_lib_python/resilience/backpressure.py +225 -225
  57. ai_lib_python/resilience/circuit_breaker.py +318 -318
  58. ai_lib_python/resilience/executor.py +344 -343
  59. ai_lib_python/resilience/fallback.py +341 -341
  60. ai_lib_python/resilience/preflight.py +413 -413
  61. ai_lib_python/resilience/rate_limiter.py +291 -291
  62. ai_lib_python/resilience/retry.py +299 -299
  63. ai_lib_python/resilience/signals.py +283 -283
  64. ai_lib_python/routing/__init__.py +118 -118
  65. ai_lib_python/routing/manager.py +593 -593
  66. ai_lib_python/routing/strategy.py +345 -345
  67. ai_lib_python/routing/types.py +397 -397
  68. ai_lib_python/structured/__init__.py +33 -33
  69. ai_lib_python/structured/json_mode.py +281 -281
  70. ai_lib_python/structured/schema.py +316 -316
  71. ai_lib_python/structured/validator.py +334 -334
  72. ai_lib_python/telemetry/__init__.py +127 -127
  73. ai_lib_python/telemetry/exporters/__init__.py +9 -9
  74. ai_lib_python/telemetry/exporters/prometheus.py +111 -111
  75. ai_lib_python/telemetry/feedback.py +446 -446
  76. ai_lib_python/telemetry/health.py +409 -409
  77. ai_lib_python/telemetry/logger.py +389 -389
  78. ai_lib_python/telemetry/metrics.py +496 -496
  79. ai_lib_python/telemetry/tracer.py +473 -473
  80. ai_lib_python/tokens/__init__.py +25 -25
  81. ai_lib_python/tokens/counter.py +282 -282
  82. ai_lib_python/tokens/estimator.py +286 -286
  83. ai_lib_python/transport/__init__.py +34 -34
  84. ai_lib_python/transport/auth.py +141 -141
  85. ai_lib_python/transport/http.py +381 -364
  86. ai_lib_python/transport/pool.py +425 -425
  87. ai_lib_python/types/__init__.py +41 -41
  88. ai_lib_python/types/events.py +343 -343
  89. ai_lib_python/types/message.py +332 -332
  90. ai_lib_python/types/tool.py +191 -191
  91. ai_lib_python/utils/__init__.py +21 -21
  92. ai_lib_python/utils/tool_call_assembler.py +317 -317
  93. {ai_lib_python-0.5.0.dist-info → ai_lib_python-0.7.0.dist-info}/METADATA +211 -52
  94. ai_lib_python-0.7.0.dist-info/RECORD +97 -0
  95. {ai_lib_python-0.5.0.dist-info → ai_lib_python-0.7.0.dist-info}/licenses/LICENSE-APACHE +201 -201
  96. {ai_lib_python-0.5.0.dist-info → ai_lib_python-0.7.0.dist-info}/licenses/LICENSE-MIT +21 -21
  97. ai_lib_python-0.5.0.dist-info/RECORD +0 -84
  98. {ai_lib_python-0.5.0.dist-info → ai_lib_python-0.7.0.dist-info}/WHEEL +0 -0
ai_lib_python/__init__.py CHANGED
@@ -1,43 +1,62 @@
1
- """
2
- ai-lib-python: Official Python Runtime for AI-Protocol
3
-
4
- The canonical Pythonic implementation for unified AI model interaction.
5
- Core principle: All logic is operators, all configuration is protocol.
6
- """
7
-
8
- from ai_lib_python.client import AiClient, AiClientBuilder, CallStats, ChatResponse
9
- from ai_lib_python.errors import AiLibError, ProtocolError, TransportError
10
- from ai_lib_python.types.events import StreamingEvent
11
- from ai_lib_python.types.message import (
12
- ContentBlock,
13
- Message,
14
- MessageContent,
15
- MessageRole,
16
- )
17
- from ai_lib_python.types.tool import ToolCall, ToolDefinition
18
-
19
- __version__ = "0.5.0"
20
-
21
- __all__ = [
22
- # Client
23
- "AiClient",
24
- "AiClientBuilder",
25
- # Errors
26
- "AiLibError",
27
- "CallStats",
28
- "ChatResponse",
29
- "ContentBlock",
30
- # Types - Message
31
- "Message",
32
- "MessageContent",
33
- "MessageRole",
34
- "ProtocolError",
35
- # Types - Events
36
- "StreamingEvent",
37
- "ToolCall",
38
- # Types - Tool
39
- "ToolDefinition",
40
- "TransportError",
41
- # Version
42
- "__version__",
43
- ]
1
+ """AI-Protocol 官方 Python 运行时:提供统一的多厂商 AI 模型交互接口。
2
+
3
+ ai-lib-python: Official Python Runtime for AI-Protocol.
4
+
5
+ The canonical Pythonic implementation for unified AI model interaction.
6
+ Core principle: All logic is operators, all configuration is protocol.
7
+ """
8
+ from __future__ import annotations
9
+
10
+ from ai_lib_python._features import (
11
+ HAS_AUDIO,
12
+ HAS_KEYRING,
13
+ HAS_TELEMETRY,
14
+ HAS_TOKENIZER,
15
+ HAS_VISION,
16
+ HAS_WATCHDOG,
17
+ require_extra,
18
+ )
19
+ from ai_lib_python.client import AiClient, AiClientBuilder, CallStats, ChatResponse
20
+ from ai_lib_python.errors import AiLibError, ProtocolError, TransportError
21
+ from ai_lib_python.types.events import StreamingEvent
22
+ from ai_lib_python.types.message import (
23
+ ContentBlock,
24
+ Message,
25
+ MessageContent,
26
+ MessageRole,
27
+ )
28
+ from ai_lib_python.types.tool import ToolCall, ToolDefinition
29
+
30
+ __version__ = "0.7.0"
31
+
32
+ __all__ = [
33
+ # Client
34
+ "AiClient",
35
+ "AiClientBuilder",
36
+ # Feature flags
37
+ "HAS_AUDIO",
38
+ "HAS_KEYRING",
39
+ "HAS_TELEMETRY",
40
+ "HAS_TOKENIZER",
41
+ "HAS_VISION",
42
+ "HAS_WATCHDOG",
43
+ "require_extra",
44
+ # Errors
45
+ "AiLibError",
46
+ "CallStats",
47
+ "ChatResponse",
48
+ "ContentBlock",
49
+ # Types - Message
50
+ "Message",
51
+ "MessageContent",
52
+ "MessageRole",
53
+ "ProtocolError",
54
+ # Types - Events
55
+ "StreamingEvent",
56
+ "ToolCall",
57
+ # Types - Tool
58
+ "ToolDefinition",
59
+ "TransportError",
60
+ # Version
61
+ "__version__",
62
+ ]
@@ -0,0 +1,51 @@
1
+ """运行时特性检测:检查可选依赖以确定可用的能力模块。
2
+
3
+ Runtime feature detection for optional extras.
4
+
5
+ Checks availability of optional dependencies to determine which
6
+ capability modules can be used.
7
+ """
8
+ from __future__ import annotations
9
+
10
+
11
+ def _check_import(module_name: str) -> bool:
12
+ """Check if a module is importable."""
13
+ try:
14
+ __import__(module_name)
15
+ return True
16
+ except ImportError:
17
+ return False
18
+
19
+
20
+ # Capability feature flags
21
+ HAS_VISION: bool = _check_import("PIL")
22
+ HAS_AUDIO: bool = _check_import("soundfile")
23
+ HAS_TELEMETRY: bool = _check_import("opentelemetry")
24
+ HAS_TOKENIZER: bool = _check_import("tiktoken")
25
+ HAS_KEYRING: bool = _check_import("keyring")
26
+ HAS_WATCHDOG: bool = _check_import("watchdog")
27
+
28
+
29
+ # Map pip package names to import module names (when they differ)
30
+ _PACKAGE_TO_MODULE: dict[str, str] = {
31
+ "pillow": "PIL",
32
+ }
33
+
34
+
35
+ def require_extra(extra_name: str, package_name: str) -> None:
36
+ """Raise ImportError with installation hint if extra is not available.
37
+
38
+ Args:
39
+ extra_name: Name of the pip extra (e.g., 'vision')
40
+ package_name: Name of the required package (e.g., 'pillow')
41
+
42
+ Raises:
43
+ ImportError: With installation instructions when package is not available.
44
+ """
45
+ module_name = _PACKAGE_TO_MODULE.get(package_name, package_name)
46
+ if _check_import(module_name):
47
+ return
48
+ raise ImportError(
49
+ f"The '{extra_name}' extra is required for this feature. "
50
+ f"Install it with: pip install ai-lib-python[{extra_name}]"
51
+ )
@@ -1,15 +1,15 @@
1
- """
2
- Batch processing module for ai-lib-python.
3
-
4
- Provides request batching and batch execution utilities.
5
- """
6
-
7
- from ai_lib_python.batch.collector import BatchCollector, BatchConfig
8
- from ai_lib_python.batch.executor import BatchExecutor, BatchResult
9
-
10
- __all__ = [
11
- "BatchCollector",
12
- "BatchConfig",
13
- "BatchExecutor",
14
- "BatchResult",
15
- ]
1
+ """
2
+ Batch processing module for ai-lib-python.
3
+
4
+ Provides request batching and batch execution utilities.
5
+ """
6
+
7
+ from ai_lib_python.batch.collector import BatchCollector, BatchConfig
8
+ from ai_lib_python.batch.executor import BatchExecutor, BatchResult
9
+
10
+ __all__ = [
11
+ "BatchCollector",
12
+ "BatchConfig",
13
+ "BatchExecutor",
14
+ "BatchResult",
15
+ ]
@@ -1,244 +1,244 @@
1
- """
2
- Batch collector for grouping requests.
3
-
4
- Collects requests and groups them for batch execution.
5
- """
6
-
7
- from __future__ import annotations
8
-
9
- import asyncio
10
- import time
11
- from dataclasses import dataclass, field
12
- from typing import TYPE_CHECKING, Any, Generic, TypeVar
13
-
14
- if TYPE_CHECKING:
15
- from collections.abc import Awaitable, Callable
16
-
17
- T = TypeVar("T")
18
- R = TypeVar("R")
19
-
20
-
21
- @dataclass
22
- class BatchConfig:
23
- """Configuration for batch collection.
24
-
25
- Attributes:
26
- max_batch_size: Maximum requests per batch
27
- max_wait_ms: Maximum wait time before flushing batch
28
- group_by: Function to determine grouping key
29
- """
30
-
31
- max_batch_size: int = 10
32
- max_wait_ms: float = 100.0
33
- group_by: Callable[[Any], str] | None = None
34
-
35
- @classmethod
36
- def default(cls) -> BatchConfig:
37
- """Create default configuration."""
38
- return cls()
39
-
40
- @classmethod
41
- def for_embeddings(cls) -> BatchConfig:
42
- """Create configuration optimized for embeddings."""
43
- return cls(
44
- max_batch_size=100,
45
- max_wait_ms=50.0,
46
- )
47
-
48
- @classmethod
49
- def for_chat(cls) -> BatchConfig:
50
- """Create configuration for chat completions."""
51
- return cls(
52
- max_batch_size=5,
53
- max_wait_ms=10.0,
54
- )
55
-
56
-
57
- @dataclass
58
- class PendingRequest(Generic[T]):
59
- """A pending request waiting for batch execution.
60
-
61
- Attributes:
62
- data: Request data
63
- future: Future to resolve with result
64
- added_at: Timestamp when added
65
- group_key: Grouping key
66
- """
67
-
68
- data: T
69
- future: asyncio.Future[Any]
70
- added_at: float = field(default_factory=time.time)
71
- group_key: str = "_default_"
72
-
73
-
74
- class BatchCollector(Generic[T, R]):
75
- """Collects requests for batch processing.
76
-
77
- Accumulates requests until batch size or time limit is reached,
78
- then triggers batch execution.
79
-
80
- Example:
81
- >>> async def process_batch(items):
82
- ... return [f"result_{i}" for i in range(len(items))]
83
- ...
84
- >>> collector = BatchCollector(
85
- ... config=BatchConfig(max_batch_size=5),
86
- ... executor=process_batch,
87
- ... )
88
- >>> result = await collector.add("request1")
89
- """
90
-
91
- def __init__(
92
- self,
93
- config: BatchConfig | None = None,
94
- executor: Callable[[list[T]], Awaitable[list[R]]] | None = None,
95
- ) -> None:
96
- """Initialize batch collector.
97
-
98
- Args:
99
- config: Batch configuration
100
- executor: Function to execute batches
101
- """
102
- self._config = config or BatchConfig.default()
103
- self._executor = executor
104
- self._pending: dict[str, list[PendingRequest[T]]] = {}
105
- self._lock = asyncio.Lock()
106
- self._timers: dict[str, asyncio.Task[None]] = {}
107
- self._running = True
108
-
109
- def set_executor(
110
- self, executor: Callable[[list[T]], Awaitable[list[R]]]
111
- ) -> None:
112
- """Set the batch executor function.
113
-
114
- Args:
115
- executor: Function to execute batches
116
- """
117
- self._executor = executor
118
-
119
- async def add(self, data: T) -> R:
120
- """Add a request to the batch.
121
-
122
- Args:
123
- data: Request data
124
-
125
- Returns:
126
- Result from batch execution
127
- """
128
- if not self._running:
129
- raise RuntimeError("Batch collector is stopped")
130
-
131
- if self._executor is None:
132
- raise RuntimeError("No executor set")
133
-
134
- # Determine group key
135
- group_key = (
136
- self._config.group_by(data) if self._config.group_by else "_default_"
137
- )
138
-
139
- # Create future for result
140
- loop = asyncio.get_event_loop()
141
- future: asyncio.Future[R] = loop.create_future()
142
-
143
- async with self._lock:
144
- # Add to pending
145
- if group_key not in self._pending:
146
- self._pending[group_key] = []
147
-
148
- self._pending[group_key].append(
149
- PendingRequest(data=data, future=future, group_key=group_key)
150
- )
151
-
152
- # Check if batch is full
153
- if len(self._pending[group_key]) >= self._config.max_batch_size:
154
- await self._flush_group(group_key)
155
- else:
156
- # Start timer if not already running
157
- if group_key not in self._timers or self._timers[group_key].done():
158
- self._timers[group_key] = asyncio.create_task(
159
- self._timer_flush(group_key)
160
- )
161
-
162
- return await future
163
-
164
- async def _timer_flush(self, group_key: str) -> None:
165
- """Flush group after timeout."""
166
- await asyncio.sleep(self._config.max_wait_ms / 1000.0)
167
-
168
- async with self._lock:
169
- if self._pending.get(group_key):
170
- await self._flush_group(group_key)
171
-
172
- async def _flush_group(self, group_key: str) -> None:
173
- """Flush a specific group.
174
-
175
- Args:
176
- group_key: Group to flush
177
- """
178
- if group_key not in self._pending or not self._pending[group_key]:
179
- return
180
-
181
- # Get pending requests
182
- requests = self._pending.pop(group_key)
183
-
184
- # Cancel timer
185
- if group_key in self._timers:
186
- self._timers[group_key].cancel()
187
- del self._timers[group_key]
188
-
189
- # Extract data
190
- data_list = [r.data for r in requests]
191
-
192
- try:
193
- # Execute batch
194
- results = await self._executor(data_list)
195
-
196
- # Resolve futures
197
- for request, result in zip(requests, results, strict=False):
198
- if not request.future.done():
199
- request.future.set_result(result)
200
-
201
- except Exception as e:
202
- # Reject all futures
203
- for request in requests:
204
- if not request.future.done():
205
- request.future.set_exception(e)
206
-
207
- async def flush(self) -> None:
208
- """Flush all pending batches."""
209
- async with self._lock:
210
- for group_key in list(self._pending.keys()):
211
- await self._flush_group(group_key)
212
-
213
- async def stop(self) -> None:
214
- """Stop the collector and flush pending requests."""
215
- self._running = False
216
- await self.flush()
217
-
218
- # Cancel all timers
219
- for timer in self._timers.values():
220
- timer.cancel()
221
- self._timers.clear()
222
-
223
- def get_pending_count(self, group_key: str | None = None) -> int:
224
- """Get count of pending requests.
225
-
226
- Args:
227
- group_key: Group to count (all if None)
228
-
229
- Returns:
230
- Number of pending requests
231
- """
232
- if group_key:
233
- return len(self._pending.get(group_key, []))
234
- return sum(len(p) for p in self._pending.values())
235
-
236
- @property
237
- def config(self) -> BatchConfig:
238
- """Get batch configuration."""
239
- return self._config
240
-
241
- @property
242
- def is_running(self) -> bool:
243
- """Check if collector is running."""
244
- return self._running
1
+ """
2
+ Batch collector for grouping requests.
3
+
4
+ Collects requests and groups them for batch execution.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import asyncio
10
+ import time
11
+ from dataclasses import dataclass, field
12
+ from typing import TYPE_CHECKING, Any, Generic, TypeVar
13
+
14
+ if TYPE_CHECKING:
15
+ from collections.abc import Awaitable, Callable
16
+
17
+ T = TypeVar("T")
18
+ R = TypeVar("R")
19
+
20
+
21
+ @dataclass
22
+ class BatchConfig:
23
+ """Configuration for batch collection.
24
+
25
+ Attributes:
26
+ max_batch_size: Maximum requests per batch
27
+ max_wait_ms: Maximum wait time before flushing batch
28
+ group_by: Function to determine grouping key
29
+ """
30
+
31
+ max_batch_size: int = 10
32
+ max_wait_ms: float = 100.0
33
+ group_by: Callable[[Any], str] | None = None
34
+
35
+ @classmethod
36
+ def default(cls) -> BatchConfig:
37
+ """Create default configuration."""
38
+ return cls()
39
+
40
+ @classmethod
41
+ def for_embeddings(cls) -> BatchConfig:
42
+ """Create configuration optimized for embeddings."""
43
+ return cls(
44
+ max_batch_size=100,
45
+ max_wait_ms=50.0,
46
+ )
47
+
48
+ @classmethod
49
+ def for_chat(cls) -> BatchConfig:
50
+ """Create configuration for chat completions."""
51
+ return cls(
52
+ max_batch_size=5,
53
+ max_wait_ms=10.0,
54
+ )
55
+
56
+
57
+ @dataclass
58
+ class PendingRequest(Generic[T]):
59
+ """A pending request waiting for batch execution.
60
+
61
+ Attributes:
62
+ data: Request data
63
+ future: Future to resolve with result
64
+ added_at: Timestamp when added
65
+ group_key: Grouping key
66
+ """
67
+
68
+ data: T
69
+ future: asyncio.Future[Any]
70
+ added_at: float = field(default_factory=time.time)
71
+ group_key: str = "_default_"
72
+
73
+
74
+ class BatchCollector(Generic[T, R]):
75
+ """Collects requests for batch processing.
76
+
77
+ Accumulates requests until batch size or time limit is reached,
78
+ then triggers batch execution.
79
+
80
+ Example:
81
+ >>> async def process_batch(items):
82
+ ... return [f"result_{i}" for i in range(len(items))]
83
+ ...
84
+ >>> collector = BatchCollector(
85
+ ... config=BatchConfig(max_batch_size=5),
86
+ ... executor=process_batch,
87
+ ... )
88
+ >>> result = await collector.add("request1")
89
+ """
90
+
91
+ def __init__(
92
+ self,
93
+ config: BatchConfig | None = None,
94
+ executor: Callable[[list[T]], Awaitable[list[R]]] | None = None,
95
+ ) -> None:
96
+ """Initialize batch collector.
97
+
98
+ Args:
99
+ config: Batch configuration
100
+ executor: Function to execute batches
101
+ """
102
+ self._config = config or BatchConfig.default()
103
+ self._executor = executor
104
+ self._pending: dict[str, list[PendingRequest[T]]] = {}
105
+ self._lock = asyncio.Lock()
106
+ self._timers: dict[str, asyncio.Task[None]] = {}
107
+ self._running = True
108
+
109
+ def set_executor(
110
+ self, executor: Callable[[list[T]], Awaitable[list[R]]]
111
+ ) -> None:
112
+ """Set the batch executor function.
113
+
114
+ Args:
115
+ executor: Function to execute batches
116
+ """
117
+ self._executor = executor
118
+
119
+ async def add(self, data: T) -> R:
120
+ """Add a request to the batch.
121
+
122
+ Args:
123
+ data: Request data
124
+
125
+ Returns:
126
+ Result from batch execution
127
+ """
128
+ if not self._running:
129
+ raise RuntimeError("Batch collector is stopped")
130
+
131
+ if self._executor is None:
132
+ raise RuntimeError("No executor set")
133
+
134
+ # Determine group key
135
+ group_key = (
136
+ self._config.group_by(data) if self._config.group_by else "_default_"
137
+ )
138
+
139
+ # Create future for result
140
+ loop = asyncio.get_event_loop()
141
+ future: asyncio.Future[R] = loop.create_future()
142
+
143
+ async with self._lock:
144
+ # Add to pending
145
+ if group_key not in self._pending:
146
+ self._pending[group_key] = []
147
+
148
+ self._pending[group_key].append(
149
+ PendingRequest(data=data, future=future, group_key=group_key)
150
+ )
151
+
152
+ # Check if batch is full
153
+ if len(self._pending[group_key]) >= self._config.max_batch_size:
154
+ await self._flush_group(group_key)
155
+ else:
156
+ # Start timer if not already running
157
+ if group_key not in self._timers or self._timers[group_key].done():
158
+ self._timers[group_key] = asyncio.create_task(
159
+ self._timer_flush(group_key)
160
+ )
161
+
162
+ return await future
163
+
164
+ async def _timer_flush(self, group_key: str) -> None:
165
+ """Flush group after timeout."""
166
+ await asyncio.sleep(self._config.max_wait_ms / 1000.0)
167
+
168
+ async with self._lock:
169
+ if self._pending.get(group_key):
170
+ await self._flush_group(group_key)
171
+
172
+ async def _flush_group(self, group_key: str) -> None:
173
+ """Flush a specific group.
174
+
175
+ Args:
176
+ group_key: Group to flush
177
+ """
178
+ if group_key not in self._pending or not self._pending[group_key]:
179
+ return
180
+
181
+ # Get pending requests
182
+ requests = self._pending.pop(group_key)
183
+
184
+ # Cancel timer
185
+ if group_key in self._timers:
186
+ self._timers[group_key].cancel()
187
+ del self._timers[group_key]
188
+
189
+ # Extract data
190
+ data_list = [r.data for r in requests]
191
+
192
+ try:
193
+ # Execute batch
194
+ results = await self._executor(data_list)
195
+
196
+ # Resolve futures
197
+ for request, result in zip(requests, results, strict=False):
198
+ if not request.future.done():
199
+ request.future.set_result(result)
200
+
201
+ except Exception as e:
202
+ # Reject all futures
203
+ for request in requests:
204
+ if not request.future.done():
205
+ request.future.set_exception(e)
206
+
207
+ async def flush(self) -> None:
208
+ """Flush all pending batches."""
209
+ async with self._lock:
210
+ for group_key in list(self._pending.keys()):
211
+ await self._flush_group(group_key)
212
+
213
+ async def stop(self) -> None:
214
+ """Stop the collector and flush pending requests."""
215
+ self._running = False
216
+ await self.flush()
217
+
218
+ # Cancel all timers
219
+ for timer in self._timers.values():
220
+ timer.cancel()
221
+ self._timers.clear()
222
+
223
+ def get_pending_count(self, group_key: str | None = None) -> int:
224
+ """Get count of pending requests.
225
+
226
+ Args:
227
+ group_key: Group to count (all if None)
228
+
229
+ Returns:
230
+ Number of pending requests
231
+ """
232
+ if group_key:
233
+ return len(self._pending.get(group_key, []))
234
+ return sum(len(p) for p in self._pending.values())
235
+
236
+ @property
237
+ def config(self) -> BatchConfig:
238
+ """Get batch configuration."""
239
+ return self._config
240
+
241
+ @property
242
+ def is_running(self) -> bool:
243
+ """Check if collector is running."""
244
+ return self._running