stage back models api

This commit is contained in:
Xi Yan 2024-09-16 13:00:39 -07:00
parent 82b5c0460e
commit 5839c61002

View file

@ -11,31 +11,4 @@ from llama_models.schema_utils import webmethod # noqa: F401
from pydantic import BaseModel # noqa: F401
@json_schema_type
class ModelSpec(BaseModel):
model_name: str = Field(description="Name of the model")
providers_spec: Dict[Api, Dict[str, str]] = Field(
default_factory=dict,
description="Map of API to the concrete provider specs. E.g. {}".format(
{
"inference": {
"provider_type": "remote::tgi",
"url": "localhost::5555",
"api_token": "hf_xxx",
}
}
),
)
class Models(Protocol):
@webmethod(route="/models/list", method="GET")
async def list_models(self) -> List[ModelSpec]: ...
@webmethod(route="/models/get", method="GET")
async def get_model(self, model_name: str) -> ModelSpec: ...
@webmethod(route="/models/register")
async def register_model(
self, name: str, provider: Api, provider_spec: Dict[str, str]
) -> ModelSpec: ...
class Models(Protocol): ...