forked from retoor/devplacepy
Update
This commit is contained in:
@@ -331,6 +331,140 @@ def test_embeddings_disabled_returns_503(local_db, monkeypatch):
|
||||
assert resp.status_code == 503
|
||||
|
||||
|
||||
class FakeImageClient_openai_gateway:
|
||||
def __init__(self, *a, **k):
|
||||
self.calls = []
|
||||
|
||||
def build_request(self, method, url, headers=None, json=None, content=None):
|
||||
return FakeRequest(method, url, json)
|
||||
|
||||
async def send(self, request):
|
||||
self.calls.append((request.url, request.json_body))
|
||||
body = request.json_body or {}
|
||||
return FakeResp_openai_gateway(
|
||||
payload={
|
||||
"created": 1,
|
||||
"model": body.get("model"),
|
||||
"data": [{"b64_json": "aGVsbG8="}],
|
||||
"usage": {"cost": 0.05},
|
||||
}
|
||||
)
|
||||
|
||||
async def aclose(self):
|
||||
pass
|
||||
|
||||
|
||||
def test_images_remaps_alias_to_configured_model(local_db, monkeypatch):
|
||||
monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
|
||||
svc = GatewayService()
|
||||
cfg = svc.effective_config()
|
||||
cfg["gateway_force_model"] = False
|
||||
cfg["gateway_image_enabled"] = True
|
||||
cfg["gateway_image_model"] = "black-forest-labs/flux.2-pro"
|
||||
rt = svc.runtime()
|
||||
run_async(
|
||||
rt.handle_images(
|
||||
{
|
||||
"model": "molodetz-img-small",
|
||||
"prompt": "award emblem",
|
||||
"size": "512x512",
|
||||
},
|
||||
cfg,
|
||||
("guest", "test"),
|
||||
"test",
|
||||
)
|
||||
)
|
||||
assert rt._client.calls[-1][1]["model"] == "black-forest-labs/flux.2-pro"
|
||||
|
||||
|
||||
def test_image_route_overrides_upstream(local_db, monkeypatch):
|
||||
from devplacepy.services.openai_gateway import routing
|
||||
|
||||
monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
|
||||
routing.provider_store.set(
|
||||
routing.ProviderIn(
|
||||
name="gwimg",
|
||||
base_url="https://routed.example/v1/chat/completions",
|
||||
api_key="img-key",
|
||||
)
|
||||
)
|
||||
routing.model_store.set(
|
||||
routing.ModelRouteIn(
|
||||
source_model="custom-img",
|
||||
provider="gwimg",
|
||||
target_model="vendor/flux-pro",
|
||||
kind="image",
|
||||
price_input_per_m=0.06,
|
||||
)
|
||||
)
|
||||
try:
|
||||
svc = GatewayService()
|
||||
cfg = svc.effective_config()
|
||||
rt = svc.runtime()
|
||||
run_async(
|
||||
rt.handle_images(
|
||||
{
|
||||
"model": "custom-img",
|
||||
"prompt": "trophy",
|
||||
"response_format": "b64_json",
|
||||
},
|
||||
cfg,
|
||||
("guest", "test"),
|
||||
"test",
|
||||
)
|
||||
)
|
||||
url, body = rt._client.calls[-1]
|
||||
assert str(url) == "https://routed.example/v1/images"
|
||||
assert body["model"] == "vendor/flux-pro"
|
||||
finally:
|
||||
routing.model_store.remove("custom-img")
|
||||
routing.provider_store.remove("gwimg")
|
||||
|
||||
|
||||
def test_images_success_records_ledger(local_db, monkeypatch):
|
||||
monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
|
||||
svc = GatewayService()
|
||||
cfg = svc.effective_config()
|
||||
cfg["gateway_image_enabled"] = True
|
||||
cfg["gateway_image_model"] = "black-forest-labs/flux.2-pro"
|
||||
rt = svc.runtime()
|
||||
before = rt.image_calls
|
||||
resp = run_async(
|
||||
rt.handle_images(
|
||||
{"model": "molodetz-img-small", "prompt": "badge"},
|
||||
cfg,
|
||||
("guest", "img_ledger"),
|
||||
"test",
|
||||
)
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
payload = json.loads(bytes(resp.body).decode())
|
||||
assert payload["data"][0]["b64_json"] == "aGVsbG8="
|
||||
assert rt.image_calls == before + 1
|
||||
row = get_table("gateway_usage_ledger").find_one(owner_id="img_ledger")
|
||||
assert row is not None
|
||||
assert row["backend"] == "image"
|
||||
assert row["endpoint"] == "images/generations"
|
||||
assert row["requested_model"] == "molodetz-img-small"
|
||||
assert row["model"] == "black-forest-labs/flux.2-pro"
|
||||
assert row["success"] == 1
|
||||
assert float(row["cost_usd"]) == 0.05
|
||||
|
||||
|
||||
def test_images_disabled_returns_503(local_db, monkeypatch):
|
||||
monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
|
||||
svc = GatewayService()
|
||||
cfg = svc.effective_config()
|
||||
cfg["gateway_image_enabled"] = False
|
||||
rt = svc.runtime()
|
||||
resp = run_async(
|
||||
rt.handle_images(
|
||||
{"prompt": "badge"}, cfg, ("guest", "test"), "test"
|
||||
)
|
||||
)
|
||||
assert resp.status_code == 503
|
||||
|
||||
|
||||
def test_compute_cost_embed_branch():
|
||||
from devplacepy.services.openai_gateway.usage import (
|
||||
Pricing,
|
||||
|
||||
Reference in New Issue
Block a user