llama-stack-mirror/llama_stack/apis/models/models.py
Ashwin Bharambe 7f1160296c Updates to server.py to clean up streaming vs non-streaming stuff
Also make sure agent turn create is correctly marked
2024-10-08 17:23:42 -07:00

37 lines
1.3 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 List, Optional, Protocol
from llama_models.schema_utils import json_schema_type, webmethod
from pydantic import BaseModel, Field
@json_schema_type
class ModelDef(BaseModel):
identifier: str = Field(
description="A unique identifier for the model type",
)
llama_model: str = Field(
description="Pointer to the core Llama family model",
)
provider_id: Optional[str] = Field(
default=None, description="The provider instance which serves this model"
)
# For now, we are only supporting core llama models but as soon as finetuned
# and other custom models (for example various quantizations) are allowed, there
# will be more metadata fields here
class Models(Protocol):
@webmethod(route="/models/list", method="GET")
async def list_models(self) -> List[ModelDef]: ...
@webmethod(route="/models/get", method="GET")
async def get_model(self, identifier: str) -> Optional[ModelDef]: ...
@webmethod(route="/models/register", method="POST")
async def register_model(self, model: ModelDef) -> None: ...