forked from phoenix-oss/llama-stack-mirror
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
from typing import Protocol
|
|
|
|
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: ...
|