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.
- ai_lib_python/__init__.py +62 -43
- ai_lib_python/_features.py +51 -0
- ai_lib_python/batch/__init__.py +15 -15
- ai_lib_python/batch/collector.py +244 -244
- ai_lib_python/batch/executor.py +224 -224
- ai_lib_python/cache/__init__.py +26 -26
- ai_lib_python/cache/backends.py +380 -380
- ai_lib_python/cache/key.py +237 -237
- ai_lib_python/cache/manager.py +332 -332
- ai_lib_python/client/__init__.py +37 -37
- ai_lib_python/client/builder.py +528 -528
- ai_lib_python/client/cancel.py +368 -368
- ai_lib_python/client/core.py +434 -433
- ai_lib_python/client/response.py +134 -134
- ai_lib_python/computer_use/__init__.py +228 -0
- ai_lib_python/drivers/__init__.py +140 -0
- ai_lib_python/drivers/anthropic.py +173 -0
- ai_lib_python/drivers/gemini.py +177 -0
- ai_lib_python/drivers/openai.py +133 -0
- ai_lib_python/embeddings/__init__.py +36 -36
- ai_lib_python/embeddings/client.py +339 -339
- ai_lib_python/embeddings/types.py +234 -234
- ai_lib_python/embeddings/vectors.py +246 -246
- ai_lib_python/errors/__init__.py +54 -41
- ai_lib_python/errors/base.py +325 -316
- ai_lib_python/errors/classification.py +240 -210
- ai_lib_python/errors/standard_codes.py +280 -0
- ai_lib_python/guardrails/__init__.py +35 -35
- ai_lib_python/guardrails/base.py +336 -336
- ai_lib_python/guardrails/filters.py +583 -583
- ai_lib_python/guardrails/validators.py +475 -475
- ai_lib_python/mcp/__init__.py +181 -0
- ai_lib_python/multimodal/__init__.py +138 -0
- ai_lib_python/pipeline/__init__.py +55 -55
- ai_lib_python/pipeline/accumulate.py +248 -248
- ai_lib_python/pipeline/base.py +240 -240
- ai_lib_python/pipeline/decode.py +281 -281
- ai_lib_python/pipeline/event_map.py +507 -506
- ai_lib_python/pipeline/fan_out.py +284 -284
- ai_lib_python/pipeline/select.py +297 -297
- ai_lib_python/plugins/__init__.py +32 -32
- ai_lib_python/plugins/base.py +294 -294
- ai_lib_python/plugins/hooks.py +296 -296
- ai_lib_python/plugins/middleware.py +285 -285
- ai_lib_python/plugins/registry.py +294 -294
- ai_lib_python/protocol/__init__.py +71 -71
- ai_lib_python/protocol/loader.py +317 -317
- ai_lib_python/protocol/manifest.py +385 -385
- ai_lib_python/protocol/v2/__init__.py +22 -0
- ai_lib_python/protocol/v2/capabilities.py +198 -0
- ai_lib_python/protocol/v2/manifest.py +256 -0
- ai_lib_python/protocol/validator.py +460 -460
- ai_lib_python/py.typed +1 -1
- ai_lib_python/registry/__init__.py +174 -0
- ai_lib_python/resilience/__init__.py +102 -102
- ai_lib_python/resilience/backpressure.py +225 -225
- ai_lib_python/resilience/circuit_breaker.py +318 -318
- ai_lib_python/resilience/executor.py +344 -343
- ai_lib_python/resilience/fallback.py +341 -341
- ai_lib_python/resilience/preflight.py +413 -413
- ai_lib_python/resilience/rate_limiter.py +291 -291
- ai_lib_python/resilience/retry.py +299 -299
- ai_lib_python/resilience/signals.py +283 -283
- ai_lib_python/routing/__init__.py +118 -118
- ai_lib_python/routing/manager.py +593 -593
- ai_lib_python/routing/strategy.py +345 -345
- ai_lib_python/routing/types.py +397 -397
- ai_lib_python/structured/__init__.py +33 -33
- ai_lib_python/structured/json_mode.py +281 -281
- ai_lib_python/structured/schema.py +316 -316
- ai_lib_python/structured/validator.py +334 -334
- ai_lib_python/telemetry/__init__.py +127 -127
- ai_lib_python/telemetry/exporters/__init__.py +9 -9
- ai_lib_python/telemetry/exporters/prometheus.py +111 -111
- ai_lib_python/telemetry/feedback.py +446 -446
- ai_lib_python/telemetry/health.py +409 -409
- ai_lib_python/telemetry/logger.py +389 -389
- ai_lib_python/telemetry/metrics.py +496 -496
- ai_lib_python/telemetry/tracer.py +473 -473
- ai_lib_python/tokens/__init__.py +25 -25
- ai_lib_python/tokens/counter.py +282 -282
- ai_lib_python/tokens/estimator.py +286 -286
- ai_lib_python/transport/__init__.py +34 -34
- ai_lib_python/transport/auth.py +141 -141
- ai_lib_python/transport/http.py +381 -364
- ai_lib_python/transport/pool.py +425 -425
- ai_lib_python/types/__init__.py +41 -41
- ai_lib_python/types/events.py +343 -343
- ai_lib_python/types/message.py +332 -332
- ai_lib_python/types/tool.py +191 -191
- ai_lib_python/utils/__init__.py +21 -21
- ai_lib_python/utils/tool_call_assembler.py +317 -317
- {ai_lib_python-0.5.0.dist-info → ai_lib_python-0.7.0.dist-info}/METADATA +211 -52
- ai_lib_python-0.7.0.dist-info/RECORD +97 -0
- {ai_lib_python-0.5.0.dist-info → ai_lib_python-0.7.0.dist-info}/licenses/LICENSE-APACHE +201 -201
- {ai_lib_python-0.5.0.dist-info → ai_lib_python-0.7.0.dist-info}/licenses/LICENSE-MIT +21 -21
- ai_lib_python-0.5.0.dist-info/RECORD +0 -84
- {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
|
-
|
|
3
|
-
|
|
4
|
-
|
|
5
|
-
|
|
6
|
-
|
|
7
|
-
|
|
8
|
-
from
|
|
9
|
-
|
|
10
|
-
from ai_lib_python.
|
|
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
|
-
|
|
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
|
+
)
|
ai_lib_python/batch/__init__.py
CHANGED
|
@@ -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
|
+
]
|
ai_lib_python/batch/collector.py
CHANGED
|
@@ -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
|