diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 0913a70..3ba988e 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -66,9 +66,9 @@ Transport settings belong to the HTTP client. Setting `httpx2.AsyncClient(trust_ ## HTTP clients -Without `http_client=`, Republic creates and closes a client for each call. Reusing a model alone does not share a connection pool across calls. +Without `http_client=`, a Provider lazily creates one client on its first request. Models created by the same Provider and model listing share its connection pool. Separate calls to the top-level model factories create separate Providers. -Supply a client to share connections and control its lifetime. Configure timeouts on that client; the provider's `timeout=` option only configures clients Republic creates. +Supply a client to share connections across Providers and control its lifetime yourself. Configure timeouts on that client; the Provider's `timeout=` option only configures clients Republic creates. ```python import asyncio @@ -88,8 +88,6 @@ async def main(): asyncio.run(main()) ``` -Republic leaves a supplied client open. In the example, leaving the application's `async with` block closes it. Exiting a model stream closes that response while leaving the supplied client available for later requests. - This client is used for inference requests and Copilot Plugin token exchange. Codex and Grok token refresh, Copilot device login, and OpenRouter authorization-code exchange create their own auth clients; an injected inference client's transport options do not configure those clients. CLI login methods use the selected CLI's network configuration. ## Request body merging diff --git a/docs/reference/provider-api.md b/docs/reference/provider-api.md index 4a4aa40..5530f14 100644 --- a/docs/reference/provider-api.md +++ b/docs/reference/provider-api.md @@ -53,6 +53,68 @@ model = provider.get_model("gpt-6-sol") Here, `api_base`, `api_key`, and `MyAuth` are supplied by the application. A top-level constructor accepts `provider:model`; a provider instance accepts only the model name. See [configuration](configuration.md) for options and credential precedence, and [authentication](../guides/authentication.md) for API keys and account logins. +### HTTP client lifecycle + +A Provider lazily creates one HTTP client and reuses it for requests, streams, model listing, and every model created by that Provider. Separate calls to the top-level model factories create separate Providers. + +For applications that manage a long-lived Provider, call `await provider.close()` at shutdown: + +```python +provider = republic.get_provider("openai") +model = provider.get_model("gpt-6-sol") +try: + response = await model.chat("Hello") +finally: + await provider.close() +``` + +Providers also support async context management. Context entries are released on exit, including when the body raises an exception or is cancelled: + +```python +async with republic.get_provider("openai") as provider: + first = provider.get_model("gpt-6-sol") + second = provider.get_model("gpt-6-luna") + await first.chat("Hello") + await second.chat("Hello") +``` + +Chat, embedding, and decision models support the same pattern. The Provider tracks entered Provider and model contexts and closes only when all of them have exited: + +```python +async with republic.get_model("openai:gpt-6-sol") as model: + response = await model.chat("Hello") +``` + +Use the same pattern with `get_embedding_model()` and `get_decision_model()`, or enter an existing model with `async with model:`. Entering a Provider or model does not create the HTTP client before its first request. Creating a model without entering its context does not register it as open. Nested contexts for the same model remain open until all its entries have exited. + +```python +provider = republic.get_provider("openai") +async with provider.get_model("gpt-6-sol") as first: + async with provider.get_embedding_model("text-embedding-3-small") as embeddings: + vector = await embeddings.embed("Hello") + # The Provider remains open while first's context is active. + response = await first.chat("Hello") +# The last model context has exited, so the Provider is now closed. +``` + +A Provider context keeps the client open while model contexts enter and exit, so sequential model contexts can share it: + +```python +async with republic.get_provider("openai") as provider: + for name in ["gpt-6-sol", "gpt-6-luna"]: + async with provider.get_model(name) as model: + response = await model.chat("Hello") + # The Provider context still keeps the client open here. +``` + +Nested Provider contexts follow the same rule: exiting an inner context does not close the client while another Provider or model context is active. Explicit `await provider.close()` closes it immediately, regardless of open contexts. + +`provider.close()` is idempotent; `provider.is_closed` reports whether it has been closed. Once closed, all models sharing that Provider reject further requests, and the Provider cannot be reopened. Finish outstanding requests and exit stream contexts before closing it. Exiting a stream context closes only that response and keeps the HTTP client available for subsequent requests. + +An internally created client must be used and closed on the event loop that first used it. Constructing the Provider outside an event loop is supported; reusing an active Provider across multiple `asyncio.run()` calls raises `RuntimeError`. + +An external client passed through `http_client=` remains caller-owned: closing the Provider leaves that client open. See [HTTP client configuration](configuration.md#http-clients) for transport settings. + ### Tool calls `tool1` and `tool2` below are `republic.Tool` schemas. diff --git a/examples/inference.py b/examples/inference.py index 2d359fc..2ce810b 100644 --- a/examples/inference.py +++ b/examples/inference.py @@ -12,96 +12,94 @@ class Profile(pydantic.BaseModel): async def test_decision(): - model = republic.get_decision_model("openrouter:~typesafe/jev-latest") - - questions: dict[str, republic.decisions.Question] = { - "is_urgent": republic.decisions.Noul( - instructions="Does this message convey urgency?", - criteria={"true": "Explicitly time-sensitive", "false": "No urgency expressed"}, - ), - "department": republic.decisions.Choice( - instructions="Which team should handle this?", - criteria={ - "billing": "Payments, invoicing, refunds", - "technical": "Bugs, outages, integrations", - "sales": "Pricing, upgrades, new accounts", - }, - ), - "frustration": republic.decisions.Score( - instructions="How frustrated is the customer?", - criteria=["Calm", "Frustrated", "Very angry"], - ), - } - - response = await model.decide( - "Help! My payouts have been failing for 3 days.", - questions=questions, - ) - print(response.answers) + async with republic.get_decision_model("openrouter:~typesafe/jev-latest") as model: + questions: dict[str, republic.decisions.Question] = { + "is_urgent": republic.decisions.Noul( + instructions="Does this message convey urgency?", + criteria={"true": "Explicitly time-sensitive", "false": "No urgency expressed"}, + ), + "department": republic.decisions.Choice( + instructions="Which team should handle this?", + criteria={ + "billing": "Payments, invoicing, refunds", + "technical": "Bugs, outages, integrations", + "sales": "Pricing, upgrades, new accounts", + }, + ), + "frustration": republic.decisions.Score( + instructions="How frustrated is the customer?", + criteria=["Calm", "Frustrated", "Very angry"], + ), + } + + response = await model.decide( + "Help! My payouts have been failing for 3 days.", + questions=questions, + ) + print(response.answers) async def test_chat(): - model = republic.get_model("google:gemini-flash-latest") + async with republic.get_model("google:gemini-flash-latest") as model: + response = await model.chat("What is the capital of France?") + print(response.text) + print("--------") - response = await model.chat("What is the capital of France?") - print(response.text) - print("--------") - - async with model.stream("给出五条今天的新闻", tools=[republic.tools.WebSearch()]) as stream: - async for event in stream: - if isinstance(event, republic.events.TextDelta): - print(event.chunk, end="", flush=True) - print() # Ensure a newline after streaming output + async with model.stream("给出五条今天的新闻", tools=[republic.tools.WebSearch()]) as stream: + async for event in stream: + if isinstance(event, republic.events.TextDelta): + print(event.chunk, end="", flush=True) + print() # Ensure a newline after streaming output async def test_reasoning(): - model = republic.get_model("deepseek:deepseek-flash") - response = await model.chat("Explain the implications of quantum computing.", reasoning_effort="none") - print(response.text) - assert not response.reasoning, "Expected no reasoning when reasoning_effort is 'none'" # noqa: S101 - - reasoning_head = False - answer_head = False - async with model.stream("Explain the implications of quantum computing.", reasoning_effort="high") as stream: - async for event in stream: - if isinstance(event, republic.events.TextDelta): - if not answer_head: - print("\n-----------------") - answer_head = True - print(event.chunk, end="", flush=True) - elif isinstance(event, republic.events.ReasoningDelta): - if not reasoning_head: - print("Thinking >") - reasoning_head = True - print(event.chunk, end="", flush=True) - print() # Ensure a newline after streaming output + async with republic.get_model("deepseek:deepseek-flash") as model: + response = await model.chat("Explain the implications of quantum computing.", reasoning_effort="none") + print(response.text) + assert not response.reasoning, "Expected no reasoning when reasoning_effort is 'none'" # noqa: S101 + + reasoning_head = False + answer_head = False + async with model.stream("Explain the implications of quantum computing.", reasoning_effort="high") as stream: + async for event in stream: + if isinstance(event, republic.events.TextDelta): + if not answer_head: + print("\n-----------------") + answer_head = True + print(event.chunk, end="", flush=True) + elif isinstance(event, republic.events.ReasoningDelta): + if not reasoning_head: + print("Thinking >") + reasoning_head = True + print(event.chunk, end="", flush=True) + print() # Ensure a newline after streaming output async def test_structured_output(): - model = republic.get_model("magpie:openrouter/moonshotai/kimi-k3") - response = await model.chat("Provide 5 fake users for testing.", output_schema=list[Profile]) - print(response.output) - assert all(isinstance(item, Profile) for item in response.output), ( # noqa: S101 - "Expected output to be a list of Profile instances" - ) + async with republic.get_model("magpie:openrouter/moonshotai/kimi-k3") as model: + response = await model.chat("Provide 5 fake users for testing.", output_schema=list[Profile]) + print(response.output) + assert all(isinstance(item, Profile) for item in response.output), ( # noqa: S101 + "Expected output to be a list of Profile instances" + ) async def test_image_generation(): - model = republic.get_model("google:gemini-3.1-flash-image-preview") - response = await model.chat("Generate an image of a futuristic city skyline at sunset.") - print(response.text) - for i, image in enumerate(response.image_parts): - if image.data is None: - raise ValueError(f"Image data for image_{i}.png is None") - with open(f"image_{i}.png", "wb") as f: - f.write(image.data) + async with republic.get_model("google:gemini-3.1-flash-image-preview") as model: + response = await model.chat("Generate an image of a futuristic city skyline at sunset.") + print(response.text) + for i, image in enumerate(response.image_parts): + if image.data is None: + raise ValueError(f"Image data for image_{i}.png is None") + with open(f"image_{i}.png", "wb") as f: + f.write(image.data) async def test_list_models(): for provider_name in ["google", "deepseek", "magpie"]: - provider = republic.get_provider(provider_name) - models = await provider.list_models() - print(f"{provider_name.capitalize()} provider has {len(models)} models") + async with republic.get_provider(provider_name) as provider: + models = await provider.list_models() + print(f"{provider_name.capitalize()} provider has {len(models)} models") async def run_all(): diff --git a/src/republic/_models.py b/src/republic/_models.py index b8232a1..01eda78 100644 --- a/src/republic/_models.py +++ b/src/republic/_models.py @@ -45,6 +45,15 @@ def __init__( self.api_format = api_format self.history = history + async def __aenter__(self) -> Self: + self.provider._open_context(self) + return self + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, traceback: TracebackType | None + ) -> None: + await self.provider._close_context(self) + @overload async def chat( self, prompt: Input, *, output_schema: None = None, **options: Unpack[ChatOptions] @@ -211,6 +220,15 @@ def __init__(self, provider: Provider, name: str, api_format: EmbeddingApiFormat self.name = name self.api_format = api_format + async def __aenter__(self) -> Self: + self.provider._open_context(self) + return self + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, traceback: TracebackType | None + ) -> None: + await self.provider._close_context(self) + async def embed(self, text: str, *, dimensions: int | None = None) -> EmbeddingResponse: """Embed one text; the vector is ``response.vector``.""" return await self.embed_many([text], dimensions=dimensions) @@ -238,6 +256,15 @@ def __init__(self, provider: Provider, name: str, api_format: DecisionApiFormat) self.name = name self.api_format = api_format + async def __aenter__(self) -> Self: + self.provider._open_context(self) + return self + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, traceback: TracebackType | None + ) -> None: + await self.provider._close_context(self) + async def decide(self, state: JSONValue, *, questions: Mapping[str, Question]) -> DecisionResponse: """Answer every question about ``state``. Answers come back under the same ids.""" request = self.api_format.decision_request(self.name, state, questions) diff --git a/src/republic/providers/base.py b/src/republic/providers/base.py index ca11136..d553a58 100644 --- a/src/republic/providers/base.py +++ b/src/republic/providers/base.py @@ -11,7 +11,8 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from email.utils import parsedate_to_datetime -from typing import TYPE_CHECKING, Any, ClassVar, TypeVar +from types import TracebackType +from typing import TYPE_CHECKING, Any, ClassVar, Self, TypeVar import httpx2 @@ -104,6 +105,60 @@ def __init__( self.retry_delay = retry_delay self.max_retry_delay = max_retry_delay self._http_client = http_client + self._owns_http_client = http_client is None + self._client_loop: asyncio.AbstractEventLoop | None = None + self._closed = False + self._open_contexts: list[Provider | ChatModel | EmbeddingModel | DecisionModel] = [] + + @property + def is_closed(self) -> bool: + """Whether this provider has been closed, independently of an external client.""" + return self._closed + + async def __aenter__(self) -> Self: + self._open_context(self) + return self + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, traceback: TracebackType | None + ) -> None: + await self._close_context(self) + + async def close(self) -> None: + """Close the provider and its internally created client; leave supplied clients open. + + Finish requests and exit stream contexts before closing. Close an owned + client on the event loop that first used it. Repeated closes are harmless. + """ + if self._closed: + return + self._check_client_loop() + self._closed = True + self._open_contexts.clear() + if self._owns_http_client and self._http_client is not None: + await self._http_client.aclose() + + def _open_context(self, context: Provider | ChatModel | EmbeddingModel | DecisionModel) -> None: + self._ensure_open() + self._check_client_loop() + # Record each entry so nested contexts for the same object stay open. + self._open_contexts.append(context) + + async def _close_context(self, context: Provider | ChatModel | EmbeddingModel | DecisionModel) -> None: + self._check_client_loop() + if context not in self._open_contexts: + return + self._open_contexts.remove(context) + if not self._open_contexts: + await self.close() + + def _ensure_open(self) -> None: + if self._closed: + raise RuntimeError("Provider is closed") + + def _check_client_loop(self) -> None: + if self._client_loop is not None and self._client_loop is not asyncio.get_running_loop(): + raise RuntimeError("Use and close the provider on the event loop that created its HTTP client") def get_model(self, name: str, *, history: HistoryProtocol | None = None) -> ChatModel: from republic._models import ChatModel @@ -129,18 +184,18 @@ async def list_models(self) -> list[ModelInfo]: auth = httpx2.USE_CLIENT_DEFAULT if self.auth is None else self.auth models: list[ModelInfo] = [] params = self._models_params(None) - async with self._client() as client: - while params is not None: - response = await self._retry( - lambda params=params: client.get( - f"{self.api_base}{self.MODELS_PATH}", params=params, headers=headers, auth=auth - ) + client = self._client() + while params is not None: + response = await self._retry( + lambda params=params: client.get( + f"{self.api_base}{self.MODELS_PATH}", params=params, headers=headers, auth=auth ) - if response.is_error: - raise APIStatusError(response.status_code, response.text, headers=response.headers) - data = response.json() - models.extend(self._parse_models(data)) - params = self._models_params(data) + ) + if response.is_error: + raise APIStatusError(response.status_code, response.text, headers=response.headers) + data = response.json() + models.extend(self._parse_models(data)) + params = self._models_params(data) return models def _models_params(self, previous: Any) -> dict[str, str] | None: @@ -179,26 +234,26 @@ def select_api_format(self, format_kind: type[_FormatT], model: str) -> _FormatT raise UnsupportedApiFormatError(f"Provider {self.name!r} supports no {format_kind.kind} API format") async def _post(self, api_format: ApiFormat, request: HttpRequest) -> httpx2.Response: - async with self._client() as client: - response = await self._retry(lambda: self._send(client, api_format, request, stream=False)) - if response.is_error: - raise APIStatusError(response.status_code, response.text, headers=response.headers) - return response + client = self._client() + response = await self._retry(lambda: self._send(client, api_format, request, stream=False)) + if response.is_error: + raise APIStatusError(response.status_code, response.text, headers=response.headers) + return response @asynccontextmanager async def _stream(self, api_format: ApiFormat, request: HttpRequest) -> AsyncGenerator[HttpStream]: - async with self._client() as client: - response = await self._retry(lambda: self._send(client, api_format, request, stream=True)) - try: - if response.is_error: - await response.aread() - raise APIStatusError(response.status_code, response.text, headers=response.headers) - end_marker = "[DONE]" if api_format.name in {"chat", "responses"} else None - yield HttpStream(_response_events(response, end_marker), dict(response.headers)) - except httpx2.RequestError as exc: - raise _connection_error(exc, headers=response.headers) from exc - finally: - await response.aclose() + client = self._client() + response = await self._retry(lambda: self._send(client, api_format, request, stream=True)) + try: + if response.is_error: + await response.aread() + raise APIStatusError(response.status_code, response.text, headers=response.headers) + end_marker = "[DONE]" if api_format.name in {"chat", "responses"} else None + yield HttpStream(_response_events(response, end_marker), dict(response.headers)) + except httpx2.RequestError as exc: + raise _connection_error(exc, headers=response.headers) from exc + finally: + await response.aclose() async def _retry(self, send: Callable[[], Awaitable[httpx2.Response]]) -> httpx2.Response: """Retry only the request, never replay a successful streaming response.""" @@ -225,14 +280,15 @@ async def _retry(self, send: Callable[[], Awaitable[httpx2.Response]]) -> httpx2 await asyncio.sleep(min(delay, self.max_retry_delay)) raise AssertionError("Unreachable retry state") - @asynccontextmanager - async def _client(self) -> AsyncGenerator[httpx2.AsyncClient]: - if self._http_client is not None: - yield self._http_client - return - # A client per call keeps providers usable across event loops. - async with httpx2.AsyncClient(timeout=self.timeout) as client: - yield client + def _client(self) -> httpx2.AsyncClient: + self._ensure_open() + self._check_client_loop() + if self._http_client is None: + # No await between checking and creating: concurrent tasks on this + # event loop share the same client, including their first request. + self._http_client = httpx2.AsyncClient(timeout=self.timeout) + self._client_loop = asyncio.get_running_loop() + return self._http_client async def _send( self, client: httpx2.AsyncClient, api_format: ApiFormat, request: HttpRequest, *, stream: bool diff --git a/tests/test_provider_lifecycle.py b/tests/test_provider_lifecycle.py new file mode 100644 index 0000000..a65e68d --- /dev/null +++ b/tests/test_provider_lifecycle.py @@ -0,0 +1,532 @@ +from __future__ import annotations + +import asyncio +from typing import Any + +import httpx2 +import pytest + +import republic +from republic.providers import base +from tests.conftest import FakeService + + +@pytest.fixture(params=["embedding", "decision"]) +def other_model(request: pytest.FixtureRequest) -> republic.EmbeddingModel | republic.DecisionModel: + if request.param == "embedding": + return republic.get_embedding_model("openai:test") + return republic.get_decision_model("typesafe:test", api_key="key") + + +async def _call_other_model(model: republic.EmbeddingModel | republic.DecisionModel, service: FakeService) -> None: + if isinstance(model, republic.EmbeddingModel): + service.reply_json({"data": [{"index": 0, "embedding": [0.1]}]}) + assert (await model.embed("Hi")).vector == [0.1] + else: + service.reply_json({"answers": {}, "usage": {}}) + assert (await model.decide("Hi", questions={})).answers == {} + + +async def test_other_model_context_reuses_client_and_closes_provider( + other_model: republic.EmbeddingModel | republic.DecisionModel, + clients: list[httpx2.AsyncClient], + service: FakeService, +) -> None: + async with other_model as entered: + assert entered is other_model + assert clients == [] + await _call_other_model(other_model, service) + await _call_other_model(other_model, service) + assert len(clients) == 1 + assert not clients[0].is_closed + assert other_model.provider.is_closed + assert clients[0].is_closed + with pytest.raises(RuntimeError, match="Provider is closed"): + async with other_model: + pytest.fail("A model with a closed provider must not be entered") + + +@pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) +async def test_other_model_context_closes_on_failure( + other_model: republic.EmbeddingModel | republic.DecisionModel, + failure: type[BaseException], + clients: list[httpx2.AsyncClient], + service: FakeService, +) -> None: + with pytest.raises(failure): + async with other_model: + await _call_other_model(other_model, service) + raise failure + assert other_model.provider.is_closed + assert clients[0].is_closed + + +async def test_unused_other_model_context_does_not_create_client( + other_model: republic.EmbeddingModel | republic.DecisionModel, clients: list[httpx2.AsyncClient] +) -> None: + async with other_model: + assert clients == [] + assert other_model.provider.is_closed + assert clients == [] + + +async def test_other_model_context_leaves_external_client_open( + other_model: republic.EmbeddingModel | republic.DecisionModel, service: FakeService +) -> None: + async with service.client() as client: + provider = republic.get_provider(other_model.provider.name, api_key="key", http_client=client) + model = ( + provider.get_embedding_model("test") + if isinstance(other_model, republic.EmbeddingModel) + else provider.get_decision_model("test") + ) + async with model: + await _call_other_model(model, service) + assert provider.is_closed + assert not client.is_closed + assert client.is_closed + + +async def test_nested_contexts_for_same_model_keep_provider_open( + other_model: republic.EmbeddingModel | republic.DecisionModel, + clients: list[httpx2.AsyncClient], + service: FakeService, +) -> None: + async with other_model: + async with other_model: + await _call_other_model(other_model, service) + assert not other_model.provider.is_closed + assert not clients[0].is_closed + await _call_other_model(other_model, service) + assert other_model.provider.is_closed + assert clients[0].is_closed + + +async def test_provider_context_keeps_other_model_client_open_after_model_exit( + other_model: republic.EmbeddingModel | republic.DecisionModel, + clients: list[httpx2.AsyncClient], + service: FakeService, +) -> None: + async with other_model.provider as provider: + async with other_model: + await _call_other_model(other_model, service) + assert not provider.is_closed + assert not clients[0].is_closed + await _call_other_model(other_model, service) + assert provider.is_closed + assert clients[0].is_closed + + +async def test_provider_context_allows_sequential_chat_model_contexts( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + async with republic.get_provider("openai") as provider: + for name in ["first", "second"]: + service.reply_json({"output": []}) + async with provider.get_model(name) as model: + await model.chat("Hi") + assert not provider.is_closed + assert not clients[0].is_closed + service.reply_json({"data": [{"id": "first"}]}) + assert (await provider.list_models())[0].id == "first" + assert len(clients) == 1 + assert provider.is_closed + assert clients[0].is_closed + + +@pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) +async def test_nested_provider_context_failure_keeps_outer_context_open( + failure: type[BaseException], clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + async with republic.get_provider("openai") as provider: + model = provider.get_model("test") + with pytest.raises(failure): + async with provider: + service.reply_json({"output": []}) + await model.chat("Hi") + raise failure + assert not provider.is_closed + assert not clients[0].is_closed + service.reply_json({"output": []}) + await model.chat("Still open") + assert provider.is_closed + assert clients[0].is_closed + + +async def test_provider_context_exit_keeps_outer_model_context_open( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + async with republic.get_model("openai:test") as model: + async with model.provider as provider: + service.reply_json({"output": []}) + await model.chat("Hi") + assert not provider.is_closed + assert not clients[0].is_closed + service.reply_json({"output": []}) + await model.chat("Still open") + assert provider.is_closed + assert clients[0].is_closed + + +async def test_explicit_close_overrides_nested_provider_and_model_contexts( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + async with republic.get_provider("openai") as provider, provider, provider.get_model("test") as model: + service.reply_json({"output": []}) + await model.chat("Hi") + await provider.close() + assert provider.is_closed + assert clients[0].is_closed + with pytest.raises(RuntimeError, match="Provider is closed"): + await model.chat("Again") + await provider.close() + assert provider.is_closed + + +@pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) +async def test_other_model_failure_keeps_sibling_context_open( + other_model: republic.EmbeddingModel | republic.DecisionModel, + failure: type[BaseException], + clients: list[httpx2.AsyncClient], + service: FakeService, +) -> None: + provider = other_model.provider + sibling = ( + provider.get_embedding_model("sibling") + if isinstance(other_model, republic.EmbeddingModel) + else provider.get_decision_model("sibling") + ) + async with sibling: + with pytest.raises(failure): + async with other_model: + await _call_other_model(other_model, service) + raise failure + assert not provider.is_closed + assert not clients[0].is_closed + await _call_other_model(sibling, service) + assert provider.is_closed + assert clients[0].is_closed + + +async def test_chat_and_embedding_contexts_share_provider_until_last_exit( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + provider = republic.get_provider("openai") + async with provider.get_model("chat") as chat: + async with provider.get_embedding_model("embedding") as embeddings: + await _call_other_model(embeddings, service) + assert not provider.is_closed + assert not clients[0].is_closed + # Re-enter an exited model while another model keeps its Provider open. + async with embeddings: + await _call_other_model(embeddings, service) + service.reply_json({"output": []}) + await chat.chat("Still open") + assert len(clients) == 1 + assert clients[0].is_closed + assert provider.is_closed + + +async def test_concurrent_model_contexts_close_only_after_both_exit( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + provider = republic.get_provider("openai") + entered = asyncio.Event() + leave = asyncio.Event() + + async def hold() -> None: + async with provider.get_model("first"): + entered.set() + await leave.wait() + + task = asyncio.create_task(hold()) + try: + await entered.wait() + async with provider.get_model("second") as model: + service.reply_json({"output": []}) + await model.chat("Hi") + assert not provider.is_closed + assert not clients[0].is_closed + finally: + leave.set() + await task + assert provider.is_closed + assert clients[0].is_closed + + +async def test_explicit_close_overrides_open_model_contexts( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + provider = republic.get_provider("openai") + async with provider.get_model("first") as first, provider.get_model("second"): + service.reply_json({"output": []}) + await first.chat("Hi") + await provider.close() + assert provider.is_closed + assert clients[0].is_closed + with pytest.raises(RuntimeError, match="Provider is closed"): + await first.chat("Again") + assert provider.is_closed + + +@pytest.fixture +def clients(monkeypatch: pytest.MonkeyPatch, service: FakeService) -> list[httpx2.AsyncClient]: + client_type = httpx2.AsyncClient + created: list[httpx2.AsyncClient] = [] + + def create(**options: Any) -> httpx2.AsyncClient: + options.setdefault("transport", httpx2.MockTransport(service._handle)) + client = client_type(**options) + created.append(client) + return client + + monkeypatch.setattr(base.httpx2, "AsyncClient", create) + return created + + +async def test_lazy_client_shared_by_models_listing_and_streams( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + provider = republic.get_provider("openai", timeout=17) + first = provider.get_model("first") + second = provider.get_model("second") + embeddings = provider.get_embedding_model("embedding") + assert clients == [] + + async with provider as entered: + assert entered is provider + assert clients == [] + service.reply_json({"output": []}) + await first.chat("Hi") + assert len(clients) == 1 + assert clients[0].timeout.read == 17 + assert not clients[0].is_closed + + service.reply_json({"output": []}) + await second.chat("Hi") + service.reply_json({"data": [{"index": 0, "embedding": [0.1]}]}) + assert (await embeddings.embed("Hi")).vector == [0.1] + service.reply_json({"data": [{"id": "first"}]}) + assert (await provider.list_models())[0].id == "first" + service.reply_events([{"type": "response.output_text.delta", "delta": "hello"}, "[DONE]"]) + async with first.stream("Hi") as stream: + [event async for event in stream] + assert stream.text == "hello" + assert not clients[0].is_closed + assert len(clients) == 1 + assert not provider.is_closed + + assert clients[0].is_closed + assert provider.is_closed + + +async def test_model_context_reuses_client_and_closes_shared_provider( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + model = republic.get_model("openai:first") + sibling = model.provider.get_model("second") + async with model as entered: + assert entered is model + assert clients == [] + service.reply_json({"output": []}) + await model.chat("Hi") + service.reply_json({"output": []}) + await sibling.chat("Hi") + service.reply_events([{"type": "response.output_text.delta", "delta": "hello"}, "[DONE]"]) + async with model.stream("Hi") as stream: + [event async for event in stream] + assert stream.text == "hello" + assert len(clients) == 1 + assert not clients[0].is_closed + assert model.provider.is_closed + assert clients[0].is_closed + with pytest.raises(RuntimeError, match="Provider is closed"): + await sibling.chat("Again") + with pytest.raises(RuntimeError, match="Provider is closed"): + async with model: + pytest.fail("A model with a closed provider must not be entered") + + +@pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) +async def test_model_context_closes_client_when_body_fails( + failure: type[BaseException], clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + service.reply_json({"output": []}) + with pytest.raises(failure): + async with republic.get_model("openai:test") as model: + await model.chat("Hi") + raise failure + assert model.provider.is_closed + assert clients[0].is_closed + + +async def test_unused_model_context_does_not_create_client(clients: list[httpx2.AsyncClient]) -> None: + async with republic.get_model("openai:test") as model: + assert clients == [] + assert model.provider.is_closed + assert clients == [] + + +async def test_model_context_leaves_external_client_open(service: FakeService) -> None: + async with service.client() as client: + service.reply_json({"output": []}) + async with republic.get_model("openai:test", http_client=client) as model: + await model.chat("Hi") + assert model.provider.is_closed + assert not client.is_closed + assert client.is_closed + + +async def test_concurrent_first_requests_create_one_client( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + async with republic.get_provider("openai") as provider: + for _ in range(8): + service.reply_json({"output": []}) + await asyncio.gather(*(provider.get_model(str(index)).chat("Hi") for index in range(8))) + assert len(clients) == 1 + assert len(service.requests) == 8 + assert not clients[0].is_closed + assert clients[0].is_closed + + +async def test_manual_close_is_idempotent_and_prevents_reopening( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + service.reply_json({"output": []}) + model = republic.get_model("openai:test") + await model.chat("Hi") + await model.provider.close() + await model.provider.close() + + assert model.provider.is_closed + assert clients[0].is_closed + with pytest.raises(RuntimeError, match="Provider is closed"): + await model.chat("Again") + with pytest.raises(RuntimeError, match="Provider is closed"): + await model.provider.list_models() + with pytest.raises(RuntimeError, match="Provider is closed"): + async with model.provider: + pytest.fail("A closed provider must not be entered") + assert len(clients) == 1 + assert len(service.requests) == 1 + + +async def test_close_unused_provider_does_not_create_client(clients: list[httpx2.AsyncClient]) -> None: + provider = republic.get_provider("openai") + await provider.close() + await provider.close() + assert clients == [] + assert provider.is_closed + with pytest.raises(RuntimeError, match="Provider is closed"): + await provider.get_model("test").chat("Hi") + assert clients == [] + + +@pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) +async def test_provider_context_closes_client_when_body_fails( + failure: type[BaseException], clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + service.reply_json({"output": []}) + provider = republic.get_provider("openai") + with pytest.raises(failure): + async with provider: + await provider.get_model("test").chat("Hi") + raise failure + assert provider.is_closed + assert clients[0].is_closed + + +async def test_request_failure_leaves_owned_client_available_until_provider_closes( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + async with republic.get_provider("openai") as provider: + model = provider.get_model("test") + service.reply_json({"error": "bad request"}, status_code=400) + with pytest.raises(republic.errors.APIStatusError): + await model.chat("Hi") + assert not clients[0].is_closed + service.reply_json({"output": []}) + await model.chat("Try again") + assert len(clients) == 1 + assert clients[0].is_closed + + +async def test_early_stream_exit_closes_response_but_keeps_client( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + service.reply_events([{"type": "response.output_text.delta", "delta": "partial"}]) + response = service._responses[0] + async with republic.get_provider("openai") as provider: + model = provider.get_model("test") + async with model.stream("Hi") as stream: + async for event in stream: + assert isinstance(event, republic.events.TextDelta) + break + assert response.is_closed + assert not clients[0].is_closed + service.reply_json({"output": []}) + await model.chat("Again") + assert len(clients) == 1 + assert clients[0].is_closed + + +async def test_external_client_survives_provider_close_and_can_be_shared(service: FakeService) -> None: + async with service.client() as client: + async with republic.get_provider("openai", http_client=client, timeout=17) as first: + service.reply_json({"output": []}) + await first.get_model("one").chat("Hi") + await first.close() + assert first.is_closed + assert not client.is_closed + assert client.timeout.read != 17 + with pytest.raises(RuntimeError, match="Provider is closed"): + await first.get_model("one").chat("Again") + async with republic.get_provider("openai", http_client=client) as second: + service.reply_json({"output": []}) + await second.get_model("two").chat("Hi") + assert not client.is_closed + assert client.is_closed + + +async def test_owned_client_rejects_other_event_loop_without_closing( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + service.reply_json({"output": []}) + async with republic.get_provider("openai") as provider: + model = provider.get_model("test") + await model.chat("Hi") + + async def elsewhere() -> None: + with pytest.raises(RuntimeError, match="event loop"): + await model.chat("Wrong loop") + with pytest.raises(RuntimeError, match="event loop"): + await provider.close() + with pytest.raises(RuntimeError, match="event loop"): + async with provider: + pytest.fail("A used provider must not be entered on another loop") + + await asyncio.to_thread(lambda: asyncio.run(elsewhere())) + assert not provider.is_closed + assert not clients[0].is_closed + service.reply_json({"output": []}) + await model.chat("Original loop") + assert len(clients) == 1 + assert len(service.requests) == 2 + assert clients[0].is_closed + + +def test_provider_can_be_constructed_before_event_loop_starts( + clients: list[httpx2.AsyncClient], service: FakeService +) -> None: + provider = republic.get_provider("openai") + assert clients == [] + service.reply_json({"output": []}) + + async def use() -> None: + async with provider: + await provider.get_model("test").chat("Hi") + + asyncio.run(use()) + assert provider.is_closed + assert clients[0].is_closed