embed-client 0.0.1__tar.gz → 1.0.0.1__tar.gz
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.
- {embed_client-0.0.1 → embed_client-1.0.0.1}/PKG-INFO +1 -1
- {embed_client-0.0.1 → embed_client-1.0.0.1}/embed_client/async_client.py +101 -15
- embed_client-1.0.0.1/embed_client/example_async_usage.py +69 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/embed_client.egg-info/PKG-INFO +1 -1
- {embed_client-0.0.1 → embed_client-1.0.0.1}/embed_client.egg-info/SOURCES.txt +1 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/pyproject.toml +1 -1
- {embed_client-0.0.1 → embed_client-1.0.0.1}/tests/test_async_client.py +104 -28
- {embed_client-0.0.1 → embed_client-1.0.0.1}/tests/test_async_client_real.py +15 -6
- embed_client-1.0.0.1/tests/test_async_client_stress.py +303 -0
- embed_client-0.0.1/embed_client/example_async_usage.py +0 -94
- {embed_client-0.0.1 → embed_client-1.0.0.1}/README.md +0 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/embed_client/__init__.py +0 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/embed_client/example_async_usage_ru.py +0 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/embed_client.egg-info/dependency_links.txt +0 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/embed_client.egg-info/requires.txt +0 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/embed_client.egg-info/top_level.txt +0 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/setup.cfg +0 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/tests/test_example_async_usage.py +0 -0
- {embed_client-0.0.1 → embed_client-1.0.0.1}/tests/test_example_async_usage_ru.py +0 -0
@@ -8,6 +8,7 @@ Async client for Embedding Service API (OpenAPI 3.0.2)
|
|
8
8
|
|
9
9
|
from typing import Any, Dict, List, Optional, Union
|
10
10
|
import aiohttp
|
11
|
+
import os
|
11
12
|
|
12
13
|
class EmbeddingServiceError(Exception):
|
13
14
|
"""Base exception for EmbeddingServiceAsyncClient."""
|
@@ -37,13 +38,13 @@ class EmbeddingServiceAsyncClient:
|
|
37
38
|
Raises:
|
38
39
|
ValueError: If base_url or port is not provided.
|
39
40
|
"""
|
40
|
-
def __init__(self, base_url: str, port: int):
|
41
|
-
|
41
|
+
def __init__(self, base_url: Optional[str] = None, port: Optional[int] = None):
|
42
|
+
self.base_url = base_url or os.getenv("EMBEDDING_SERVICE_BASE_URL", "http://localhost")
|
43
|
+
if not self.base_url:
|
42
44
|
raise ValueError("base_url must be provided.")
|
43
|
-
|
45
|
+
self.port = port or int(os.getenv("EMBEDDING_SERVICE_PORT", "8001"))
|
46
|
+
if self.port is None:
|
44
47
|
raise ValueError("port must be provided.")
|
45
|
-
self.base_url = base_url.rstrip("/")
|
46
|
-
self.port = port
|
47
48
|
self._session: Optional[aiohttp.ClientSession] = None
|
48
49
|
|
49
50
|
def _make_url(self, path: str, base_url: Optional[str] = None, port: Optional[int] = None) -> str:
|
@@ -51,6 +52,23 @@ class EmbeddingServiceAsyncClient:
|
|
51
52
|
port_val = port if port is not None else self.port
|
52
53
|
return f"{url}:{port_val}{path}"
|
53
54
|
|
55
|
+
def _format_error_response(self, error: str, lang: Optional[str] = None, text: Optional[str] = None) -> Dict[str, Any]:
|
56
|
+
"""
|
57
|
+
Format error response in a standard way.
|
58
|
+
Args:
|
59
|
+
error (str): Error message
|
60
|
+
lang (str, optional): Language of the text that caused the error
|
61
|
+
text (str, optional): Text that caused the error
|
62
|
+
Returns:
|
63
|
+
dict: Formatted error response
|
64
|
+
"""
|
65
|
+
response = {"error": f"Embedding service error: {error}"}
|
66
|
+
if lang is not None:
|
67
|
+
response["lang"] = lang
|
68
|
+
if text is not None:
|
69
|
+
response["text"] = text
|
70
|
+
return response
|
71
|
+
|
54
72
|
async def __aenter__(self):
|
55
73
|
self._session = aiohttp.ClientSession()
|
56
74
|
return self
|
@@ -135,6 +153,37 @@ class EmbeddingServiceAsyncClient:
|
|
135
153
|
except Exception as e:
|
136
154
|
raise EmbeddingServiceError(f"Unexpected error: {e}") from e
|
137
155
|
|
156
|
+
def _validate_texts(self, texts: List[str]) -> None:
|
157
|
+
"""
|
158
|
+
Validate input texts before sending to the API.
|
159
|
+
Args:
|
160
|
+
texts (List[str]): List of texts to validate
|
161
|
+
Raises:
|
162
|
+
EmbeddingServiceAPIError: If texts are invalid
|
163
|
+
"""
|
164
|
+
if not texts:
|
165
|
+
raise EmbeddingServiceAPIError({
|
166
|
+
"code": -32602,
|
167
|
+
"message": "Empty texts list provided"
|
168
|
+
})
|
169
|
+
|
170
|
+
invalid_texts = []
|
171
|
+
for i, text in enumerate(texts):
|
172
|
+
if not isinstance(text, str):
|
173
|
+
invalid_texts.append(f"Text at index {i} is not a string")
|
174
|
+
continue
|
175
|
+
if not text or not text.strip():
|
176
|
+
invalid_texts.append(f"Text at index {i} is empty or contains only whitespace")
|
177
|
+
elif len(text.strip()) < 2: # Минимальная длина текста
|
178
|
+
invalid_texts.append(f"Text at index {i} is too short (minimum 2 characters)")
|
179
|
+
|
180
|
+
if invalid_texts:
|
181
|
+
raise EmbeddingServiceAPIError({
|
182
|
+
"code": -32602,
|
183
|
+
"message": "Invalid input texts",
|
184
|
+
"details": invalid_texts
|
185
|
+
})
|
186
|
+
|
138
187
|
async def cmd(self, command: str, params: Optional[Dict[str, Any]] = None, base_url: Optional[str] = None, port: Optional[int] = None) -> Dict[str, Any]:
|
139
188
|
"""
|
140
189
|
Execute a command via JSON-RPC protocol.
|
@@ -144,32 +193,69 @@ class EmbeddingServiceAsyncClient:
|
|
144
193
|
base_url (str, optional): Override base URL.
|
145
194
|
port (int, optional): Override port.
|
146
195
|
Returns:
|
147
|
-
dict: Command execution result
|
196
|
+
dict: Command execution result or error response in format:
|
197
|
+
{
|
198
|
+
"error": {
|
199
|
+
"code": <код ошибки>,
|
200
|
+
"message": <сообщение об ошибке>,
|
201
|
+
"details": <опциональные детали ошибки>
|
202
|
+
}
|
203
|
+
}
|
204
|
+
или
|
205
|
+
{
|
206
|
+
"result": {
|
207
|
+
"success": true,
|
208
|
+
"data": {
|
209
|
+
"embeddings": [[...], ...]
|
210
|
+
}
|
211
|
+
}
|
212
|
+
}
|
148
213
|
"""
|
214
|
+
if not command:
|
215
|
+
raise EmbeddingServiceAPIError({
|
216
|
+
"code": -32602,
|
217
|
+
"message": "Command is required"
|
218
|
+
})
|
219
|
+
|
220
|
+
# Валидация текстов для команды embed
|
221
|
+
if command == "embed" and params and "texts" in params:
|
222
|
+
self._validate_texts(params["texts"])
|
223
|
+
|
149
224
|
url = self._make_url("/cmd", base_url, port)
|
150
225
|
payload = {"command": command}
|
151
226
|
if params is not None:
|
152
227
|
payload["params"] = params
|
228
|
+
|
153
229
|
try:
|
154
230
|
async with self._session.post(url, json=payload) as resp:
|
155
231
|
await self._raise_for_status(resp)
|
156
232
|
data = await resp.json()
|
157
|
-
|
233
|
+
|
158
234
|
if "error" in data:
|
159
235
|
raise EmbeddingServiceAPIError(data["error"])
|
236
|
+
|
237
|
+
if "result" in data:
|
238
|
+
res = data["result"]
|
239
|
+
if isinstance(res, dict) and "success" in res and res["success"] is False:
|
240
|
+
if "error" in res:
|
241
|
+
raise EmbeddingServiceAPIError(res["error"])
|
242
|
+
|
160
243
|
return data
|
161
|
-
|
162
|
-
raise
|
163
|
-
except EmbeddingServiceHTTPError:
|
164
|
-
raise
|
165
|
-
except EmbeddingServiceConnectionError:
|
166
|
-
raise
|
244
|
+
|
167
245
|
except aiohttp.ClientConnectionError as e:
|
168
|
-
raise
|
246
|
+
raise EmbeddingServiceAPIError({
|
247
|
+
"code": -32000,
|
248
|
+
"message": f"Connection error: {e}"
|
249
|
+
}) from e
|
169
250
|
except aiohttp.ClientResponseError as e:
|
170
251
|
raise EmbeddingServiceHTTPError(e.status, e.message) from e
|
252
|
+
except EmbeddingServiceHTTPError:
|
253
|
+
raise
|
171
254
|
except Exception as e:
|
172
|
-
raise
|
255
|
+
raise EmbeddingServiceAPIError({
|
256
|
+
"code": -32000,
|
257
|
+
"message": f"Unexpected error: {e}"
|
258
|
+
}) from e
|
173
259
|
|
174
260
|
async def _raise_for_status(self, resp: aiohttp.ClientResponse):
|
175
261
|
try:
|
@@ -0,0 +1,69 @@
|
|
1
|
+
"""
|
2
|
+
Example usage of EmbeddingServiceAsyncClient.
|
3
|
+
|
4
|
+
This example demonstrates how to use the async client to check the health of the embedding service,
|
5
|
+
request embeddings, and handle all possible errors.
|
6
|
+
|
7
|
+
Run this script with:
|
8
|
+
python -m asyncio embed_client/example_async_usage.py --base-url http://localhost --port 8001
|
9
|
+
|
10
|
+
You can also set EMBED_CLIENT_BASE_URL and EMBED_CLIENT_PORT environment variables.
|
11
|
+
"""
|
12
|
+
|
13
|
+
import asyncio
|
14
|
+
import sys
|
15
|
+
import os
|
16
|
+
from embed_client.async_client import EmbeddingServiceAsyncClient
|
17
|
+
|
18
|
+
def get_params():
|
19
|
+
base_url = None
|
20
|
+
port = None
|
21
|
+
for i, arg in enumerate(sys.argv):
|
22
|
+
if arg in ("--base-url", "-b") and i + 1 < len(sys.argv):
|
23
|
+
base_url = sys.argv[i + 1]
|
24
|
+
if arg in ("--port", "-p") and i + 1 < len(sys.argv):
|
25
|
+
port = sys.argv[i + 1]
|
26
|
+
if not base_url:
|
27
|
+
base_url = os.environ.get("EMBED_CLIENT_BASE_URL")
|
28
|
+
if not port:
|
29
|
+
port = os.environ.get("EMBED_CLIENT_PORT")
|
30
|
+
if not base_url or not port:
|
31
|
+
print("Error: base_url and port must be provided via --base-url/--port arguments or EMBED_CLIENT_BASE_URL/EMBED_CLIENT_PORT environment variables.")
|
32
|
+
sys.exit(1)
|
33
|
+
return None, None
|
34
|
+
return base_url, int(port)
|
35
|
+
|
36
|
+
async def main():
|
37
|
+
base_url, port = get_params()
|
38
|
+
async with EmbeddingServiceAsyncClient(base_url=base_url, port=port) as client:
|
39
|
+
# Check health
|
40
|
+
health = await client.health()
|
41
|
+
print("Service health:", health)
|
42
|
+
|
43
|
+
# Request embeddings for a list of texts
|
44
|
+
texts = ["hello world", "test embedding"]
|
45
|
+
result = await client.cmd("embed", params={"texts": texts})
|
46
|
+
|
47
|
+
if "error" in result:
|
48
|
+
print(f"Error occurred: {result['error']}")
|
49
|
+
if "lang" in result:
|
50
|
+
print(f"Language: {result['lang']}")
|
51
|
+
if "text" in result:
|
52
|
+
print(f"Text: {result['text']}")
|
53
|
+
else:
|
54
|
+
vectors = result["result"]
|
55
|
+
print(f"Embeddings for {len(texts)} texts:")
|
56
|
+
for i, vec in enumerate(vectors):
|
57
|
+
print(f" Text: {texts[i]!r}\n Vector: {vec[:5]}... (total {len(vec)} dims)")
|
58
|
+
|
59
|
+
# Example: error handling for invalid command
|
60
|
+
result = await client.cmd("health")
|
61
|
+
print("Health check result:", result)
|
62
|
+
|
63
|
+
# Example: error handling for empty command
|
64
|
+
# result = await client.cmd("")
|
65
|
+
# if "error" in result:
|
66
|
+
# print(f"Error for empty command: {result['error']}")
|
67
|
+
|
68
|
+
if __name__ == "__main__":
|
69
|
+
asyncio.run(main())
|
@@ -2,6 +2,7 @@ import pytest
|
|
2
2
|
import pytest_asyncio
|
3
3
|
from unittest.mock import patch, MagicMock
|
4
4
|
from embed_client.async_client import EmbeddingServiceAsyncClient, EmbeddingServiceAPIError, EmbeddingServiceHTTPError, EmbeddingServiceError, EmbeddingServiceConnectionError
|
5
|
+
import aiohttp
|
5
6
|
|
6
7
|
BASE_URL = "http://testserver"
|
7
8
|
PORT = 1234
|
@@ -67,18 +68,80 @@ async def test_cmd(client):
|
|
67
68
|
mock_post.assert_called_with(make_url("/cmd"), json={"command": "embed", "params": {"texts": ["abc"]}})
|
68
69
|
|
69
70
|
@pytest.mark.asyncio
|
70
|
-
async def test_init_requires_base_url_and_port():
|
71
|
-
|
72
|
-
|
73
|
-
|
74
|
-
|
71
|
+
async def test_init_requires_base_url_and_port(monkeypatch):
|
72
|
+
# Сохраняем и очищаем переменные окружения
|
73
|
+
monkeypatch.delenv("EMBEDDING_SERVICE_BASE_URL", raising=False)
|
74
|
+
monkeypatch.delenv("EMBEDDING_SERVICE_PORT", raising=False)
|
75
|
+
# Если не передано ничего и нет переменных окружения, будет дефолт
|
76
|
+
client = EmbeddingServiceAsyncClient()
|
77
|
+
assert client.base_url == "http://localhost"
|
78
|
+
assert client.port == 8001
|
79
|
+
# Если явно передан base_url и port
|
80
|
+
client2 = EmbeddingServiceAsyncClient(base_url="http://test", port=1234)
|
81
|
+
assert client2.base_url == "http://test"
|
82
|
+
assert client2.port == 1234
|
83
|
+
|
84
|
+
@pytest.mark.asyncio
|
85
|
+
async def test_cmd_empty_command(client):
|
86
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
87
|
+
await client.cmd("")
|
88
|
+
assert "Command is required" in str(excinfo.value)
|
89
|
+
|
90
|
+
@pytest.mark.asyncio
|
91
|
+
async def test_cmd_connection_error(client):
|
92
|
+
with patch.object(client._session, 'post', side_effect=aiohttp.ClientConnectionError("Connection failed")):
|
93
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
94
|
+
await client.cmd("embed", params={"texts": ["abc"]})
|
95
|
+
assert "Connection error" in str(excinfo.value)
|
96
|
+
|
97
|
+
@pytest.mark.asyncio
|
98
|
+
async def test_cmd_http_error(client):
|
99
|
+
with patch.object(client._session, 'post', side_effect=aiohttp.ClientResponseError(
|
100
|
+
request_info=MagicMock(),
|
101
|
+
history=(),
|
102
|
+
status=500,
|
103
|
+
message="Internal Server Error"
|
104
|
+
)):
|
105
|
+
with pytest.raises(EmbeddingServiceHTTPError) as excinfo:
|
106
|
+
await client.cmd("embed", params={"texts": ["abc"]})
|
107
|
+
assert "HTTP 500" in str(excinfo.value)
|
108
|
+
|
109
|
+
@pytest.mark.asyncio
|
110
|
+
async def test_cmd_api_error(client):
|
111
|
+
mock_response = MockAiohttpResponse(json_data={"error": "Invalid command"})
|
112
|
+
with patch.object(client._session, 'post', return_value=mock_response):
|
113
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
114
|
+
await client.cmd("invalid_command")
|
115
|
+
assert "Invalid command" in str(excinfo.value)
|
116
|
+
|
117
|
+
@pytest.mark.asyncio
|
118
|
+
async def test_cmd_with_lang_and_text(client):
|
119
|
+
mock_response = MockAiohttpResponse(json_data={"error": "Invalid text"})
|
120
|
+
with patch.object(client._session, 'post', return_value=mock_response):
|
121
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
122
|
+
await client.cmd("embed", params={
|
123
|
+
"texts": ["test"],
|
124
|
+
"lang": "en",
|
125
|
+
"text": "test text"
|
126
|
+
})
|
127
|
+
assert "Invalid text" in str(excinfo.value)
|
128
|
+
|
129
|
+
@pytest.mark.asyncio
|
130
|
+
async def test_cmd_success(client):
|
131
|
+
mock_response = MockAiohttpResponse(json_data={"result": [[1.0, 2.0, 3.0]]})
|
132
|
+
with patch.object(client._session, 'post', return_value=mock_response):
|
133
|
+
result = await client.cmd("embed", params={"texts": ["test"]})
|
134
|
+
assert "result" in result
|
135
|
+
assert result["result"] == [[1.0, 2.0, 3.0]]
|
75
136
|
|
76
137
|
# Некорректные параметры: не-строка в texts
|
77
138
|
@pytest.mark.asyncio
|
78
139
|
async def test_embed_non_string_text(client):
|
79
|
-
|
80
|
-
|
81
|
-
|
140
|
+
"""Test validation of non-string values in texts list."""
|
141
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
142
|
+
await client.cmd("embed", params={"texts": [123, "ok"]})
|
143
|
+
assert "Invalid input texts" in str(excinfo.value)
|
144
|
+
assert "Text at index 0 is not a string" in str(excinfo.value)
|
82
145
|
|
83
146
|
# Некорректные параметры: невалидный params
|
84
147
|
@pytest.mark.asyncio
|
@@ -183,25 +246,6 @@ async def test_get_commands_unexpected_error(client):
|
|
183
246
|
with pytest.raises(EmbeddingServiceError):
|
184
247
|
await client.get_commands()
|
185
248
|
|
186
|
-
# Аналогично для cmd
|
187
|
-
@pytest.mark.asyncio
|
188
|
-
async def test_cmd_http_error(client):
|
189
|
-
with patch.object(client._session, 'post', side_effect=EmbeddingServiceHTTPError(500, "fail")):
|
190
|
-
with pytest.raises(EmbeddingServiceHTTPError):
|
191
|
-
await client.cmd("embed", params={"texts": ["abc"]})
|
192
|
-
|
193
|
-
@pytest.mark.asyncio
|
194
|
-
async def test_cmd_connection_error(client):
|
195
|
-
with patch.object(client._session, 'post', side_effect=EmbeddingServiceConnectionError("fail")):
|
196
|
-
with pytest.raises(EmbeddingServiceConnectionError):
|
197
|
-
await client.cmd("embed", params={"texts": ["abc"]})
|
198
|
-
|
199
|
-
@pytest.mark.asyncio
|
200
|
-
async def test_cmd_unexpected_error(client):
|
201
|
-
with patch.object(client._session, 'post', side_effect=ValueError("fail")):
|
202
|
-
with pytest.raises(EmbeddingServiceError):
|
203
|
-
await client.cmd("embed", params={"texts": ["abc"]})
|
204
|
-
|
205
249
|
# Покрытие: _raise_for_status - ClientResponseError
|
206
250
|
@pytest.mark.asyncio
|
207
251
|
async def test_raise_for_status_http_error():
|
@@ -238,4 +282,36 @@ async def test_aenter_aexit_exceptions():
|
|
238
282
|
raise RuntimeError("fail")
|
239
283
|
client._session = BadSession()
|
240
284
|
with pytest.raises(RuntimeError):
|
241
|
-
await client.__aexit__(None, None, None)
|
285
|
+
await client.__aexit__(None, None, None)
|
286
|
+
|
287
|
+
@pytest.mark.asyncio
|
288
|
+
async def test_embed_validation():
|
289
|
+
"""Test validation of input texts for embed command."""
|
290
|
+
client = EmbeddingServiceAsyncClient()
|
291
|
+
|
292
|
+
# Test empty texts list
|
293
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
294
|
+
await client.cmd("embed", params={"texts": []})
|
295
|
+
assert "Empty texts list provided" in str(excinfo.value)
|
296
|
+
|
297
|
+
# Test empty strings
|
298
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
299
|
+
await client.cmd("embed", params={"texts": ["", " "]})
|
300
|
+
assert "Invalid input texts" in str(excinfo.value)
|
301
|
+
assert "Text at index 0 is empty" in str(excinfo.value)
|
302
|
+
assert "Text at index 1 is empty" in str(excinfo.value)
|
303
|
+
|
304
|
+
# Test too short texts
|
305
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
306
|
+
await client.cmd("embed", params={"texts": ["a", "b"]})
|
307
|
+
assert "Invalid input texts" in str(excinfo.value)
|
308
|
+
assert "Text at index 0 is too short" in str(excinfo.value)
|
309
|
+
assert "Text at index 1 is too short" in str(excinfo.value)
|
310
|
+
|
311
|
+
# Test mixed valid and invalid texts
|
312
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
313
|
+
await client.cmd("embed", params={"texts": ["valid text", "", " ", "a"]})
|
314
|
+
assert "Invalid input texts" in str(excinfo.value)
|
315
|
+
assert "Text at index 1 is empty" in str(excinfo.value)
|
316
|
+
assert "Text at index 2 is empty" in str(excinfo.value)
|
317
|
+
assert "Text at index 3 is too short" in str(excinfo.value)
|
@@ -58,12 +58,22 @@ def extract_vectors(result):
|
|
58
58
|
elif "result" in result:
|
59
59
|
if isinstance(result["result"], list):
|
60
60
|
return result["result"]
|
61
|
-
elif
|
61
|
+
elif (
|
62
|
+
isinstance(result["result"], dict)
|
63
|
+
and "embeddings" in result["result"]
|
64
|
+
):
|
62
65
|
return result["result"]["embeddings"]
|
66
|
+
elif (
|
67
|
+
isinstance(result["result"], dict)
|
68
|
+
and "data" in result["result"]
|
69
|
+
and isinstance(result["result"]["data"], dict)
|
70
|
+
and "embeddings" in result["result"]["data"]
|
71
|
+
):
|
72
|
+
return result["result"]["data"]["embeddings"]
|
63
73
|
else:
|
64
74
|
pytest.fail("No embeddings in result['result']")
|
65
75
|
else:
|
66
|
-
pytest.fail("No embeddings
|
76
|
+
pytest.fail("No embeddings in result")
|
67
77
|
|
68
78
|
@pytest.mark.asyncio
|
69
79
|
@pytest.mark.integration
|
@@ -84,10 +94,9 @@ async def test_real_embed_vector(real_client):
|
|
84
94
|
async def test_real_embed_empty_texts(real_client):
|
85
95
|
if not await is_service_available():
|
86
96
|
pytest.skip("Real service on localhost:8001 is not available.")
|
87
|
-
|
88
|
-
|
89
|
-
assert
|
90
|
-
assert len(vectors) == 0
|
97
|
+
with pytest.raises(EmbeddingServiceAPIError) as excinfo:
|
98
|
+
await real_client.cmd("embed", params={"texts": []})
|
99
|
+
assert "Empty texts list provided" in str(excinfo.value)
|
91
100
|
|
92
101
|
@pytest.mark.asyncio
|
93
102
|
@pytest.mark.integration
|
@@ -0,0 +1,303 @@
|
|
1
|
+
"""
|
2
|
+
Stress tests for EmbeddingServiceAsyncClient.
|
3
|
+
|
4
|
+
These tests verify the client's behavior under heavy load with parallel processing
|
5
|
+
of large text batches.
|
6
|
+
"""
|
7
|
+
|
8
|
+
import pytest
|
9
|
+
import pytest_asyncio
|
10
|
+
import asyncio
|
11
|
+
import time
|
12
|
+
import random
|
13
|
+
from typing import List, Dict, Any, Optional
|
14
|
+
from embed_client.async_client import EmbeddingServiceAsyncClient
|
15
|
+
|
16
|
+
BASE_URL = "http://localhost"
|
17
|
+
PORT = 8001
|
18
|
+
|
19
|
+
# Constants for retry logic
|
20
|
+
MAX_RETRIES = 3
|
21
|
+
RETRY_DELAY = 1.0 # seconds
|
22
|
+
MAX_CONCURRENT_TASKS = 100 # Limit concurrent tasks to prevent system overload
|
23
|
+
|
24
|
+
def generate_test_texts(count: int) -> List[str]:
|
25
|
+
"""Generate test texts with unique content."""
|
26
|
+
return [f"Test text {i} for stress testing" for i in range(count)]
|
27
|
+
|
28
|
+
async def process_with_retry(
|
29
|
+
client: EmbeddingServiceAsyncClient,
|
30
|
+
texts: List[str],
|
31
|
+
max_retries: int = MAX_RETRIES,
|
32
|
+
retry_delay: float = RETRY_DELAY
|
33
|
+
) -> Dict[str, Any]:
|
34
|
+
"""
|
35
|
+
Process texts with retry logic for handling runtime exceptions.
|
36
|
+
|
37
|
+
Args:
|
38
|
+
client: EmbeddingServiceAsyncClient instance
|
39
|
+
texts: List of texts to process
|
40
|
+
max_retries: Maximum number of retry attempts
|
41
|
+
retry_delay: Delay between retries in seconds
|
42
|
+
|
43
|
+
Returns:
|
44
|
+
Dict containing the result or error information
|
45
|
+
"""
|
46
|
+
for attempt in range(max_retries):
|
47
|
+
try:
|
48
|
+
return await client.cmd("embed", params={"texts": texts})
|
49
|
+
except Exception as e:
|
50
|
+
if attempt == max_retries - 1: # Last attempt
|
51
|
+
return {
|
52
|
+
"error": f"Failed after {max_retries} attempts: {str(e)}",
|
53
|
+
"exception": str(e),
|
54
|
+
"attempts": attempt + 1
|
55
|
+
}
|
56
|
+
await asyncio.sleep(retry_delay * (attempt + 1)) # Exponential backoff
|
57
|
+
|
58
|
+
async def process_batch(
|
59
|
+
client: EmbeddingServiceAsyncClient,
|
60
|
+
texts: List[str],
|
61
|
+
batch_size: int,
|
62
|
+
max_concurrent: int = MAX_CONCURRENT_TASKS
|
63
|
+
) -> List[Dict[str, Any]]:
|
64
|
+
"""
|
65
|
+
Process a batch of texts in parallel with controlled concurrency.
|
66
|
+
|
67
|
+
Args:
|
68
|
+
client: EmbeddingServiceAsyncClient instance
|
69
|
+
texts: List of texts to process
|
70
|
+
batch_size: Size of each batch
|
71
|
+
max_concurrent: Maximum number of concurrent tasks
|
72
|
+
|
73
|
+
Returns:
|
74
|
+
List of results for each batch
|
75
|
+
"""
|
76
|
+
semaphore = asyncio.Semaphore(max_concurrent)
|
77
|
+
|
78
|
+
async def process_with_semaphore(batch: List[str]) -> Dict[str, Any]:
|
79
|
+
async with semaphore:
|
80
|
+
return await process_with_retry(client, batch)
|
81
|
+
|
82
|
+
tasks = []
|
83
|
+
for i in range(0, len(texts), batch_size):
|
84
|
+
batch = texts[i:i + batch_size]
|
85
|
+
task = asyncio.create_task(process_with_semaphore(batch))
|
86
|
+
tasks.append(task)
|
87
|
+
|
88
|
+
return await asyncio.gather(*tasks, return_exceptions=True)
|
89
|
+
|
90
|
+
@pytest_asyncio.fixture
|
91
|
+
async def stress_client():
|
92
|
+
"""Create a client instance for stress testing."""
|
93
|
+
async with EmbeddingServiceAsyncClient(base_url=BASE_URL, port=PORT) as client:
|
94
|
+
yield client
|
95
|
+
|
96
|
+
@pytest.mark.asyncio
|
97
|
+
@pytest.mark.stress
|
98
|
+
async def test_parallel_processing(stress_client):
|
99
|
+
"""Test parallel processing of 10K texts."""
|
100
|
+
total_texts = 10_000
|
101
|
+
batch_size = 100 # Process 100 texts at a time
|
102
|
+
texts = generate_test_texts(total_texts)
|
103
|
+
|
104
|
+
start_time = time.time()
|
105
|
+
results = await process_batch(stress_client, texts, batch_size)
|
106
|
+
end_time = time.time()
|
107
|
+
|
108
|
+
# Verify results
|
109
|
+
success_count = 0
|
110
|
+
error_count = 0
|
111
|
+
exception_count = 0
|
112
|
+
|
113
|
+
for result in results:
|
114
|
+
if isinstance(result, Exception):
|
115
|
+
exception_count += 1
|
116
|
+
continue
|
117
|
+
# Сначала assert на error/result
|
118
|
+
assert ("error" in result) or ("result" in result), f"Neither 'error' nor 'result' in response: {result}"
|
119
|
+
if "error" in result:
|
120
|
+
error_count += 1
|
121
|
+
if "exception" in result:
|
122
|
+
print(f"Error with exception: {result['error']}")
|
123
|
+
elif "result" in result:
|
124
|
+
res = result["result"]
|
125
|
+
assert isinstance(res, dict), f"result is not a dict: {res}"
|
126
|
+
if "success" in res and res["success"] is False:
|
127
|
+
error_count += 1
|
128
|
+
if "error" in res:
|
129
|
+
print(f"Error in result: {res['error']}")
|
130
|
+
else:
|
131
|
+
success_count += 1
|
132
|
+
assert "data" in res, f"No 'data' in result: {res}"
|
133
|
+
data = res["data"]
|
134
|
+
assert isinstance(data, dict), f"data is not a dict: {data}"
|
135
|
+
assert "embeddings" in data, f"No 'embeddings' in data: {data}"
|
136
|
+
assert isinstance(data["embeddings"], list), f"'embeddings' is not a list: {data}"
|
137
|
+
|
138
|
+
# Calculate statistics
|
139
|
+
total_time = end_time - start_time
|
140
|
+
texts_per_second = total_texts / total_time
|
141
|
+
|
142
|
+
print(f"\nStress Test Results:")
|
143
|
+
print(f"Total texts processed: {total_texts}")
|
144
|
+
print(f"Successful batches: {success_count}")
|
145
|
+
print(f"Failed batches: {error_count}")
|
146
|
+
print(f"Exception batches: {exception_count}")
|
147
|
+
print(f"Total time: {total_time:.2f} seconds")
|
148
|
+
print(f"Processing speed: {texts_per_second:.2f} texts/second")
|
149
|
+
|
150
|
+
# Assertions
|
151
|
+
assert success_count > 0, "No successful batches processed"
|
152
|
+
assert total_time > 0, "Invalid processing time"
|
153
|
+
# Allow some errors and exceptions under stress
|
154
|
+
assert error_count + exception_count < total_texts * 0.1, "Too many errors/exceptions"
|
155
|
+
|
156
|
+
@pytest.mark.asyncio
|
157
|
+
@pytest.mark.stress
|
158
|
+
async def test_concurrent_connections(stress_client):
|
159
|
+
"""Test multiple concurrent connections to the service."""
|
160
|
+
total_connections = 50
|
161
|
+
texts_per_connection = 200
|
162
|
+
texts = generate_test_texts(texts_per_connection)
|
163
|
+
|
164
|
+
async def single_connection() -> Dict[str, Any]:
|
165
|
+
try:
|
166
|
+
async with EmbeddingServiceAsyncClient(base_url=BASE_URL, port=PORT) as client:
|
167
|
+
return await process_with_retry(client, texts)
|
168
|
+
except Exception as e:
|
169
|
+
return {
|
170
|
+
"error": f"Connection failed: {str(e)}",
|
171
|
+
"exception": str(e)
|
172
|
+
}
|
173
|
+
|
174
|
+
start_time = time.time()
|
175
|
+
tasks = [asyncio.create_task(single_connection()) for _ in range(total_connections)]
|
176
|
+
results = await asyncio.gather(*tasks, return_exceptions=True)
|
177
|
+
end_time = time.time()
|
178
|
+
|
179
|
+
# Verify results
|
180
|
+
success_count = 0
|
181
|
+
error_count = 0
|
182
|
+
exception_count = 0
|
183
|
+
|
184
|
+
for result in results:
|
185
|
+
if isinstance(result, Exception):
|
186
|
+
exception_count += 1
|
187
|
+
continue
|
188
|
+
# Сначала assert на error/result
|
189
|
+
assert ("error" in result) or ("result" in result), f"Neither 'error' nor 'result' in response: {result}"
|
190
|
+
if "error" in result:
|
191
|
+
error_count += 1
|
192
|
+
if "exception" in result:
|
193
|
+
print(f"Error with exception: {result['error']}")
|
194
|
+
elif "result" in result:
|
195
|
+
res = result["result"]
|
196
|
+
assert isinstance(res, dict), f"result is not a dict: {res}"
|
197
|
+
if "success" in res and res["success"] is False:
|
198
|
+
error_count += 1
|
199
|
+
if "error" in res:
|
200
|
+
print(f"Error in result: {res['error']}")
|
201
|
+
else:
|
202
|
+
success_count += 1
|
203
|
+
assert "data" in res, f"No 'data' in result: {res}"
|
204
|
+
data = res["data"]
|
205
|
+
assert isinstance(data, dict), f"data is not a dict: {data}"
|
206
|
+
assert "embeddings" in data, f"No 'embeddings' in data: {data}"
|
207
|
+
assert isinstance(data["embeddings"], list), f"'embeddings' is not a list: {data}"
|
208
|
+
|
209
|
+
# Calculate statistics
|
210
|
+
total_time = end_time - start_time
|
211
|
+
total_texts = total_connections * texts_per_connection
|
212
|
+
texts_per_second = total_texts / total_time
|
213
|
+
|
214
|
+
print(f"\nConcurrent Connections Test Results:")
|
215
|
+
print(f"Total connections: {total_connections}")
|
216
|
+
print(f"Texts per connection: {texts_per_connection}")
|
217
|
+
print(f"Total texts processed: {total_texts}")
|
218
|
+
print(f"Successful connections: {success_count}")
|
219
|
+
print(f"Failed connections: {error_count}")
|
220
|
+
print(f"Exception connections: {exception_count}")
|
221
|
+
print(f"Total time: {total_time:.2f} seconds")
|
222
|
+
print(f"Processing speed: {texts_per_second:.2f} texts/second")
|
223
|
+
|
224
|
+
# Assertions
|
225
|
+
assert success_count > 0, "No successful connections"
|
226
|
+
assert total_time > 0, "Invalid processing time"
|
227
|
+
# Allow some errors and exceptions under stress
|
228
|
+
assert error_count + exception_count < total_connections * 0.2, "Too many errors/exceptions"
|
229
|
+
|
230
|
+
@pytest.mark.asyncio
|
231
|
+
@pytest.mark.stress
|
232
|
+
async def test_error_handling_under_load(stress_client):
|
233
|
+
"""Test error handling under load with mixed valid and invalid inputs."""
|
234
|
+
total_texts = 5_000
|
235
|
+
batch_size = 50
|
236
|
+
texts = generate_test_texts(total_texts)
|
237
|
+
|
238
|
+
# Add some invalid inputs
|
239
|
+
invalid_texts = [None, "", " " * 1000] # Empty, None, and very long text
|
240
|
+
texts.extend(invalid_texts)
|
241
|
+
|
242
|
+
# Add some random delays to simulate real-world conditions
|
243
|
+
async def process_with_delay(batch: List[str]) -> Dict[str, Any]:
|
244
|
+
await asyncio.sleep(random.uniform(0.1, 0.5))
|
245
|
+
return await process_with_retry(stress_client, batch)
|
246
|
+
|
247
|
+
start_time = time.time()
|
248
|
+
tasks = []
|
249
|
+
for i in range(0, len(texts), batch_size):
|
250
|
+
batch = texts[i:i + batch_size]
|
251
|
+
task = asyncio.create_task(process_with_delay(batch))
|
252
|
+
tasks.append(task)
|
253
|
+
results = await asyncio.gather(*tasks, return_exceptions=True)
|
254
|
+
end_time = time.time()
|
255
|
+
|
256
|
+
# Verify results
|
257
|
+
success_count = 0
|
258
|
+
error_count = 0
|
259
|
+
exception_count = 0
|
260
|
+
|
261
|
+
for result in results:
|
262
|
+
if isinstance(result, Exception):
|
263
|
+
exception_count += 1
|
264
|
+
continue
|
265
|
+
# Сначала assert на error/result
|
266
|
+
assert ("error" in result) or ("result" in result), f"Neither 'error' nor 'result' in response: {result}"
|
267
|
+
if "error" in result:
|
268
|
+
error_count += 1
|
269
|
+
if "exception" in result:
|
270
|
+
print(f"Error with exception: {result['error']}")
|
271
|
+
elif "result" in result:
|
272
|
+
res = result["result"]
|
273
|
+
assert isinstance(res, dict), f"result is not a dict: {res}"
|
274
|
+
if "success" in res and res["success"] is False:
|
275
|
+
error_count += 1
|
276
|
+
if "error" in res:
|
277
|
+
print(f"Error in result: {res['error']}")
|
278
|
+
else:
|
279
|
+
success_count += 1
|
280
|
+
assert "data" in res, f"No 'data' in result: {res}"
|
281
|
+
data = res["data"]
|
282
|
+
assert isinstance(data, dict), f"data is not a dict: {data}"
|
283
|
+
assert "embeddings" in data, f"No 'embeddings' in data: {data}"
|
284
|
+
assert isinstance(data["embeddings"], list), f"'embeddings' is not a list: {data}"
|
285
|
+
|
286
|
+
# Calculate statistics
|
287
|
+
total_time = end_time - start_time
|
288
|
+
texts_per_second = total_texts / total_time
|
289
|
+
|
290
|
+
print(f"\nError Handling Under Load Test Results:")
|
291
|
+
print(f"Total texts processed: {total_texts}")
|
292
|
+
print(f"Successful batches: {success_count}")
|
293
|
+
print(f"Failed batches: {error_count}")
|
294
|
+
print(f"Exception batches: {exception_count}")
|
295
|
+
print(f"Total time: {total_time:.2f} seconds")
|
296
|
+
print(f"Processing speed: {texts_per_second:.2f} texts/second")
|
297
|
+
|
298
|
+
# Assertions
|
299
|
+
assert success_count > 0, "No successful batches processed"
|
300
|
+
assert error_count > 0, "No errors detected with invalid inputs"
|
301
|
+
assert total_time > 0, "Invalid processing time"
|
302
|
+
# Allow some exceptions under stress
|
303
|
+
assert exception_count < total_texts * 0.05, "Too many exceptions"
|
@@ -1,94 +0,0 @@
|
|
1
|
-
"""
|
2
|
-
Example usage of EmbeddingServiceAsyncClient.
|
3
|
-
|
4
|
-
This example demonstrates how to use the async client to check the health of the embedding service,
|
5
|
-
request embeddings, and handle all possible exceptions.
|
6
|
-
|
7
|
-
Run this script with:
|
8
|
-
python -m asyncio embed_client/example_async_usage.py --base-url http://localhost --port 8001
|
9
|
-
|
10
|
-
You can also set EMBED_CLIENT_BASE_URL and EMBED_CLIENT_PORT environment variables.
|
11
|
-
"""
|
12
|
-
|
13
|
-
import asyncio
|
14
|
-
import sys
|
15
|
-
import os
|
16
|
-
from embed_client.async_client import (
|
17
|
-
EmbeddingServiceAsyncClient,
|
18
|
-
EmbeddingServiceConnectionError,
|
19
|
-
EmbeddingServiceHTTPError,
|
20
|
-
EmbeddingServiceAPIError,
|
21
|
-
EmbeddingServiceError,
|
22
|
-
)
|
23
|
-
|
24
|
-
def get_params():
|
25
|
-
base_url = None
|
26
|
-
port = None
|
27
|
-
for i, arg in enumerate(sys.argv):
|
28
|
-
if arg in ("--base-url", "-b") and i + 1 < len(sys.argv):
|
29
|
-
base_url = sys.argv[i + 1]
|
30
|
-
if arg in ("--port", "-p") and i + 1 < len(sys.argv):
|
31
|
-
port = sys.argv[i + 1]
|
32
|
-
if not base_url:
|
33
|
-
base_url = os.environ.get("EMBED_CLIENT_BASE_URL")
|
34
|
-
if not port:
|
35
|
-
port = os.environ.get("EMBED_CLIENT_PORT")
|
36
|
-
if not base_url or not port:
|
37
|
-
print("Error: base_url and port must be provided via --base-url/--port arguments or EMBED_CLIENT_BASE_URL/EMBED_CLIENT_PORT environment variables.")
|
38
|
-
sys.exit(1)
|
39
|
-
return None, None
|
40
|
-
return base_url, int(port)
|
41
|
-
|
42
|
-
async def main():
|
43
|
-
base_url, port = get_params()
|
44
|
-
# Always use try/except to handle all possible errors
|
45
|
-
try:
|
46
|
-
async with EmbeddingServiceAsyncClient(base_url=base_url, port=port) as client:
|
47
|
-
# Check health
|
48
|
-
try:
|
49
|
-
health = await client.health()
|
50
|
-
print("Service health:", health)
|
51
|
-
except EmbeddingServiceConnectionError as e:
|
52
|
-
print("[Connection error]", e)
|
53
|
-
return
|
54
|
-
except EmbeddingServiceHTTPError as e:
|
55
|
-
print(f"[HTTP error] {e.status}: {e.message}")
|
56
|
-
return
|
57
|
-
except EmbeddingServiceError as e:
|
58
|
-
print("[Other error]", e)
|
59
|
-
return
|
60
|
-
|
61
|
-
# Request embeddings for a list of texts
|
62
|
-
texts = ["hello world", "test embedding"]
|
63
|
-
try:
|
64
|
-
result = await client.cmd("embed", params={"texts": texts})
|
65
|
-
vectors = result["result"]
|
66
|
-
print(f"Embeddings for {len(texts)} texts:")
|
67
|
-
for i, vec in enumerate(vectors):
|
68
|
-
print(f" Text: {texts[i]!r}\n Vector: {vec[:5]}... (total {len(vec)} dims)")
|
69
|
-
except EmbeddingServiceAPIError as e:
|
70
|
-
print("[API error]", e.error)
|
71
|
-
except EmbeddingServiceHTTPError as e:
|
72
|
-
print(f"[HTTP error] {e.status}: {e.message}")
|
73
|
-
except EmbeddingServiceConnectionError as e:
|
74
|
-
print("[Connection error]", e)
|
75
|
-
except EmbeddingServiceError as e:
|
76
|
-
print("[Other error]", e)
|
77
|
-
|
78
|
-
# Example: error handling for invalid command
|
79
|
-
try:
|
80
|
-
await client.cmd("not_a_command")
|
81
|
-
except EmbeddingServiceAPIError as e:
|
82
|
-
print("[API error for invalid command]", e.error)
|
83
|
-
|
84
|
-
# Example: error handling for empty texts
|
85
|
-
try:
|
86
|
-
await client.cmd("embed", params={"texts": []})
|
87
|
-
except EmbeddingServiceAPIError as e:
|
88
|
-
print("[API error for empty texts]", e.error)
|
89
|
-
|
90
|
-
except Exception as e:
|
91
|
-
print("[Unexpected error]", e)
|
92
|
-
|
93
|
-
if __name__ == "__main__":
|
94
|
-
asyncio.run(main())
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|