mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
de32a2a06e
fix(count_tokens): backport October 19 fixes to stable/1.104.x
(cherry picked from commit 5e4d160423)
49 lines
1.8 KiB
Python
49 lines
1.8 KiB
Python
"""The fake provider shared by every integration test that takes the `provider` fixture.
|
|
|
|
Deployments in `proxy_config.yaml` point at `PROVIDER_URL`, so one server answers for all of them. A test
|
|
queues the replies it expects with `expect` and reads what the proxy sent with `received`. Tests run one at a
|
|
time against it; the `provider` fixture checks nothing is left over between tests.
|
|
The proxy's model-info refresh is answered here.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from collections import deque
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass, field
|
|
from typing import Final
|
|
|
|
from tests.integration._support.wire import Reply, Request, Wire, wire_server
|
|
|
|
PROVIDER_PORT: Final = int(os.environ.get("INTEGRATION_PROVIDER_PORT", "8191"))
|
|
PROVIDER_URL: Final = f"http://127.0.0.1:{PROVIDER_PORT}"
|
|
MODEL_DISCOVERY: Final = ("GET", "/v1/models")
|
|
_UNQUEUED: Final = Reply(status=500, body=b'{"error": "the shared fake provider has no reply queued for this request"}')
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class SharedProvider:
|
|
wire: Wire
|
|
replies: deque[Reply]
|
|
last_test: str | None = field(default=None)
|
|
|
|
def expect(self, *replies: Reply) -> None:
|
|
self.replies.extend(replies)
|
|
|
|
def received(self) -> tuple[Request, ...]:
|
|
return tuple(request for request in self.wire.drain() if (request.method, request.target) != MODEL_DISCOVERY)
|
|
|
|
|
|
@contextmanager
|
|
def shared_provider() -> Iterator[SharedProvider]:
|
|
replies: Final[deque[Reply]] = deque()
|
|
|
|
def respond(request: Request) -> Reply:
|
|
if (request.method, request.target) == MODEL_DISCOVERY:
|
|
return Reply(body=b'{"object":"list","data":[]}')
|
|
return replies.popleft() if replies else _UNQUEUED
|
|
|
|
with wire_server(respond, port=PROVIDER_PORT) as wire:
|
|
yield SharedProvider(wire, replies)
|