mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/lmproto: add OpenAI-compatible models() list API
Decode GET /models into ModelsResponse and return sorted unique ids.
This commit is contained in:
@@ -21,6 +21,8 @@ from .model import FinishReason as FinishReason
|
||||
from .model import GrammarSyntax as GrammarSyntax
|
||||
from .model import JsonObjectResponseFormat as JsonObjectResponseFormat
|
||||
from .model import Modality as Modality
|
||||
from .model import ModelObject as ModelObject
|
||||
from .model import ModelsResponse as ModelsResponse
|
||||
from .model import NamedChatMessage as NamedChatMessage
|
||||
from .model import NamedRole as NamedRole
|
||||
from .model import ReasoningEffort as ReasoningEffort
|
||||
|
||||
@@ -15,6 +15,7 @@ from .model import (
|
||||
ChatCompletionChunk,
|
||||
ChatCompletionResponse,
|
||||
ChatCompletionsParam,
|
||||
ModelsResponse,
|
||||
StreamToolCallDelta,
|
||||
)
|
||||
|
||||
@@ -95,6 +96,7 @@ class BaseOpenAIClient:
|
||||
encoder: msgspec.json.Encoder
|
||||
decoder: msgspec.json.Decoder[ChatCompletionResponse]
|
||||
chunk_decoder: msgspec.json.Decoder[ChatCompletionChunk]
|
||||
models_decoder: msgspec.json.Decoder[ModelsResponse]
|
||||
|
||||
def __init__(self, config: OpenAIConfig) -> None:
|
||||
self.session = niquests.AsyncSession(
|
||||
@@ -104,6 +106,25 @@ class BaseOpenAIClient:
|
||||
self.encoder = msgspec.json.Encoder()
|
||||
self.decoder = msgspec.json.Decoder(ChatCompletionResponse)
|
||||
self.chunk_decoder = msgspec.json.Decoder(ChatCompletionChunk)
|
||||
self.models_decoder = msgspec.json.Decoder(ModelsResponse)
|
||||
|
||||
async def models(self) -> list[str]:
|
||||
"""List model ids via OpenAI-compatible ``GET /models``.
|
||||
|
||||
Returns sorted unique ``id`` values from the response ``data`` array.
|
||||
"""
|
||||
resp = await self.session.get("/models", stream=False)
|
||||
await self._ensure_ok(resp)
|
||||
body = await read_response_body(resp)
|
||||
if body is None:
|
||||
msg = "models response body is empty"
|
||||
raise RuntimeError(msg)
|
||||
if not isinstance(body, (bytes, bytearray)):
|
||||
msg = f"models response body has unexpected type {type(body)!r}"
|
||||
raise TypeError(msg)
|
||||
parsed = self.models_decoder.decode(bytes(body))
|
||||
ids = {item.id for item in parsed.data if item.id}
|
||||
return sorted(ids)
|
||||
|
||||
async def _ensure_ok(self, resp: object) -> None:
|
||||
"""Raise if the HTTP response is an error (stream or non-stream)."""
|
||||
|
||||
@@ -258,3 +258,19 @@ class ChatCompletionChunk(Struct):
|
||||
choices: list[ChunkChoice]
|
||||
# Final usage chunk may omit or null usage depending on provider.
|
||||
usage: dict[str, Any] | None | Unset = UNSET
|
||||
|
||||
|
||||
class ModelObject(Struct):
|
||||
"""One entry from OpenAI-compatible ``GET /models``."""
|
||||
|
||||
id: str
|
||||
object: Literal["model"] | Unset = UNSET
|
||||
created: int | Unset = UNSET
|
||||
owned_by: str | Unset = UNSET
|
||||
|
||||
|
||||
class ModelsResponse(Struct):
|
||||
"""OpenAI-compatible ``GET /models`` list body."""
|
||||
|
||||
data: list[ModelObject]
|
||||
object: Literal["list"] | Unset = UNSET
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import msgspec
|
||||
import pytest
|
||||
|
||||
from plyngent.lmproto.openai_compatible.client import OpenAIClient
|
||||
from plyngent.lmproto.openai_compatible.config import OpenAIConfig
|
||||
from plyngent.lmproto.openai_compatible.model import ModelObject, ModelsResponse
|
||||
|
||||
|
||||
def test_models_response_decode() -> None:
|
||||
raw = b'{"object":"list","data":[{"id":"b","object":"model"},{"id":"a","owned_by":"x"}]}'
|
||||
parsed = msgspec.json.decode(raw, type=ModelsResponse)
|
||||
assert [m.id for m in parsed.data] == ["b", "a"]
|
||||
assert isinstance(parsed.data[0], ModelObject)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_models_sorted_unique(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
client = OpenAIClient(OpenAIConfig(access_key_or_token="sk", base_url="https://example/v1"))
|
||||
|
||||
class _Resp:
|
||||
status_code = 200
|
||||
|
||||
@property
|
||||
def content(self) -> bytes:
|
||||
return (
|
||||
b'{"object":"list","data":['
|
||||
b'{"id":"m2"},{"id":"m1"},{"id":"m2"},{"id":""}'
|
||||
b"]}"
|
||||
)
|
||||
|
||||
async def fake_get(path: str, **kwargs: object) -> _Resp:
|
||||
assert path == "/models"
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(client.session, "get", fake_get)
|
||||
ids = await client.models()
|
||||
assert ids == ["m1", "m2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_models_http_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
client = OpenAIClient(OpenAIConfig(access_key_or_token="sk", base_url="https://example/v1"))
|
||||
|
||||
class _Resp:
|
||||
status_code = 401
|
||||
|
||||
@property
|
||||
def content(self) -> bytes:
|
||||
return b'{"error":"nope"}'
|
||||
|
||||
def close(self) -> None:
|
||||
return None
|
||||
|
||||
async def fake_get(path: str, **kwargs: object) -> _Resp:
|
||||
del path, kwargs
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(client.session, "get", fake_get)
|
||||
with pytest.raises(RuntimeError, match="HTTP 401"):
|
||||
_ = await client.models()
|
||||
Reference in New Issue
Block a user