File size: 4,544 Bytes
e1dced2 a77c60e e1dced2 a77c60e e1dced2 a77c60e ee7e79e a77c60e ee7e79e e1dced2 a77c60e e1dced2 ee7e79e a77c60e ee7e79e a77c60e e1dced2 | 1 2 3 4 5 6 7 8 9 10 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | import numpy as np
import pytest
from openai import APIStatusError
from sync.embeddings_client import EmbeddingsClient
class _FakeEmbeddingDatum:
def __init__(self, embedding, index):
self.embedding = embedding
self.index = index
class _FakeEmbeddingsResponse:
def __init__(self, data):
self.data = data
def test_embed_batch_parses_response(monkeypatch):
captured = {}
def fake_create(self, model, input):
captured["model"] = model
captured["input"] = input
return _FakeEmbeddingsResponse([
_FakeEmbeddingDatum([0.2] * 1024, 1),
_FakeEmbeddingDatum([0.1] * 1024, 0),
])
monkeypatch.setattr("openai.resources.embeddings.Embeddings.create", fake_create)
client = EmbeddingsClient(base_url="http://fake", api_key="k")
result = client.embed_batch(["a", "b"])
assert captured["input"] == ["a", "b"]
assert len(result) == 2
assert result[0].shape == (1024,)
assert result[0].dtype == np.float32
assert np.isclose(result[0][0], 0.1)
assert np.isclose(result[1][0], 0.2)
def test_embed_batch_empty_input_makes_no_requests(monkeypatch):
called = {"n": 0}
def fake_create(self, model, input):
called["n"] += 1
return _FakeEmbeddingsResponse([])
monkeypatch.setattr("openai.resources.embeddings.Embeddings.create", fake_create)
client = EmbeddingsClient(base_url="http://fake", api_key="k")
assert client.embed_batch([]) == []
assert called["n"] == 0
def test_embed_batch_chunks_large_input(monkeypatch):
captured = []
def fake_create(self, model, input):
captured.append(list(input))
# index resets to 0 within each request, as on the real API
return _FakeEmbeddingsResponse([
_FakeEmbeddingDatum([float(i)] * 1024, i) for i in range(len(input))
])
monkeypatch.setattr("openai.resources.embeddings.Embeddings.create", fake_create)
monkeypatch.setattr("sync.embeddings_client._BATCH_SIZE", 32)
client = EmbeddingsClient(base_url="http://fake", api_key="k")
texts = [f"t{i}" for i in range(80)]
result = client.embed_batch(texts)
assert len(captured) == 3 # 32 + 32 + 16
assert captured[0] == texts[:32]
assert captured[1] == texts[32:64]
assert captured[2] == texts[64:]
assert len(result) == 80
# order preserved within and across chunks (each chunk's vectors come back 0..n-1)
assert [round(v[0]) for v in result] == list(range(32)) + list(range(32)) + list(range(16))
def test_embed_batch_retries_on_timeout_then_succeeds(monkeypatch):
import httpx
from openai import APITimeoutError
calls = {"n": 0}
def fake_create(self, model, input):
calls["n"] += 1
if calls["n"] == 1:
raise APITimeoutError(request=httpx.Request("POST", "http://fake/embeddings"))
return _FakeEmbeddingsResponse([_FakeEmbeddingDatum([0.5] * 1024, 0)])
monkeypatch.setattr("openai.resources.embeddings.Embeddings.create", fake_create)
monkeypatch.setattr("sync.embeddings_client.time.sleep", lambda _: None)
client = EmbeddingsClient(base_url="http://fake", api_key="k")
result = client.embed_batch(["a"])
assert calls["n"] == 2
assert np.isclose(result[0][0], 0.5)
def test_embed_batch_raises_after_exhausted_retries(monkeypatch):
import httpx
from openai import APITimeoutError
calls = {"n": 0}
def fake_create(self, model, input):
calls["n"] += 1
raise APITimeoutError(request=httpx.Request("POST", "http://fake/embeddings"))
monkeypatch.setattr("openai.resources.embeddings.Embeddings.create", fake_create)
monkeypatch.setattr("sync.embeddings_client.time.sleep", lambda _: None)
client = EmbeddingsClient(base_url="http://fake", api_key="k")
with pytest.raises(RuntimeError):
client.embed_batch(["a"])
assert calls["n"] == 3 # _MAX_RETRIES attempts
def test_embed_batch_raises_runtime_error_on_api_error(monkeypatch):
import httpx
def fake_create(self, model, input):
request = httpx.Request("POST", "http://fake/embeddings")
response = httpx.Response(500, request=request, json={"error": "boom"})
raise APIStatusError("boom", response=response, body={"error": "boom"})
monkeypatch.setattr("openai.resources.embeddings.Embeddings.create", fake_create)
client = EmbeddingsClient(base_url="http://fake", api_key="k")
with pytest.raises(RuntimeError):
client.embed_batch(["a"])
|